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

rubin_sim/maf/slicers/time_interval_slicers.py

@@ -0,0 +1,283 @@
+"""MAF slicers to slice into time intervals based on MJD
+
+Primarily intended for hourglass plots.
+"""
+
+# pylint: disable=too-many-arguments
+__all__ = (
+    "TimeIntervalSlicer",
+    "BlockIntervalSlicer",
+    "VisitIntervalSlicer",
+    "SlicerNotSetupError",
+)
+
+from collections import defaultdict
+
+# imports
+from functools import wraps
+
+import numpy as np
+import pandas as pd
+
+from .base_slicer import BaseSlicer
+
+
+class SlicerNotSetupError(Exception):
+    """Thrown when a slicer is not setup for the method called."""
+
+
+# interface functions
+
+# classes
+
+
+class TimeIntervalSlicer(BaseSlicer):
+    """Base for all time interval slicers.
+
+    Slices in constant time intervals.
+
+    Parameters
+    ----------
+    interval_seconds : `int`
+        Duration of slice time intervals, in seconds
+    mjd_column_name : `str`
+        Name of column on which to slice, must be in units of days
+    badval : `float`
+        Value to use for bad values in slice
+    verbose : `bool`
+        Print extra information?
+    """
+
+    def __init__(
+        self,
+        interval_seconds=90,
+        mjd_column_name="observationStartMJD",
+        badval=np.nan,
+        verbose=False,
+    ):
+        super().__init__(verbose=verbose, badval=badval)
+        self.interval_seconds = interval_seconds
+        self.mjd_column_name = mjd_column_name
+        self.columns_needed = [mjd_column_name]
+        self.sim_idxs = defaultdict(list)  # pylint: disable=invalid-name
+
+    def setup_slicer(self, sim_data, maps=None):
+        visit_mjds = sim_data[self.mjd_column_name]
+        start_mjd = np.floor(np.min(visit_mjds)).astype(int)
+        end_mjd = np.ceil(np.max(visit_mjds)).astype(int)
+        interval_days = self.interval_seconds / (24 * 60 * 60.0)
+
+        mjd_bin_edges = np.arange(start_mjd, end_mjd + interval_days, interval_days)
+
+        self.sim_idxs.update(
+            pd.DataFrame(
+                {
+                    "visit_idx": np.arange(len(visit_mjds)),
+                    "sid": np.digitize(visit_mjds, mjd_bin_edges),
+                }
+            )
+            .groupby("sid")
+            .agg(list)
+            .to_dict()["visit_idx"]
+        )
+
+        mjds = np.arange(start_mjd, end_mjd, interval_days)
+        self.slice_points["sid"] = np.arange(len(mjds))
+        self.slice_points["mjd"] = mjds
+        self.slice_points["duration"] = np.full_like(mjds, self.interval_seconds)
+        self.nslice = len(mjds)
+        self.shape = self.nslice
+        self._run_maps(maps)
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):  # pylint: disable=invalid-name
+            idxs = self.sim_idxs[islice]
+
+            try:
+                _ = idxs[0]
+            except (TypeError, IndexError):
+                idxs = [idxs]
+
+            slice_points = {
+                "sid": islice,
+                "mjd": self.slice_points["mjd"][islice],
+                "duration": self.slice_points["duration"][islice],
+            }
+
+            return {"idxs": idxs, "slice_point": slice_points}
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def __eq__(self, other_slicer):
+        """Evaluate if slicers are equivalent."""
+        if not isinstance(other_slicer, self.__class__):
+            return False
+
+        for key in ["sid", "mjd", "duration"]:
+            if not np.array_equal(other_slicer.slice_points[key], self.slice_points[key]):
+                return False
+
+        return True
+
+    def _slice_sim_data(self, *args, **kwargs):
+        raise SlicerNotSetupError()
+
+
+class BlockIntervalSlicer(TimeIntervalSlicer):
+    """Slices into intervals with common "note" values and no long gaps.
+
+    Parameters
+    ----------
+    mjd_column_name : `str`
+        Name of column on which to slice, must be in units of days
+    duration_column_name : `str`
+        Name of column with the duration of each visit (in seconds)
+    note_column_name : `str`
+        Name of column with the visit note.
+    badval : `float`
+        Value to use for bad values in slice
+    verbose : `bool`
+        Print extra information?
+    """
+
+    # Gap between visits in the same block, in hours
+    gap_tolerance = 0.1
+
+    def __init__(
+        self,
+        mjd_column_name="observationStartMJD",
+        duration_column_name="visitTime",
+        note_column_name="scheduler_note",
+        badval=np.nan,
+        verbose=False,
+    ):
+        super().__init__(verbose=verbose, badval=badval)
+        self.mjd_column_name = mjd_column_name
+        self.duration_column_name = duration_column_name
+        self.note_column_name = note_column_name
+        self.columns_needed = [
+            mjd_column_name,
+            duration_column_name,
+            note_column_name,
+        ]
+        self.sim_idxs = defaultdict(list)  # pylint: disable=invalid-name
+
+    def setup_slicer(self, sim_data, maps=None):
+        visits = pd.DataFrame(
+            sim_data,
+            index=pd.Index(np.arange(len(sim_data[self.mjd_column_name])), name="visit_idx"),
+        )
+        visits.rename(
+            columns={
+                self.mjd_column_name: "mjd",
+                self.duration_column_name: "duration",
+                self.note_column_name: "scheduler_note",
+            },
+            inplace=True,
+        )
+        # convert to hours
+        visits.sort_values("mjd", inplace=True)
+        visits["end_mjd"] = visits.mjd + visits.duration / (60 * 60 * 24.0)
+
+        same_note = visits.scheduler_note == visits.scheduler_note.shift(-1)
+        adjacent_times = visits.end_mjd + self.gap_tolerance / 24.0 > visits.mjd.shift(-1)
+        visits["sid"] = np.logical_not(np.logical_and(same_note, adjacent_times)).cumsum().shift()
+        visits.fillna({"sid": 0}, inplace=True)
+        visits["sid"] = visits["sid"].astype(int)
+
+        blocks = visits.groupby("sid").agg(
+            mjd=pd.NamedAgg(column="mjd", aggfunc="min"),
+            end_mjd=pd.NamedAgg(column="end_mjd", aggfunc="max"),
+        )
+        blocks["duration"] = (blocks.end_mjd - blocks.mjd) * 24 * 60 * 60
+
+        self.nslice = len(blocks)
+        self.shape = self.nslice
+        self.sim_idxs.update(
+            visits.reset_index()[["sid", "visit_idx"]].groupby("sid").agg(list).to_dict()["visit_idx"]
+        )
+
+        self.slice_points["sid"] = blocks.reset_index().sid.values
+        self.slice_points["mjd"] = blocks.mjd.values
+        self.slice_points["duration"] = blocks.duration.values
+        self._run_maps(maps)
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):  # pylint: disable=invalid-name
+            idxs = self.sim_idxs[islice]
+
+            try:
+                _ = idxs[0]
+            except TypeError:
+                idxs = [idxs]
+
+            slice_points = {
+                "sid": islice,
+                "mjd": self.slice_points["mjd"][islice],
+                "duration": self.slice_points["duration"][islice],
+            }
+
+            return {"idxs": idxs, "slice_point": slice_points}
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+
+class VisitIntervalSlicer(TimeIntervalSlicer):
+    """Slices into intervals each of which contain one visit
+
+    Parameters
+    ----------
+    mjd_column_name : `str`
+        Name of column on which to slice, must be in units of days
+    duration_column_name : `str`
+        Name of column with the duration of each visit (in seconds)
+    badval : `float`
+        Value to use for bad values in slice
+    verbose : `bool`
+        Print extra information?
+    """
+
+    def __init__(
+        self,
+        mjd_column_name="observationStartMJD",
+        duration_column_name="visitTime",
+        extra_column_names=tuple(),
+        badval=np.nan,
+        verbose=False,
+    ):
+        super().__init__(verbose=verbose, badval=badval)
+        self.mjd_column_name = mjd_column_name
+        self.duration_column_name = duration_column_name
+        self.extra_column_names = extra_column_names
+        self.columns_needed = [mjd_column_name, duration_column_name]
+        self.sim_idxs = None  # pylint: disable=invalid-name
+
+    def setup_slicer(self, sim_data, maps=None):
+        self.nslice = len(sim_data[self.mjd_column_name])
+        self.shape = self.nslice
+
+        self.sim_idxs = np.argsort(sim_data[self.mjd_column_name])
+        self.slice_points["sid"] = np.arange(self.nslice)
+        self.slice_points["mjd"] = sim_data[self.mjd_column_name]
+        self.slice_points["duration"] = sim_data[self.duration_column_name]
+        for column_name in self.extra_column_names:
+            self.slice_points[column_name] = sim_data[column_name]
+        self._run_maps(maps)
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):  # pylint: disable=invalid-name
+            idxs = self.sim_idxs[islice]
+            slice_points = {
+                "sid": [idxs],
+                "mjd": self.slice_points["mjd"][islice],
+                "duration": self.slice_points["duration"][islice],
+            }
+            for column_name in self.extra_column_names:
+                slice_points[column_name] = self.slice_points[column_name]
+
+            return {"idxs": idxs, "slice_point": slice_points}
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+
+# internal functions & classes

rubin_sim/maf/slicers/uni_slicer.py

@@ -0,0 +1,57 @@
+"""UniSlicer - no slicing at all, simply return all data points."""
+
+__all__ = ("UniSlicer",)
+
+from functools import wraps
+
+import numpy as np
+
+from .base_slicer import BaseSlicer
+
+
+class UniSlicer(BaseSlicer):
+    """UniSlicer."""
+
+    def __init__(self, verbose=True, badval=-666):
+        """Instantiate unislicer."""
+        super(UniSlicer, self).__init__(verbose=verbose, badval=badval)
+        self.nslice = 1
+        self.shape = self.nslice
+        self.slice_points["sid"] = np.array(
+            [
+                0,
+            ],
+            int,
+        )
+        self.plot_funcs = []
+
+    def setup_slicer(self, sim_data, maps=None):
+        """Use sim_data to set indexes to return."""
+        self._run_maps(maps)
+        sim_dataCol = sim_data.dtype.names[0]
+        self.indices = np.ones(len(sim_data[sim_dataCol]), dtype="bool")
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):
+            """Return all indexes in sim_data."""
+            idxs = self.indices
+            slice_point = {"sid": islice}
+            for key in self.slice_points:
+                if len(np.shape(self.slice_points[key])) == 0:
+                    keyShape = 0
+                else:
+                    keyShape = np.shape(self.slice_points[key])[0]
+                if keyShape == self.nslice:
+                    slice_point[key] = self.slice_points[key][islice]
+                else:
+                    slice_point[key] = self.slice_points[key]
+            return {"idxs": idxs, "slice_point": slice_point}
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def __eq__(self, other_slicer):
+        """Evaluate if slicers are equivalent."""
+        if isinstance(other_slicer, UniSlicer):
+            return True
+        else:
+            return False

rubin_sim/maf/slicers/user_points_slicer.py

@@ -0,0 +1,141 @@
+__all__ = ("UserPointsSlicer",)
+
+import numpy as np
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+
+from rubin_sim.maf.plots.spatial_plotters import BaseHistogram, BaseSkyMap
+
+from .base_spatial_slicer import BaseSpatialSlicer
+
+
+class UserPointsSlicer(BaseSpatialSlicer):
+    """A spatial slicer based on a user-provided list of points.
+    The data_slices returned are the visit pointings which overlap
+    each of these points.
+
+    Parameters
+    ----------
+    ra : `list` or `numpy.ndarray`
+        User-selected RA points, in degrees. Stored internally in radians.
+    dec : `list` or `numpy.ndarray`
+        User-selected Dec points, in degrees. Stored internally in radians.
+    lon_col : `str`, optional
+        Name of the longitude (RA equivalent) column in the input data.
+        Default fieldRA
+    lat_col : `str`, optional
+        Name of the latitude (Dec equivalent) column in the input data.
+        Default fieldDec
+    latLonDeg : `bool`, optional
+        Flag indicating whether the lon and lat values will be in
+        degrees (True) or radians (False).
+        Default True.
+    verbose : `bool`, optional
+        Flag to indicate whether or not to write additional
+        information to stdout during runtime.
+        Default True.
+    badval : `float`, optional
+        Bad value flag, relevant for plotting. Default -666.
+    leafsize : `int`, optional
+        Leafsize value for kdtree. Default 100.
+    radius : `float`, optional
+        Radius for matching in the kdtree.
+        Equivalent to the radius of the FOV. Degrees.
+        Default 2.45.
+    use_camera : `bool`, optional
+        Flag to indicate whether to use the LSST camera footprint or not.
+        Default True.
+    camera_radius : `float`, optional
+        max_radius for the LsstCameraFootprint (degrees).
+    camera_footprint_file : `str`, optional
+        Name of the camera footprint map to use.
+        Can be None, which will use the default.
+    rotSkyPosColName : `str`, optional
+        Name of the rotSkyPos column in the input  data.
+        Only used if use_camera is True.
+        Describes the camera orientation compared to the sky.
+        Default rotSkyPos.
+    """
+
+    def __init__(
+        self,
+        ra,
+        dec,
+        lon_col="fieldRA",
+        lat_col="fieldDec",
+        lat_lon_deg=True,
+        verbose=True,
+        badval=-666,
+        leafsize=100,
+        radius=2.45,
+        use_camera=True,
+        camera_radius=1.94,
+        camera_footprint_file=None,
+        rot_sky_pos_col_name="rotSkyPos",
+    ):
+        super().__init__(
+            lon_col=lon_col,
+            lat_col=lat_col,
+            lat_lon_deg=lat_lon_deg,
+            verbose=verbose,
+            badval=badval,
+            radius=radius,
+            leafsize=leafsize,
+            use_camera=use_camera,
+            camera_radius=camera_radius,
+            camera_footprint_file=camera_footprint_file,
+            rot_sky_pos_col_name=rot_sky_pos_col_name,
+        )
+        # check that ra and dec are iterable,
+        # if not, they are probably naked numbers, wrap in list
+        if not hasattr(ra, "__iter__"):
+            ra = [ra]
+        if not hasattr(dec, "__iter__"):
+            dec = [dec]
+        if len(ra) != len(dec):
+            raise ValueError("RA and Dec must be the same length")
+        ra = np.radians(ra)
+        dec = np.radians(dec)
+        self.slice_points["sid"] = np.arange(np.size(ra))
+        self.slice_points["ra"] = np.array(ra)
+        self.slice_points["dec"] = np.array(dec)
+        c = SkyCoord(ra=self.slice_points["ra"] * u.rad, dec=self.slice_points["dec"] * u.rad).transform_to(
+            "galactic"
+        )
+        gall, galb = c.l.rad, c.b.rad
+        self.slice_points["gall"] = gall
+        self.slice_points["galb"] = galb
+
+        self.nslice = np.size(ra)
+        self.shape = self.nslice
+        self.slicer_init = {
+            "ra": ra,
+            "dec": dec,
+            "lon_col": lon_col,
+            "lat_col": lat_col,
+            "radius": radius,
+        }
+        self.plot_funcs = [BaseSkyMap, BaseHistogram]
+
+    def __eq__(self, other_slicer):
+        """Evaluate if two slicers are equivalent."""
+        result = False
+        # check the slice_points
+        for key in other_slicer.slice_points:
+            if key in self.slice_points.keys():
+                if not np.array_equal(other_slicer.slice_points[key], self.slice_points[key]):
+                    return False
+            else:
+                return False
+        if isinstance(other_slicer, UserPointsSlicer):
+            if other_slicer.nslice == self.nslice:
+                if np.array_equal(
+                    other_slicer.slice_points["ra"], self.slice_points["ra"]
+                ) and np.array_equal(other_slicer.slice_points["dec"], self.slice_points["dec"]):
+                    if other_slicer.lon_col == self.lon_col and other_slicer.lat_col == self.lat_col:
+                        if other_slicer.radius == self.radius:
+                            if other_slicer.use_camera == self.use_camera:
+                                if other_slicer.rotSkyPosColName == self.rotSkyPosColName:
+                                    if np.array_equal(other_slicer.shape, self.shape):
+                                        result = True
+        return result

rubin_sim/maf/stackers/__init__.py

@@ -0,0 +1,13 @@
+from .base_stacker import *
+from .coord_stackers import *
+from .date_stackers import *
+from .general_stackers import *
+from .get_col_info import *
+from .label_stackers import *
+from .m5_optimal_stacker import *
+from .mo_stackers import *
+from .n_follow_stacker import *
+from .neo_dist_stacker import *
+from .sdss_stackers import *
+from .sn_stacker import *
+from .teff_stacker import *