support custom groups for zproc, e.g. 1x_depth

py/desispec/io/meta.py

@@ -69,7 +69,7 @@ def get_readonly_filepath(filepath):
     return filepath
 
 def findfile(filetype, night=None, expid=None, camera=None,
-        tile=None, groupname=None,
+        tile=None, groupname=None, subgroup=None,
         healpix=None, nside=64,
         band=None, spectrograph=None,
         survey=None, faprogram=None,
@@ -88,6 +88,7 @@ def findfile(filetype, night=None, expid=None, camera=None,
         camera : 'b0' 'r1' .. 'z9'
         tile : integer tile (pointing) number
         groupname : spectral grouping name (e.g. "healpix", "cumulative", "pernight")
+        subgroup : (str) subgrouping name for non-standard groupnames
         healpix : healpix pixel number
         nside : healpix nside
         band : one of 'b','r','z' identifying the camera band
@@ -203,16 +204,16 @@ def findfile(filetype, night=None, expid=None, camera=None,
         #
         # spectra- tile based
         #
-        coadd_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/coadd-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
-        rrdetails_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/rrdetails-{spectrograph:d}-{tile:d}-{nightprefix}{night}.h5',
-        rrmodel_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/rrmodel-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
-        spectra_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/spectra-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits.gz',
-        redrock_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/redrock-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
-        qso_mgii_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/qso_mgii-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
-        qso_qn_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/qso_qn-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
-        emline_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{night}/emline-{spectrograph:d}-{tile:d}-{nightprefix}{night}.fits',
+        coadd_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/coadd-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
+        rrdetails_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/rrdetails-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.h5',
+        rrmodel_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/rrmodel-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
+        spectra_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/spectra-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits.gz',
+        redrock_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/redrock-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
+        qso_mgii_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/qso_mgii-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
+        qso_qn_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/qso_qn-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
+        emline_tile='{specprod_dir}/tiles/{groupname}/{tile:d}/{subgroup}/emline-{spectrograph:d}-{tile:d}-{nightprefix}{subgroup}.fits',
         #
-        # spectra- single exp tile based
+        # spectra- single exp tile based requires cusom formatting for expid:08d
         #
         coadd_single='{specprod_dir}/tiles/perexp/{tile:d}/{expid:08d}/coadd-{spectrograph:d}-{tile:d}-exp{expid:08d}.fits',
         rrdetails_single='{specprod_dir}/tiles/perexp/{tile:d}/{expid:08d}/rrdetails-{spectrograph:d}-{tile:d}-exp{expid:08d}.h5',
@@ -267,12 +268,18 @@ def findfile(filetype, night=None, expid=None, camera=None,
             ):
         groupname = 'cumulative'
 
-    if str(groupname) == "cumulative":
+    if groupname == "cumulative":
         nightprefix = "thru"
-    elif groupname == 'perexp':
+        subgroup = str(night)
+    elif groupname == "pernight":
+        nightprefix = ""
+        subgroup = str(night)
+    elif groupname == "perexp":
         nightprefix = "exp"
-    else:
+    elif groupname == "healpix":
         nightprefix = ""
+    else:
+        nightprefix = str(groupname)+'-'
 
     #- backwards compatibility: try interpreting groupname as a healpix number
     if healpix is None and tile is None and groupname is not None:
@@ -297,6 +304,9 @@ def findfile(filetype, night=None, expid=None, camera=None,
     if isinstance(tile, str):    tile = int(tile)
     if isinstance(spectrograph, str): spectrograph = int(spectrograph)
 
+    #- Determine if this is healpix-based or tile-based objects, and update
+    #- location dict for which flavor of coadd/spectra/redrock/etc is needed,
+    #- removing the _hp, _single, _tile suffixes from the keys
     loc_copy = location.copy()
     if tile is not None:
         log.debug("Tile-based files selected; healpix-based files and input will be ignored.")
@@ -314,6 +324,11 @@ def findfile(filetype, night=None, expid=None, camera=None,
                 if key.endswith('_tile'):
                     root_key = key.removesuffix('_tile')
                     location[root_key] = val
+
+        ## cumulative and pernight use night as subgroup
+        if groupname in ('cumulative', 'pernight'):
+            subgroup = night
+
     else:
         ## If not tile based then use the hp naming scheme
         ## Do loop to improve scaling with additional file types
@@ -372,7 +387,8 @@ def findfile(filetype, night=None, expid=None, camera=None,
 
     actual_inputs = {
         'specprod_dir':specprod_dir, 'specprod':specprod, 'qaprod_dir':qaprod_dir, 'tiles_dir':tiles_dir,
-        'night':night, 'expid':expid, 'tile':tile, 'camera':camera, 'groupname':groupname,
+        'night':night, 'expid':expid, 'tile':tile, 'camera':camera,
+        'groupname':groupname, 'subgroup':subgroup,
         'healpix':healpix, 'nside':nside, 'hpixdir':hpixdir, 'band':band,
         'spectrograph':spectrograph, 'nightprefix':nightprefix, 'month':month
         }
@@ -385,9 +401,17 @@ def findfile(filetype, night=None, expid=None, camera=None,
     if 'rawdata_dir' in required_inputs:
         actual_inputs['rawdata_dir'] = rawdata_dir
 
+    #- If any inputs missing, print all missing inputs, then raise single ValueError
+    missing_inputs = False
     for i in required_inputs:
         if actual_inputs[i] is None:
-            raise ValueError("Required input '{0}' is not set for type '{1}'!".format(i,filetype))
+            log.error("Required input '{0}' is not set for type '{1}'!".format(i,filetype))
+            missing_inputs = True
+
+    if missing_inputs:
+        msg = f"Missing inputs for {location[filetype]}"
+        log.critical(msg)
+        raise ValueError(msg)
 
     #- normpath to remove extraneous double slashes /a/b//c/d
     filepath = os.path.normpath(location[filetype].format(**actual_inputs))

py/desispec/scripts/zproc.py

@@ -58,7 +58,7 @@ def parse(options=None):
     parser = argparse.ArgumentParser(usage="{prog} [options]")
 
     parser.add_argument("-g", "--groupname", type=str,
-                        help="Redshift grouping type: cumulative, perexp, pernight, healpix")
+                        help="Redshift grouping type: cumulative, perexp, pernight, healpix, or custom name")
 
     #- Options for tile-based redshifts
     tiles_options = parser.add_argument_group("tile-based options (--groupname perexp, pernight, or cumulative)")
@@ -74,6 +74,8 @@ def parse(options=None):
     tiles_options.add_argument("-c", "--cameras", type=str,
                         help="Subset of cameras to process, either as a camword (e.g. a012)" +
                              "Or a comma separated list (e.g. b0,r0,z0).")
+    parser.add_argument("--subgroup", type=str,
+                        help="subgroup to use for non-standard groupname values")
 
     #- Options for healpix-based redshifts
     healpix_options = parser.add_argument_group("healpix-based options (--groupname healpix)")
@@ -225,6 +227,21 @@ def main(args=None, comm=None):
 
             raise ValueError(msg)
 
+    ## Unpack arguments for shorter names (tileid might be None, ok)
+    tileid, groupname, subgroup = args.tileid, args.groupname, args.subgroup
+
+    known_groups = ['cumulative', 'pernight', 'perexp', 'healpix']
+    if groupname not in known_groups:
+        if subgroup is None:
+            msg = f'Non-standard --groupname={groupname} requires --subgroup'
+            if rank == 0:
+                log.critical(msg)
+            raise ValueError(msg)
+        else:
+            msg = f'Non-standard {groupname=} not in {known_groups}; using {subgroup=}'
+            if rank == 0:
+                log.warning(msg)
+
     #- redrock non-MPI mode isn't compatible with GPUs,
     #- so if zproc is running in non-MPI mode, force --no-gpu
     #- https://github.com/desihub/redrock/issues/223
@@ -276,15 +293,6 @@ def main(args=None, comm=None):
     else:
         camword = create_camword(args.cameras)
 
-    ## Unpack arguments for shorter names (tileid might be None, ok)
-    tileid, groupname = args.tileid, args.groupname
-
-    known_groups = ['cumulative', 'pernight', 'perexp', 'healpix']
-    if groupname not in known_groups:
-        msg = 'obstype {} not in {}'.format(groupname, known_groups)
-        log.error(msg)
-        raise ValueError(msg)
-
     if args.batch:
         err = 0
         #-------------------------------------------------------------------------
@@ -293,6 +301,7 @@ def main(args=None, comm=None):
             ## create the batch script
             cmdline = list(sys.argv).copy()
             scriptfile = create_desi_zproc_batch_script(group=groupname,
+                                                        subgroup=subgroup,
                                                         tileid=tileid,
                                                         cameras=camword,
                                                         thrunight=args.thrunight,
@@ -395,6 +404,9 @@ def main(args=None, comm=None):
     if rank == 0:
         log.info('------------------------------')
         log.info('Groupname {}'.format(groupname))
+        if subgroup is not None:
+            log.info('Subgroup {}'.format(subgroup))
+
         if args.healpix is not None:
             log.info(f'Healpixels {args.healpix}')
         else:
@@ -443,10 +455,14 @@ def main(args=None, comm=None):
     if groupname == 'healpix':
         findfileopts = dict(groupname=groupname, survey=args.survey, faprogram=args.program)
     else:
-        findfileopts = dict(night=thrunight, tile=tileid, groupname=groupname)
-        if groupname == 'perexp':
+        findfileopts = dict(tile=tileid, groupname=groupname, subgroup=subgroup)
+        if groupname in ('cumulative', 'pernight'):
+            findfileopts['night'] = thrunight
+        elif groupname == 'perexp':
             assert len(expids) == 1
             findfileopts['expid'] = expids[0]
+        elif subgroup is not None:
+            findfileopts['subgroup'] = subgroup
 
     timer.stop('preflight')
 

py/desispec/test/test_io.py

@@ -844,11 +844,12 @@ def test_findfile(self):
         with self.assertRaises(ValueError) as cm:
             foo = findfile('stdstars',expid=2,spectrograph=0)
         the_exception = cm.exception
-        self.assertEqual(str(the_exception), "Required input 'night' is not set for type 'stdstars'!")
+        self.assertTrue(str(the_exception), "Missing inputs for")
+
         with self.assertRaises(ValueError) as cm:
             foo = findfile('spectra', survey='main', groupname=123)
         the_exception = cm.exception
-        self.assertEqual(str(the_exception), "Required input 'faprogram' is not set for type 'spectra'!")
+        self.assertTrue(str(the_exception), "Missing inputs for")
 
         #- Some findfile calls require $DESI_SPECTRO_DATA; others do not
         del os.environ['DESI_SPECTRO_DATA']
@@ -891,20 +892,59 @@ def test_findfile(self):
         with self.assertRaises(ValueError):
             a = findfile('cframe', night=20200317, expid=18, camera='Hasselblad')
 
-        # Test healpix versus tiles
+        # Test healpix versus tiles for various groupings
         a = findfile('spectra', groupname='5286', survey='main', faprogram='BRIGHT')
         b = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
                          os.environ['SPECPROD'],
                          'healpix', 'main', 'bright', '52', '5286',
                          'spectra-main-bright-5286.fits.gz')
         self.assertEqual(a, b)
+
         a = findfile('spectra', tile=68000, night=20200314, spectrograph=2)
         b = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
                          os.environ['SPECPROD'], 'tiles', 'cumulative',
                          '68000', '20200314',
                          'spectra-2-68000-thru20200314.fits.gz')
         self.assertEqual(a, b)
 
+        a = findfile('coadd', tile=68000, groupname='perexp', expid=1234,
+                     spectrograph=2)
+        b = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
+                         os.environ['SPECPROD'], 'tiles', 'perexp',
+                         '68000', '00001234',
+                         'coadd-2-68000-exp00001234.fits')
+        self.assertEqual(a, b)
+
+        a = findfile('coadd', tile=68000, groupname='1x_depth', subgroup=42,
+                     spectrograph=2)
+        b = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
+                         os.environ['SPECPROD'], 'tiles', '1x_depth',
+                         '68000', '42',
+                         'coadd-2-68000-1x_depth-42.fits')
+        self.assertEqual(a, b)
+
+        a = findfile('redrock', tile=68000, groupname='coffeeboba', subgroup=13,
+                     spectrograph=2)
+        b = os.path.join(os.environ['DESI_SPECTRO_REDUX'],
+                         os.environ['SPECPROD'], 'tiles', 'coffeeboba',
+                         '68000', '13',
+                         'redrock-2-68000-coffeeboba-13.fits')
+        self.assertEqual(a, b)
+
+        #- groupname shouldn't impact non-tile non-healpix files
+        night = 20201010
+        expid = 1234
+        tileid = 8888
+        refpath = os.path.expandvars(f'$DESI_SPECTRO_DATA/{night}/{expid:08d}/fiberassign-{tileid:06d}.fits.gz')
+        for groupname in ('healpix', 'pernight', 'cumulative', 'perexp', 'blatfoo'):
+            testpath = findfile('fiberassign', night=night, expid=expid, tile=tileid, groupname=groupname)
+            self.assertEqual(testpath, refpath)
+
+        refpath = os.path.expandvars(f'$DESI_SPECTRO_REDUX/$SPECPROD/preproc/{night}/{expid:08d}/tilepix-{tileid}.json')
+        for groupname in ('healpix', 'pernight', 'cumulative', 'perexp', 'blatfoo'):
+            testpath = findfile('tilepix', night=night, expid=expid, tile=tileid, groupname=groupname)
+            self.assertEqual(testpath, refpath)
+
         #- Can't set both tile and healpix
         with self.assertRaises(ValueError):
             findfile('redrock', tile=1234, healpix=1234, survey='main', faprogram='dark')

py/desispec/workflow/desi_proc_funcs.py

@@ -455,7 +455,7 @@ def determine_resources(ncameras, jobdesc, nexps=1, forced_runtime=None, queue=N
     elif jobdesc == 'NIGHTLYBIAS':
         ncores, runtime = 15, 5
         nodes = 2
-    elif jobdesc in ['PEREXP', 'PERNIGHT', 'CUMULATIVE']:
+    elif jobdesc in ['PEREXP', 'PERNIGHT', 'CUMULATIVE', 'CUSTOMZTILE']:
         if system_name.startswith('perlmutter'):
             nodes, runtime = 1, 50  #- timefactor will bring time back down
         else: