desihub · akremin · Apr 1, 2025 · Mar 28, 2025 · Mar 28, 2025 · Mar 31, 2025
diff --git a/py/redrock/results.py b/py/redrock/results.py
@@ -13,7 +13,7 @@
 from astropy.table import Table
 
 from .utils import encode_column
-from .templates import parse_fulltype
+from .templates import parse_fulltype, make_fulltype
 
 def write_zscan(filename, zscan, zfit, clobber=False):
     """Writes redrock.zfind results to a file.
@@ -52,24 +52,34 @@ def write_zscan(filename, zscan, zfit, clobber=False):
             if zfit[colname].dtype.kind == 'U':
                 zfit.replace_column(colname, np.char.encode(zfit[colname], 'ascii'))
 
+    #- convert columns to lower precision for final output
+    for col in ['zerr', 'chi2', 'zz', 'zzchi2', 'coeff', 'deltachi2']:
+        zfit[col] = zfit[col].astype(np.float32)
+
+    zfit['zwarn'] = zfit['zwarn'].astype(np.int32)
+    zfit['npixels'] = zfit['npixels'].astype(np.int32)
+    zfit['znum'] = zfit['znum'].astype(np.int8)
+    zfit['ncoeff'] = zfit['ncoeff'].astype(np.int8)
+
     zbest = zfit[zfit['znum'] == 0]
     zbest.remove_column('znum')
 
-    zbest.write(filename, path='zbest', format='hdf5')
+    # zbest.write(tempfile, path='zbest', format='hdf5')
 
     targetids = np.asarray(zbest['targetid'])
     spectypes = list(zscan[targetids[0]].keys())
 
     tempfile = filename + '.tmp'
     fx = h5py.File(tempfile, mode='w')
     fx['targetids'] = targetids
+    fx['zbest'] = zbest.as_array()
 
     for spectype in spectypes:
-        zchi2 = np.vstack([zscan[t][spectype]['zchi2'] for t in targetids])
-        penalty = np.vstack([zscan[t][spectype]['penalty'] for t in targetids])
+        zchi2 = np.vstack([zscan[t][spectype]['zchi2'].astype(np.float32) for t in targetids])
+        penalty = np.vstack([zscan[t][spectype]['penalty'].astype(np.float32) for t in targetids])
         zcoeff = list()
         for t in targetids:
-            tmp = zscan[t][spectype]['zcoeff']
+            tmp = zscan[t][spectype]['zcoeff'].astype(np.float32)
             tmp = tmp.reshape((1,)+tmp.shape)
             zcoeff.append(tmp)
         zcoeff = np.vstack(zcoeff)
@@ -82,34 +92,38 @@ def write_zscan(filename, zscan, zfit, clobber=False):
     for targetid in targetids:
         ii = np.where(zfit['targetid'] == targetid)[0]
         fx['zfit/{}/zfit'.format(targetid)] = zfit[ii].as_array()
-        #- TODO: fx['zfit/{}/model']
 
     fx.close()
     os.rename(tempfile, filename)
 
 
-def read_zscan(filename, select_targetids=None):
+def read_zscan(filename, select_targetids=None, upper=False, nozfit=False):
     """Read redrock.zfind results from a file.
 
     Args:
         filename (str): redrock details (.h5) filename
-        select_targetids (list): array of TARGETIDs to read
+
+    Options:
+        select_targetids (int or array of int): TARGETID(s) to read
+        upper (bool): convert table column names to UPPERCASE
+        nozfit (bool): do not include zfit information
 
     Returns:
-        tuple: ``(zscan, zfit)`` where
-        ``zfit`` is a Table with the N best fits for each target
-        per spectype and subtype; and
-        ``zscan`` is a nested dictionary ``zscan[targetid][templatetype]``
-        with keys:
+        ``zscan``, or ``(zscan, zfit)`` tuple where
+        ``zscan`` is a nested dictionary ``zscan[targetid][templatetype]`` with keys:
 
                 - redshifts: array of redshifts scanned
                 - zchi2: array of chi2 of fit vs. redshifts
                 - penalty: array of chi2 penalties for unphysical fits at each z
                 - zcoeff: array of coefficients fit per redshifts
                 - zfit: table of N best-fit redshifts for just this target and templatetype
+
+        ``zfit`` is a Table with the N best fits for each target per spectype and subtype
     """
     import h5py
-    # zbest = Table.read(filename, format='hdf5', path='zbest')
+    if np.isscalar(select_targetids):
+        select_targetids = [select_targetids,]
+
     with h5py.File(os.path.expandvars(filename), mode='r') as fx:
         targetids = fx['targetids'][()]  # .value
         fulltypes = list(fx['zscan'].keys())
@@ -137,23 +151,62 @@ def read_zscan(filename, select_targetids=None):
                 zscan[targetid][fulltype]['zchi2'] = zchi2[i]
                 zscan[targetid][fulltype]['penalty'] = penalty[i]
                 zscan[targetid][fulltype]['zcoeff'] = zcoeff[i]
-                thiszfit = fx['/zfit/{}/zfit'.format(targetid)][()]
-                ii = (thiszfit['spectype'].astype('U') == spectype)
-                ii &= (thiszfit['subtype'].astype('U') == subtype)
-                thiszfit = Table(thiszfit[ii])
-                thiszfit.remove_columns(['targetid', 'znum', 'deltachi2'])
-                thiszfit.replace_column('spectype',
-                    encode_column(thiszfit['spectype']))
-                thiszfit.replace_column('subtype',
-                    encode_column(thiszfit['subtype']))
-                zscan[targetid][fulltype]['zfit'] = thiszfit
-
-        zfit = [fx['zfit/{}/zfit'.format(tid)][()] for tid in targetids[indx]]
+
+        if nozfit:
+            return zscan
+
+        #- Add zfit information too
+        zfit = read_zfit(filename, select_targetids=select_targetids, upper=upper)
+        if upper:
+            targetid_col = 'TARGETID'
+            spectype_col = 'SPECTYPE'
+            subtype_col = 'SUBTYPE'
+        else:
+            targetid_col = 'targetid'
+            spectype_col = 'spectype'
+            subtype_col = 'subtype'
+
+        for zfit_tid in zfit.group_by( (targetid_col, spectype_col, subtype_col) ).groups:
+            targetid = zfit_tid[targetid_col][0]
+            spectype = zfit_tid[spectype_col][0]
+            subtype = zfit_tid[subtype_col][0]
+            fulltype = make_fulltype(spectype, subtype)
+            zscan[targetid][fulltype]['zfit'] = zfit_tid
+
+        return zscan, zfit
+
+def read_zfit(filename, select_targetids=None, upper=False):
+    """
+    Read and stack `zfit` tables from filename with N best fits
+
+    Args:
+        filename (str): path to Redrock details hdf5 file
+
+    Options:
+        select_targetids (int or array of int): read only these targetids
+        upper (bool): convert table column names to UPPERCASE
+
+    Returns: astropy Table of zfit results, or dict of Tables keyed by targetid
+    """
+    import h5py
+    if np.isscalar(select_targetids):
+        select_targetids = [select_targetids,]
+
+    zfit = dict()
+    with h5py.File(os.path.expandvars(filename), mode='r') as fx:
+        if select_targetids is None:
+            select_targetids = fx['targetids'][:]
+
+        zfit = [fx[f'zfit/{tid}/zfit'][:] for tid in select_targetids]
         zfit = Table(np.hstack(zfit))
         zfit.replace_column('spectype', encode_column(zfit['spectype']))
         zfit.replace_column('subtype', encode_column(zfit['subtype']))
 
-    return zscan, zfit
+    if upper:
+        for colname in list(zfit.colnames):
+            zfit.rename_column(colname, colname.upper())
+
+    return zfit
 
 
 def read_zscan_redrock(filename):

diff --git a/py/redrock/test/test_io.py b/py/redrock/test/test_io.py
@@ -7,7 +7,7 @@
 import numpy as np
 
 from .. import utils as rrutils
-from ..results import read_zscan, write_zscan
+from ..results import read_zscan, write_zscan, read_zfit
 from ..zfind import zfind
 from ..templates import DistTemplate
 
@@ -61,13 +61,22 @@ def test_zscan_io(self):
         zscan2, zfit2 = read_zscan(self.testfile)
 
         self.assertEqual(zfit1.colnames, zfit2.colnames)
+
+        #- z and targetid retain original precision/size; others may have lower precision
+        self.assertEqual(zfit1['z'].dtype, zfit2['z'].dtype)
+        self.assertEqual(zfit1['targetid'].dtype, zfit2['targetid'].dtype)
+
         for cn in zfit1.colnames:
-            np.testing.assert_equal(zfit1[cn], zfit2[cn])
+            d2 = zfit2[cn]
+            d1 = zfit1[cn].astype(d2.dtype)  #- e.g. float64 -> float32
+            np.testing.assert_equal(d1, d2)
 
         for targetid in zscan1:
             for spectype in zscan1[targetid]:
                 for key in zscan1[targetid][spectype]:
                     d1 = zscan1[targetid][spectype][key]
+                    if key != 'redshifts':
+                        d1 = d1.astype(np.float32)
                     d2 = zscan2[targetid][spectype][key]
                     self.assertTrue(np.all(d1==d2), 'data mismatch {}/{}/{}'.format(targetid, spectype, key))
 
@@ -83,5 +92,50 @@ def test_read_zscan(self):
                 self.assertEqual(len(zz['redshifts']), len(zz['zchi2']))
                 self.assertGreater(len(zz['zfit']), 0, f'Empty zfit Table for {targetid=} {fulltype=}')
 
+        #- targetid subsets
+        targetids = np.unique(zfit['targetid'])
+
+        zscan, zfit = read_zscan(zscanfile, select_targetids=targetids[1])
+        self.assertEqual(len(np.unique(zfit['targetid'])), 1)
+        self.assertEqual(np.unique(zfit['targetid'])[0], targetids[1])
+
+        zscan, zfit = read_zscan(zscanfile, select_targetids=targetids[0:2])
+        self.assertEqual(len(np.unique(zfit['targetid'])), 2)
+        self.assertTrue(np.all(np.unique(zfit['targetid']) == targetids[0:2]))
+
+        #- UPPERCASE column names
+        zscan, zfit = read_zscan(zscanfile, upper=True)
+        self.assertIn('TARGETID', zfit.colnames)
+
+        #- only zscan; not (zscan,zfit)
+        results = read_zscan(zscanfile, nozfit=True)
+        self.assertTrue(isinstance(results, dict))  #- not tuple
+        self.assertIn(targetids[0], results.keys())
+
+    def test_read_zfit(self):
+        """Test read_zfit with pre-generated data"""
+        import importlib
+        filename = importlib.resources.files('redrock.test').joinpath('data/rrdetails-test.h5')
+        zfit = read_zfit(filename)
+        targetids = np.unique(zfit['targetid'])
+        self.assertEqual(len(targetids), 3)
+
+        #- Select two targetids
+        zfit2 = read_zfit(filename, select_targetids=targetids[0:2])
+        targetids2 = np.unique(zfit2['targetid'])
+        self.assertEqual(len(targetids2), 2)
+        self.assertTrue(np.all(targetids2 == targetids[0:2]))
+
+        #- Select a single targetid as integer
+        zfit1 = read_zfit(filename, select_targetids=targetids[1])
+        targetids1 = np.unique(zfit1['targetid'])
+        self.assertEqual(len(targetids1), 1)
+        self.assertEqual(targetids1[0], targetids[1])
+
+        #- force column names to UPPERCASE
+        zfit = read_zfit(filename, upper=True)
+        self.assertIn('TARGETID', zfit.colnames)
+        self.assertIn('Z', zfit.colnames)
+        self.assertNotIn('z', zfit.colnames)
 
 
diff --git a/py/redrock/zscan.py b/py/redrock/zscan.py
@@ -13,6 +13,10 @@
 import numpy as np
 from scipy.optimize import lsq_linear, nnls
 
+#- Largest all-nines number that works for both float32 and float64,
+#- to be used for chi2 of failed fits, etc.
+HUGE_CHI2=9999999.0
+
 #try:
 #    import cupy as cp
 #    import cupyx.scipy
@@ -175,9 +179,9 @@ def _zchi2_one(Tb, weights, flux, wflux, zcoeff, solve_matrices_algorithm):
     try:
         zcoeff[:] = solve_matrices(M, y, solve_algorithm=solve_matrices_algorithm, use_gpu=False)
     except np.linalg.LinAlgError:
-        return 9e99
+        return HUGE_CHI2
     except NotImplementedError:
-        return 9e99
+        return HUGE_CHI2
 
     model = Tb.dot(zcoeff)
 
@@ -645,7 +649,7 @@ def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_mat
     """
     zchi2 = np.zeros(nz)
     if (weights.sum() == 0):
-        zchi2[:] = 9e99
+        zchi2[:] = HUGE_CHI2
         zcoeff = np.zeros((nz, nbasis))
         return (zchi2, zcoeff)
     if (use_gpu):
@@ -709,7 +713,7 @@ def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_mat
             zcoeff = solve_matrices(all_M, all_y, solve_algorithm=solve_matrices_algorithm,
                                     solver_args=solver_args, use_gpu=True)
         except NotImplementedError:
-            zchi2[:] = 9e99
+            zchi2[:] = HUGE_CHI2
             zcoeff = np.zeros((nz, nbasis))
             return (zchi2, zcoeff)
 
@@ -757,10 +761,10 @@ def calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_mat
             try:
                 zcoeff[i,:] = solve_matrices(M, y, solve_algorithm=solve_matrices_algorithm, solver_args=solver_args, use_gpu=False)
             except np.linalg.LinAlgError:
-                zchi2[i] = 9e99
+                zchi2[i] = HUGE_CHI2
                 continue
             except NotImplementedError:
-                zchi2[i] = 9e99
+                zchi2[i] = HUGE_CHI2
                 continue
 
             #4) Calculate dot products individually for each template
@@ -850,7 +854,7 @@ def calc_zchi2(target_ids, target_data, dtemplate, progress=None, use_gpu=False)
             #Use spectral_data() to get numpy arrays
             (weights, flux, wflux) = spectral_data(target_data[j].spectra)
         if np.sum(weights) == 0:
-            zchi2[j,:] = 9e99
+            zchi2[j,:] = HUGE_CHI2
             #Update progress for multiprocessing!!
             if dtemplate.comm is None and progress is not None:
                 progress.put(1)
@@ -1141,7 +1145,7 @@ def solve_matrices(M, y, solve_algorithm="PCA", solver_args=None, use_gpu=False)
                     res = nnls(Mcpu[j,:,:], ycpu[j,:])
                     zcoeff[j,:] = res[0]
                 except Exception:
-                    zcoeff[j,:] = 9e99
+                    zcoeff[j,:] = HUGE_CHI2
             return cp.asarray(zcoeff)
         else:
             try:
@@ -1169,7 +1173,7 @@ def solve_matrices(M, y, solve_algorithm="PCA", solver_args=None, use_gpu=False)
                     res = lsq_linear(Mcpu[j,:,:], ycpu[j,:], bounds=bounds, method='bvls')
                     zcoeff[j,:] = res.x
                 except np.linalg.LinAlgError:
-                    zcoeff[j,:] = 9e99
+                    zcoeff[j,:] = HUGE_CHI2
             return cp.asarray(zcoeff)
         else:
             try: