rubin-sim 2.6.1.dev2__py3-none-any.whl

rubin_sim/maf/stackers/base_stacker.py

@@ -0,0 +1,221 @@
+__all__ = ("StackerRegistry", "BaseStacker")
+
+import inspect
+import warnings
+
+import numpy as np
+
+
+class StackerRegistry(type):
+    """
+    Meta class for Stackers, to build a registry of stacker classes.
+    """
+
+    def __init__(cls, name, bases, dict):
+        super(StackerRegistry, cls).__init__(name, bases, dict)
+        if not hasattr(cls, "registry"):
+            cls.registry = {}
+        if not hasattr(cls, "source_dict"):
+            cls.source_dict = {}
+        modname = inspect.getmodule(cls).__name__
+        if modname.startswith("rubin_sim.maf.stackers"):
+            modname = ""
+        else:
+            if len(modname.split(".")) > 1:
+                modname = ".".join(modname.split(".")[:-1]) + "."
+            else:
+                modname = modname + "."
+        stackername = modname + name
+        if stackername in cls.registry:
+            raise Exception(
+                "Redefining stacker %s! (there are >1 stackers with the same name)" % (stackername)
+            )
+        if stackername != "BaseStacker":
+            cls.registry[stackername] = cls
+        cols_added = cls.cols_added
+        for col in cols_added:
+            cls.source_dict[col] = cls
+
+    def get_class(cls, stackername):
+        return cls.registry[stackername]
+
+    def help(cls, doc=False):
+        for stackername in sorted(cls.registry):
+            if not doc:
+                print(stackername)
+            if doc:
+                print("---- ", stackername, " ----")
+                print(cls.registry[stackername].__doc__)
+                stacker = cls.registry[stackername]()
+                print(" Columns added to SimData: ", ",".join(stacker.cols_added))
+                print(" Default columns required: ", ",".join(stacker.cols_req))
+
+
+class BaseStacker(metaclass=StackerRegistry):
+    """Base MAF Stacker.
+
+    Stackers add columns generated at run-time to the simdata array.
+    """
+
+    # List of the names of the columns generated by the Stacker.
+    cols_added = []
+
+    def __init__(self):
+        """
+        Instantiate the stacker.
+        This method should be overriden by the user.
+        This serves to define the API required by MAF.
+        """
+        # Add the list of new columns generated by the stacker
+        # as class attributes (colsAdded - above).
+
+        # List of the names of the columns required from the database
+        # (to generate the Stacker columns).
+        self.cols_req = []
+        # Optional: specify the new column types.
+        self.cols_added_dtypes = None
+        # Optional: provide a list of units for the columns
+        # defined in colsAdded.
+        self.units = [None]
+
+    def __hash__(self):
+        return None
+
+    def __eq__(self, other_stacker):
+        """
+        Evaluate if two stackers are equivalent.
+
+        This method is required to determine if metric_bundles may be
+        evaluated at the same time, on the same data.
+        """
+        # If the class names are different, they are not 'the same'.
+        if self.__class__.__name__ != other_stacker.__class__.__name__:
+            return False
+        # Otherwise, this is the same stacker class,
+        # but may be instantiated differently.
+        # We have to delve a little further,
+        # and compare the kwargs & attributes for each stacker.
+        state_now = dir(self)
+        for key in state_now:
+            if not key.startswith("_") and key != "registry" and key != "run" and key != "next":
+                if not hasattr(other_stacker, key):
+                    return False
+                # If the attribute is from numpy, assume it's an array
+                # and test it
+                if type(getattr(self, key)).__module__ == np.__name__:
+                    if not np.array_equal(getattr(self, key), getattr(other_stacker, key)):
+                        return False
+                else:
+                    if getattr(self, key) != getattr(other_stacker, key):
+                        return False
+        return True
+
+    def __ne__(self, other_stacker):
+        """
+        Evaluate if two stackers are not equal.
+        """
+        if self == other_stacker:
+            return False
+        else:
+            return True
+
+    def _add_stacker_cols(self, sim_data):
+        """
+        Add the new Stacker columns to the sim_data array.
+        If columns already present in sim_data,
+        just allows 'run' method to overwrite.
+        Returns sim_data array with these columns added
+        (so 'run' method can set their values).
+        """
+        if not hasattr(self, "cols_added_dtypes") or self.cols_added_dtypes is None:
+            self.cols_added_dtypes = [float for col in self.cols_added]
+        # Create description of new recarray.
+        newdtype = sim_data.dtype.descr
+        cols_present = [False] * len(self.cols_added)
+        for i, (col, dtype) in enumerate(zip(self.cols_added, self.cols_added_dtypes)):
+            if col in sim_data.dtype.names:
+                if sim_data[col][0] is not None:
+                    cols_present[i] = True
+                    warnings.warn(
+                        "Warning - column %s already present in sim_data, may be overwritten "
+                        "(depending on stacker)." % (col)
+                    )
+            else:
+                newdtype += [(col, dtype)]
+        new_data = np.empty(sim_data.shape, dtype=newdtype)
+        # Add references to old data.
+        for col in sim_data.dtype.names:
+            new_data[col] = sim_data[col]
+        # Were all columns present and populated with something not None?
+        # If so, then consider 'all there'.
+        if sum(cols_present) == len(self.cols_added):
+            cols_present = True
+        else:
+            cols_present = False
+        return new_data, cols_present
+
+    def run(self, sim_data, override=False):
+        """
+        Run the stacker, adding new columns.
+
+        Parameters
+        ----------
+        sim_data : `np.ndarray`, (N, M)
+            The data to be used to evaluate metrics.
+        override : `bool`, optional
+            If True, recalculate new (stacker) columns even if present.
+            If False, calculates stacker columns only if they are not present.
+
+
+        Returns
+        -------
+        sim_data : `np.ndarray`, (N, MM)
+            The input data, plus additional stacker columns.
+        """
+        # Add new columns
+        if len(sim_data) == 0:
+            return sim_data
+        sim_data, cols_present = self._add_stacker_cols(sim_data)
+        # If override is set, it means go ahead and recalculate stacker values.
+        if override:
+            cols_present = False
+        # Run the method to calculate/add new data.
+        try:
+            return self._run(sim_data, cols_present)
+        except TypeError:
+            warnings.warn(
+                "Please update the stacker %s so that the _run method matches the current API. "
+                "This will give you the option to skip re-running stackers if the columns are "
+                "already present." % (self.__class__.__name__)
+            )
+            return self._run(sim_data)
+
+    def _run(self, sim_data, cols_present=False):
+        """Do the work to add the new columns. This method should be overriden
+        in subclasses.
+
+        Parameters
+        ----------
+        sim_data: `np.ndarray`, (N, M)
+            The observation data, provided by the MAF framework.
+        cols_present: `bool`, optional
+            Flag to indicate what to do with the columns already present in
+            the data.  This will also be provided by the MAF framework --
+            but your _run method can use the value. If it is 'True' and you
+            do trust the existing value, the _run method can simply
+            return sim_data without additional calculations.
+
+        Returns
+        -------
+        sim_data : `np.ndarray`, (N, MM)
+            The simdata, with the columns added or updated
+            (or simply already present).
+        """
+        # By moving the calculation of these columns to a separate method,
+        # we add the possibility of using stackers with pandas dataframes.
+        # The _addStackerCols method won't work with dataframes, but the
+        #  _run methods are quite likely to (depending on their details),
+        #  as they are just populating columns.
+        raise NotImplementedError(
+            "Not Implemented: " "the child stackers should implement their own _run methods"
+        )

rubin_sim/maf/stackers/coord_stackers.py

@@ -0,0 +1,224 @@
+__all__ = ("ra_dec2_alt_az", "GalacticStacker", "EclipticStacker", "HealpixStacker")
+
+import healpy as hp
+import numpy as np
+from astropy import units as u
+from astropy.coordinates import SkyCoord, get_sun
+from astropy.time import Time
+from rubin_scheduler.utils import calc_lmst
+
+from .base_stacker import BaseStacker
+
+
+def wrap_ra(ra):
+    """
+    Wrap only RA values into 0-2pi (using mod).
+
+    Parameters
+    ----------
+    ra : numpy.ndarray
+        RA in radians
+
+    Returns
+    -------
+    numpy.ndarray
+        Wrapped RA values, in radians.
+    """
+    ra = ra % (2.0 * np.pi)
+    return ra
+
+
+def ra_dec2_alt_az(ra, dec, lat, lon, mjd, altonly=False):
+    """Convert RA/Dec (and telescope site lat/lon) to alt/az.
+
+    This uses simple equations and ignores aberation, precession, nutation.
+
+    Parameters
+    ----------
+    ra : `np.ndarray`, (N,)
+        RA, in radians.
+    dec : `np.ndarray`, (N,)
+        Dec, in radians. Must be same length as `ra`.
+    lat : `float`
+        Latitude of the observatory in radians.
+    lon : `float`
+        Longitude of the observatory in radians.
+    mjd : `float`
+        Modified Julian Date.
+    altonly : `bool`, optional
+        Calculate altitude only.
+
+    Returns
+    -------
+    alt : `np.ndarray`, (N,)
+        Altitude, same length as `ra` and `dec`. Radians.
+    az : `np.ndarray`, (N,)
+        Azimuth, same length as `ra` and `dec`. Radians.
+    """
+    lmst = calc_lmst(mjd, lon)
+    lmst = lmst / 12.0 * np.pi  # convert to rad
+    ha = lmst - ra
+    sindec = np.sin(dec)
+    sinlat = np.sin(lat)
+    coslat = np.cos(lat)
+    sinalt = sindec * sinlat + np.cos(dec) * coslat * np.cos(ha)
+    # make sure sinalt is in the expected range.
+    sinalt = np.where(sinalt < -1, -1, sinalt)
+    sinalt = np.where(sinalt > 1, 1, sinalt)
+    alt = np.arcsin(sinalt)
+    if altonly:
+        az = None
+    else:
+        cosaz = (sindec - np.sin(alt) * sinlat) / (np.cos(alt) * coslat)
+        cosaz = np.where(cosaz < -1, -1, cosaz)
+        cosaz = np.where(cosaz > 1, 1, cosaz)
+        az = np.arccos(cosaz)
+        signflip = np.where(np.sin(ha) > 0)
+        az[signflip] = 2.0 * np.pi - az[signflip]
+    return alt, az
+
+
+class GalacticStacker(BaseStacker):
+    """Add the galactic coordinates of each RA/Dec pointing: gall, galb
+
+    Parameters
+    ----------
+    ra_col : str, optional
+        Name of the RA column. Default fieldRA.
+    dec_col : str, optional
+        Name of the Dec column. Default fieldDec.
+    """
+
+    cols_added = ["gall", "galb"]
+
+    def __init__(self, ra_col="fieldRA", dec_col="fieldDec", degrees=True):
+        self.cols_req = [ra_col, dec_col]
+        self.ra_col = ra_col
+        self.dec_col = dec_col
+        self.degrees = degrees
+        if self.degrees:
+            self.units = ["degrees", "degrees"]
+        else:
+            self.units = ["radians", "radians"]
+
+    def _run(self, sim_data, cols_present=False):
+        # raCol and DecCol in radians, gall/b in radians.
+        if cols_present:
+            # Column already present in data;
+            # assume it is correct and does not need recalculating.
+            return sim_data
+        if self.degrees:
+            c = SkyCoord(ra=sim_data[self.ra_col] * u.deg, dec=sim_data[self.dec_col] * u.deg).transform_to(
+                "galactic"
+            )
+        else:
+            c = SkyCoord(ra=sim_data[self.ra_col] * u.rad, dec=sim_data[self.dec_col] * u.rad).transform_to(
+                "galactic"
+            )
+        sim_data["gall"] = c.l.rad
+        sim_data["galb"] = c.b.rad
+        return sim_data
+
+
+class EclipticStacker(BaseStacker):
+    """Add the ecliptic coordinates of each RA/Dec pointing:
+    eclipLat, eclipLon
+
+    Optionally subtract off the sun's ecliptic longitude and wrap.
+
+    Parameters
+    ----------
+    mjd_col : str, optional
+        Name of the MJD column. Default expMJD.
+    ra_col : str, optional
+        Name of the RA column. Default fieldRA.
+    dec_col : str, optional
+        Name of the Dec column. Default fieldDec.
+    subtract_sun_lon : bool, optional
+        Flag to subtract the sun's ecliptic longitude. Default False.
+    """
+
+    cols_added = ["eclipLat", "eclipLon"]
+
+    def __init__(
+        self,
+        mjd_col="observationStartMJD",
+        ra_col="fieldRA",
+        dec_col="fieldDec",
+        degrees=True,
+        subtract_sun_lon=False,
+    ):
+        self.cols_req = [mjd_col, ra_col, dec_col]
+        self.subtract_sun_lon = subtract_sun_lon
+        self.degrees = degrees
+        if self.degrees:
+            self.units = ["degrees", "degrees"]
+        else:
+            self.units = ["radians", "radians"]
+        self.mjd_col = mjd_col
+        self.ra_col = ra_col
+        self.dec_col = dec_col
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data;
+            # assume it is correct and does not need recalculating.
+            return sim_data
+        for i in np.arange(sim_data.size):
+            if self.degrees:
+                coord = SkyCoord(
+                    ra=sim_data[self.ra_col] * u.degree,
+                    dec=sim_data[self.dec_col] * u.degree,
+                )
+            else:
+                coord = SkyCoord(ra=sim_data[self.ra_col] * u.rad, dec=sim_data[self.dec_col] * u.rad)
+            coord_ecl = coord.geocentricmeanecliptic
+            sim_data["eclipLat"] = coord_ecl.lat.rad
+
+            if self.subtract_sun_lon:
+                times = Time(sim_data[self.mjd_col])
+                sun = get_sun(times)
+                sun_ecl = sun.geocentricmeanecliptic
+                lon = wrap_ra(coord_ecl.lon.rad - sun_ecl.lon.rad)
+                sim_data["eclipLon"] = lon
+            else:
+                sim_data["eclipLon"] = coord_ecl.lon.rad
+        if self.degrees:
+            sim_data["eclipLon"] = np.degrees(sim_data["eclipLon"])
+            sim_data["eclipLat"] = np.degrees(sim_data["eclipLat"])
+        return sim_data
+
+
+class HealpixStacker(BaseStacker):
+    """Add healpix id
+
+    Parameters
+    ----------
+    ra_col : str, optional
+        Name of the RA column. Default fieldRA.
+    dec_col : str, optional
+        Name of the Dec column. Default fieldDec.
+    nside : int
+        nside healpix to use. Default to 128.
+    """
+
+    cols_added = ["ring_healpix", "nest_healpix"]
+
+    def __init__(self, ra_col: str = "fieldRA", dec_col: str = "fieldDec", nside: int = 128):
+        self.cols_req = [ra_col, dec_col]
+        self.ra_col = ra_col
+        self.dec_col = dec_col
+        self.nside = nside
+        self.units = [f"healpix (ring, nside={nside})", f"healpix (nest, nside={nside})"]
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data;
+            # assume it is correct and does not need recalculating.
+            return sim_data
+
+        ra = sim_data[self.ra_col]
+        decl = sim_data[self.dec_col]
+        sim_data["ring_healpix"] = hp.ang2pix(self.nside, ra, decl, nest=False, lonlat=True)
+        sim_data["nest_healpix"] = hp.ang2pix(self.nside, ra, decl, nest=True, lonlat=True)
+        return sim_data

rubin_sim/maf/stackers/date_stackers.py

@@ -0,0 +1,192 @@
+__all__ = (
+    "ObservationStartDatetime64Stacker",
+    "ObservationStartTimestampStacker",
+    "DayObsStacker",
+    "DayObsMJDStacker",
+    "DayObsISOStacker",
+)
+
+import numpy as np
+import pandas as pd
+from astropy.time import Time
+
+from .base_stacker import BaseStacker
+
+
+class ObservationStartDatetime64Stacker(BaseStacker):
+    """Add the observation start time as a numpy.datetime64."""
+
+    cols_added = ["observationStartDatetime64"]
+
+    def __init__(
+        self,
+        mjd_col="observationStartMJD",
+    ):
+        self.mjd_col = mjd_col
+        self.cols_req = [self.mjd_col]
+        self.units = [None]
+        self.cols_added_dtypes = ["datetime64[ns]"]
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data; assume it is correct and does not
+            # need recalculating.
+            return sim_data
+
+        sim_data["observationStartDatetime64"] = Time(sim_data[self.mjd_col], format="mjd").datetime64
+
+        return sim_data
+
+
+class ObservationStartTimestampStacker(BaseStacker):
+    """Add the observation start time as a pandas.Timestamp."""
+
+    cols_added = ["start_timestamp"]
+
+    def __init__(
+        self,
+        mjd_col="observationStartMJD",
+    ):
+        self.mjd_col = mjd_col
+        self.cols_req = [self.mjd_col]
+        self.units = [None]
+        self.cols_added_dtypes = ["O"]
+
+    def run(self, sim_data, override=False):
+        # Override the run from the base class, not _run,
+        # because the implementation
+        # of _add_stackers_cols in run in the base closs
+        # clobbers the type for the new column,
+        # which we really need.
+
+        visits = sim_data if isinstance(sim_data, pd.DataFrame) else pd.DataFrame(sim_data)
+
+        if "start_timestamp" in visits.columns and not override:
+            return sim_data
+
+        if len(visits[self.mjd_col]) > 0:
+            visits["start_timestamp"] = pd.to_datetime(
+                visits[self.mjd_col] + 2400000.5, origin="julian", unit="D", utc=True
+            )
+        else:
+            # If we are passed an empty series, be sure to return add an empty
+            # series of the correct type back.
+            # This is handy if the result is being passed to bokeh
+            # for plotting.
+            visits["start_timestamp"] = pd.to_datetime(2460000.5, origin="julian", unit="D", utc=True)
+
+        match sim_data:
+            case pd.DataFrame():
+                return visits
+            case dict():
+                return visits.to_dict()
+            case _:
+                return visits.to_records(index=False)
+
+
+def _compute_day_obs_mjd(mjd):
+    day_obs_mjd = np.floor(mjd - 0.5).astype("int")
+    return day_obs_mjd
+
+
+def _compute_day_obs_astropy_time(mjd):
+    day_obs_time = Time(_compute_day_obs_mjd(mjd), format="mjd")
+    return day_obs_time
+
+
+def _compute_day_obs_iso8601(mjd):
+    iso_times = _compute_day_obs_astropy_time(mjd).iso
+
+    # Work both for mjd as a scalar and a numpy array
+    if isinstance(iso_times, str):
+        day_obs_iso = iso_times[:10]
+    else:
+        day_obs_iso = np.array([d[:10] for d in iso_times])
+
+    return day_obs_iso
+
+
+def _compute_day_obs_int(mjd):
+    day_obs_int = np.array([d.replace("-", "") for d in _compute_day_obs_iso8601(mjd)])
+
+    return day_obs_int
+
+
+class DayObsStacker(BaseStacker):
+    """Add dayObs as, as defined by SITCOMTN-32.
+
+    Parameters
+    ----------
+    mjd_col : `str`
+        The column with the observatin start MJD.
+    """
+
+    cols_added = ["dayObs"]
+
+    def __init__(self, mjd_col="observationStartMJD"):
+        self.mjd_col = mjd_col
+        self.cols_req = [self.mjd_col]
+        self.units = ["days"]
+        self.cols_added_dtypes = [int]
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data; assume it is correct and does not
+            # need recalculating.
+            return sim_data
+
+        sim_data[self.cols_added[0]] = _compute_day_obs_int(sim_data[self.mjd_col])
+        return sim_data
+
+
+class DayObsMJDStacker(BaseStacker):
+    """Add dayObs defined by SITCOMTN-32, as an MJD.
+
+    Parameters
+    ----------
+    mjd_col : `str`
+        The column with the observatin start MJD.
+    """
+
+    cols_added = ["day_obs_mjd"]
+
+    def __init__(self, mjd_col="observationStartMJD"):
+        self.mjd_col = mjd_col
+        self.cols_req = [self.mjd_col]
+        self.units = ["days"]
+        self.cols_added_dtypes = [int]
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data; assume it is correct and does not
+            # need recalculating.
+            return sim_data
+
+        sim_data[self.cols_added[0]] = _compute_day_obs_mjd(sim_data[self.mjd_col])
+        return sim_data
+
+
+class DayObsISOStacker(BaseStacker):
+    """Add dayObs as defined by SITCOMTN-32, in ISO 8601 format.
+
+    Parameters
+    ----------
+    mjd_col : `str`
+        The column with the observatin start MJD."""
+
+    cols_added = ["day_obs_iso8601"]
+
+    def __init__(self, mjd_col="observationStartMJD"):
+        self.mjd_col = mjd_col
+        self.cols_req = [self.mjd_col]
+        self.units = [None]
+        self.cols_added_dtypes = [(str, 10)]
+
+    def _run(self, sim_data, cols_present=False):
+        if cols_present:
+            # Column already present in data; assume it is correct and does not
+            # need recalculating.
+            return sim_data
+
+        sim_data[self.cols_added[0]] = _compute_day_obs_iso8601(sim_data[self.mjd_col])
+        return sim_data