Merge pull request #39 from p-slash/save-sm-weights

enable saving smooth weights in delta files

py/qsonic/io.py

@@ -27,46 +27,39 @@ def add_io_parser(parser=None):
         parser = argparse.ArgumentParser(
             formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-    iogroup = parser.add_argument_group(
-        'Input/output parameters and selections')
-    iogroup.add_argument(
+    ingroup = parser.add_argument_group('Input options')
+    ingroup.add_argument(
         "--input-dir", '-i', required=True,
         help="Input directory to healpix")
-    iogroup.add_argument(
+    ingroup.add_argument(
         "--catalog", required=True,
         help="Catalog filename")
-    iogroup.add_argument(
-        "--outdir", '-o',
-        help="Output directory to save deltas.")
-    iogroup.add_argument(
+    ingroup.add_argument(
         "--mock-analysis", action="store_true",
         help="Input folder is mock. Uses nside=16")
-    iogroup.add_argument(
+    ingroup.add_argument(
         "--keep-surveys", nargs='+', default=['sv3', 'main'],
         help="Surveys to keep.")
-    iogroup.add_argument(
-        "--coadd-arms", action="store_true",
-        help="Coadds arms when saving.")
-
-    iogroup.add_argument(
+    ingroup.add_argument(
         "--skip-resomat", action="store_true",
         help="Skip reading resolution matrix for 3D.")
-    iogroup.add_argument(
+    ingroup.add_argument(
         "--arms", default=['B', 'R'], choices=['B', 'R', 'Z'], nargs='+',
         help="Arms to read.")
-    iogroup.add_argument(
-        "--min-rsnr", type=float, default=0.,
-        help="Minium SNR <F/sigma> above Lya.")
-    iogroup.add_argument(
-        "--skip", type=_float_range(0, 1), default=0.,
-        help="Skip short spectra lower than given ratio.")
-    iogroup.add_argument(
+
+    outgroup = parser.add_argument_group('Output options')
+    outgroup.add_argument(
+        "--outdir", '-o',
+        help="Output directory to save deltas.")
+    outgroup.add_argument(
+        "--coadd-arms", action="store_true",
+        help="Coadds arms when saving.")
+    outgroup.add_argument(
+        "--save-smooth-weights", action="store_true",
+        help="Save smoothed weights instead.")
+    outgroup.add_argument(
         "--save-by-hpx", action="store_true",
         help="Save by healpix. If not, saves by MPI rank.")
-    iogroup.add_argument(
-        "--keep-nonforest-pixels", action="store_true",
-        help="Keeps non forest wavelengths. Memory intensive!")
-
     return parser
 
 
@@ -141,7 +134,9 @@ def read_spectra_onehealpix(
 
 
 def save_deltas(
-        spectra_list, outdir, varlss_interp, save_by_hpx=False, mpi_rank=None):
+        spectra_list, outdir, varlss_interp,
+        save_by_hpx=False, mpi_rank=None, use_ivar_sm=False
+):
     """ Saves given list of spectra as deltas. NO coaddition of arms.
     Each arm is saved separately. Only valid spectra are saved.
 
@@ -153,10 +148,12 @@ def save_deltas(
         Output directory. Does not save if empty of None
     varlss_interp: Interpolator
         Interpolator for LSS variance
-    save_by_hpx: bool
+    save_by_hpx: bool, default: False
         Saves by healpix if True. Has priority over mpi_rank
-    mpi_rank: int
+    mpi_rank: int, default: None
         Rank of the MPI process. Save by `mpi_rank` if passed.
+    use_ivar_sm: bool, default: False
+        Use :attr:`Spectrum.forestivar_sm` in weights instead.
 
     Raises
     ---------
@@ -188,7 +185,7 @@ def save_deltas(
             f"{outdir}/delta-{healpix}.fits", 'rw', clobber=True)
 
         for spec in qsonic.spectrum.valid_spectra(hp_specs):
-            spec.write(results, varlss_interp)
+            spec.write(results, varlss_interp, use_ivar_sm)
 
         results.close()
 

py/qsonic/scripts/qsonic_fit.py

@@ -32,6 +32,18 @@ def get_parser(add_help=True):
         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
     parser = qsonic.io.add_io_parser(parser)
+
+    analysis_group = parser.add_argument_group('Analysis options')
+    analysis_group.add_argument(
+        "--min-rsnr", type=float, default=0.,
+        help="Minium SNR <F/sigma> above Lya.")
+    analysis_group.add_argument(
+        "--skip", type=qsonic.io._float_range(0, 1), default=0.,
+        help="Skip short spectra lower than given ratio.")
+    analysis_group.add_argument(
+        "--keep-nonforest-pixels", action="store_true",
+        help="Keeps non forest wavelengths. Memory intensive!")
+
     parser = qsonic.spectrum.add_wave_region_parser(parser)
     parser = qsonic.masks.add_mask_parser(parser)
     parser = add_picca_continuum_parser(parser)

py/qsonic/spectrum.py

@@ -398,19 +398,24 @@ def mean_snr(self):
             npix += armpix
         return snr / npix
 
-    def write(self, fts_file, varlss_interp):
+    def write(self, fts_file, varlss_interp, use_ivar_sm=False):
         """Writes each arm to FITS file separately.
 
         Writes 'LAMBDA', 'DELTA', 'IVAR', 'WEIGHT', 'CONT' columns and
         'RESOMAT' column if resolution matrix is present to extention name
-        'targetid-arm'. FITS file must be initialized before. Each arm has its
-        own `MEANSNR`. `weights` in the file are **not** smoothed.
+        ``targetid-arm``. FITS file must be initialized before.
+        Each arm has its own `MEANSNR`. "WEIGHT" in the file do not use
+        smoothed inverse variance by default, which can be changed by
+        ``use_ivar_sm=True``.
 
         Arguments
         ---------
         fts_file: FITS file
+            The file handler, not filename.
         varlss_interp: Fast1DInterpolator or any other interpolator.
             LSS variance interpolator.
+        use_ivar_sm: bool, default: False
+            Use :attr:`forestivar_sm` in weights instead.
         """
         hdr_dict = {
             'LOS_ID': self.targetid,
@@ -423,6 +428,7 @@ def write(self, fts_file, varlss_interp):
             'MEANSNR': 0.,
             'RSNR': self.rsnr,
             'DELTA_LAMBDA': self.dwave,
+            'SMWEIGHT': use_ivar_sm,
         }
 
         for arm, wave_arm in self.forestwave.items():
@@ -439,7 +445,9 @@ def write(self, fts_file, varlss_interp):
             delta = self.forestflux[arm] / cont_est - 1
             ivar = self.forestivar[arm] * cont_est**2
             var_lss = varlss_interp(wave_arm)
-            weight = ivar / (1 + ivar * var_lss)
+            ivar_w = (
+                self.forestivar_sm[arm] * cont_est**2 if use_ivar_sm else ivar)
+            weight = ivar_w / (1 + ivar_w * var_lss)
 
             cols = [wave_arm, delta, ivar, weight, cont_est]
             if self.forestreso: