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

rubin_sim/maf/slicers/base_spatial_slicer.py

@@ -0,0 +1,202 @@
+# The base class for all spatial slicers.
+# Slicers are 'data slicers' at heart; spatial slicers slice data by RA/Dec and
+# return the relevant indices in the sim_data to the metric.
+# The primary things added here are the methods to slice the data
+# (for any spatial slicer) as this uses a KD-tree built on spatial
+# (RA/Dec type) indexes.
+
+__all__ = ("BaseSpatialSlicer",)
+
+import warnings
+from functools import wraps
+
+import numpy as np
+import rubin_scheduler.utils as simsUtils
+
+from rubin_sim.maf.plots.spatial_plotters import BaseHistogram, BaseSkyMap
+
+from .base_slicer import BaseSlicer
+
+
+def rotate(x, y, rotation_angle_rad):
+    """rotate 2d points around the origin (0,0)"""
+    cos_rad = np.cos(rotation_angle_rad)
+    sin_rad = np.sin(rotation_angle_rad)
+    qx = cos_rad * x - sin_rad * y
+    qy = sin_rad * x + cos_rad * y
+    return qx, qy
+
+
+class BaseSpatialSlicer(BaseSlicer):
+    """Base spatial slicer object, contains additional functionality
+    for spatial slicing, including setting up and traversing a kdtree
+    containing the simulated data points.
+
+    Parameters
+    ----------
+    lon_col : `str`, optional
+        Name of the longitude (RA equivalent) column.
+    lat_col : `str`, optional
+        Name of the latitude (Dec equivalent) column.
+    rot_sky_pos_col_name : `str`, optional
+        Name of the rotSkyPos column in the input  data.
+        Only used if use_camera is True.
+        Describes the orientation of the camera orientation on the sky.
+    lat_lon_deg : `bool`, optional
+        Flag indicating whether lat and lon values from input data are
+        in degrees (True) or radians (False).
+    verbose : `bool`, optional
+        Verbosity flag - noisy or quiet.
+    badval : `float`, optional
+        Bad value flag, relevant for plotting.
+    leafsize : `int`, optional
+        Leafsize value for kdtree.
+    radius : `float`, optional
+        Radius for matching in the kdtree.
+        Equivalent to the radius of the FOV, in degrees.
+    use_camera : `bool`, optional
+        Flag to indicate whether to use the LSST camera footprint or not.
+    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 file.
+    """
+
+    def __init__(
+        self,
+        lon_col="fieldRA",
+        lat_col="fieldDec",
+        rot_sky_pos_col_name="rotSkyPos",
+        lat_lon_deg=True,
+        verbose=True,
+        badval=-666,
+        leafsize=100,
+        radius=2.45,
+        use_camera=True,
+        camera_radius=1.94,
+        camera_footprint_file=None,
+    ):
+        super().__init__(verbose=verbose, badval=badval)
+        self.lon_col = lon_col
+        self.lat_col = lat_col
+        self.latLonDeg = lat_lon_deg
+        self.rotSkyPosColName = rot_sky_pos_col_name
+        self.columns_needed = [lon_col, lat_col, rot_sky_pos_col_name]
+        self.use_camera = use_camera
+        self.camera_radius = camera_radius
+        self.camera_footprint_file = camera_footprint_file
+        self.radius = radius
+        self.leafsize = leafsize
+        self.slicer_init = {
+            "lon_col": self.lon_col,
+            "lat_col": self.lat_col,
+            "radius": self.radius,
+            "badval": self.badval,
+            "use_camera": self.use_camera,
+        }
+        # RA and Dec are required slice_point info for any spatial slicer.
+        # slice_point RA/Dec are in radians.
+        self.slice_points["sid"] = None
+        self.slice_points["ra"] = None
+        self.slice_points["dec"] = None
+        self.nslice = None
+        self.shape = None
+        self.plot_funcs = [BaseHistogram, BaseSkyMap]
+
+    def setup_slicer(self, sim_data, maps=None):
+        """Use sim_data[self.lon_col] and sim_data[self.lat_col]
+        (in radians) to set up KDTree.
+
+        Parameters
+        -----------
+        sim_data : `numpy.ndarray`
+            The simulated data, including the location of each pointing.
+        maps : `list` of `rubin_sim.maf.maps` objects, optional
+            List of maps (such as dust extinction) that will run to build up
+            additional data at each slice_point. This additional data
+            is available to metrics via the slice_point dictionary.
+        """
+        if maps is not None:
+            if self.cache_size != 0 and len(maps) > 0:
+                warnings.warn(
+                    "Warning:  Loading maps but cache on." "Should probably set use_cache=False in slicer."
+                )
+            self._run_maps(maps)
+        self._set_rad(self.radius)
+
+        if self.latLonDeg:
+            self.data_ra = np.radians(sim_data[self.lon_col])
+            self.data_dec = np.radians(sim_data[self.lat_col])
+            self.data_rot = np.radians(sim_data[self.rotSkyPosColName])
+        else:
+            self.data_ra = sim_data[self.lon_col]
+            self.data_dec = sim_data[self.lat_col]
+            self.data_rot = sim_data[self.rotSkyPosColName]
+        if self.use_camera:
+            self._setupLSSTCamera()
+
+        self._build_tree(self.data_ra, self.data_dec, self.leafsize)
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):
+            """Return indexes for relevant opsim data at slice_point
+            (slice_point=lon_col/lat_col value .. usually ra/dec).
+            """
+
+            # Build dict for slice_point info
+            slice_point = {"sid": islice}
+            sx, sy, sz = simsUtils._xyz_from_ra_dec(
+                self.slice_points["ra"][islice], self.slice_points["dec"][islice]
+            )
+            # Query against tree.
+            indices = self.opsimtree.query_ball_point((sx, sy, sz), self.rad)
+
+            if (self.use_camera) & (len(indices) > 0):
+                # Find the indices *of those indices*
+                # which fall in the camera footprint
+                camera_idx = self.camera(
+                    self.slice_points["ra"][islice],
+                    self.slice_points["dec"][islice],
+                    self.data_ra[indices],
+                    self.data_dec[indices],
+                    self.data_rot[indices],
+                )
+                indices = np.array(indices)[camera_idx]
+
+            # Loop through all the slice_point keys.
+            # If the first dimension of slice_point[key] has the same shape
+            # as the slicer, assume it is information per slice_point.
+            # Otherwise, pass the whole slice_point[key] information.
+            # Useful for stellar LF maps where we want to pass only the
+            # relevant LF and the bins that go with it.
+            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": indices, "slice_point": slice_point}
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)
+
+    def _setupLSSTCamera(self):
+        """If we want to include the camera chip gaps, etc."""
+        self.camera = simsUtils.LsstCameraFootprint(
+            units="radians", footprint_file=self.camera_footprint_file, max_radius=self.camera_radius
+        )
+
+    def _build_tree(self, sim_dataRa, sim_dataDec, leafsize=100):
+        """Build KD tree on sim_dataRA/Dec.
+
+        sim_dataRA, sim_dataDec = RA and Dec values (in radians).
+        leafsize = the number of Ra/Dec pointings in each leaf node.
+        """
+        self.opsimtree = simsUtils._build_tree(sim_dataRa, sim_dataDec, leafsize)
+
+    def _set_rad(self, radius=1.75):
+        """Set radius (in degrees) for kdtree search."""
+        self.rad = simsUtils.xyz_angular_radius(radius)

rubin_sim/maf/slicers/healpix_sdss_slicer.py

@@ -0,0 +1,118 @@
+__all__ = ("HealpixSDSSSlicer",)
+
+from functools import wraps
+
+import matplotlib.path as mplPath
+import numpy as np
+from rubin_scheduler.utils import _xyz_from_ra_dec, gnomonic_project_toxy
+
+from rubin_sim.maf.plots import HealpixSDSSSkyMap
+
+from .healpix_slicer import HealpixSlicer
+
+
+class HealpixSDSSSlicer(HealpixSlicer):
+    """For use with SDSS stripe 82 square images"""
+
+    def __init__(
+        self,
+        nside=128,
+        lon_col="RA1",
+        lat_col="Dec1",
+        verbose=True,
+        use_cache=True,
+        radius=17.0 / 60.0,
+        leafsize=100,
+        **kwargs,
+    ):
+        """Using one corner of the chip as the spatial key and the
+        diagonal as the radius."""
+        super().__init__(
+            verbose=verbose,
+            lon_col=lon_col,
+            lat_col=lat_col,
+            radius=radius,
+            leafsize=leafsize,
+            use_cache=use_cache,
+            nside=nside,
+        )
+        self.cornerLables = ["RA1", "Dec1", "RA2", "Dec2", "RA3", "Dec3", "RA4", "Dec4"]
+        self.plot_funcs = [
+            HealpixSDSSSkyMap,
+        ]
+
+    def setup_slicer(self, sim_data, maps=None):
+        """
+        Use sim_data[self.lon_col] and sim_data[self.lat_col]
+        (in radians) to set up KDTree.
+        """
+        self._run_maps(maps)
+        self._build_tree(sim_data[self.lon_col], sim_data[self.lat_col], self.leafsize)
+        self._set_rad(self.radius)
+        self.corners = sim_data[self.cornerLables]
+
+        @wraps(self._slice_sim_data)
+        def _slice_sim_data(islice):
+            """Return indexes for relevant opsim data at slice_point
+            (slice_point=lon_col/lat_col value .. usually ra/dec)."""
+            sx, sy, sz = _xyz_from_ra_dec(self.slice_points["ra"][islice], self.slice_points["dec"][islice])
+            # Query against tree.
+            initIndices = self.opsimtree.query_ball_point((sx, sy, sz), self.rad)
+            # Loop through all the images and check if the slice_point
+            # is inside the corners of the chip
+            indices = []
+            # Gnomic project all the corners that are near the slice point,
+            # centered on slice point
+            x1, y1 = gnomonic_project_toxy(
+                self.corners["RA1"][initIndices],
+                self.corners["Dec1"][initIndices],
+                self.slice_points["ra"][islice],
+                self.slice_points["dec"][islice],
+            )
+            x2, y2 = gnomonic_project_toxy(
+                self.corners["RA2"][initIndices],
+                self.corners["Dec2"][initIndices],
+                self.slice_points["ra"][islice],
+                self.slice_points["dec"][islice],
+            )
+            x3, y3 = gnomonic_project_toxy(
+                self.corners["RA3"][initIndices],
+                self.corners["Dec3"][initIndices],
+                self.slice_points["ra"][islice],
+                self.slice_points["dec"][islice],
+            )
+            x4, y4 = gnomonic_project_toxy(
+                self.corners["RA4"][initIndices],
+                self.corners["Dec4"][initIndices],
+                self.slice_points["ra"][islice],
+                self.slice_points["dec"][islice],
+            )
+
+            for i, ind in enumerate(initIndices):
+                # Use matplotlib to make a polygon on
+                bbPath = mplPath.Path(
+                    np.array(
+                        [
+                            [x1[i], y1[i]],
+                            [x2[i], y2[i]],
+                            [x3[i], y3[i]],
+                            [x4[i], y4[i]],
+                            [x1[i], y1[i]],
+                        ]
+                    )
+                )
+                # Check if the slice_point is inside the image corners
+                # and append to list if it is
+                if bbPath.contains_point((0.0, 0.0)) == 1:
+                    indices.append(ind)
+
+            return {
+                "idxs": indices,
+                "slice_point": {
+                    "sid": self.slice_points["sid"][islice],
+                    "ra": self.slice_points["ra"][islice],
+                    "dec": self.slice_points["dec"][islice],
+                },
+            }
+
+        setattr(self, "_slice_sim_data", _slice_sim_data)

rubin_sim/maf/slicers/healpix_slicer.py

@@ -0,0 +1,189 @@
+"""A slicer that uses a Healpix grid to calculate metric values
+(at the center of each healpixel)."""
+
+# Requires healpy
+# See more documentation on healpy here
+# http://healpy.readthedocs.org/en/latest/tutorial.html
+# Also requires numpy (for histogram and power spectrum plotting)
+
+__all__ = ("HealpixSlicer",)
+
+import healpy as hp
+import numpy as np
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+
+from rubin_sim.maf.plots.spatial_plotters import HealpixHistogram, HealpixSkyMap
+
+from .base_spatial_slicer import BaseSpatialSlicer
+
+
+class HealpixSlicer(BaseSpatialSlicer):
+    """
+    A spatial slicer that evaluates pointings on a healpix-based grid.
+
+    Note that a healpix slicer is intended to evaluate the sky
+    on a spatial grid. Usually this grid will be something like RA as
+    the lon_col and Dec as the lat_col.
+    However, it could also be altitude and azimuth,
+    in which case altitude as lat_col, and azimuth as lon_col.
+    An additional alternative is to use HA/Dec, with lon_col=HA, lat_col=Dec.
+
+    When plotting with RA/Dec, the default HealpixSkyMap can be used,
+    corresponding to
+    {'rot': (0, 0, 0), 'flip': 'astro'}.
+    When plotting with alt/az, either the LambertSkyMap can be used
+    (if Basemap is installed)
+    or the HealpixSkyMap can be used with a modified plot_dict,
+    {'rot': (90, 90, 90), 'flip': 'geo'}.
+    When plotting with HA/Dec, only the HealpixSkyMap can be used,
+    with a modified plot_dict of
+    {'rot': (0, -30, 0), 'flip': 'geo'}.
+
+    Parameters
+    ----------
+    nside : `int`, optional
+        The nside parameter of the healpix grid. Must be a power of 2.
+    lon_col : `str`, optional
+        Name of the longitude (RA equivalent) column to use from the
+        input data.
+    lat_col : `str`, optional
+        Name of the latitude (Dec equivalent) column to use from
+        the input data.
+    lat_lon_deg : `bool`, optional
+        Flag indicating whether the lat and lon values in the input data are
+        in degrees (True) or radians (False).
+    verbose : `bool`, optional
+        Flag to indicate whether or not to write additional information to
+        stdout during runtime.
+    use_cache : `bool`, optional
+        Flag allowing the user to indicate whether or not to cache
+        (and reuse) metric results calculated with the same set of
+        simulated data pointings.
+        This can be safely set to True for slicers not using maps and
+        will result in increased speed.
+        When calculating metric results using maps, the map data at each
+        individual ra/dec point may influence the metric results and so
+        use_cache should be set to False.
+    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.
+    use_camera : `bool`, optional
+        Flag to indicate whether to use the LSST camera footprint or not.
+    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 footprint map.
+    rot_sky_pos_col_name : `str`, optional
+        Name of the rotSkyPos column in the input  data.
+        Only used if use_camera is True.
+        Describes the orientation of the camera orientation
+        compared to the sky.
+    badval: `float`, optional
+        The value to place use in place of bad data values.
+        Use np.nan unless specifically needing other values.
+    """
+
+    def __init__(
+        self,
+        nside=128,
+        lon_col="fieldRA",
+        lat_col="fieldDec",
+        lat_lon_deg=True,
+        verbose=True,
+        use_cache=True,
+        leafsize=100,
+        radius=2.45,
+        use_camera=True,
+        camera_radius=1.94,
+        camera_footprint_file=None,
+        rot_sky_pos_col_name="rotSkyPos",
+        badval=np.nan,
+    ):
+        super().__init__(
+            verbose=verbose,
+            lon_col=lon_col,
+            lat_col=lat_col,
+            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,
+            lat_lon_deg=lat_lon_deg,
+        )
+        # Valid values of nside are powers of 2.
+        # nside=64 gives about 1 deg resolution
+        # nside=256 gives about 13' resolution (~1 CCD)
+        # nside=1024 gives about 3' resolution
+        # Check validity of nside:
+        if not (hp.isnsideok(nside)):
+            raise ValueError("Valid values of nside are powers of 2.")
+        self.nside = int(nside)
+        self.pix_area = hp.nside2pixarea(self.nside)
+        self.nslice = hp.nside2npix(self.nside)
+        self.shape = self.nslice
+        if self.verbose:
+            print(
+                "Healpix slicer using NSIDE=%d, " % (self.nside)
+                + "approximate resolution %f arcminutes" % (hp.nside2resol(self.nside, arcmin=True))
+            )
+        # Set variables so slicer can be re-constructed
+        self.slicer_init = {
+            "nside": nside,
+            "lon_col": lon_col,
+            "lat_col": lat_col,
+            "radius": radius,
+        }
+        self.use_cache = use_cache
+        if use_cache:
+            # use_cache set the size of the cache for the memoize function in
+            # sliceMetric.
+            bin_res = hp.nside2resol(nside)  # Pixel size in radians
+            # Set the cache size to be ~2x the circumference
+            self.cache_size = int(np.round(4.0 * np.pi / bin_res))
+        # Set up slice_point metadata.
+        self.slice_points["nside"] = nside
+        self.slice_points["sid"] = np.arange(self.nslice)
+        self.slice_points["ra"], self.slice_points["dec"] = self._pix2radec(self.slice_points["sid"])
+        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
+
+        # Set the default plotting functions.
+        self.plot_funcs = [HealpixSkyMap, HealpixHistogram]
+
+    def __eq__(self, other_slicer):
+        """Evaluate if two slicers are equivalent."""
+        # If the two slicers are both healpix slicers, check nsides value.
+        result = False
+        if isinstance(other_slicer, HealpixSlicer):
+            if other_slicer.nside == self.nside:
+                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.all(other_slicer.shape == self.shape):
+                                    if other_slicer.use_cache == self.use_cache:
+                                        result = True
+        return result
+
+    def _pix2radec(self, islice):
+        """Given the pixel number / sliceID, return the RA/Dec of the
+        pointing, in radians.
+        """
+        # Calculate RA/Dec in RADIANS of pixel in this healpix slicer.
+        # Note that ipix could be an array,
+        # in which case RA/Dec values will be an array also.
+        lat, ra = hp.pix2ang(self.nside, islice)
+        # Move dec to +/- 90 degrees
+        dec = np.pi / 2.0 - lat
+        return ra, dec