psfellipse_srcs

#!/usr/bin/env python


# Copyright (C) 2022 Smithsonian Astrophysical Observatory
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#

'Generate ellipse regions that enclose ~90% of the PSF'


import sys
import os
from tempfile import NamedTemporaryFile

import numpy as np

import ciao_contrib.logger_wrapper as lw
from pycrates import read_file

AcisEdge = 32    # pixels, dither pattern


toolname = "psfellipse_srcs"
__revision__ = "07 July 2022"

lw.initialize_logger(toolname)
lgr = lw.get_logger(toolname)
verb0 = lgr.verbose0
verb1 = lgr.verbose1


def lookup_psffile_in_caldb(calfile):
    """
    Lookup psffile -- a noop right now.  Must supply filename
    """

    return calfile


def _make_not_coords_table(theta, phi):
    'Create temp file with new coords'

    notdetx = np.array(theta) * np.cos(np.radians(phi))
    notdety = np.array(theta) * np.sin(np.radians(phi))

    tmproot = NamedTemporaryFile(dir=os.environ["ASCDS_WORK_PATH"],
                                 delete=False)
    tmproot.close()

    outfile = tmproot.name

    with open(outfile, "w") as fp:
        fp.write("#not_detx\tnot_dety\n")
        for x, y in zip(notdetx, notdety):
            fp.write(f"{x}\t{y}\n")

    return outfile


def get_psf_size(psffile, theta, phi, pixsize, keywords):
    """
    Get PSF size for given fration/theta/phi
    """
    from ciao_contrib.runtool import dmimgpick

    tmpfile = _make_not_coords_table(theta, phi)

    def _run_dmimgpick(block_name):
        dmimgpick(tmpfile, psffile+f"[{block_name}]", tmpout,
                  method="weight", clobber=True)
        vals = read_file(tmpout).get_column(block_name).values
        vals[~np.isfinite(vals)] = 1

        return vals

    try:
        tmproot = NamedTemporaryFile(dir=os.environ["ASCDS_WORK_PATH"],
                                     delete=False)
        tmproot.close()
        tmpout = tmproot.name

        major = _run_dmimgpick("MAJOR_AXIS")
        minor = _run_dmimgpick("MINOR_AXIS")
        angle = _run_dmimgpick("ANGLE")

    finally:
        if os.path.exists(tmpfile):
            os.unlink(tmpfile)
        if os.path.exists(tmpout):
            os.unlink(tmpout)

    # Convert from arcsec to pixels
    cnv = 1.0/(pixsize)
    major *= cnv
    minor *= cnv

    # Correct angle for ROLL
    if "ROLL_PNT" not in keywords:
        verb0("WARNING: could not find ROLL_PNT keyword, " +
              "rotation angle may be wrong")
    else:
        angle += 90 - float(keywords["ROLL_PNT"])

    radii = [np.array(xx) for xx in zip(major, minor)]

    return np.array(radii), np.array(angle)


def get_keys_from_crate(myfile):
    """
    Get all keys into dictionary
    """
    kk = myfile.get_keynames()
    mykeys = {k: myfile.get_key_value(k) for k in kk}
    return mykeys


def write_output(outfile, radii, angle, ra_vals, dec_vals, coords, mykeys):
    """
    Write output file
    """
    shape = ['ellipse'] * len(angle)
    component = range(1, len(angle) + 1)

    print(radii)

    cnam = [x.upper() for x in ['shape', 'x', 'y', 'r', 'rotang',
                                'component', 'ra', 'dec', 'theta',
                                'phi', 'detx', 'dety', "chip_id",
                                "chipx", "chipy", "near_chip_edge"]]
    cval = [shape, coords["x"], coords["y"],
            radii, angle, component, ra_vals, dec_vals,
            coords["theta"], coords["phi"], coords["detx"],
            coords["dety"], coords["chip_id"], coords["chipx"],
            coords["chipy"], coords["offchip"]]
    units = ['', 'pixel', 'pixel', 'pixel', 'deg', '', 'deg', 'deg',
             'arcmin', 'deg', "pixel", "pixel", "", "pixel", "pixel",
             ""]
    desc = ['region geometry',
            'X center',
            'Y center',
            'Radii',
            'Rotation Angle',
            'Region Number',
            'Right Ascencion',
            'Declination',
            'Off axis angle',
            'Azimuth angle',
            'Detector X',
            'Detector Y',
            'Chip ID',
            'Chip X',
            'Chip Y',
            "Is POS near or off any chip edge?"]

    import pycrates as pc
    from pytransform import LINEARTransform

    tab = pc.TABLECrate()

    tab.name = 'REGION'
    for ii in range(len(cnam)):
        cr = pc.CrateData()
        cr.name = cnam[ii]
        cr.desc = desc[ii]
        cr.unit = units[ii]
        cr.values = cval[ii]
        tab.add_column(cr)

    lt = LINEARTransform()
    lt.get_parameter("SCALE").set_value(coords["pixsize"])
    lt.get_parameter("OFFSET").set_value(0.0)
    lt.name = "CEL_R"
    rc = tab.get_column("r")
    # Crates bug won't allow transform to be set? Error on write()
    # rc._set_transform(lt)

    for k in mykeys:
        ck = pc.CrateKey()
        ck.name = k
        ck.value = mykeys[k]
        tab.add_key(ck)

    tab.write(outfile=outfile, clobber=True)


def check_chip_edge(coords, mykeys, edge):
    """
    Check if within dither size of chip edge (for now)
    """

    if mykeys["INSTRUME"].lower() == "hrc":
        coords["offchip"] = [False]*len(coords["chipx"])
        return

    coords["offchip"] = [False]*len(coords["chipx"])

    xmin = 1 + edge
    xmax = 1024 - edge

    if "FIRSTROW" not in mykeys or "NROWS" not in mykeys:
        verb0("Missing 'FIRSTROW' or 'NROWS' keyword, cannot " +
              "determine if this is a subarray dataset")

    # Handle subarray, (doesn't do windows)
    ylo = mykeys["FIRSTROW"] if "FIRSTROW" in mykeys else 1
    nrow = mykeys["NROWS"] if "NROWS" in mykeys else 1024

    ymin = ylo + edge
    ymax = (nrow-1) - edge

    def offchip(pos):
        x = pos[0]
        y = pos[1]
        if xmin < x < xmax and ymin < y < ymax:
            return False
        return True

    coords["offchip"] = list(map(offchip, zip(coords["chipx"],
                                              coords["chipy"])))


#
# Main Routine
#
@lw.handle_ciao_errors(toolname, __revision__)
def main():
    'Main routine'

    from ciao_contrib.runtool import add_tool_history
    from ciao_contrib._tools.fileio import outfile_clobber_checks
    from ciao_contrib.param_soaker import get_params
    from ciao_contrib.parse_pos import get_radec_from_pos
    from coords.chandra import cel_to_chandra

    # get parameters
    pars = get_params(toolname, "rw", sys.argv,
                      verbose={"set": lw.set_verbosity, "cmd": verb1})

    outfile_clobber_checks(pars["clobber"], pars["outfile"])

    # Parse coords
    ra_vals, dec_vals = get_radec_from_pos(pars["pos"])

    # Convert ra/dec to theta/phi and x/y
    myfile = read_file(pars["infile"])
    mykeys = get_keys_from_crate(myfile)
    coords = cel_to_chandra(mykeys, ra_vals, dec_vals)

    check_chip_edge(coords, mykeys, AcisEdge)

    # Get PSF size
    radii, angle = get_psf_size(pars["psffile"], coords["theta"],
                                coords["phi"], coords["pixsize"], mykeys)

    # Write output
    write_output(pars["outfile"], radii, angle, ra_vals, dec_vals,
                 coords, mykeys)
    # add_tool_history(pars["outfile"], toolname, pars,
    #                 toolversion=__revision__)


if __name__ == "__main__":
    main()
    sys.exit(0)