Merge pull request #296 from desihub/negoii-final

include negative [OII] flux penalty in final minimization

py/redrock/fitz.py

@@ -16,6 +16,7 @@
 from .rebin import rebin_template
 
 from .zscan import calc_zchi2_one, spectral_data, calc_zchi2_batch
+from .zscan import calc_negOII_penalty
 
 from .zwarning import ZWarningMask as ZW
 
@@ -126,7 +127,8 @@ def prior_on_coeffs(n_nbh, deg_legendre, sigma, ncamera):
     return prior
 
 
-def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False, deg_legendre=None, zminfit_npoints=15, per_camera=False, n_nearest=None, prior_sigma=None):
+def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False, deg_legendre=None,
+         zminfit_npoints=15, per_camera=False, n_nearest=None, prior_sigma=None):
     """Refines redshift measurement around up to nminima minima.
 
     TODO:
@@ -164,7 +166,7 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
     dwave = { s.wavehash:s.wave for s in spectra }
 
     (weights, flux, wflux) = spectral_data(spectra)
-    
+
     if (use_gpu):
         #Get CuPy arrays of weights, flux, wflux
         #These are created on the first call of gpu_spectral_data() for a
@@ -242,6 +244,10 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
                                                  solve_matrices_algorithm=template.solve_matrices_algorithm,
                                                  use_gpu=use_gpu)
 
+        #- Penalize chi2 for negative [OII] flux; ad-hoc
+        if hasattr(template, 'OIItemplate'):
+            zzchi2 += calc_negOII_penalty(template.OIItemplate, zzcoeff)
+
         #- fit parabola to 3 points around minimum
         i = min(max(np.argmin(zzchi2),1), len(zz)-2)
         zmin, sigma, chi2min, zwarn = minfit(zz[i-1:i+2], zzchi2[i-1:i+2])

py/redrock/templates.py

@@ -161,6 +161,13 @@ def __init__(self, filename=None, spectype=None, redshifts=None,
             raise ValueError(f'Template method {self._method} unrecognized; '
                               f'should be one of {valid_template_methods}')
 
+        # Special case for GALAXY templates and PCA fitting, which can include
+        # negative templates or coefficients: Cache a template that covers just
+        # [OII].
+        if self._rrtype == 'GALAXY' and self._method == "PCA":
+            isOII = (3724 <= self.wave) & (self.wave <= 3733)
+            self.OIItemplate = self.flux[:,isOII]
+
         if filename is not None:
             print(f'INFO: {os.path.basename(filename)} {self}')
         else:

py/redrock/zscan.py

@@ -609,7 +609,8 @@ def calc_batch_dot_product_3d3d_gpu(a, b, transpose_a=False, fullprecision=True)
 ###    are very obviously analagous though and should be highly
 ###    maintainable.  The main difference is the extra loop on the CPU version
 
-def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_matrices_algorithm=None, solver_args=None, use_gpu=False, fullprecision=True, prior=None):
+def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_matrices_algorithm=None,
+                     solver_args=None, use_gpu=False, fullprecision=True, prior=None):
 
     """Calculate a batch of chi2.
     For many redshifts and a set of spectral data, compute the chi2 for
@@ -765,6 +766,24 @@ def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_mat
             zchi2[i] = np.dot( (flux - model)**2, weights )
     return (zchi2, zcoeff)
 
+def calc_negOII_penalty(OIItemplate, coeff):
+    """return penalty term for model having negative [OII] flux
+
+    Args:
+        OIItemplate[nwave, nbasis]: portion of template around [OII] doublet
+        coeff[nz, nbasis]: coefficients from full template fit to data
+
+    Returns:
+        penalty[nz]: penalty per redshift bin
+    """
+    #- evaluate template and sum over [OII] doublet
+    OIIflux = np.sum(coeff @ OIItemplate, axis=1)
+
+    #- Linear penalty with negative [OII]; no mathematical basis but
+    #- pragmatically works ok
+    penalty = -OIIflux * (OIIflux < 0)
+
+    return penalty
 
 ###!!! NOTE - this is the main method for the v3 algorithm
 ###    In this version, everything is done in batch on the GPU but the
@@ -813,12 +832,6 @@ def calc_zchi2(target_ids, target_data, dtemplate, progress=None, use_gpu=False)
     zchi2penalty = np.zeros( (ntargets, nz) )
     zcoeff = np.zeros( (ntargets, nz, nbasis) )
 
-    # Redshifts near [OII]; used only for galaxy templates
-    if dtemplate.template.template_type == 'GALAXY':
-        isOII = (3724 <= dtemplate.template.wave) & \
-            (dtemplate.template.wave <= 3733)
-        OIItemplate = dtemplate.template.flux[:,isOII].T
-
     ## Redshifted templates are now already in format needed - dict of 3d
     # arrays (CUPY or numpy).
     tdata = dtemplate.local.data
@@ -847,12 +860,14 @@ def calc_zchi2(target_ids, target_data, dtemplate, progress=None, use_gpu=False)
         # for all templates for all three spectra for this target
 
         # For coarse z scan, use fullprecision = False to maximize speed
-        (zchi2[j,:], zcoeff[j,:,:]) = calc_zchi2_batch(target_data[j].spectra, tdata, weights, flux, wflux, nz, nbasis, dtemplate.template.solve_matrices_algorithm, use_gpu=use_gpu, fullprecision=False)
-
-        #- Penalize chi2 for negative [OII] flux; ad-hoc
-        if dtemplate.template.template_type == 'GALAXY':
-            OIIflux = np.sum(zcoeff[j] @ OIItemplate.T, axis=1)
-            zchi2penalty[j][OIIflux < 0] = -OIIflux[OIIflux < 0]
+        (zchi2[j,:], zcoeff[j,:,:]) = calc_zchi2_batch(
+            target_data[j].spectra, tdata, weights, flux, wflux, nz, nbasis,
+            dtemplate.template.solve_matrices_algorithm, use_gpu=use_gpu,
+            fullprecision=False)
+
+        #- Galaxy templates penalize chi2 for negative [OII] flux; ad-hoc
+        if hasattr(dtemplate.template, 'OIItemplate'):
+            zchi2penalty[j] = calc_negOII_penalty(dtemplate.template.OIItemplate, zcoeff[j])
 
         if dtemplate.comm is None and progress is not None:
             progress.put(1)