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

rubin_sim/maf/maps/galactic_plane_priority_maps.py

@@ -0,0 +1,210 @@
+__all__ = (
+    "gp_priority_map_components_to_keys",
+    "galplane_priority_map",
+    "GalacticPlanePriorityMap",
+)
+
+import os
+import warnings
+
+import healpy as hp
+import numpy as np
+from astropy.io import fits
+from rubin_scheduler.data import get_data_dir
+
+from rubin_sim.maf.maps import BaseMap
+from rubin_sim.maf.utils import radec2pix
+
+
+def gp_priority_map_components_to_keys(filtername, science_map):
+    """A convenience function to help keep the map key
+    formats in sync in various places.
+    """
+    return f"galplane_priority_{science_map}:{filtername}"
+
+
+def gp_priority_map_keys_to_components(mapname):
+    s, f = mapname.replace("galplane_priority_", "").split(":")
+    return f, s
+
+
+def galplane_priority_map(
+    nside=64,
+    get_keys=False,
+    ra=None,
+    dec=None,
+    pixels=None,
+    interp=False,
+    map_path=None,
+    use_alt_maps=False,
+):
+    """Reads and saves the galactic plane priority maps.
+
+    Parameters
+    ----------
+    nside : `int`
+        Healpixel resolution (2^x). At present, this must be 64.
+    get_keys : `bool`, opt
+        Set this to True to retrieve *only* the keys
+        (such as the science map names) for the maps.
+        Default False.
+    ra : `np.ndarray` or `float`, opt
+        RA (can take numpy array).
+        Default None sets up healpix array of nside. Radians.
+    dec : `np.ndarray` or `float`, opt
+        Dec (can take numpy array).
+        Default None set up healpix array of nside. Radians.
+    pixels : `np.ndarray`, opt
+        Healpixel IDs, to sub-select particular healpix points.
+        Default uses all points.
+        Easiest way to access healpix values.
+        Note that the pixels in the healpix array MUST come from a
+        healpix grid with the same nside as the galactic plane priority map.
+        Using different nsides can potentially fail silently.
+    interp : `bool`, opt
+        Should returned values be interpolated (True)
+        or just nearest neighbor (False).
+        Default False.
+    map_path : `str`, opt
+        Path to directory containing dust map files.
+        Default None, uses $RUBIN_SIM_DATA_DIR/maps.
+    use_alt_maps : `bool`, opt
+        Use the priority_GalPlane_footprint_alt_map_data_{ugrizysum}.fits
+        files instead of the default
+        priority_galPlane_footprint_map_data_{ugrizysum}.fits files.
+    """
+    # This is a function that will read the galactic plane priority map data
+    # and hold onto it indefinitely
+    # this also lets us use a range of slicers, as it will set the slice_point
+    # data appropriately.
+
+    # This function's primary goal is to return this information to the map,
+    # to use for the slicer.
+    # So you MUST specify ra/dec or pixels -- or only retireve the keys
+    if get_keys is False:
+        if (ra is None) & (dec is None) & (pixels is None):
+            raise RuntimeError("Need to set ra,dec or pixels.")
+
+    # This reads and stores the galactic plane priority maps
+    # The galactic plane priority maps are only available in nside 64
+    # There are several different versions of the map -
+    # but we will almost always run all of the galactic plane metrics
+    # together, so we'll just read them all at once here
+    if nside != 64:
+        raise RuntimeError("Currently only available with nside=64")
+
+    # Do we need to read from disk?
+    if not hasattr(galplane_priority_map, "maps"):
+        galplane_priority_map.nside = 64
+        galplane_priority_map.filterlist = ["u", "g", "r", "i", "z", "y", "sum"]
+        if map_path is not None:
+            data_dir = map_path
+        else:
+            data_dir = os.path.join(get_data_dir(), "maps", "GalacticPlanePriorityMaps")
+        galplane_priority_map.maps = {}
+        science_maps = []
+        for f in galplane_priority_map.filterlist:
+            if use_alt_maps:
+                filename = f"priority_GalPlane_footprint_alt_map_data_{f}.fits"
+            else:
+                filename = f"priority_GalPlane_footprint_map_data_{f}.fits"
+            mapfile = os.path.join(data_dir, filename)
+            with fits.open(mapfile) as hdul:
+                galplane_priority_map.maps[f] = hdul[1].data
+                print(f"Read galplane priority map from {mapfile}")
+                science_maps += list(galplane_priority_map.maps[f].dtype.names)
+            science_maps = list(set(science_maps))
+            galplane_priority_map.science_maps = science_maps
+
+    galplane_priority_map.keys = [
+        gp_priority_map_components_to_keys(f, s)
+        for f in galplane_priority_map.filterlist
+        for s in galplane_priority_map.science_maps
+    ]
+
+    if get_keys:
+        return galplane_priority_map.keys
+
+    if pixels is not None:
+        # Assume if pixels is set, then interp is irrelevant
+        maps = {}
+        for k in galplane_priority_map.keys:
+            f, s = gp_priority_map_keys_to_components(k)
+            maps[k] = galplane_priority_map.maps[f][s][pixels]
+    # BUT, if we were provided ra/dec then have to see if we should interpolate
+    else:
+        if interp:  # find interpolated values at given ra/dec
+            maps = {}
+            for k in galplane_priority_map.keys:
+                f, s = gp_priority_map_keys_to_components(k)
+                maps[k] = hp.get_interp_val(galplane_priority_map.maps[f][s], np.pi / 2.0 - dec, ra)
+        else:  # just find nearest neighbor pixel value and return those
+            pixels = radec2pix(nside, ra, dec)
+            maps = {}
+            for k in galplane_priority_map.keys:
+                f, s = gp_priority_map_keys_to_components(k)
+                maps[k] = galplane_priority_map.maps[f][s][pixels]
+    return maps
+
+
+class GalacticPlanePriorityMap(BaseMap):
+    """Add the galactic plane priority map data to the slice points.
+
+    This calls galactic_plane_priority_map to read the map data, and then
+    assigns the appropriate values to each slice_point.
+    If the slicer is an nside=64 healpix slicer, this is trivial.
+    (other use-cases currently experimental and not supported).
+
+    Add keys corresponding to each of the galplane priority map elements.
+
+    Parameters
+    ----------
+    interp : `bool`, opt
+        Interpolate the dust map at each slice_point (True)
+        or just use the nearest value (False).
+        Default is False.
+    nside : `int`, opt
+        Default nside value to read the dust map from disk.
+        Primarily useful if the slicer is not a healpix slicer.
+        Default 64.
+    map_path : `str`, opt
+        Define a path to the directory holding the dust map files.
+        Default None, which uses RUBIN_SIM_DATA_DIR.
+    """
+
+    def __init__(self, interp=False, nside=64, map_path=None):
+        self.keynames = galplane_priority_map(get_keys=True)
+        self.interp = interp
+        self.nside = nside
+        self.map_path = map_path
+
+    def run(self, slice_points):
+        # If the slicer has nside,
+        # it's a healpix slicer so we can read the map directly
+        if "nside" in slice_points:
+            if slice_points["nside"] != self.nside:
+                warnings.warn(
+                    f"Slicer value of nside {slice_points['nside']} different "
+                    f"from map value {self.nside}, using slicer value"
+                )
+            maps = galplane_priority_map(
+                slice_points["nside"],
+                pixels=slice_points["sid"],
+                map_path=self.map_path,
+            )
+            for key in self.keynames:
+                slice_points[key] = maps[key]
+        # Not a healpix slicer, look up values based on RA,dec
+        # with possible interpolation
+        else:
+            maps = galplane_priority_map(
+                self.nside,
+                ra=slice_points["ra"],
+                dec=slice_points["dec"],
+                interp=self.interp,
+                map_path=self.map_path,
+            )
+            for key in self.keynames:
+                slice_points[key] = maps[key]
+
+        return slice_points

rubin_sim/maf/maps/stellar_density_map.py

@@ -0,0 +1,73 @@
+__all__ = ("StellarDensityMap",)
+
+import os
+
+import healpy as hp
+import numpy as np
+from rubin_scheduler.data import get_data_dir
+
+from rubin_sim.maf.utils import radec2pix
+
+from . import BaseMap
+
+
+class StellarDensityMap(BaseMap):
+    """Read and hold the cumulative stellar luminosity function for
+    each slice point.
+
+    The underlying stellar luminosity function map is nside = 64, and contains
+    stars per sq degree at a series of magnitudes (the map contains
+    `starLumFunc_<filter>` and `starMapBins_<filter>`).
+    For slice points which do not match nside=64, the map uses the nearest
+    healpix point on the nside=64 grid.
+
+    The stellar luminosity function comes from the GalFast model.
+
+    Parameters
+    ----------
+    startype : `str` ('allstars', 'wdstars')
+        Load the luminosity function for all stars ('allstars'),
+        which includes main-sequence stars
+        white dwarfs, blue horozontal branch, RR Lyrae, and Cepheids.
+        The 'wdstars' option only includes white dwarf stars.
+    filtername : `str`
+        Filter to use. Options of u,g,r,i,z,y
+    """
+
+    def __init__(self, startype="allstars", filtername="r", map_dir=None):
+        if map_dir is not None:
+            self.map_dir = map_dir
+        else:
+            self.map_dir = os.path.join(get_data_dir(), "maps", "StarMaps")
+        self.filtername = filtername
+        self.keynames = [
+            f"starLumFunc_{self.filtername}",
+            f"starMapBins_{self.filtername}",
+        ]
+        if startype == "allstars":
+            self.startype = ""
+        else:
+            self.startype = startype + "_"
+
+    def _read_map(self):
+        filename = "starDensity_%s_%snside_64.npz" % (self.filtername, self.startype)
+        star_map = np.load(os.path.join(self.map_dir, filename))
+        self.star_map = star_map["starDensity"].copy()
+        self.star_map_bins = star_map["bins"].copy()
+        self.starmap_nside = hp.npix2nside(np.size(self.star_map[:, 0]))
+
+    def run(self, slice_points):
+        self._read_map()
+
+        nside_match = False
+        if "nside" in slice_points:
+            if slice_points["nside"] == self.starmap_nside:
+                slice_points[f"starLumFunc_{self.filtername}"] = self.star_map
+                nside_match = True
+        if not nside_match:
+            # Compute the healpix for each slice_point on the nside=64 grid
+            indx = radec2pix(self.starmap_nside, slice_points["ra"], slice_points["dec"])
+            slice_points[f"starLumFunc_{self.filtername}"] = self.star_map[indx, :]
+
+        slice_points[f"starMapBins_{self.filtername}"] = self.star_map_bins
+        return slice_points

rubin_sim/maf/maps/trilegal_map.py

@@ -0,0 +1,82 @@
+__all__ = ("TrilegalDensityMap",)
+
+import os
+
+import healpy as hp
+import numpy as np
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+from rubin_scheduler.data import get_data_dir
+from rubin_scheduler.utils import _build_tree, _hpid2_ra_dec, _xyz_from_ra_dec
+
+from . import BaseMap
+
+
+class TrilegalDensityMap(BaseMap):
+    """Read and hold the cumulative stellar luminosity function for
+    each slice point.
+
+    The stellar luminosity function comes from the TRILEGAL model.
+
+    Parameters
+    ----------
+    filtername : `str`, opt
+        Filter to use. Options of u,g,r,i,z,y. Default r.
+    nside : `int`, opt
+        The HEALpix nside (can be 64 or 128). Default 64.
+    ext : `bool`, opt
+        Use the full sky maps. Default True.
+
+    Notes
+    -----
+    The underlying stellar luminosity function map is available in a
+    variety of nsides, and contains
+    stars per sq degree at a series of magnitudes (the map contains
+    `starLumFunc_<filter>` and `starMapBins_<filter>`).
+    For slice points which do not match one of the native nside options,
+    the map uses the nearest healpix point on the specified nside grid.
+    """
+
+    def __init__(self, filtername="r", nside=64, ext=True):
+        self.map_dir = os.path.join(get_data_dir(), "maps", "TriMaps")
+        self.filtername = filtername
+        self.keynames = [
+            f"starLumFunc_{self.filtername}",
+            f"starMapBins_{self.filtername}",
+        ]
+        self.nside = nside
+        self.ext = ext
+
+    def _read_map(self):
+        if self.ext:
+            filename = "TRIstarDensity_%s_nside_%i_ext.npz" % (
+                self.filtername,
+                self.nside,
+            )
+        else:
+            filename = "TRIstarDensity_%s_nside_%i.npz" % (self.filtername, self.nside)
+        star_map = np.load(os.path.join(self.map_dir, filename))
+        self.star_map = star_map["starDensity"].copy()
+        self.star_map_bins = star_map["bins"].copy()
+        self.starmap_nside = hp.npix2nside(np.size(self.star_map[:, 0]))
+        # note, the trilegal maps are in galactic coordinates
+        # and use nested healpix.
+        gal_l, gal_b = _hpid2_ra_dec(self.nside, np.arange(hp.nside2npix(self.nside)), nest=True)
+
+        # Convert that to RA,dec. Then do nearest neighbor lookup.
+        c = SkyCoord(l=gal_l * u.rad, b=gal_b * u.rad, frame="galactic").transform_to("icrs")
+        ra = c.ra.rad
+        dec = c.dec.rad
+
+        self.tree = _build_tree(ra, dec)
+
+    def run(self, slice_points):
+        self._read_map()
+
+        x, y, z = _xyz_from_ra_dec(slice_points["ra"], slice_points["dec"])
+
+        dist, indices = self.tree.query(list(zip(x, y, z)))
+
+        slice_points["starLumFunc_%s" % self.filtername] = self.star_map[indices, :]
+        slice_points["starMapBins_%s" % self.filtername] = self.star_map_bins
+        return slice_points

rubin_sim/maf/metadata_dir.py

@@ -0,0 +1,90 @@
+from . import batches as batches
+from .db import ResultsDb
+from .metric_bundles import MetricBundleGroup
+from .slicers import HealpixSlicer, make_wfd_subset_slicer
+
+__all__ = ("metadata_dir",)
+
+import argparse
+import glob
+import os
+import shutil
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+
+def metadata_dir():
+    """
+    Run the metadata batch on all .db files in a directory.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--db", type=str, default=None)
+    parser.add_argument(
+        "--nside",
+        type=float,
+        default=64,
+        help="nside to use for the healpix slicer and subsetslicer for metrics.",
+    )
+    parser.add_argument(
+        "--wfd_threshold",
+        type=float,
+        default=750,
+        help="Threshold number of visits per pointing to use to define the WFD footprint."
+        "Default value of 750 corresponds to the minimum median value per pointing from the SRD.",
+    )
+    parser.add_argument(
+        "--no_clobber",
+        dest="no_clobber",
+        action="store_true",
+        default=False,
+        help="Do not remove existing directory outputs",
+    )
+    args = parser.parse_args()
+
+    # If runNames not given, scan for sim_name databases in current
+    # directory and use those
+    # Note that 'runNames' can be full path to directories
+
+    if args.db is None:
+        # Just look for any .db files in this directory
+        db_files = glob.glob("*.db")
+        # But remove trackingDb and results_db if they're there
+        try:
+            db_files.remove("trackingDb_sqlite.db")
+        except ValueError:
+            pass
+        try:
+            db_files.remove("resultsDb_sqlite.db")
+        except ValueError:
+            pass
+    elif isinstance(args.db, str):
+        db_files = [args.db]
+    else:
+        db_files = args.db
+
+    sim_names = [os.path.basename(name).replace(".db", "") for name in db_files]
+
+    for filename, sim_name in zip(db_files, sim_names):
+        # Connect to the database
+        colmap = batches.col_map_dict()
+
+        # Set and create if needed the output directory
+        out_dir = sim_name + "_meta"
+        if not args.no_clobber:
+            if os.path.isdir(out_dir):
+                shutil.rmtree(out_dir)
+
+        # Find the 'wfd' footprint - use the scheduler footprint.
+        allsky_slicer = HealpixSlicer(nside=args.nside)
+        wfd_slicer = make_wfd_subset_slicer(nside=args.nside)
+
+        bdict = batches.info_bundle_dicts(allsky_slicer, wfd_slicer, sim_name, colmap)
+
+        # Set up the resultsDB
+        results_db = ResultsDb(out_dir=out_dir)
+        # Go and run it
+        group = MetricBundleGroup(bdict, filename, out_dir=out_dir, results_db=results_db, save_early=False)
+        group.run_all(clear_memory=True, plot_now=True)
+        results_db.close()

rubin_sim/maf/metric_bundles/__init__.py

@@ -0,0 +1,3 @@
+from .metric_bundle import *
+from .metric_bundle_group import *
+from .mo_metric_bundle import *