Skip to content

Commit

Permalink
support custom groups for zproc
Browse files Browse the repository at this point in the history
  • Loading branch information
Stephen Bailey authored and Stephen Bailey committed Jul 25, 2024
1 parent fe6cbb9 commit 96922c7
Show file tree
Hide file tree
Showing 5 changed files with 142 additions and 39 deletions.
54 changes: 39 additions & 15 deletions py/desispec/io/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -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:
Expand All @@ -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.")
Expand All @@ -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
Expand Down Expand Up @@ -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
}
Expand All @@ -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))
Expand Down
40 changes: 28 additions & 12 deletions py/desispec/scripts/zproc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)")
Expand All @@ -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)")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
#-------------------------------------------------------------------------
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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')

Expand Down
46 changes: 43 additions & 3 deletions py/desispec/test/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Expand Down Expand Up @@ -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='healpix')
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='healpix')
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')
Expand Down
2 changes: 1 addition & 1 deletion py/desispec/workflow/desi_proc_funcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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', 'CUSTOM_TILE']:
if system_name.startswith('perlmutter'):
nodes, runtime = 1, 50 #- timefactor will bring time back down
else:
Expand Down
Loading

0 comments on commit 96922c7

Please sign in to comment.