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

rubin_sim/maf/metrics/transient_metrics.py

@@ -0,0 +1,268 @@
+__all__ = ("TransientMetric",)
+
+
+import numpy as np
+
+from .base_metric import BaseMetric
+
+
+class TransientMetric(BaseMetric):
+    """
+    Calculate what fraction of the transients would be detected. Best paired
+    with a spatial slicer.
+    We are assuming simple light curves with no color evolution.
+
+    Parameters
+    ----------
+    trans_duration : float, optional
+        How long the transient lasts (days). Default 10.
+    peak_time : float, optional
+        How long it takes to reach the peak magnitude (days). Default 5.
+    rise_slope : float, optional
+        Slope of the light curve before peak time (mags/day).
+        This should be negative since mags are backwards (magnitudes decrease
+        towards brighter fluxes).
+        Default 0.
+    decline_slope : float, optional
+        Slope of the light curve after peak time (mags/day).
+        This should be positive since mags are backwards. Default 0.
+    uPeak : float, optional
+        Peak magnitude in u band. Default 20.
+    gPeak : float, optional
+        Peak magnitude in g band. Default 20.
+    rPeak : float, optional
+        Peak magnitude in r band. Default 20.
+    iPeak : float, optional
+        Peak magnitude in i band. Default 20.
+    zPeak : float, optional
+        Peak magnitude in z band. Default 20.
+    yPeak : float, optional
+        Peak magnitude in y band. Default 20.
+    survey_duration : float, optional
+        Length of survey (years).
+        Default 10.
+    survey_start : float, optional
+        MJD for the survey start date.
+        Default None (uses the time of the first observation).
+    detect_m5_plus : float, optional
+        An observation will be used if the light curve magnitude is brighter
+        than m5+detect_m5_plus.
+        Default 0.
+    n_pre_peak : int, optional
+        Number of observations (in any filter(s)) to demand before peak_time,
+        before saying a transient has been detected.
+        Default 0.
+    n_per_lc : int, optional
+        Number of sections of the light curve that must be sampled above the
+        detect_m5_plus theshold
+        (in a single filter) for the light curve to be counted.
+        For example, setting n_per_lc = 2 means a light curve  is only
+        considered detected if there is at least 1 observation in the first
+        half of the LC, and at least one in the second half of the LC.
+        n_per_lc = 4 means each quarter of the light curve must be detected to
+        count.
+        Default 1.
+    n_filters : int, optional
+        Number of filters that need to be observed for an object to be counted
+        as detected.
+        Default 1.
+    n_phase_check : int, optional
+        Sets the number of phases that should be checked.
+        One can imagine pathological cadences where many objects pass the
+        detection criteria,
+        but would not if the observations were offset by a phase-shift.
+        Default 1.
+    count_method : {'full' 'partialLC'}, defaults to 'full'
+        Sets the method of counting max number of transients. if 'full', the
+        only full light curves that fit the survey duration are counted. If
+        'partialLC', then the max number of possible transients is taken to be
+        the integer floor
+    """
+
+    def __init__(
+        self,
+        metric_name="TransientDetectMetric",
+        mjd_col="observationStartMJD",
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        trans_duration=10.0,
+        peak_time=5.0,
+        rise_slope=0.0,
+        decline_slope=0.0,
+        survey_duration=10.0,
+        survey_start=None,
+        detect_m5_plus=0.0,
+        u_peak=20,
+        g_peak=20,
+        r_peak=20,
+        i_peak=20,
+        z_peak=20,
+        y_peak=20,
+        n_pre_peak=0,
+        n_per_lc=1,
+        n_filters=1,
+        n_phase_check=1,
+        count_method="full",
+        **kwargs,
+    ):
+        self.mjd_col = mjd_col
+        self.m5_col = m5_col
+        self.filter_col = filter_col
+        super(TransientMetric, self).__init__(
+            col=[self.mjd_col, self.m5_col, self.filter_col],
+            units="Fraction Detected",
+            metric_name=metric_name,
+            **kwargs,
+        )
+        self.peaks = {
+            "u": u_peak,
+            "g": g_peak,
+            "r": r_peak,
+            "i": i_peak,
+            "z": z_peak,
+            "y": y_peak,
+        }
+        self.trans_duration = trans_duration
+        self.peak_time = peak_time
+        self.rise_slope = rise_slope
+        self.decline_slope = decline_slope
+        self.survey_duration = survey_duration
+        self.survey_start = survey_start
+        self.detect_m5_plus = detect_m5_plus
+        self.n_pre_peak = n_pre_peak
+        self.n_per_lc = n_per_lc
+        self.n_filters = n_filters
+        self.n_phase_check = n_phase_check
+        self.count_method = count_method
+
+    def light_curve(self, time, filters):
+        """
+        Calculate the magnitude of the object at each time, in each filter.
+
+        Parameters
+        ----------
+        time : numpy.ndarray
+            The times of the observations.
+        filters : numpy.ndarray
+            The filters of the observations.
+
+        Returns
+        -------
+        numpy.ndarray
+            The magnitudes of the object at each time, in each filter.
+        """
+        lc_mags = np.zeros(time.size, dtype=float)
+        rise = np.where(time <= self.peak_time)
+        lc_mags[rise] += self.rise_slope * time[rise] - self.rise_slope * self.peak_time
+        decline = np.where(time > self.peak_time)
+        lc_mags[decline] += self.decline_slope * (time[decline] - self.peak_time)
+        for key in self.peaks:
+            f_match = np.where(filters == key)
+            lc_mags[f_match] += self.peaks[key]
+        return lc_mags
+
+    def run(self, data_slice, slice_point=None):
+        """
+        Calculate the detectability of a transient with the specified
+        lightcurve.
+
+        Parameters
+        ----------
+        data_slice : numpy.array
+            Numpy structured array containing the data related to the visits
+            provided by the slicer.
+        slice_point : dict, optional
+            Dictionary containing information about the slice_point currently
+            active in the slicer.
+
+        Returns
+        -------
+        float
+            The total number of transients that could be detected.
+        """
+        # Total number of transients that could go off back-to-back
+        if self.count_method == "partialLC":
+            _n_trans_max = np.ceil(self.survey_duration / (self.trans_duration / 365.25))
+        else:
+            _n_trans_max = np.floor(self.survey_duration / (self.trans_duration / 365.25))
+        tshifts = np.arange(self.n_phase_check) * self.trans_duration / float(self.n_phase_check)
+        n_detected = 0
+        n_trans_max = 0
+        for tshift in tshifts:
+            # Compute the total number of back-to-back transients are possible
+            # to detect given the survey duration and the transient duration.
+            n_trans_max += _n_trans_max
+            if tshift != 0:
+                n_trans_max -= 1
+            if self.survey_start is None:
+                survey_start = data_slice[self.mjd_col].min()
+            time = (data_slice[self.mjd_col] - survey_start + tshift) % self.trans_duration
+
+            # Which lightcurve does each point belong to
+            lc_number = np.floor((data_slice[self.mjd_col] - survey_start) / self.trans_duration)
+
+            lc_mags = self.light_curve(time, data_slice[self.filter_col])
+
+            # How many criteria needs to be passed
+            detect_thresh = 0
+
+            # Flag points that are above the SNR limit
+            detected = np.zeros(data_slice.size, dtype=int)
+            detected[np.where(lc_mags < data_slice[self.m5_col] + self.detect_m5_plus)] += 1
+            detect_thresh += 1
+
+            # If we demand points on the rise
+            if self.n_pre_peak > 0:
+                detect_thresh += 1
+                ord = np.argsort(data_slice[self.mjd_col])
+                data_slice = data_slice[ord]
+                detected = detected[ord]
+                lc_number = lc_number[ord]
+                time = time[ord]
+                ulc_number = np.unique(lc_number)
+                left = np.searchsorted(lc_number, ulc_number)
+                right = np.searchsorted(lc_number, ulc_number, side="right")
+                # Note here I'm using np.searchsorted to basically do a
+                # 'group by' might be clearer to use
+                # scipy.ndimage.measurements.find_objects or pandas, but this
+                # numpy function is known for being efficient.
+                for le, ri in zip(left, right):
+                    # Number of points where there are a detection
+                    good = np.where(time[le:ri] < self.peak_time)
+                    nd = np.sum(detected[le:ri][good])
+                    if nd >= self.n_pre_peak:
+                        detected[le:ri] += 1
+
+            # Check if we need multiple points per light curve
+            # or multiple filters
+            if (self.n_per_lc > 1) | (self.n_filters > 1):
+                # make sure things are sorted by time
+                ord = np.argsort(data_slice[self.mjd_col])
+                data_slice = data_slice[ord]
+                detected = detected[ord]
+                lc_number = lc_number[ord]
+                time = time[ord]
+                ulc_number = np.unique(lc_number)
+                left = np.searchsorted(lc_number, ulc_number)
+                right = np.searchsorted(lc_number, ulc_number, side="right")
+                detect_thresh += self.n_filters
+
+                for le, ri in zip(left, right):
+                    points = np.where(detected[le:ri] > 0)
+                    ufilters = np.unique(data_slice[self.filter_col][le:ri][points])
+                    phase_sections = np.floor(time[le:ri][points] / self.trans_duration * self.n_per_lc)
+                    for filt_name in ufilters:
+                        good = np.where(data_slice[self.filter_col][le:ri][points] == filt_name)
+                        if np.size(np.unique(phase_sections[good])) >= self.n_per_lc:
+                            detected[le:ri] += 1
+
+            # Find the unique number of light curves that passed the required
+            # number of conditions
+            n_detected += np.size(np.unique(lc_number[np.where(detected >= detect_thresh)]))
+
+        # Rather than keeping a single "detected" variable, maybe make a mask
+        # for each criteria, then reduce functions like: reduce_singleDetect,
+        # reduce_NDetect, reduce_PerLC, reduce_perFilter.
+        # The way I'm running now it would speed things up.
+
+        return float(n_detected) / n_trans_max

rubin_sim/maf/metrics/uniformity_metrics.py

@@ -0,0 +1,418 @@
+__all__ = ("SingleLinearMultibandModelMetric", "NestedLinearMultibandModelMetric")
+
+import healpy as hp
+import numpy as np
+
+from .base_metric import BaseMetric
+from .exgal_m5 import ExgalM5
+from .weak_lensing_systematics_metric import ExgalM5WithCuts, RIZDetectionCoaddExposureTime
+
+
+class SingleLinearMultibandModelMetric(BaseMetric):
+    """Calculate a single linear combination of depths.
+
+    This is useful for calculating density or redshift fluctuations
+    from depth fluctuations on the sky,
+    for a single redshift bins (thus it is a single metric).
+    For multiple bins, see NestedLinearMultibandModelMetric.
+
+    Parameters
+    ----------
+    arr_of_model_dicts : `dict`
+        Linear coefficients to be applied to the M5 depth values.
+        Keys should be filter names + 'cst' if there is a constant offset.
+    post_processing_fn : lambda
+        Post processing function to apply to the linear combination of inputs.
+        Default `lambda x: x` simply returns the output.
+    extinction_cut : `float`, optional
+        E(B-V) cut on extinction (0.2 by default).
+    n_filters : `int`, optional
+        Cut on the number of filters required (6 by default).
+    min_depth_cut : `dict`, optional
+        Cut the areas of low depths, based on cuts in this dict.
+        Keys should be filter names. Default is no cuts.
+    max_depth_cut : `dict`, optional
+        Cut the areas of high depths, based on cuts in this dict.
+        Keys should be filter names. Default is no cuts.
+    mean_depth : `dict`, optional
+        Mean depths, which will be subtracted from the inputs
+        before applying model.
+        Keys should be filter names. Default is zero in each band.
+        In a lot of cases, instead of feeding mean_depths,
+        one can remove the monopole in the constructed healpix maps.
+    m5_col : `str`, optional
+        Column name for the m5 depth.
+    filter_col : `str`, optional
+        Column name for the filter.
+    units : `str`, optional
+        Label for "units" in the output, for use in plots.
+    badval : `float`, optional
+        Value to return for metric failure.
+        Specified here so that exgalm5 uses the same value.
+
+    Returns
+    -------
+    result : `float`
+        Linear combination of the M5 depths (minus mean depth, if provided).
+        result = np.sum([
+                model_dict[band] * M5_depth[band]
+                for band in ["u", "g", "r", "i", "z", "y"]
+                ])
+        if post_processing_fn is provided:
+        result => post_processing_fn(result)
+    """
+
+    def __init__(
+        self,
+        model_dict,
+        post_processing_fn=lambda x: x,
+        extinction_cut=0.2,
+        n_filters=6,
+        min_depth_cut=None,
+        max_depth_cut=None,
+        mean_depth=None,
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        units="mag",
+        badval=np.nan,
+        **kwargs,
+    ):
+        if min_depth_cut is None:
+            min_depth_cut = {}
+        self.min_depth_cut = min_depth_cut
+        if max_depth_cut is None:
+            max_depth_cut = {}
+        self.max_depth_cut = max_depth_cut
+        if mean_depth is None:
+            # Set mean_depth to 0 if not specified
+            mean_depth = dict([(f, 0) for f in "ugrizy"])
+        self.mean_depth = mean_depth
+
+        self.n_filters = n_filters
+        self.extinction_cut = extinction_cut
+        if "cst" in model_dict:
+            self.cst = model_dict["cst"]
+        else:
+            self.cst = 0.0
+        lsst_filters = ["u", "g", "r", "i", "z", "y"]
+        self.bands = [x for x in model_dict.keys() if x in lsst_filters]
+        self.model_arr = np.array([model_dict[x] for x in self.bands])[:, None]
+
+        self.m5_col = m5_col
+        self.filter_col = filter_col
+        self.post_processing_fn = post_processing_fn
+
+        self.exgal_m5 = ExgalM5(m5_col=m5_col, filter_col=filter_col, badval=badval)
+        super().__init__(
+            col=[m5_col, filter_col], units=units, maps=self.exgal_m5.maps, badval=badval, **kwargs
+        )
+
+    def run(self, data_slice, slice_point=None):
+        """ """
+
+        if slice_point["ebv"] > self.extinction_cut:
+            return self.badval
+
+        n_filters = len(set(data_slice[self.filter_col]))
+        if n_filters < self.n_filters:
+            return self.badval
+
+        # add extinction_cut and depth cuts? and n_filters cuts
+        depths = np.vstack(
+            [
+                self.exgal_m5.run(data_slice[data_slice[self.filter_col] == lsst_filter], slice_point)
+                for lsst_filter in self.bands
+            ]
+        ).T
+        for i, lsst_filter in enumerate(self.bands):
+            if lsst_filter in self.mean_depth:
+                depths[i] -= self.mean_depth[lsst_filter]
+        # Don't raise Exceptions in the middle of metrics -
+        # assert depths.shape[0] >= 1
+        if depths.shape[0] < 1:
+            return self.badval
+
+        if np.any(depths == self.badval):
+            return self.badval
+
+        for i, lsst_filter in enumerate(self.bands):
+            if (
+                lsst_filter in self.min_depth_cut
+                and depths[0, i] < self.min_depth_cut[lsst_filter] - self.mean_depth[lsst_filter]
+            ):
+                return self.badval
+            if (
+                lsst_filter in self.max_depth_cut
+                and depths[0, i] > self.max_depth_cut[lsst_filter] - self.mean_depth[lsst_filter]
+            ):
+                return self.badval
+
+        val = self.post_processing_fn(self.cst + np.dot(depths, self.model_arr))
+        return val
+
+
+class NestedLinearMultibandModelMetric(BaseMetric):
+    """Calculate multiple linear combinations of depths.
+
+    This is useful for calculating density or redshift fluctuations
+    from depth fluctuations on the sky,
+    for multiple redshift bins (thus it is a nested metric).
+    For a single bin, see LinearMultibandModelMetric.
+
+    Points of contact / contributors: Boris Leistedt
+
+    Parameters
+    ----------
+    arr_of_model_dicts : `list` [ `dicts` ]
+        Array of linear coefficients to be applied to the M5 depth values.
+        Keys should be filter names + 'cst' if there is a constant offset.
+    post_processing_fn : lambda
+        Post processing function to apply to the linear combination of inputs.
+        Default `lambda x: x` simply returns the output.
+    extinction_cut : `float`, optional
+        E(B-V) cut on extinction (0.2 by default).
+    n_filters : `int`, optional
+        Cut on the number of filters required (6 by default).
+    min_depth_cut : `dict`, optional
+        Cut the areas of low depths, based on cuts in this dict.
+        Keys should be filter names. Default is no cuts.
+    max_depth_cut : `dict`, optional
+        Cut the areas of high depths, based on cuts in this dict.
+        Keys should be filter names. Default is no cuts.
+    mean_depth : `dict`, optional
+        Mean depths, which will be subtracted from the inputs
+        before applying model.
+        Keys should be filter names. Default is zero in each band.
+        In a lot of cases, instead of feeding mean_depths,
+        one can remove the monopole in the constructed healpix maps.
+    m5_col : `str`, optional
+        Column name for the m5 depth.
+    filter_col : `str`, optional
+        Column name for the filter.
+    units : `str`, optional
+        Label for "units" in the output, for use in plots.
+    badval : `float`, optional
+        Value to return for metric failure.
+        Specified here so that exgalm5 uses the same value.
+
+    Returns
+    -------
+    result : `list` [ `float` ]
+        List is the same size as arr_of_model_dicts.
+        For each element, return a linear combination of the M5 depths
+        (minus mean depth, if provided).
+        for i, model_dict in enumerate(arr_of_model_dicts):
+            result[i] = np.sum([
+                model_dict[band] * M5_depth[band]
+                for band in ["u", "g", "r", "i", "z", "y"]
+                ])
+        if post_processing_fn is provided:
+        result[i] => post_processing_fn(result[i])
+    """
+
+    def __init__(
+        self,
+        arr_of_model_dicts,
+        post_processing_fn=lambda x: x,
+        extinction_cut=0.2,
+        n_filters=6,
+        min_depth_cut=None,
+        max_depth_cut=None,
+        mean_depth=None,
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        units="mag",
+        badval=np.nan,
+        **kwargs,
+    ):
+        self.arr_of_model_dicts = arr_of_model_dicts
+        self.n_filters = n_filters
+        self.extinction_cut = extinction_cut
+        if min_depth_cut is None:
+            min_depth_cut = {}
+        self.min_depth_cut = min_depth_cut
+        if max_depth_cut is None:
+            max_depth_cut = {}
+        self.max_depth_cut = max_depth_cut
+        if mean_depth is None:
+            mean_depth = {}
+        self.mean_depth = mean_depth
+
+        self.m5_col = m5_col
+        self.badval = badval
+        self.mask_val = hp.UNSEEN
+        self.filter_col = filter_col
+        self.post_processing_fn = post_processing_fn
+
+        self.n_bins = len(self.arr_of_model_dicts)
+        self.bad_val_arr = np.repeat(self.badval, self.n_bins)
+
+        self.exgal_m5 = ExgalM5(
+            m5_col=m5_col,
+            filter_col=filter_col,
+            badval=self.badval,
+        )
+        super().__init__(
+            col=[m5_col, filter_col],
+            badval=self.badval,
+            units=units,
+            maps=self.exgal_m5.maps,
+            metric_dtype="object",
+            **kwargs,
+        )
+
+    def run(self, data_slice, slice_point=None):
+        # apply extinction and n_filters cut
+        if slice_point["ebv"] > self.extinction_cut:
+            return self.bad_val_arr
+
+        n_filters = len(set(data_slice[self.filter_col]))
+        if n_filters < self.n_filters:
+            return self.bad_val_arr
+
+        lsst_filters = ["u", "g", "r", "i", "z", "y"]
+
+        # initialize dictionary of outputs
+        # types = [float]*n_bins
+        result_arr = np.zeros((self.n_bins,), dtype=float)
+
+        for binnumber, model_dict in enumerate(self.arr_of_model_dicts):
+
+            # put together a vector of depths
+            depths = np.vstack(
+                [
+                    self.exgal_m5.run(data_slice[data_slice[self.filter_col] == lsst_filter], slice_point)
+                    for lsst_filter in lsst_filters
+                ]
+            ).T
+
+            # if there are depth cuts, apply them
+            for i, lsst_filter in enumerate(lsst_filters):
+                if lsst_filter in self.min_depth_cut and depths[0, i] < self.min_depth_cut[lsst_filter]:
+                    depths[0, i] = self.badval
+                if lsst_filter in self.max_depth_cut and depths[0, i] > self.max_depth_cut[lsst_filter]:
+                    depths[0, i] = self.badval
+
+            # if provided, subtract the mean
+            for i, lsst_filter in enumerate(lsst_filters):
+                if lsst_filter in self.mean_depth:
+                    depths[i] -= self.mean_depth[lsst_filter]
+
+            # now assemble the results
+            model_arr = np.array([model_dict[x] if x in model_dict else 0.0 for x in lsst_filters])
+            cst = model_dict["cst"] if "cst" in model_dict else 0.0
+
+            result_arr[binnumber] = self.post_processing_fn(cst + np.dot(depths, model_arr))
+
+        # returns array where each element of the original
+        # input arr_of_model_dicts contains a scalar value resulting
+        # from a linear combination of the six depths
+        return result_arr
+
+
+class NestedRIZExptimeExgalM5Metric(BaseMetric):
+    """TODO"""
+
+    def __init__(
+        self,
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        exptime_col="visitExposureTime",
+        extinction_cut=0.2,
+        n_filters=6,
+        depth_cut=24,
+        metric_name="new_nested_FoM",
+        lsst_filter="i",
+        badval=np.nan,
+        **kwargs,
+    ):
+        maps = ["DustMap"]
+        self.m5_col = m5_col
+        self.filter_col = filter_col
+        self.exptime_col = exptime_col
+        self.lsst_filter = lsst_filter
+        self.n_filters = n_filters
+
+        cols = [self.m5_col, self.filter_col, self.exptime_col]
+        super().__init__(cols, metric_name=metric_name, maps=maps, badval=badval, **kwargs)
+        self.riz_exptime_metric = RIZDetectionCoaddExposureTime(ebvlim=extinction_cut)
+        self.exgalm5_metric = ExgalM5WithCuts(
+            m5_col=m5_col,
+            filter_col=filter_col,
+            extinction_cut=extinction_cut,
+            depth_cut=depth_cut,
+            lsst_filter=lsst_filter,
+            badval=badval,
+            n_filters=n_filters,
+        )
+
+        self.metric_dtype = "object"
+
+    def run(self, data_slice, slice_point):
+
+        # set up array to hold the results
+        names = ["exgal_m5", "riz_exptime"]
+        types = [float] * 2
+        result = np.zeros(1, dtype=list(zip(names, types)))
+        result["exgal_m5"] = self.exgalm5_metric.run(data_slice, slice_point)
+        result["riz_exptime"] = self.riz_exptime_metric.run(
+            data_slice[data_slice[self.filter_col] == "i"], slice_point
+        )
+
+        return result
+
+
+class MultibandExgalM5(BaseMetric):
+    """Calculate multiple linear combinations of depths."""
+
+    def __init__(
+        self,
+        extinction_cut=0.2,
+        n_filters=6,
+        m5_col="fiveSigmaDepth",
+        filter_col="band",
+        units="mag",
+        badval=np.nan,
+        **kwargs,
+    ):
+        self.n_filters = n_filters
+        self.extinction_cut = extinction_cut
+
+        self.m5_col = m5_col
+        self.badval = badval
+        self.filter_col = filter_col
+
+        self.lsst_filters = ["u", "g", "r", "i", "z", "y"]
+        self.bad_val_arr = np.repeat(self.badval, len(self.lsst_filters))
+
+        self.exgal_m5 = ExgalM5(
+            m5_col=m5_col,
+            filter_col=filter_col,
+            badval=self.badval,
+        )
+        super().__init__(
+            col=[m5_col, filter_col],
+            badval=self.badval,
+            units=units,
+            maps=self.exgal_m5.maps,
+            metric_dtype="object",
+            **kwargs,
+        )
+
+    def run(self, data_slice, slice_point=None):
+        # apply extinction and n_filters cut
+        if slice_point["ebv"] > self.extinction_cut:
+            return self.bad_val_arr
+
+        n_filters = len(set(data_slice[self.filter_col]))
+        if n_filters < self.n_filters:
+            return self.bad_val_arr
+
+        depths = np.vstack(
+            [
+                self.exgal_m5.run(data_slice[data_slice[self.filter_col] == lsst_filter], slice_point)
+                for lsst_filter in self.lsst_filters
+            ]
+        ).T
+
+        return depths.ravel()