Merge pull request #39 from /issues/38/reddener

Issues/38/reddener

src/rail/tools/photometry_tools.py

@@ -20,8 +20,6 @@
 dustmaps_sfd = tables_io.lazy_modules.lazyImport('dustmaps.sfd')
 
 
-
-
 # default column names in DC2
 LSST_BANDS = 'ugrizy'
 DEFAULT_MAG_COLS = [f"mag_{band}_lsst" for band in LSST_BANDS]
@@ -412,13 +410,13 @@ def __call__(self, data):
         return self.get_handle('output')
 
 
-class Dereddener(RailStage):
+class DustMapBase(RailStage):
     """Utility stage that does dereddening
     
     Note: set copy_all_cols=True to copy all 
     columns in data, copy_cols will be ignored
     """
-    name = 'Dereddener'
+    name = 'DustMapBase'
 
     config_options = RailStage.config_options.copy()
     config_options.update(bands='ugrizy')
@@ -429,8 +427,8 @@ class Dereddener(RailStage):
     config_options.update(copy_cols=[])
     config_options.update(copy_all_cols=False)
 
-    inputs = [('input', Hdf5Handle)]
-    outputs = [('output', Hdf5Handle)]
+    inputs = [('input', PqHandle)]
+    outputs = [('output', PqHandle)]
 
     def fetch_map(self):
         dust_map_dict = dict(sfd=dustmaps_sfd)
@@ -439,12 +437,15 @@ def fetch_map(self):
         except KeyError as msg:  # pragma: no cover
             raise KeyError(f"Unknown dustmap {self.config.dustmap_name}, options are {list(dust_map_dict.keys())}") from msg
 
-        if os.path.exists(os.path.join(self.config.dustmap_dir, self.config.dustmap_name)):  # pragma: no cover
+        dustmap_dir = os.path.expandvars(self.config.dustmap_dir)
+        dustmap_path = os.path.join(dustmap_dir, self.config.dustmap_name)
+        if os.path.exists(dustmap_path):  # pragma: no cover
             # already downloaded, return
             return
 
         dust_map_config = dustmaps_config.config
-        dust_map_config['data_dir'] = self.config.dustmap_dir
+        # dust_map_config['data_dir'] = self.config.dustmap_dir
+        dust_map_config['data_dir'] = dustmap_dir
         fetch_func = dust_map_submod.fetch
         fetch_func()
 
@@ -455,7 +456,7 @@ def __init__(self, args, comm=None):
     def run(self):
         data = self.get_data('input', allow_missing=True)
         out_data = {}
-        coords = SkyCoord(data['ra'], data['decl'], unit = 'deg',frame='fk5')
+        coords = SkyCoord(data['ra'], data['dec'], unit = 'deg',frame='fk5')
         dust_map_dict = dict(sfd=dustmaps_sfd.SFDQuery)
         try:
             dust_map_class = dust_map_dict[self.config.dustmap_name]
@@ -469,7 +470,7 @@ def run(self):
         for i, band_ in enumerate(self.config.bands):
             band_mag_name = self.config.mag_name.format(band=band_)
             mag_vals = data[band_mag_name]
-            out_data[band_mag_name] = mag_vals - ebvvec*self.config.band_a_env[i]
+            out_data[band_mag_name] = self._calc_values(mag_vals, ebvvec, self.config.band_a_env[i])
             band_mag_name_list.append(band_mag_name)
 
         # check if copy_all_cols set to true:
@@ -482,7 +483,8 @@ def run(self):
                 if col_ not in band_mag_name_list:
                     out_data[col_] = data[col_]
 
-        self.add_data('output', out_data)
+        out_data_pd = pd.DataFrame(out_data)
+        self.add_data('output', out_data_pd)
 
     def __call__(self, data):
         """Return a converted table
@@ -500,3 +502,23 @@ def __call__(self, data):
         self.set_data('input', data)
         self.run()
         return self.get_handle('output')
+
+
+class Dereddener(DustMapBase):
+    """Utility stage that does dereddening
+    
+    """
+    name = 'Dereddener'
+
+    def _calc_values(self, mag_vals, ebvvec, band_a_env):
+        return mag_vals - ebvvec*band_a_env
+
+
+class Reddener(DustMapBase):
+    """Utility stage that does reddening
+    
+    """
+    name = 'Reddener'
+
+    def _calc_values(self, mag_vals, ebvvec, band_a_env):
+        return mag_vals + ebvvec*band_a_env

tests/astro_tools/test_core.py

@@ -20,8 +20,8 @@
     TableHandle,
 )
 from rail.core.stage import RailStage
-from rail.tools.photometry_tools import HyperbolicMagnitudes, HyperbolicSmoothing, PhotometryManipulator, LSSTFluxToMagConverter, Dereddener
-from rail.utils.path_utils import RAILDIR
+from rail.tools.photometry_tools import HyperbolicMagnitudes, HyperbolicSmoothing, PhotometryManipulator, LSSTFluxToMagConverter, Dereddener, Reddener
+from rail.utils.path_utils import RAILDIR, find_rail_file
 #from rail.tools.util_stages import (
 #    LSSTFluxToMagConverter,
 #    Dereddener,
@@ -37,7 +37,7 @@ def test_flux2mag():
     DS = RailStage.data_store
     DS.clear()
 
-    testFile = os.path.join(RAILDIR, "rail", "examples_data", "testdata", "rubin_dm_dc2_example.pq")
+    testFile = find_rail_file(os.path.join("examples_data", "testdata", "rubin_dm_dc2_example2.pq"))
     test_data = DS.read_file("test_data", TableHandle, testFile)
 
     fluxToMag = LSSTFluxToMagConverter.make_stage(name='flux2mag')
@@ -49,10 +49,10 @@ def test_dereddener():
     DS = RailStage.data_store
     DS.clear()
 
-    testFile = os.path.join(RAILDIR, "rail", "examples_data", "testdata", "rubin_dm_dc2_example.pq")
+    testFile = find_rail_file(os.path.join("examples_data", "testdata", "rubin_dm_dc2_example2.pq"))
     test_data = DS.read_file("test_data", TableHandle, testFile)
 
-    fluxToMag = LSSTFluxToMagConverter.make_stage(name='flux2mag', copy_cols=dict(ra='ra', decl='decl'))
+    fluxToMag = LSSTFluxToMagConverter.make_stage(name='flux2mag', copy_cols=dict(ra='ra', dec='decl'))
 
     is_temp_dir = False
     dustmap_dir = os.environ.get('RAIL_DUSTMAP_DIR')
@@ -62,10 +62,12 @@ def test_dereddener():
         is_temp_dir = True
 
     dereddener = Dereddener.make_stage(name='dereddner', dustmap_dir=dustmap_dir)
+    reddener = Reddener.make_stage(name='reddner', dustmap_dir=dustmap_dir)
     dereddener.fetch_map()
 
     flux_data = fluxToMag(test_data)
     dered_data = dereddener(flux_data)
+    red_data = reddener(flux_data)
 
     if is_temp_dir:
         tmp_dustmap_dir.cleanup()