Merge pull request #2111 from desihub/zcat_units

desi_zcatalog --add-units option for DR1 patching

bin/desi_zcatalog

@@ -17,6 +17,7 @@ from __future__ import absolute_import, division, print_function
 
 import sys, os, glob
 import argparse
+import importlib.resources
 
 import numpy as np
 from numpy.lib.recfunctions import append_fields
@@ -27,10 +28,11 @@ from astropy.table import Table, hstack, vstack
 from desiutil.log import get_logger
 from desispec import io
 from desispec.zcatalog import find_primary_spectra
-from desispec.io.util import get_tempfilename, checkgzip, replace_prefix
+from desispec.io.util import get_tempfilename, checkgzip, replace_prefix, write_bintable
 from desispec.io.table import read_table
 from desispec.coaddition import coadd_fibermap
 from desispec.util import parse_keyval
+from desiutil.annotate import load_csv_units
 
 def match(table1,table2,key="TARGETID") :
     """
@@ -161,6 +163,9 @@ parser.add_argument('--recoadd-fibermap', action='store_true',
 parser.add_argument('--ztile', action='store_true',
         help="Used with --recoadd-fibermap, this is a tile-based recoadd "
              "not a healpix-based recoadd")
+parser.add_argument('--add-units', action='store_true',
+        help="Add units to output catalog from desidatamodel "
+             "column descriptions")
 
 # parser.add_argument("--match", type=str, nargs="*",
 #         help="match other tables (targets,truth...)")
@@ -176,6 +181,14 @@ if args.indir is None:
 if args.outfile is None:
     args.outfile = io.findfile('zcatalog')
 
+#- If adding units, check dependencies before doing a lot of work
+if args.add_units:
+    try:
+        import desidatamodel
+    except ImportError:
+        log.critical('Unable to import desidatamodel, required to add units (try "module load desidatamodel" first)')
+        sys.exit(1)
+
 #- Get redrock*.fits files in subdirs, excluding e.g. redrock*.log
 
 log.info(f'Looking for redrock files in subdirectories of {args.indir}')
@@ -372,7 +385,8 @@ zcat = np.array(zcat)
 #- Inherit header from first input, but remove keywords that don't apply
 #- across multiple files
 header = fitsio.read_header(redrockfiles[0], 0)
-for key in ['SPGRPVAL', 'TILEID', 'SPECTRO', 'PETAL', 'NIGHT', 'EXPID', 'HPXPIXEL']:
+for key in ['SPGRPVAL', 'TILEID', 'SPECTRO', 'PETAL', 'NIGHT', 'EXPID', 'HPXPIXEL',
+            'NAXIS', 'BITPIX', 'SIMPLE', 'EXTEND']:
     if key in header:
         header.delete(key)
 
@@ -392,12 +406,24 @@ if args.header is not None:
         key, value = parse_keyval(keyval)
         header[key] = value
 
+#- Add units if requested
+if args.add_units:
+    datamodeldir = str(importlib.resources.files('desidatamodel'))
+    unitsfile = os.path.join(datamodeldir, 'data', 'column_descriptions.csv')
+    log.info(f'Adding units from {unitsfile}')
+    units, comments = load_csv_units(unitsfile)
+else:
+    units = dict()
+    comments = dict()
+
 log.info(f'Writing {args.outfile}')
 tmpfile = get_tempfilename(args.outfile)
-fitsio.write(tmpfile, zcat, header=header, extname='ZCATALOG', clobber=True)
+
+write_bintable(tmpfile, zcat, header=header, extname='ZCATALOG',
+               units=units, clobber=True)
 
 if not args.minimal and expfm is not None:
-    fitsio.write(tmpfile, expfm, extname='EXP_FIBERMAP')
+    write_bintable(tmpfile, expfm, extname='EXP_FIBERMAP', units=units)
 
 os.rename(tmpfile, args.outfile)
 

py/desispec/io/util.py

@@ -187,6 +187,18 @@ def write_bintable(filename, data, header=None, comments=None, units=None,
     """Utility function to write a fits binary table complete with
     comments and units in the FITS header too.  DATA can either be
     dictionary, an Astropy Table, a numpy.recarray or a numpy.ndarray.
+
+    Args:
+        filename: full path to filename to write
+        data: dict or table-like data to write
+
+    Options:
+        header: dict-like header key/value pairs to propagate
+        comments (dict): comments[COLNAME] per column
+        units (dict): units[COLNAME] per column
+        extname (str): extension name for EXTNAME header
+        clobber (bool): if True, overwrite pre-existing file
+        primary_extname (str): EXTNAME to use for primary HDU=0
     """
     from astropy.table import Table
 
@@ -221,16 +233,22 @@ def write_bintable(filename, data, header=None, comments=None, units=None,
     #
     # Add comments and units to the *columns* of the table.
     #
-    for i in range(1, 999):
+    for i in range(1, 1000):
         key = 'TTYPE'+str(i)
         if key not in hdu.header:
             break
         else:
-            value = hdu.header[key]
-            if value in comments:
-                hdu.header[key] = (value, comments[value])
-            if value in units:
-                hdu.header['TUNIT'+str(i)] = (units[value], value+' units')
+            colname = hdu.header[key]
+            if colname in comments:
+                hdu.header[key] = (colname, comments[colname])
+            if colname in units and units[colname].strip() != '':
+                tunit_key = 'TUNIT'+str(i)
+                if tunit_key in hdu.header and hdu.header[tunit_key] != units[colname]:
+                    log.warning(f'Overriding {colname} units {hdu.header[tunit_key]} -> {units[colname]}')
+
+                # Add TUNITnn key after TFORMnn key (which is right after TTYPEnn)
+                tform_key = 'TFORM'+str(i)
+                hdu.header.insert(tform_key, (tunit_key, units[colname], colname+' units'), after=True)
     #
     # Add checksum cards.
     #