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

rubin_sim/maf/maf_contrib/kne_metrics.py

@@ -0,0 +1,473 @@
+__all__ = ("get_kne_filename", "KnLc", "KNePopMetric", "generate_kn_pop_slicer")
+
+import glob
+import os
+
+import numpy as np
+from rubin_scheduler.data import get_data_dir
+from rubin_scheduler.utils import SURVEY_START_MJD, uniform_sphere
+
+from rubin_sim.maf.metrics import BaseMetric
+from rubin_sim.maf.slicers import UserPointsSlicer
+from rubin_sim.maf.utils import m52snr
+from rubin_sim.phot_utils import DustValues
+
+
+def get_kne_filename(inj_params_list=None):
+    """Given kilonova parameters, get the filename from the grid of models
+    developed by M. Bulla
+
+    Parameters
+    ----------
+    inj_params_list : `list` [`dict`]
+        parameters for the kilonova model such as
+        mass of the dynamical ejecta (mej_dyn), mass of the disk wind ejecta
+        (mej_wind), semi opening angle of the cylindrically-symmetric ejecta
+        fan ('phi'), and viewing angle ('theta'). For example
+        inj_params_list =
+        [{'mej_dyn': 0.005, 'mej_wind': 0.050, 'phi': 30, 'theta': 25.8}]
+    """
+    # Get files, model grid developed by M. Bulla
+    datadir = get_data_dir()
+    file_list = glob.glob(os.path.join(datadir, "maf", "bns", "*.dat"))
+
+    # If no specific parameters passed - return everything.
+    if inj_params_list is None or len(inj_params_list) == 0:
+        return file_list
+
+    # Otherwise find the parameters for each file and
+    # then find the relevant matches.
+    params = {}
+    matched_files = []
+    for filename in file_list:
+        key = filename.replace(".dat", "").split("/")[-1]
+        params[key] = {}
+        params[key]["filename"] = filename
+        key_split = key.split("_")
+        # Binary neutron star merger models
+        if key_split[0] == "nsns":
+            mejdyn = float(key_split[2].replace("mejdyn", ""))
+            mejwind = float(key_split[3].replace("mejwind", ""))
+            phi0 = float(key_split[4].replace("phi", ""))
+            theta = float(key_split[5])
+            params[key]["mej_dyn"] = mejdyn
+            params[key]["mej_wind"] = mejwind
+            params[key]["phi"] = phi0
+            params[key]["theta"] = theta
+        # Neutron star--black hole merger models
+        elif key_split[0] == "nsbh":
+            mej_dyn = float(key_split[2].replace("mejdyn", ""))
+            mej_wind = float(key_split[3].replace("mejwind", ""))
+            phi = float(key_split[4].replace("phi", ""))
+            theta = float(key_split[5])
+            params[key]["mej_dyn"] = mej_dyn
+            params[key]["mej_wind"] = mej_wind
+            params[key]["phi"] = phi
+            params[key]["theta"] = theta
+    for key in params.keys():
+        for inj_params in inj_params_list:
+            match = all([np.isclose(params[key][var], inj_params[var]) for var in inj_params.keys()])
+            if match:
+                matched_files.append(params[key]["filename"])
+                print(f"Found match for {inj_params}")
+    print(
+        f"Found matches for {len(matched_files)}/{len(inj_params_list)} \
+          sets of parameters"
+    )
+
+    return matched_files
+
+
+class KnLc:
+    """Read in some KNe lightcurves
+
+    Parameters
+    ----------
+    file_list : `list` [`str`] or None
+        List of file paths to load. If None, loads up all the files
+        from data/bns/
+    """
+
+    def __init__(self, file_list=None):
+        if file_list is None:
+            datadir = get_data_dir()
+            # Get files, model grid developed by M. Bulla
+            file_list = glob.glob(os.path.join(datadir, "maf", "bns", "*.dat"))
+
+        filts = ["u", "g", "r", "i", "z", "y"]
+        magidxs = [1, 2, 3, 4, 5, 6]
+
+        # Let's organize the data in to a list of dicts for easy lookup
+        self.data = []
+        for filename in file_list:
+            mag_ds = np.loadtxt(filename)
+            t = mag_ds[:, 0]
+            new_dict = {}
+            for ii, (filt, magidx) in enumerate(zip(filts, magidxs)):
+                new_dict[filt] = {"ph": t, "mag": mag_ds[:, magidx]}
+            self.data.append(new_dict)
+
+    def interp(self, t, filtername, lc_indx=0):
+        """Do the interpolation of the lightcurve for a given time and filter.
+
+        Parameters
+        ----------
+        t : `np.ndarray`, (N,)
+            The times to interpolate the light curve to.
+        filtername : `str`
+            The filter. one of ugrizy
+        lc_index : `int`, optional
+            Which file to use.
+
+        Returns
+        -------
+        result : `np.ndarray`, (N,)
+            Array of lightcurve brightnesses at the times of t.=
+        """
+
+        result = np.interp(
+            t,
+            self.data[lc_indx][filtername]["ph"],
+            self.data[lc_indx][filtername]["mag"],
+            left=99,
+            right=99,
+        )
+        return result
+
+
+class KNePopMetric(BaseMetric):
+    def __init__(
+        self,
+        metric_name="KNePopMetric",
+        mjd_col="observationStartMJD",
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        night_col="night",
+        pts_needed=2,
+        file_list=None,
+        mjd0=SURVEY_START_MJD,
+        output_lc=False,
+        badval=-666,
+        **kwargs,
+    ):
+        maps = ["DustMap"]
+        self.mjd_col = mjd_col
+        self.m5_col = m5_col
+        self.filter_col = filter_col
+        self.night_col = night_col
+        self.pts_needed = pts_needed
+        # `bool` variable, if True the light curve will be exported
+        self.output_lc = output_lc
+
+        self.lightcurves = KnLc(file_list=file_list)
+
+        self.mjd0 = mjd0
+
+        dust_properties = DustValues()
+        self.ax1 = dust_properties.ax1
+
+        cols = [self.mjd_col, self.m5_col, self.filter_col, self.night_col]
+        super(KNePopMetric, self).__init__(
+            col=cols,
+            units="Detected, 0 or 1",
+            metric_name=metric_name,
+            maps=maps,
+            badval=badval,
+            **kwargs,
+        )
+
+    def _multi_detect(self, around_peak):
+        """Simple detection criteria:
+        detect at least a certain number of times
+        """
+        result = 1
+        # Detected data points
+        if np.size(around_peak) < self.pts_needed:
+            return 0
+
+        return result
+
+    def _ztfrest_simple(
+        self,
+        around_peak,
+        mags,
+        mags_unc,
+        t,
+        filters,
+        min_dt=0.125,
+        min_fade=0.3,
+        max_rise=-1.0,
+        select_red=False,
+        select_blue=False,
+    ):
+        """Selection criteria based on rise or decay rate;
+        simplified version of the methods employed by the ZTFReST project
+        (Andreoni & Coughlin et al., 2021)
+
+        Parameters
+        ----------
+        around_peak : `np.ndarray`, (N,)
+            indexes corresponding to 5sigma detections
+        mags : `np.ndarray`, (N,)
+            magnitudes obtained interpolating models on the data_slice
+        t : `np.ndarray`, (N,)
+            relative times
+        filters : `np.ndarray`, (N,)
+            filters in which detections happened
+        min_dt : `float`
+            minimum time gap between first and last detection in a given band
+        min_fade : `float`
+            fade rate threshold (positive, mag/day)
+        max_rise : `float`
+            rise rate threshold (negative, mag/day)
+        select_red : `bool`
+            if True, only red 'izy' filters will be considered
+        select_blue : `bool`
+            if True, only blue 'ugr' filters will be considered
+
+        Examples
+        ----------
+        A transient:
+            rising by 0.74 mag/day will pass a threshold max_rise=-0.5
+            rising by 0.74 mag/day will not pass a threshold max_rise=-1.0
+            fading by 0.6 mag/day will pass a threshold min_fade=0.3
+            fading by 0.2 mag/day will not pass a threshold min_fade=0.3
+        """
+        result = 1
+
+        # Quick check on the number of detected points
+        if np.size(around_peak) < self.pts_needed:
+            return 0
+        # Quick check on the time gap between first and last detection
+        elif np.max(t[around_peak]) - np.min(t[around_peak]) < min_dt:
+            return 0
+        else:
+            evol_rate = []
+            fil = []
+            # Check time gaps and rise or fade rate for each band
+            for f in set(filters):
+                if select_red is True and f not in "izy":
+                    continue
+                elif select_blue is True and f not in "ugr":
+                    continue
+                times_f = t[around_peak][np.where(filters == f)[0]]
+                mags_f = mags[around_peak][np.where(filters == f)[0]]
+                mags_unc_f = mags_unc[around_peak][np.where(filters == f)[0]]
+
+                # Check if the evolution is significant enough
+                idx_max = np.argmax(mags_f)
+                idx_min = np.argmin(mags_f)
+                if mags_f[idx_min] + mags_unc_f[idx_min] < mags_f[idx_max] - mags_unc_f[idx_max]:
+                    signif = True
+                else:
+                    signif = False
+
+                # Time difference between max and min
+                dt_f = np.abs(times_f[idx_max] - times_f[idx_min])
+
+                # Get the evolution rate, if the time gap condition is met
+                if dt_f > min_dt and signif is True:
+                    # Calculate evolution rate
+                    evol_rate_f = (np.max(mags_f) - np.min(mags_f)) / (times_f[idx_max] - times_f[idx_min])
+                    evol_rate.append(evol_rate_f)
+                else:
+                    evol_rate.append(0)
+                fil.append(f)
+            if len(evol_rate) == 0:
+                return 0
+            # Check if the conditions on the evolution rate are met
+            if np.max(evol_rate) < min_fade and np.min(evol_rate) > max_rise:
+                return 0
+
+        return result
+
+    def _multi_color_detect(self, filters):
+        """Color-based simple detection criteria:
+        detect at least twice, with at least two filters
+        """
+        result = 1
+        # detected in at least two filters
+        if np.size(np.unique(filters)) < 2:
+            return 0
+
+        return result
+
+    def _red_color_detect(self, filters, min_det=4):
+        """Detected at least min_det times in either izy colors
+
+        Parameters
+        ----------
+        filters : `np.ndarray`, (N,)
+            filters in which detections happened
+        min_det : `float` or `int`
+            minimum number of detections required in izy bands
+        """
+        result = 1
+        # Number of detected points in izy bands
+        n_red_det = (
+            np.size(np.where(filters == "i")[0])
+            + np.size(np.where(filters == "z")[0])
+            + np.size(np.where(filters == "y")[0])
+        )
+        # Condition
+        if n_red_det < min_det:
+            return 0
+
+        return result
+
+    def _blue_color_detect(self, filters, min_det=4):
+        """Detected at least min_det times in either ugr colors
+
+        Parameters
+        ----------
+        filters : `np.ndarray`, (N,)
+            filters in which detections happened
+        min_det : `float` or `int`
+            minimum number of detections required in ugr bands
+        """
+        result = 1
+        # Number of detected points in ugr bands
+        n_blue_det = (
+            np.size(np.where(filters == "u")[0])
+            + np.size(np.where(filters == "g")[0])
+            + np.size(np.where(filters == "r")[0])
+        )
+        # Condition
+        if n_blue_det < min_det:
+            return 0
+
+        return result
+
+    def run(self, data_slice, slice_point=None):
+        result = {}
+        t = data_slice[self.mjd_col] - self.mjd0 - slice_point["peak_time"]
+        mags = np.zeros(t.size, dtype=float)
+
+        for filtername in np.unique(data_slice[self.filter_col]):
+            infilt = np.where(data_slice[self.filter_col] == filtername)
+            mags[infilt] = self.lightcurves.interp(t[infilt], filtername, lc_indx=slice_point["file_indx"])
+            # Apply dust extinction on the light curve
+            a_x = self.ax1[filtername] * slice_point["ebv"]
+            mags[infilt] += a_x
+
+            distmod = 5 * np.log10(slice_point["distance"] * 1e6) - 5.0
+            mags[infilt] += distmod
+
+        # Find the detected points
+        around_peak = np.where((t > 0) & (t < 30) & (mags < data_slice[self.m5_col]))[0]
+        # Filters in which the detections happened
+        filters = data_slice[self.filter_col][around_peak]
+        # Magnitude uncertainties with Gaussian approximation
+        snr = m52snr(mags, data_slice[self.m5_col])
+        mags_unc = 2.5 * np.log10(1.0 + 1.0 / snr)
+
+        result["multi_detect"] = self._multi_detect(around_peak)
+        result["ztfrest_simple"] = self._ztfrest_simple(
+            around_peak, mags, mags_unc, t, filters, select_red=False
+        )
+        result["ztfrest_simple_red"] = self._ztfrest_simple(
+            around_peak, mags, mags_unc, t, filters, select_red=True
+        )
+        result["ztfrest_simple_blue"] = self._ztfrest_simple(
+            around_peak, mags, mags_unc, t, filters, select_blue=True
+        )
+        result["multi_color_detect"] = self._multi_color_detect(filters)
+        result["red_color_detect"] = self._red_color_detect(filters)
+        result["blue_color_detect"] = self._blue_color_detect(filters)
+
+        # Export the light curve
+        if self.output_lc is True:
+            mags[np.where(mags > 50)[0]] = 99.0
+            result["lc"] = [
+                data_slice[self.mjd_col],
+                mags,
+                mags_unc,
+                data_slice[self.m5_col],
+                data_slice[self.filter_col],
+            ]
+            result["lc_colnames"] = ("t", "mag", "mag_unc", "maglim", "filter")
+
+        return result
+
+    def reduce_multi_detect(self, metric):
+        return metric["multi_detect"]
+
+    def reduce_ztfrest_simple(self, metric):
+        return metric["ztfrest_simple"]
+
+    def reduce_ztfrest_simple_red(self, metric):
+        return metric["ztfrest_simple_red"]
+
+    def reduce_ztfrest_simple_blue(self, metric):
+        return metric["ztfrest_simple_blue"]
+
+    def reduce_multi_color_detect(self, metric):
+        return metric["multi_color_detect"]
+
+    def reduce_red_color_detect(self, metric):
+        return metric["red_color_detect"]
+
+    def reduce_blue_color_detect(self, metric):
+        return metric["blue_color_detect"]
+
+
+def generate_kn_pop_slicer(
+    t_start=1,
+    t_end=3652,
+    n_events=10000,
+    seed=42,
+    n_files=308,
+    d_min=10,
+    d_max=300,
+    ra=None,
+    dec=None,
+):
+    """Generate a population of KNe events, and put the info about them
+    into a UserPointSlicer object
+
+    Parameters
+    ----------
+    t_start : `float`, optional
+        The night to start kilonova events on (days)
+    t_end : `float`, optional
+        The final night of kilonova events
+    n_events : `int`, optional
+        The number of kilonova events to generate
+    seed : `float`, optional
+        The seed passed to np.random
+    n_files : `int`, optional
+        The number of different kilonova lightcurves to use
+        This should match the length of the filenames list passed
+        to the KNePopMetric directly.
+    d_min : `float` or `int`, optional
+        Minimum luminosity distance (Mpc)
+    d_max : `float` or `int`, optional
+        Maximum luminosity distance (Mpc)
+    ra, dec : `np.ndarray`, (N,) or None
+        The ra and dec to use for event positions.
+        Generates uniformly on the spehere if None. (degrees)
+    """
+
+    def rndm(a, b, g, size=1):
+        """Power-law gen for pdf(x) propto x^{g-1} for a<=x<=b"""
+        r = np.random.random(size=size)
+        ag, bg = a**g, b**g
+        return (ag + (bg - ag) * r) ** (1.0 / g)
+
+    if ra is None:
+        ra, dec = uniform_sphere(n_events, seed=seed)
+
+    peak_times = np.random.uniform(low=t_start, high=t_end, size=n_events)
+    file_indx = np.floor(np.random.uniform(low=0, high=n_files, size=n_events)).astype(int)
+
+    # Define the distance
+    distance = rndm(d_min, d_max, 4, size=n_events)
+
+    # Set up the slicer to evaluate the catalog we just made
+    slicer = UserPointsSlicer(ra, dec, lat_lon_deg=True, badval=0)
+    # Add any additional information about each object to the slicer
+    slicer.slice_points["peak_time"] = peak_times
+    slicer.slice_points["file_indx"] = file_indx
+    slicer.slice_points["distance"] = distance
+
+    return slicer

rubin_sim/maf/maf_contrib/lss_metrics.py

@@ -0,0 +1,45 @@
+__all__ = ("GalaxyCountsMetric",)
+
+import healpy as hp
+import numpy as np
+import scipy
+
+from rubin_sim.maf.metrics.base_metric import BaseMetric
+from rubin_sim.maf.metrics.exgal_m5 import ExgalM5
+
+
+class GalaxyCountsMetric(BaseMetric):
+    """Estimate the number of galaxies expected at a particular (extragalactic)
+    coadded depth.
+    """
+
+    def __init__(self, m5_col="fiveSigmaDepth", nside=128, metric_name="GalaxyCounts", **kwargs):
+        self.m5_col = m5_col
+        super(GalaxyCountsMetric, self).__init__(col=self.m5_col, metric_name=metric_name, **kwargs)
+        # Use the extinction corrected coadded depth metric to calculate
+        # the depth at each point.
+        self.coaddmetric = ExgalM5(m5_col=self.m5_col)
+        # Total of 41253.0 galaxies across the sky (at what magnitude?).
+        # This didn't seem to work quite right for me..
+        self.scale = 41253.0 / hp.nside2npix(nside) / 5000.0
+        # Reset units (otherwise uses magnitudes).
+        self.units = "Galaxy Counts"
+
+    def _gal_count(self, apparent_mag, coaddm5):
+        # Order for galCount must be apparent mag, then coaddm5,
+        # for scipy.integrate method.
+        dn_gal = np.power(10.0, -3.52) * np.power(10.0, 0.34 * apparent_mag)
+        completeness = 0.5 * scipy.special.erfc(apparent_mag - coaddm5)
+        return dn_gal * completeness
+
+    def run(self, data_slice, slice_point=None):
+        # Calculate the coadded depth.
+        coaddm5 = self.coaddmetric.run(data_slice)
+        # Calculate the number of galaxies.
+        # From Carroll et al, 2014 SPIE (http://arxiv.org/abs/1501.04733)
+        # I'm not entirely certain this gives a properly calibrated number
+        # of galaxy counts, however it is proportional to the expected number
+        # at least (and should be within an order of magnitude)
+        num_gal, int_err = scipy.integrate.quad(self._gal_count, -np.inf, 32, args=coaddm5)
+        num_gal *= self.scale
+        return num_gal

rubin_sim/maf/maf_contrib/lss_obs_strategy/__init__.py

@@ -0,0 +1,3 @@
+from .constants_for_pipeline import *
+from .galaxy_counts_metric_extended import *
+from .galaxy_counts_with_pixel_calibration import *

rubin_sim/maf/maf_contrib/lss_obs_strategy/constants_for_pipeline.py

@@ -0,0 +1,73 @@
+# Various things declared here to be imported when running
+# artificialStructureCalculation.
+# Makes the updates easier, since the constants/objects defined here
+# are used by different
+# functions, e.g. power law constants are called by
+# GalaxyCountsMetric_extended as well as
+# GalaxyCounts_withPixelCalibration.
+#
+# Included here:
+#  * power law constants from the i-band mock catalog; based on mocks
+#  from Padilla et al.
+#  * normalization constant for galaxy counts from the mock catalogs.
+#
+# Humna Awan
+# humna.awan@rutgers.edu
+
+# Power law constants for each z-bin based on N. D. Padilla et al.'s
+# mock catalogs
+# General power law form: 10**(a*m+b)
+# Declare the dictionary for the power law constants
+from collections import OrderedDict
+
+power_law_const_a = OrderedDict()
+power_law_const_b = OrderedDict()
+# 0.<z<0.15
+power_law_const_a["0.<z<0.15"] = 0.166  # July 18 Catalog
+power_law_const_b["0.<z<0.15"] = -0.989  # July 18 Catalog
+
+# 0.15<z<0.37
+power_law_const_a["0.15<z<0.37"] = 0.212  # July 18 Catalog
+power_law_const_b["0.15<z<0.37"] = -1.540  # July 18 Catalog
+
+# 0.37<z<0.66
+power_law_const_a[" 0.37<z<0.66"] = 0.250  # July 18 Catalog
+power_law_const_b[" 0.37<z<0.66"] = -2.242  # July 18 Catalog
+
+# 0.66<z<1.0
+power_law_const_a["0.66<z<1.0"] = 0.265  # July 18 Catalog
+power_law_const_b["0.66<z<1.0"] = -2.660  # July 18 Catalog
+
+# 1.0<z<1.5
+power_law_const_a["1.0<z<1.5"] = 0.283  # July 18 Catalog
+power_law_const_b["1.0<z<1.5"] = -3.242  # July 18 Catalog
+
+# 1.5<z<2.0
+power_law_const_a["1.5<z<2.0"] = 0.308  # July 18 Catalog
+power_law_const_b["1.5<z<2.0"] = -4.169  # July 18 Catalog
+
+# 2.0<z<2.5
+power_law_const_a["2.0<z<2.5"] = 0.364  # July 18 Catalog
+power_law_const_b["2.0<z<2.5"] = -6.261  # July 18 Catalog
+
+# 2.5<z<3.0
+power_law_const_a["2.5<z<3.0"] = 0.346  # July 18 Catalog
+power_law_const_b["2.5<z<3.0"] = -5.661  # July 18 Catalog
+
+# 3.0<z<3.5
+power_law_const_a["3.0<z<3.5"] = 0.381  # July 18 Catalog
+power_law_const_b["3.0<z<3.5"] = -6.810  # July 18 Catalog
+
+# 3.5<z<4.0
+power_law_const_a["3.5<z<4.0"] = 0.424  # July 18 Catalog
+power_law_const_b["3.5<z<4.0"] = -8.362  # July 18 Catalog
+
+# 4.0<z<15.0
+power_law_const_a["4.0<z<15.0"] = 0.462  # July 18 Catalog
+power_law_const_b["4.0<z<15.0"] = -9.474  # July 18 Catalog
+
+
+# Normalization to match the total galaxies from  mock catalogs to match
+# up to CFHTLS counts for
+# i<25.5 galaxy catalog. Using July 18-19 i-band mock catalogs.
+normalization_constant = 3.70330701802