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

rubin_sim/maf/maf_contrib/presto_color_kne_pop_metric.py

@@ -0,0 +1,441 @@
+__all__ = ("PrestoColorKNePopMetric", "generate_presto_pop_slicer")
+
+import os
+import pickle
+import warnings
+from itertools import combinations
+
+import numpy as np
+import pandas as pd
+from rubin_scheduler.data import get_data_dir
+from rubin_scheduler.utils import SURVEY_START_MJD, uniform_sphere
+
+import rubin_sim.maf.metrics as metrics
+import rubin_sim.maf.slicers as slicers
+from rubin_sim.phot_utils import DustValues
+
+from .kne_metrics import KnLc
+
+
+def radec2gal(ra, dec):
+    """convert from ra/dec to galactic l/b"""
+    from astropy import units as u
+    from astropy.coordinates import SkyCoord
+
+    c = SkyCoord(ra=ra, dec=dec, unit=(u.degree, u.degree))
+
+    gal_l = c.galactic.l.degree
+    gal_b = c.galactic.b.degree
+    return gal_l, gal_b
+
+
+def _load_hash(
+    file_galactic="TotalCubeNorm_1000Obj.pkl",
+    file_extragalactic="TotalCubeNorm_1000Obj.pkl",
+    skyregion="extragalactic",
+):
+    """Helper function to load large hash table.
+
+    Because this is kept outside the metric attributes, it allows easy resuse.
+    Note that this does mean running different sky regions at the same
+    time may end up thrashing the data in the hash/hash table.
+
+    Parameters
+    ----------
+    skyregion : `str`
+        The skyregion of interst.
+        Only two options: 'galactic' and 'extragalactic'
+    filePathGalactic : `str`
+        File containing galactic Presto-Color phase space information
+    filePathExtragalactic : `str`
+        File containing galactic Presto-Color phase space information
+    """
+
+    if hasattr(_load_hash, "InfoDict"):
+        if skyregion == _load_hash.skyregion:
+            return _load_hash.InfoDict, _load_hash.HashTable
+
+    data_dir = get_data_dir()
+    if skyregion == "galactic":
+        file_path = os.path.join(data_dir, "maf", file_galactic)
+    elif skyregion == "extragalactic":
+        file_path = os.path.join(data_dir, "maf", file_extragalactic)
+
+    with open(file_path, "rb") as f:
+        _load_hash.InfoDict = pickle.load(f)
+        _load_hash.HashTable = pickle.load(f)
+    _load_hash.skyregion = skyregion
+    return _load_hash.InfoDict, _load_hash.HashTable
+
+
+def generate_presto_pop_slicer(
+    skyregion="galactic",
+    t_start=1,
+    t_end=3652,
+    n_events=10000,
+    seed=42,
+    n_files=100,
+    d_min=10,
+    d_max=300,
+    gb_cut=20,
+):
+    """Generate a population of KNe events, and put the info about them
+    into a UserPointSlicer object
+
+    Parameters
+    ----------
+    skyregion : `str`
+        The skyregion of interest.
+        Only two options: 'galactic' and 'extragalactic'
+    t_start : `float`
+        The night to start kilonova events on (days)
+    t_end : `float`
+        The final night of kilonova events
+    n_events : `int`
+        The number of kilonova events to generate
+    seed : `float`
+        The seed passed to np.random
+    n_files : `int`
+        The number of different kilonova lightcurves to use
+    d_min : `float` or `int`
+        Minimum luminosity distance (Mpc)
+    d_max : `float` or `int`
+        Maximum luminosity distance (Mpc)
+
+    Returns
+    -------
+    kne_slicer : `~.maf.UserPointsSlicer`
+    """
+
+    def rndm(a, b, g, size=1):
+        """Power-law gen for pdf(x) proportional to 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)
+
+    ra, dec = uniform_sphere(n_events, seed=seed)
+
+    # Convert ra, dec to gl, gb
+    gl, gb = radec2gal(ra, dec)
+
+    # Determine if the object is in the Galaxy plane
+    if skyregion == "galactic":  # keep the galactic events
+        ra = ra[np.abs(gb) < gb_cut]
+        dec = dec[np.abs(gb) < gb_cut]
+    elif skyregion == "extragalactic":  # keep the extragalactic events.
+        ra = ra[np.abs(gb) > gb_cut]
+        dec = dec[np.abs(gb) > gb_cut]
+    else:
+        warnings.warn("Skyregion %s not recognized, using whole sky" % skyregion)
+
+    n_events = len(ra)
+
+    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 = slicers.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
+
+
+class PrestoColorKNePopMetric(metrics.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,
+        skyregion="galactic",
+        thr=0.003,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        file_list : `str` or None, optional
+            File containing input lightcurves
+        mjd0 : `float`, optional
+            MJD of the start of the survey.
+        output_lc : `bool`, optional
+            Flag to whether or not to output lightcurve for each object.
+        skyregion : `str`, optional
+            The skyregion of interest.
+            Only two options: 'galactic' and 'extragalactic'
+        thr : `float`, optional
+            Threshold for "classification" of events via the Score_S
+        """
+        maps = ["DustMap"]
+        self.mjd_col = mjd_col
+        self.m5_col = m5_col
+        self.filter_col = filter_col
+        self.night_col = night_col
+        # Boolean variable, if True the light curve will be exported
+        self.output_lc = output_lc
+        self.thr = thr
+        self.skyregion = skyregion
+        # read in file as light curve object;
+        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().__init__(col=cols, units="Detected, 0 or 1", metric_name=metric_name, maps=maps, **kwargs)
+
+        # Unused ..
+        self.pts_needed = pts_needed
+
+    def _presto_color_detect(self, around_peak, filters):
+        """Detection criteria of presto cadence:
+        at least three detections at two filters;
+
+        Parameters
+        ----------
+        around_peak : `np.ndarray`, (N,)
+            indexes corresponding to 5sigma detections
+        filters : `np.ndarray`, (N,)
+            filters in which detections happened
+        """
+        result = 1
+
+        if np.size(around_peak) < 3:
+            result = 0
+
+        flts, flts_count = np.unique(
+            filters,
+            return_counts=True,
+        )
+        if np.size(flts) < 2:
+            result = 0
+        elif np.max(flts_count) < 2:
+            # if no filters have visits larger than 2, set detection false
+            result = 0
+
+        return result
+
+    def _enquiry(self, hash_table, info_dict, band1, band2, d_t1, d_t2, d_mag, color):
+        """
+        Return the value in the probability cube provided the coordinates
+        in the Presto-color phase space of an observation triplet.
+
+        Parameters
+        ----------
+        hash_table : `np.ndarray`, (N,)
+            Contains the values of the 6-D Presto-color phase space
+        info_dict : `dict`
+            Contains the essential information of the hash_table abobe.
+        band1, band2 : `str`, `str`
+            The two filters that comprise the Presto-color observation triplet.
+            The filters are the 6 bands of LSST: u, g, r, i, z, y.
+            Band1 and band2 should be different.
+        d_t1, d_t2 : `float`, `float`
+            The time gaps of the Presto-color observation triplet.
+        d_mag : `float`
+            The magnitude change between from the observations of the same band
+        color : `float`
+            The difference in magnitude of observations in different bands.
+
+        hash_table and info_dict have to be loaded from premade data
+        Presto-color data file.
+        """
+
+        #         if abs(d_t1) > abs(d_t1-d_t2):
+        #             d_t1, d_t2 = d_t1-d_t2, -d_t2
+
+        if not (
+            info_dict["BinMag"][0] <= d_mag < info_dict["BinMag"][-1]
+            and info_dict["BinColor"][0] <= color < info_dict["BinColor"][-1]
+        ):
+            return 0
+
+        ind1 = info_dict["BandPairs"].index(band1 + band2)
+
+        time_pair_grid = [
+            info_dict["dT1s"][abs(d_t1 - info_dict["dT1s"]).argmin()],
+            info_dict["dT2s"][abs(d_t2 - info_dict["dT2s"]).argmin()],
+        ]
+
+        ind2 = np.where((info_dict["TimePairs"] == time_pair_grid).all(axis=1))[0][0]
+        ind3 = np.where(d_mag >= info_dict["BinMag"])[0][-1]
+        ind4 = np.where(color >= info_dict["BinColor"])[0][-1]
+
+        return hash_table[ind1, ind2, ind3, ind4]
+
+    def _get_score(self, result, hash_table, info_dict, thr):
+        """Get the score of a strategy from the Presto-color perspective.
+
+        Parameters
+        ----------
+        result : `pd.DataFrame`
+            Dataframe that contains the results of the observations.
+            The columns include
+            t: the time of the observation
+            mag: the detected magnitude
+            maglim: the limit fiveSigmaDepth that can be detected
+            filter: the filter used for the observation
+        hash_table : `np.ndarray`, (N,)
+            Contains the values of the 6-D Presto-color phase space
+        info_dict : `dict`
+            Contains the essential information of the hash_table abobe.
+        scoreType : `str`
+            Two types of scores were designed:
+            'S' type involves a threshold,
+            'P' type work without a threshold.
+        thr : `float`
+            The threashold need for type 'S' score.
+            The default value is 0.003 (3-sigma)
+
+        hash_table and info_dict have to be loaded from the premade
+        Presto-color data file.
+        """
+
+        time_lim1 = 8.125 / 24  # 8 h 7.5 min
+        time_lim2 = 32.25 / 24  # 32 h 15 min
+
+        detects = result[result.mag < result.maglim]
+
+        # reset index
+        detects = detects.reset_index(drop=True)
+
+        # Times for valid detections
+        ts = detects.t.values
+        # Find out the differences between each pair
+        d_ts = ts.reshape(1, len(ts)) - ts.reshape(len(ts), 1)
+
+        # The time differences should be within 32 hours (2 nights)
+        d_tindex0, d_tindex1 = np.where(abs(d_ts) < time_lim2)
+
+        phase_space_coords = []
+
+        # loop through the rows of the matrix of valid time differences
+        for ii in range(d_ts.shape[0]):
+            groups_of_three = np.array(
+                [
+                    [ii] + list(jj)
+                    for jj in list(combinations(d_tindex1[(d_tindex0 == ii) * (d_tindex1 > ii)], 2))
+                ]
+            )
+
+            for indices in groups_of_three:
+                bands = detects["filter"][indices].values
+
+                # print('bands: ', bands)
+                if len(np.unique(bands)) != 2:
+                    continue
+
+                # The band appears once will be band2
+                occurence = np.array([np.count_nonzero(ii == bands) for ii in bands])
+                # The index of observation in band2
+                index2 = indices[occurence == 1][0]
+                # The index of the first observation in band1
+                index11 = indices[occurence == 2][0]
+                # The index of the second observation in band1
+                index12 = indices[occurence == 2][1]
+
+                if (
+                    abs(d_ts[index12, index2]) < abs(d_ts[index11, index2])
+                    and abs(d_ts[index12, index2]) < time_lim1
+                ):
+                    index11, index12 = index12, index11
+                elif abs(d_ts[index11, index2]) > time_lim1:
+                    continue
+
+                d_t1 = d_ts[index11, index2]
+                d_t2 = d_ts[index11, index12]
+
+                band1 = bands[occurence == 2][0]
+                band2 = bands[occurence == 1][0]
+
+                if band1 + band2 == "uy" or band1 + band2 == "yu":
+                    continue
+
+                d_mag = (detects.mag[index11] - detects.mag[index12]) * np.sign(d_t2)
+                color = detects.mag[index11] - detects.mag[index2]
+
+                phase_space_coords.append([band1, band2, d_t1, d_t2, d_mag, color])
+
+        score_s = 0
+        score_p = [0]
+
+        for phase_space_coord in phase_space_coords:
+            rate = self._enquiry(hash_table, info_dict, *phase_space_coord)
+
+            if score_s == 0 and rate < thr:
+                score_s = 1
+
+            score_p.append((1 - rate))
+
+        return score_s, max(score_p)
+
+    def run(self, data_slice, slice_point=None):
+        data_slice.sort(order=self.mjd_col)
+        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]
+
+        # presto color
+        result["presto_color_detect"] = self._presto_color_detect(around_peak, filters)
+
+        # Export the light curve
+        idx = np.where(mags < 100)[0]
+        lc = {
+            "t": data_slice[self.mjd_col][idx],
+            "mag": mags[idx],
+            "maglim": data_slice[self.m5_col][idx],
+            "filter": data_slice[self.filter_col][idx],
+        }
+
+        if self.output_lc is True:
+            result["lc"] = lc
+            result["slice_point"] = slice_point
+
+        if result["presto_color_detect"] == 1:
+            info_dict, hash_table = _load_hash(skyregion=self.skyregion)
+            result["scoreS"], result["scoreP"] = self._get_score(
+                pd.DataFrame(lc),
+                hash_table=hash_table,
+                info_dict=info_dict,
+                thr=self.thr,
+            )
+        else:
+            result["scoreS"] = 0
+            result["scoreP"] = 0
+        return result
+
+    def reduce_presto_color_detect(self, metric):
+        return metric["presto_color_detect"]
+
+    def reduce_score_s(self, metric):
+        return metric["scoreS"]
+
+    def reduce_score_p(self, metric):
+        return metric["scoreP"]

rubin_sim/maf/maf_contrib/selfcal_uniformity_metric.py

@@ -0,0 +1,171 @@
+__all__ = ("PhotometricSelfCalUniformityMetric",)
+
+import os
+import time
+
+import healpy as hp
+import numpy as np
+from astropy.coordinates import SkyCoord
+from astropy.table import Table
+from rubin_scheduler.data import get_data_dir
+from rubin_scheduler.utils import healbin
+from scipy.stats import median_abs_deviation
+
+from rubin_sim.maf.metrics import BaseMetric
+from rubin_sim.selfcal import LsqrSolver, OffsetSNR, generate_catalog
+from rubin_sim.selfcal.offsets import OffsetSys
+
+
+def _match(arr1, arr2):
+    st1 = np.argsort(arr1)
+    sub1 = np.searchsorted(arr1, arr2, sorter=st1)
+    if arr2.max() > arr1.max():
+        bad = sub1 == arr1.size
+        sub1[bad] = arr1.size - 1
+
+    (sub2,) = np.where(arr1[st1[sub1]] == arr2)
+    sub1 = st1[sub1[sub2]]
+
+    return sub1, sub2
+
+
+class PhotometricSelfCalUniformityMetric(BaseMetric):
+    def __init__(
+        self,
+        nside_residual=128,
+        highglat_cut=30.0,
+        outlier_nsig=4.0,
+        metric_name="PhotometricSelfCalUniformityMetric",
+        filter_name="r",
+    ):
+        cols = [
+            "observationid",
+            "fieldra",
+            "fielddec",
+            "fiveSigmaDepth",
+            "rotSkyPos",
+            "filter",
+        ]
+        super().__init__(col=cols, metric_name=metric_name)
+
+        filename = os.path.join(get_data_dir(), "maf", "monster_stars_uniformity_i15-18_sampled.parquet")
+        self.stars = Table.read(filename)
+
+        # We have to rename dec to decl for the selfcal code.
+        if "dec" in self.stars.dtype.names:
+            self.stars["decl"] = self.stars["dec"]
+
+        self.stars = self.stars.as_array()
+
+        self.nside_residual = nside_residual
+        self.highglat_cut = highglat_cut
+        self.outlier_nsig = outlier_nsig
+
+        self.units = "mmag"
+        self.filter_name = filter_name
+
+    def run(self, data_slice, slice_point=None):
+        offsets = [OffsetSys(error_sys=0.03), OffsetSNR(lsst_filter=self.filter_name)]
+
+        visits = np.zeros(
+            len(data_slice),
+            dtype=[
+                ("observationId", "i8"),
+                ("ra", "f8"),
+                ("dec", "f8"),
+                ("fiveSigmaDepth", "f8"),
+                ("rotSkyPos", "f8"),
+            ],
+        )
+        visits["observationId"] = data_slice["observationId"]
+        visits["ra"] = data_slice["fieldRA"]
+        visits["dec"] = data_slice["fieldDec"]
+        visits["fiveSigmaDepth"] = data_slice["fiveSigmaDepth"]
+        visits["rotSkyPos"] = data_slice["rotSkyPos"]
+
+        good_stars = np.isfinite(self.stars[f"{self.filter_name}mag"])
+        stars = self.stars[good_stars]
+
+        observed_stars = generate_catalog(
+            visits,
+            stars,
+            offsets=offsets,
+            lsst_filter=self.filter_name,
+            n_patches=16,
+            verbose=False,
+        )
+
+        solver = LsqrSolver(observed_stars)
+
+        print("Starting solver...")
+        t0 = time.time()
+
+        solver.run()
+        fit_patches, fit_stars = solver.return_solution()
+
+        t1 = time.time()
+        dt = (t1 - t0) / 60.0
+        print("runtime= %.1f min" % dt)
+
+        # Trim stars to the ones with solutions.
+        a, b = _match(stars["id"], fit_stars["id"])
+        stars_trimmed = stars[a]
+        fit_stars_trimmed = fit_stars[b]
+
+        # Residuals after fit, removing floating zeropoint
+        resid = stars_trimmed[f"{self.filter_name}mag"] - fit_stars_trimmed["fit_mag"]
+        resid = resid - np.median(resid)
+
+        resid_map = healbin(
+            stars_trimmed["ra"],
+            stars_trimmed["dec"],
+            resid,
+            self.nside_residual,
+            reduce_func=np.median,
+        )
+
+        (ipring,) = np.where(resid_map > hp.UNSEEN)
+
+        scatter_full = median_abs_deviation(resid_map[ipring], scale="normal")
+
+        # Fraction of (4) sigma outliers.
+        highscat = np.abs(resid_map[ipring]) > self.outlier_nsig * scatter_full
+        outlier_frac_full = highscat.sum() / len(ipring)
+
+        # Cut to high latitude
+        ra, dec = hp.pix2ang(self.nside_residual, ipring, nest=False, lonlat=True)
+        coords = SkyCoord(ra, dec, frame="icrs", unit="deg")
+        b = coords.galactic.b.value
+
+        high_glat = np.abs(b) > self.highglat_cut
+        scatter_highglat = median_abs_deviation(resid_map[ipring[high_glat]], scale="normal")
+
+        # Fraction of (4) sigma outliers.
+        highscat = np.abs(resid_map[ipring[high_glat]]) > self.outlier_nsig * scatter_highglat
+        outlier_frac_highglat = highscat.sum() / len(ipring)
+
+        # Convert to mmag
+        scatter_full = scatter_full * 1000.0
+        scatter_highglat = scatter_highglat * 1000.0
+
+        result = {
+            "scatter_full": scatter_full,
+            "scatter_highglat": scatter_highglat,
+            "outlier_frac_full": outlier_frac_full,
+            "outlier_frac_highglat": outlier_frac_highglat,
+            "uniformity_map": resid_map,
+        }
+
+        return result
+
+    def reduce_scatter_full(self, metric_value):
+        return metric_value["scatter_full"]
+
+    def reduce_scatter_highglat(self, metric_value):
+        return metric_value["scatter_highglat"]
+
+    def reduce_outlier_frac_full(self, metric_value):
+        return metric_value["outlier_frac_full"]
+
+    def reduce_outlier_frac_highglat(self, metric_value):
+        return metric_value["outlier_frac_highglat"]

rubin_sim/maf/maf_contrib/star_count_mass_metric.py

@@ -0,0 +1,64 @@
+__all__ = ("StarCountMassMetric",)
+
+import numpy as np
+
+from rubin_sim.maf.metrics import BaseMetric
+
+from .star_counts import starcount_bymass
+
+# Example for CountMassMetric
+# Mike Lund - Vanderbilt University
+# mike.lund@gmail.com
+# Last edited 8/15/2015
+# Motivation: The distances to stars in LSST will be significant enough that
+# the structure of the Galaxy will be readily apparent because of its
+# influence on the number of stars in a given field.
+# Any metric concerned with the number of potential objects to be detected
+# will need to feature not only the effects of the cadence but also the
+# number of objects per field.
+# This metric identifies the number of stars in a given field in a particular
+# mass range that will be fainter than the saturation limit of 16th magnitude
+# and still bright enough to have noise less than 0.03 mag.
+# M1 and M2 are the low and high limits of the mass range in solar masses.
+# 'band' is the band for the observations to be made in.
+# Requires StarCounts.StarCounts
+
+# NOTE
+# There are stellar luminosity function maps available within MAF
+# that may supersede these StarCount functions
+
+
+class StarCountMassMetric(BaseMetric):
+    """Find the number of stars in a given field in the mass range
+    fainter than magnitude 16 and bright enough to have noise less than
+    0.03 in a given band.
+    M1 and m2 are the upper and lower limits of the mass range.
+    'band' is the band to be observed.
+
+    This metric uses the stellar distance and luminosity equations
+    contributed by Mike Lund, which are based on the Galfast model.
+    There are some imposed limitations on the expected magnitudes
+    of the stars included for the metric, based on assuming saturation
+    at 16th magnitude and not considering stars with magnitude
+    uncertainties greater than 0.03 (based on photometry/m5 alone).
+
+    Parameters
+    ----------
+    m1 : `float`
+        Lower limit of the mass range.
+    m2 : `float`
+        Upper limit of the mass range.
+    band : `str`
+        Bandpass to consider.
+    """
+
+    def __init__(self, m1=0.9, m2=1.0, band="i", **kwargs):
+        self.m1 = m1
+        self.m2 = m2
+        self.band = band
+        super(StarCountMassMetric, self).__init__(col=[], **kwargs)
+
+    def run(self, data_slice, slice_point=None):
+        self.dec_col = np.degrees(data_slice[0][3])
+        self.ra_col = np.degrees(data_slice[0][2])
+        return starcount_bymass.starcount_bymass(self.ra_col, self.dec_col, self.m1, self.m2, self.band)