U/jrbogart/snana support

cfg/object_types.yaml

@@ -60,3 +60,10 @@ object_types :
         dstype : gaia_dr2_20200414
       area_partition : None
       sed_method : use_lut
+  snana :
+      file_template : 'snana_(?P<healpix>\d+).parquet'
+      data_file_type : parquet
+      area_partition :
+        { type: healpix, ordering : ring, nside : 32}
+      sed_model : snana
+      internal_extinction : None

skycatalogs/catalog_creator.py

@@ -18,8 +18,8 @@
 from .utils.config_utils import create_config, assemble_SED_models
 from .utils.config_utils import assemble_MW_extinction, assemble_cosmology, assemble_object_types, assemble_provenance, write_yaml
 from .utils.parquet_schema_utils import make_galaxy_schema, make_galaxy_flux_schema, make_star_flux_schema, make_pointsource_schema
+from .utils.creator_utils import make_MW_extinction_av, make_MW_extinction_rv
 from .objects.base_object import LSST_BANDS
-
 from .objects.base_object import ObjectCollection
 
 # from dm stack
@@ -29,33 +29,9 @@
 Code to create a sky catalog for particular object types
 """
 
-#####pixels = [9556, 9683, 9684, 9812, 9813, 9940]
-
 __all__ = ['CatalogCreator']
 
 _MW_rv_constant = 3.1
-_Av_adjustment = 2.742
-
-def _make_MW_extinction(ra, dec):
-    '''
-    Given arrays of ra & dec, create a MW Av column corresponding to V-band
-    correction.
-    See "Plotting Dust Maps" example in
-    https://dustmaps.readthedocs.io/en/latest/examples.html
-
-    The coefficient _Av_adjustment comes Table 6 in Schlafly & Finkbeiner (2011)
-    See http://iopscience.iop.org/0004-637X/737/2/103/article#apj398709t6
-
-    Parameters
-    ----------
-    ra, dec - arrays specifying positions where Av is to be computed
-    Return:
-    Array of Av values
-    '''
-
-    sfd = SFDQuery()
-    ebv_raw = np.array(sfd.query_equ(ra, dec))
-    return _Av_adjustment * ebv_raw
 
 def _generate_sed_path(ids, subdir, cmp):
     '''
@@ -79,11 +55,11 @@ def _get_tophat_info(columns):
     '''
     Parameters
     ----------
-    columns    list of column names including the ones with per-tophat information
+    columns    list of column names including the ones with per-tophat info
 
     Returns
     -------
-    sed_bins        List of  the tophat bins, sorted by "start" (left edge) value
+    sed_bins        List of  the tophat bins, sorted by "start" (left edge)
     sed_bulge_names To be fetched from input catalog
     sed_disk_names  To be fetched from input catalog
     '''
@@ -340,17 +316,19 @@ def __init__(self, parts, area_partition=None, skycatalog_root=None,
 
     def _make_tophat_columns(self, dat, names, cmp):
         '''
-        Create columns sed_val_cmp, cmp_magnorm where cmp is one of "disk", "bulge", "knots"
+        Create columns sed_val_cmp, cmp_magnorm where cmp is one of "disk",
+        "bulge", "knots"
+
         Parameters
         ----------
-        dat          Data read from input galaxy catalog. Includes keys for everything in names
-                     plus entry for redshiftHubble
+        dat          Data read from input galaxy catalog. Includes keys for
+                     everything in names plus entry for redshiftHubble
         names        Names of SED columns for this component
         cmp          Component name
 
         Returns
         -------
-        Add keys  sed_val_cmp, cmp_magnorm to dat.  Values are associated data. Return dat.
+        Add keys  sed_val_cmp, cmp_magnorm to input dat. Then return dat.
         '''
         sed_vals = (np.array([dat[k] for k in names]).T).tolist()
         dat['sed_val_' + cmp] = sed_vals
@@ -544,8 +522,8 @@ def create_galaxy_pixel(self, pixel, gal_cat, arrow_schema):
                                         native_filters=hp_filter,
                                         filters=mag_cut_filter)
 
-        df['MW_rv'] = np.full_like(df['sersic_bulge'], _MW_rv_constant)
-        df['MW_av'] = _make_MW_extinction(df['ra'], df['dec'])
+        df['MW_rv'] = make_MW_extinction_rv(df['ra'], df['dec'])
+        df['MW_av'] = make_MW_extinction_av(df['ra'], df['dec'])
         self._logger.debug('Made extinction')
 
         # Some columns need to be renamed
@@ -582,6 +560,7 @@ def create_galaxy_pixel(self, pixel, gal_cat, arrow_schema):
         else:
             subpixel_masks = {pixel : None}
 
+
         for p, val in subpixel_masks.items():
             output_path = os.path.join(self._output_dir, f'galaxy_{p}.parquet')
             if os.path.exists(output_path):
@@ -825,12 +804,12 @@ def create_pointsource_pixel(self, pixel, arrow_schema, star_cat=None,
 
             # NOTE MW_av calculation for stars does not use SFD dust map
             star_df['MW_av'] = star_df['ebv'] * _MW_rv_constant
+
             star_df['variability_model'] = np.full((nobj,), '')
             star_df['salt2_params'] = np.full((nobj,), None)
             out_table = pa.Table.from_pandas(star_df, schema=arrow_schema)
             self._logger.debug('Created arrow table from star dataframe')
 
-            #writer = pq.ParquetWriter(output_path, arrow_schema)
             writer.write_table(out_table)
 
         if sn_cat:
@@ -850,9 +829,8 @@ def create_pointsource_pixel(self, pixel, arrow_schema, star_cat=None,
             nobj = len(sn_df['ra'])
             sn_df['object_type'] = np.full((nobj,), self._sn_object_type)
 
-            sn_df['MW_rv'] = np.full((nobj,), _MW_rv_constant, np.float32())
-            sn_df['MW_av'] = _make_MW_extinction(np.array(sn_df['ra']),
-                                                   np.array(sn_df['dec']))
+            sn_df['MW_rv'] = make_MW_extinction_rv(sn_df['ra'], sn_df['dec'])
+            sn_df['MW_av'] = make_MW_extinction_av(sn_df['ra'], sn_df['dec'])
 
             # Add fillers for columns not relevant for sn
             sn_df['sed_filepath'] = np.full((nobj),'')

skycatalogs/objects/base_object.py

@@ -554,6 +554,8 @@ def __getitem__(self, key):
                                       self._id[key], object_type, self, key)
 
         elif type(key) == slice:
+            if key.start is None:
+                key.start = 0
             ixdata = [i for i in range(min(key.stop,len(self._ra)))]
             ixes = itertools.islice(ixdata, key.start, key.stop, key.step)
             return [self._object_class(self._ra[i], self._dec[i], self._id[i],

skycatalogs/objects/snana_object.py

@@ -0,0 +1,170 @@
+import bisect
+import galsim
+import h5py
+import numpy as np
+from .base_object import BaseObject,ObjectCollection
+from skycatalogs.utils.exceptions import SkyCatalogsRuntimeError
+
+__all__ = ['SnanaObject', 'SnanaCollection']
+
+class SnanaObject(BaseObject):
+    _type_name = 'snana'
+
+    def __init__(self, ra, dec, id, object_type, belongs_to, belongs_index,
+                 redshift=None):
+        super().__init__(ra, dec, id, object_type, belongs_to, belongs_index,
+                         redshift)
+        self._mjds = None
+        self._lambda = None
+
+        # indices of elements in mjd array bounding our mjd
+        self._mjd_ix_l = None
+        self._mjd_ix_u = None
+
+    def _get_sed(self, mjd=None):
+        if mjd is None:
+            mjd = self._belongs_to._mjd
+        mjd_start = self.get_native_attribute('start_mjd')
+        mjd_end = self.get_native_attribute('end_mjd')
+        if mjd < mjd_start or mjd > mjd_end:
+            return None, 0.0
+
+        return self._linear_interp_SED(mjd), 0.0
+
+    def get_gsobject_components(self, gsparams=None, rng=None):
+        if gsparams is not None:
+            gsparams = galsim.GSParams(**gsparams)
+        return {'this_object': galsim.DeltaFunction(gsparams=gsparams)}
+
+    def get_observer_sed_component(self, component, mjd=None):
+        sed, _ = self._get_sed(mjd=mjd)
+        if sed is not None:
+            sed = self._apply_component_extinction(sed)
+        return sed
+
+    def get_LSST_flux(self, band, sed=None, cache=True, mjd=None):
+        def _flux_ratio(mag):
+            # -0.9210340371976184 = -np.log(10)/2.5.
+            return np.exp(-0.921034037196184 * mag)
+
+        flux = super().get_LSST_flux(band, sed, mjd)
+
+        if flux < 0:
+            raise SkyCatalogsRuntimeError('Negative flux')
+
+        if flux == 0.0:
+            return flux
+
+        mjd_ix_l, mjd_ix_u, mjd_fraction = self._find_mjd_interval(mjd)
+
+        with h5py.File(self._belongs_to._SED_file, 'r') as f:
+            cors = f[self._id][f'magcor_{band}']
+
+            # interpolate corrections
+            if mjd_ix_l == mjd_ix_u:
+                mag_cor = cors[mjd_ix_l]
+            else:
+                mag_cor = cors[mjd_ix_l] + mjd_fraction *\
+                    (cors[mjd_ix_u] - cors[mjd_ix_l])
+
+        #dbg = True
+        dbg = False
+
+        # Do everything in flux units
+        flux_cor = _flux_ratio(mag_cor)
+        corrected_flux = flux * flux_cor
+
+        if dbg:
+            print(f'Band {band} uncorrected flux: {flux}')
+            print(f'                mag correction: {mag_cor}')
+            print(f' multiplicative flux correction: {flux_cor}')
+
+        if cache:
+            att = f'lsst_flux_{band}'
+            setattr(self, att, corrected_flux)
+        return corrected_flux
+
+    def _find_mjd_interval(self, mjd=None):
+        '''
+        Find indices into mjd array of elements bounding our mjd
+        Also compute and constant needed for interpolation.  If we're
+        using "standard" mjd, also store these numbers.
+
+        Parameters
+        ----------
+        mjd     float   If None use the one stored in our ObjectCollection
+
+        Returns
+        -------
+        A tuple:     index below, index above, and fraction of distance
+                     mjd is from entry below to entry above
+        '''
+        if not mjd:
+            mjd = self._belongs_to._mjd
+            store = True
+        else:
+            store = (mjd == self._belongs_to._mjd)
+
+        if store:
+            if self._mjd_ix_l is not None:
+                # just return previously-computed values
+                return self._mjd_ix_l, self._mjd_ix_u, self._mjd_fraction
+
+        if  self._mjds is None:
+            with h5py.File(self._belongs_to._SED_file, 'r') as f:
+                self._mjds = np.array(f[self._id]['mjd'])
+        mjds = self._mjds
+
+        mjd_fraction = None
+        index = bisect.bisect(mjds, mjd)
+        if index == 0:
+            mjd_ix_l =  mjd_ix_u = 0
+        elif index == len(mjds):
+            mjd_ix_l = mjd_ix_u = index - 1
+        else:
+            mjd_ix_l = index - 1
+            mjd_ix_u = index
+            mjd_fraction = (mjd - mjds[mjd_ix_l]) /\
+                (mjds[mjd_ix_u] - mjds[mjd_ix_l])
+        if store:
+            self._mjd_ix_l = mjd_ix_l
+            self._mjd_ix_u = mjd_ix_u
+            self._mjd_fraction = mjd_fraction
+
+        return mjd_ix_l, mjd_ix_u, mjd_fraction
+
+
+    def _linear_interp_SED(self, mjd=None):
+        '''
+        Return galsim SED obtained by interpolating between SEDs
+        for nearest mjds among the templates
+        '''
+        mjd_ix_l,mjd_ix_u,mjd_fraction = self._find_mjd_interval(mjd)
+
+        with h5py.File(self._belongs_to._SED_file, 'r') as f:
+            if self._mjds is None or self._lambda is None:
+                self._mjds = np.array(f[self._id]['mjd'])
+                self._lambda = np.array(f[self._id]['lambda'])
+
+            if mjd_ix_l == mjd_ix_u:
+                flambda = f[self._id]['flamba'][mjd_ix_l]
+            else:
+                mjd_ix = mjd_ix_u
+                below = f[self._id]['flambda'][mjd_ix - 1]
+                above = f[self._id]['flambda'][mjd_ix]
+                flambda = below + mjd_fraction * (above - below)
+
+            lut = galsim.LookupTable(f[self._id]['lambda'],
+                                     flambda,
+                                     interpolant='linear')
+        return galsim.SED(lut, wave_type='A', flux_type='flambda')
+
+
+class SnanaCollection(ObjectCollection):
+    '''
+    This class (so far) differs from the vanilla ObjectCollection only
+    in that it keeps track of where the file is which contains a library
+    of SEDs for each sn
+    '''
+    def set_SED_file(self, SED_file):
+        self._SED_file = SED_file

skycatalogs/scripts/adjust_snana.py

@@ -0,0 +1,28 @@
+'''
+Read in original SNANA parquet file; output file identical except for
+added columns MW_av, MW_rv
+'''
+import os
+import argparse
+from skycatalogs.utils.add_extinction import AddExtinction
+
+parser = argparse.ArgumentParser(description='''
+Make a new skyCatalogs parquet file from an old, adding columns for MW_av, MW_rv
+''')
+
+parser.add_argument('indir', help='Directory containing input file(s). Required')
+parser.add_argument('outdir', help='Directory for output. Required')
+parser.add_argument('--pixels', type=int, nargs='*', default=[],
+                    help='healpix pixels for which new files will be created. Required')
+parser.add_argument('--starts-with', help='That part of the filename preceding healpixel',
+                    default='snana_')
+
+args = parser.parse_args()
+
+if os.path.abspath(args.indir) == os.path.abspath(args.outdir):
+    raise ValueError('Input and output directories must be different')
+
+writer = AddExtinction(args.indir, args.outdir, args.starts_with)
+
+for p in args.pixels:
+    writer.write(p)