fit_source_map.py

import os
import shutil
import sys
from functools import partial
from typing import cast

import h5py
import numpy as np
import sqlalchemy as sqy
from astropy import constants as const
from astropy import units as u
from mpi4py import MPI
from pixell import enmap
from sotodlib.core import AxisManager, Context
from sotodlib.site_pipeline import jobdb
from sotodlib.site_pipeline.jobdb import Job

import lat_beams.fitting.map.bessel as fb
import lat_beams.fitting.map.gauss as fg
from lat_beams.beam_utils import (
    crop_maps,
    estimate_cent,
    get_fwhm_radial_bins,
    process_model,
    radial_profile,
)
from lat_beams.fitting.map.base import make_guess
from lat_beams.plotting import plot_map
from lat_beams.utils import (
    ErrCode,
    fail,
    get_args_cfg,
    init_log,
    set_tag,
    setup_cfg,
    setup_jobs,
    setup_paths,
    update_jobs_retry,
)

comm = MPI.COMM_WORLD
myrank = comm.Get_rank()
nproc = comm.Get_size()


def get_jobdict(jdb):
    jobdict = {
        f"{job.tags['obs_id']}-{job.tags['wafer_slot']}-{job.tags['stream_id']}-{job.tags['array']}-{job.tags['band']}-{job.tags['comps']}-{job.tags['split']}": job
        for job in jdb.get_jobs(jclass="fit_map")
    }
    return jobdict


def get_jobit(jdb, det_splits, obs_ids, start_time, stop_time):
    maplist = jdb.get_jobs(jclass="beam_map", jstate="done", locked=False)
    maplist = np.array_split(maplist, nproc)[myrank]
    timestamps = np.array(
        [float(info.tags["obs_id"].split("_")[1]) for info in maplist]
    )
    tmsk = (timestamps >= start_time) * (timestamps < stop_time)
    maplist = maplist[tmsk]
    if obs_ids is not None:
        sub_ids = [obs_id.split(":") for obs_id in obs_ids]
        obs_list = [obs_id[0] for obs_id in sub_ids]
        ws_list = []
        ba_list = []
        maplist = [info for info in maplist if info.tags["obs_id"] in obs_list]
        for info in maplist:
            idx = obs_list.index(info.tags["obs_id"])
            sub_id = sub_ids[idx]
            ws_list += [sub_id[1] if len(sub_id) > 1 else None]
            ba_list += [sub_id[2] if len(sub_id) > 2 else None]
    else:
        ws_list = [None] * len(maplist)
        ba_list = [None] * len(maplist)

    to_ret = []
    for info, wsl, bal in zip(maplist, ws_list, ba_list):
        if cfg.tube_slots and not any(
            ot in info.tags["obs_id"] for ot in cfg.tube_slots
        ):
            continue
        if wsl is not None and info.tags["wafer_slot"] not in wsl:
            continue
        if bal is not None and info.tags["band"] not in bal:
            continue
        to_ret += [
            (info, ds) for ds in info.tags["splits"].split(",") if ds in det_splits
        ]
    return to_ret


def get_jobstr(mjob):
    mjob, split = mjob
    job_str = f"{mjob.tags['obs_id']}-{mjob.tags['wafer_slot']}-{mjob.tags['stream_id']}-{mjob.tags['array']}-{mjob.tags['band']}-{mjob.tags['comps']}-{split}"
    return job_str


def get_tags(mjob):
    mjob, split = mjob
    tags = {
        "obs_id": mjob.tags["obs_id"],
        "wafer_slot": mjob.tags["wafer_slot"],
        "stream_id": mjob.tags["stream_id"],
        "array": mjob.tags["array"],
        "band": mjob.tags["band"],
        "comps": mjob.tags["comps"],
        "source": mjob.tags["source"],
        "message": "",
        "errcode": 0,
        "resid": "",
        "resid_weights": "",
        "config": "",
        "split": split,
    }
    return tags


# Setup logger
logger = init_log()
if logger.extra is None:
    raise ValueError("Logger doesn't have adapter set up!")
logger.extra = cast(dict, logger.extra)

# Get settings
args, cfg_dict = get_args_cfg()
cfg, cfg_str = setup_cfg(
    args,
    cfg_dict,
    {
        "fit_source_list": "source_list",
        "map_mask_size": "mask_size",
        "fwhm_tol_map": "fwhm_tol",
    },
)
cfg.aperature *= u.m

# Setup folders and files
plot_dir, data_dir = setup_paths(
    cfg.root_dir,
    "beams",
    cfg.tel,
    f"{cfg.pointing_type}{(cfg.append!='')*'_'}{cfg.append}{(cfg.single_det)*'_single_det'}",
)
outfile = None
if myrank == 0:
    of_path = os.path.join(data_dir, f"beam_pars{cfg.test_append}.h5")
    of_path_noa = os.path.join(data_dir, f"beam_pars.h5")
    if os.path.isfile(of_path_noa) and cfg.copy_fits_test and of_path != of_path_noa:
        shutil.copyfile(of_path_noa, of_path)
    outfile = h5py.File(of_path, "a")

# Get det splits
det_split_names = ["full"] + cfg.det_splits

# Get the jobs, make them if we need to
ctx = Context(cfg.ctx_path)
if ctx.obsdb is None:
    raise ValueError("No obsdb in context!")
jdb, all_jobs = setup_jobs(
    comm,
    data_dir,
    "fit_map",
    get_jobdict,
    partial(
        get_jobit,
        det_splits=det_split_names,
        obs_ids=args.obs_ids,
        start_time=cfg.start_time,
        stop_time=cfg.stop_time,
    ),
    get_jobstr,
    get_tags,
    cfg.source_list,
    args.overwrite,
    args.retry_failed,
    args.job_memory,
    args.job_memory_buffer,
    False,
    logger,
    cfg.test_append,
)

# Even things out
joblist = np.array_split(all_jobs, nproc)[myrank].tolist()
n_maps = comm.allgather(len(joblist))
max_maps = np.max(n_maps)
if n_maps[0] != max_maps:
    raise ValueError("Root doesn't have max maps!")
joblist += [None] * (1 + max_maps - len(joblist))

profiler = None
if args.profile:
    from pyinstrument import Profiler

    profiler = Profiler()
    profiler.start()
    logger.info("Running in profiler mode! Only one map per proc will be kept")
    joblist = [joblist[0], None]

to_save = (None, None)
map_jobs = jdb.get_jobs(jclass="beam_map", jstate="done")
map_jobdict = {
    f"{job.tags['obs_id']}-{job.tags['wafer_slot']}-{job.tags['stream_id']}-{job.tags['array']}-{job.tags['band']}": job
    for job in map_jobs
}
job = None
for i, j in enumerate(joblist):
    comm.barrier()
    sys.stdout.flush()
    to_save = comm.gather(to_save, root=0)
    if myrank == 0 and to_save is not None and outfile is not None:
        for aman, path in to_save:
            if aman is None:
                to_save = (None, None)
                continue
            aman.save(outfile, path, overwrite=True)
        outfile.flush()
    comm.barrier()

    if job is not None:
        to_save = (None, None)
        logger.debug("Writing to db")
        update_jobs_retry(jdb, [job], nproc * 10, logger)
    job = None
    if j is not None:
        with jdb.session_scope() as session:
            job = session.get(Job, j.id)
            session.expunge(job)

    if job is None:
        to_save = (None, None)
        continue

    job.mark_visited()
    obs_id = job.tags["obs_id"]
    sid = job.tags["stream_id"]
    ufm = job.tags["array"]
    ws = job.tags["wafer_slot"]
    band = job.tags["band"]
    split = job.tags["split"]
    logger.extra["extra"] = (
        f" [{obs_id} {ufm} {band} {split} ({i+1}/{len(joblist) - 1})]"
    )
    logger.log(25, "Fitting")

    # Get map job
    job_str = f"{obs_id}-{ws}-{sid}-{ufm}-{band}"
    if job_str not in map_jobdict:
        msg = "No map job"
        logger.debug("%s", msg)
        set_tag(job, "message", msg)
        job.jstate = cast(sqy.Column[str], jobdb.JState.failed)
        to_save = (None, None)
        continue
    map_job = map_jobdict[job_str]
    comps = map_job.tags["comps"]

    # Save config info
    set_tag(job, "config", cfg_str)
    set_tag(job, "comps", comps)

    # Figure out paths
    map_path = os.path.join(
        data_dir, map_job.tags["solved"].format(split=job.tags["split"])
    )
    ivar_path = os.path.join(
        data_dir, map_job.tags["weights"].format(split=job.tags["split"])
    )
    # Load the maps
    try:
        solved = enmap.read_map(map_path)[0]
        solved = cast(enmap.ndmap, solved)
        weights = enmap.read_map(ivar_path)[0][0]
        weights = cast(enmap.ndmap, weights)
    except FileNotFoundError:
        msg = "Missing map files"
        fail(job, ErrCode.MAP_MISSING, msg, logger)
        to_save = (None, None)
        continue
    if solved.wcs is None:
        raise ValueError("WCS is None")
    pixsize = 3600 * solved.wcs.wcs.cdelt[1]  # type: ignore

    # Check if this is a bogus map
    if np.sum(~(weights == 0)) == 0:
        msg = "Weights all 0"
        fail(job, ErrCode.ZERO_IVAR, msg, logger)
        to_save = (None, None)
        continue

    # Estimatr SNR
    snr_extent_pix = int(cfg.snr_extent // pixsize)
    cent = estimate_cent(solved, weights, cfg.smooth_kern / pixsize, cfg.buf)
    sig = solved[cent]
    noise = solved.copy()
    xmin = max(0, cent[0] - snr_extent_pix)
    xmax = min(solved.shape[0], cent[0] + snr_extent_pix)
    ymin = max(0, cent[1] - snr_extent_pix)
    ymax = min(solved.shape[1], cent[1] + snr_extent_pix)
    noise[xmin:xmax, ymin:ymax] = np.nan
    noise = np.nanstd(np.diff(noise))
    snr = sig / noise

    if snr < cfg.min_snr:
        msg = "Data SNR too low"
        fail(job, ErrCode.SNR_LOW, msg, logger)
        to_save = (None, None)
        continue

    # Slice things
    solved, weights = crop_maps([solved, weights], cent, int(cfg.extent // pixsize))
    posmap = enmap.posmap(solved.shape, solved.wcs)
    cent = estimate_cent(solved, weights, cfg.smooth_kern / pixsize, cfg.buf_cropped)
    fscale_fac = 90.0 / float(band[1:]) if cfg.apply_fscale else 1
    band_mask_size = np.deg2rad(fscale_fac * cfg.mask_size)

    # Make weights and zero things out
    weights[~np.isfinite(weights)] = 0
    weights[~np.isfinite(solved)] = 0
    solved[~np.isfinite(solved)] = 0

    # Setup aman for output
    aman = AxisManager()
    aman.wrap("noise", noise * u.pW)

    # Fit gaussian model
    cent, smoothed = estimate_cent(
        solved, weights, cfg.smooth_kern / pixsize, cfg.buf_cropped, True
    )
    maxval = np.max(solved)
    guess = make_guess(
        amp=np.nan_to_num(smoothed[cent].item(), True, maxval, maxval, maxval),
        fwhm_xi=np.deg2rad(cfg.nominal_fwhm[band] / 60.0),
        fwhm_eta=np.deg2rad(cfg.nominal_fwhm[band] / 60.0),
        xi0=posmap[1][cent[0], cent[1]].item(),
        eta0=posmap[0][cent[0], cent[1]].item(),
        phi=0,
        off=0,
    )
    gauss_params, model = fg.fit_gauss_map(
        solved,
        weights,
        posmap,
        guess,
        "pW",
        cfg.sym_gauss,
        -1,
    )
    if gauss_params is None or model is None:
        msg = "Gauss fit failed"
        fail(job, ErrCode.FIT_FAILED, msg, logger)
        to_save = (None, None)
        continue

    # Compute the gaussian model
    model = fg.gaussian2d_from_aman(posmap, gauss_params)

    # Check clipping
    c = np.unravel_index(np.argmax(model, axis=None), model.shape)
    min_c_dist = np.min(np.hstack((c, np.array(solved.shape) - np.array(c)))) * pixsize
    if min_c_dist < 120 * cfg.nominal_fwhm[band]:
        msg = "Source too close to edge of map"
        fail(job, ErrCode.CLOSE_TO_EDGE, msg, logger)
        to_save = (None, None)
        continue

    # Get FWHM from data
    cent = (int(c[-1]), int(c[-2]))
    rprof = radial_profile(solved - gauss_params.off.value, cent)
    r = np.linspace(0, len(rprof), len(rprof)) * pixsize
    rmsk = r < 3 * 60 * cfg.nominal_fwhm[band] / 2.355
    data_fwhm = get_fwhm_radial_bins(r[rmsk], rprof[rmsk], interpolate=True) * u.arcsec
    aman.wrap("data_fwhm", data_fwhm)
    aman.wrap("r", r * u.arcsec)
    aman.wrap("rprof", rprof * u.pW)

    # FWHM check
    if (
        np.isnan(data_fwhm)
        or abs(1 - data_fwhm.value / (60 * cfg.nominal_fwhm[band])) > cfg.fwhm_tol
    ):
        msg = "Data FWHM out of tolerance"
        fail(job, ErrCode.FWHM_TOL, msg, logger)
        to_save = (None, None)
        continue

    # Process and save fit model
    gauss_params = process_model(
        gauss_params,
        solved - gauss_params.off.value,
        model - gauss_params.off.value,
        float(noise),
        cfg.min_snr,
        (cent[-2], cent[-1]),
        u.pW,
        pixsize,
        data_fwhm,
        cfg.min_sigma,
        job,
        logger,
    )
    if gauss_params is None:
        to_save = (None, None)
        continue
    aman.wrap("gauss", gauss_params)
    off = gauss_params.off.value
    for to_parent in ["amp", "off", "xi0", "eta0"]:
        aman.wrap(to_parent, gauss_params[to_parent])
    aman.wrap("final_model", "gauss")

    # Get bessel beam if we want
    if cfg.bessel_beam:
        try:
            bessel_beam_params, model = fb.fit_bessel_map(
                solved,
                weights,
                posmap,
                gauss_params,
                "pW",
                cfg.n_bessel,
                cfg.n_multipoles,
                cfg.aperature,
                const.c / (float(band[1:]) * u.GHz),  # type: ignore
                band_mask_size,
                np.log((10**cfg.bessel_wing_n_sigma) / fscale_fac) / np.log(10),
                cfg.skip_multipoles,
            )
        except Exception as e:
            msg = f"Bessel fit failed with error {str(e)}"
            fail(job, ErrCode.FIT_FAILED, msg, logger)
            to_save = (None, None)
            continue
        if bessel_beam_params is None or model is None:
            msg = "Bessel fit failed"
            fail(job, ErrCode.FIT_FAILED, msg, logger)
            to_save = (None, None)
            continue

        off = bessel_beam_params.off.value
        bessel_beam_params = process_model(
            bessel_beam_params,
            solved - bessel_beam_params.off.value,
            model - bessel_beam_params.off.value,
            float(noise),
            cfg.min_snr,
            (cent[-2], cent[-1]),
            u.pW,
            pixsize,
            data_fwhm,
            cfg.min_sigma,
            job,
            logger,
        )
        if bessel_beam_params is None:
            to_save = (None, None)
            continue
        aman.wrap("bessel", bessel_beam_params)
        aman.final_model = "bessel"

    # Save residual
    resid = solved.copy()
    resid -= model
    fname = map_path  # map_job.tags["solved"]
    enmap.write_map(
        os.path.join(data_dir, f"{'_'.join(fname.split('_')[:-1])}_resid.fits"),
        resid,
        "fits",
        allow_modify=True,
    )
    enmap.write_map(
        os.path.join(data_dir, f"{'_'.join(fname.split('_')[:-1])}_resid_weights.fits"),
        weights,
        "fits",
        allow_modify=True,
    )
    set_tag(
        job,
        "resid",
        f"{'_'.join(fname.split('_')[:-1])}_resid.fits",
    )
    set_tag(
        job,
        "resid_weights",
        f"{'_'.join(fname.split('_')[:-1])}_resid_weights.fits",
    )

    # Plot
    obs = ctx.obsdb.get(obs_id)
    ufm_plot_dir = os.path.join(
        plot_dir,
        f'{job.tags["source"]}{cfg.test_append}',
        str(obs["timestamp"])[:5],
        obs_id,
        job.tags["array"],
        split,
    )
    os.makedirs(ufm_plot_dir, exist_ok=True)
    [[dec_min, ra_min], [dec_max, ra_max]] = 3600 * np.rad2deg(
        solved.corners(corner=False)
    )
    plt_cent = (aman.xi0.to(u.arcsec).value, aman.eta0.to(u.arcsec).value)
    norm = float(1.0 / aman.amp.value)
    posmap = cast(enmap.ndmap, np.rad2deg(posmap) * 3600)
    maps = [(solved, "data"), (model - off, "model"), (resid, "resid")]
    for log in (False, True):
        _norm = norm if log else 1
        plot_map(
            [m[0] * _norm for m in maps],
            posmap,
            pixsize,
            cfg.extent,
            plt_cent,
            ufm_plot_dir,
            f"{obs_id} {ufm} {band} {split} T fit",
            [m[1] for m in maps],
            log_thresh=cfg.log_thresh,
            append="",
            units='"',
            log=log,
        )

    # Save
    aman.wrap("data_solid_angle_corr", aman[aman.final_model].data_solid_angle_corr)
    aman_path = os.path.join(obs_id, ufm, band, split)
    to_save = (aman, aman_path)

    set_tag(job, "message", "Success")
    job.jstate = cast(sqy.Column[str], jobdb.JState.done)

comm.barrier()
logger.extra["extra"] = ""

if args.profile and profiler is not None:
    profiler.stop()
    profiler.write_html(f"profile_{myrank}.html")

if outfile is not None:
    outfile.close()
sys.stdout.flush()
logger.info("Done with all fits")