Update tpcorr creation code.

py/desispec/tpcorrparam.py

@@ -26,6 +26,7 @@
 from astropy.io import fits
 from desiutil.log import get_logger
 import desispec.calibfinder
+import desispec.io
 
 
 def cubic(p, x, y):
@@ -134,7 +135,8 @@ def gather_tpcorr(recs):
         EXPID: expid of exposure
     """
     out = np.zeros(3*len(recs), dtype=[
-        ('X', '5000f4'), ('Y', '5000f4'), ('TPCORR', '5000f4'),
+        ('X', '5000f4'), ('Y', '5000f4'),
+        ('TPCORR', '5000f4'), ('TPCORR_MODEL', '5000f4'),
         ('CAMERA', 'U1'), ('NIGHT', 'i4'), ('EXPID', 'i4')])
     count = -1
     for rec in recs:
@@ -153,12 +155,14 @@ def gather_tpcorr(recs):
                     continue
                 fiber = frame.fibermap['FIBER']
                 out['TPCORR'][count, fiber] = sky.throughput_corrections
+                out['TPCORR_MODEL'][count, fiber] = (
+                    sky.throughput_corrections_model)
                 out['X'][count, fiber] = frame.fibermap['FIBER_X']
                 out['Y'][count, fiber] = frame.fibermap['FIBER_Y']
     return out
 
 
-def gather_dark_tpcorr(specprod=None):
+def gather_dark_tpcorr(specprod=None, nproc=4):
     """Gathers TPCORR for dark exposures expected to be well behaved.
 
     This wraps gather_tpcorr, making a sensible selection of dark
@@ -174,13 +178,23 @@ def gather_dark_tpcorr(specprod=None):
     if specprod is None:
         specprod = os.environ.get('SPECPROD', 'daily')
     expfn = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
-                         specprod, 'exposures-{specprod}.fits')
+                         specprod, f'exposures-{specprod}.fits')
     exps = fits.getdata(expfn)
     m = ((exps['EXPTIME'] > 300) & (exps['SKY_MAG_R_SPEC'] > 20.5) &
          (exps['SKY_MAG_R_SPEC'] < 30) & (exps['FAPRGRM'] == 'dark'))
+    exps = exps[m]
     if specprod == 'daily':
         exps = exps[exps['NIGHT'] >= 20210901]
-    return gather_tpcorr(exps[m])
+    log = get_logger()
+    log.info(f'Gathering skies for {len(exps)} exposures...')
+    if nproc > 1:
+        from multiprocessing import pool
+        p = pool.Pool(nproc)
+        res = p.map(gather_tpcorr, [[exp] for exp in exps])
+        res = np.concatenate(res)
+    else:
+        res = gather_tpcorr(exps)
+    return res
 
 
 def pca_tpcorr(tpcorr):
@@ -321,17 +335,7 @@ def chi(param):
     return out
 
 
-# note! one worry: in the current PR, the meaning of TPCORR will change
-# moving forward, and we won't have any of the current throughput measurements
-# for computing these coefficients in the future.  We'd have to develop
-# an alternative scheme for computing these coefficients or would have
-# to change this PR so that the new and old TPCORR have similar meanings.
-# That option would mean passing calculate_throughput_corrections some
-# goofy sky model that doesn't include the new model TPCORR from this PR.
-# Or we add another model extension with the new TPCORR from this PR,
-# and then this routine would require the product of the model TPCORR and
-# the measured TPCORR.
-def make_tpcorrparam_files(tpcorr=None, spatial=None, dirname=None):
+def make_tpcorrparam_files(tpcorr=None, spatial=None, dirname=None, npca=2):
     """Gather, fit, and write tpcorrparam files to output directory.
 
     This function is intended for producing tpcorrparam files to
@@ -346,27 +350,58 @@ def make_tpcorrparam_files(tpcorr=None, spatial=None, dirname=None):
         dirname (optional): directory to which to write files.  Defaults
             to DESI_SPECTRO_CALIB.
     """
+    log = get_logger()
     if dirname is None:
         dirname = os.environ['DESI_SPECTRO_CALIB']
     if tpcorr is None:
         tpcorr = gather_dark_tpcorr()
+    if 'TPCORR_MODEL' in tpcorr.dtype.names:
+        tpcorr = tpcorr.copy()
+        tpcorr['TPCORR'] *= tpcorr['TPCORR_MODEL']
+        log.info('Multiplying TPCORR by TPCORR_MODEL')
+        # if TPCORR_MODEL is in there, that means that the measured TPCORR
+        # are measured after the model TPCORR have been removed.
+        # this puts the model back in so that the TPCORR field has the
+        # "total" TPCORR"
+
+    m = ((tpcorr['X'] == 0) | (tpcorr['Y'] == 0) |
+         ~np.isfinite(tpcorr['X']) | ~np.isfinite(tpcorr['Y']) |
+         ~np.isfinite(tpcorr['TPCORR']))
+    fiber = np.arange(5000)[None, :]
+    # ignore problematic region on r4 in device that was replaced.
+    m = m | ((tpcorr['EXPID'][:, None] > 1.07e5) &
+             (tpcorr['EXPID'][:, None] < 1.53e5) &
+             (fiber >= 2250) & (fiber < 2272) &
+             (tpcorr['CAMERA'][:, None] == 'r'))
+    # ignore problematic region on z5 in device that was replaced.
+    m = m | ((tpcorr['CAMERA'][:, None] == 'z') &
+             (tpcorr['EXPID'][:, None] < 1.55e5) &
+             (fiber > 2660) & (fiber < 2690))
+    # ignore fiber 1098 on the B camera as too noisy (parallel charge
+    # transfer issue)
+    m = m | ((fiber == 1098) & (tpcorr['CAMERA'][:, None] == 'b'))
+    mzero = m.copy()
+    tpcorr['TPCORR'][mzero] = 0
+
     if spatial is None:
         spatial = fit_tpcorr_per_fiber(tpcorr)
+
     tpcorrspatmod = np.zeros_like(tpcorr['TPCORR'])
     for f in 'brz':
         m = tpcorr['CAMERA'] == f
         tpcorrspatmod[m] = tpcorrspatialmodel(
             spatial[f], tpcorr['X'][m], tpcorr['Y'][m])
-    m = ((tpcorr['X'] == 0) | (tpcorr['Y'] == 0) |
-         ~np.isfinite(tpcorr['X']) | ~np.isfinite(tpcorr['Y']) |
-         ~np.isfinite(tpcorr['TPCORR']))
+
     tpcorrresid = tpcorr.copy()
     tpcorrresid['TPCORR'] -= tpcorrspatmod
     tpcorrresid['TPCORR'] += 1
-    tpcorrresid['TPCORR'][m] = 0  # invalid value
+    tpcorrresid['TPCORR'][mzero] = 0  # invalid value
+
     pca = pca_tpcorr(tpcorrresid)
+    constants = dict()
     for f in 'brz':
         constant = spatial[f]['PARAM'][:, 0].copy()
+        constants[f] = constant
         spatial[f]['PARAM'][:, 0] = 1
         for i in range(10):
             sm = desispec.calibfinder.sp2sm(i)
@@ -378,5 +413,6 @@ def make_tpcorrparam_files(tpcorr=None, spatial=None, dirname=None):
                          fits.Header(dict(extname='SPATIAL')))
             fits.append(fn, constant[i*500:(i+1)*500],
                         fits.Header(dict(extname='MEAN')))
-            fits.append(fn, pca[f][1]['pca'][:2, i*500:(i+1)*500],
+            fits.append(fn, pca[f][1]['pca'][:npca, i*500:(i+1)*500],
                         fits.Header(dict(extname='PCA')))
+    return tpcorrresid, spatial, pca, constants