def test_get_next_expid(self):
a = obs.get_next_expid()
b = obs.get_next_expid()
c = obs.get_next_expid()
self.assertNotEqual(a, b)
self.assertNotEqual(b, c)
def batch_newexp(batchfile,
flavors,
nspec=5000,
night=None,
expids=None,
nodes=None,
pixprod=None,
desi_spectro_sim=None,
tileids=None,
seed=None):
'''
Write a slurm batch script for run newexp-desi for the list of flavors
'''
nexp = len(flavors)
timestr = '00:30:00'
logfile = '{}.%j.log'.format(batchfile)
np.random.seed(seed)
seeds = np.random.randint(2**32, size=nexp)
if night is None:
night = obs.get_night()
if expids is None:
expids = obs.get_next_expid(nexp)
if nodes is None:
nodes = calc_nodes(nexp, tasktime=1.5, maxtime=20)
if tileids is None:
tileids = list()
for flavor in flavors:
flavor = flavor.lower()
t = obs.get_next_tileid(program=flavor)
tileids.append(t)
if flavor in ('arc', 'flat'):
obs.update_obslog(obstype=flavor, program='calib', tileid=t)
elif flavor in ('bright', 'bgs', 'mws'):
obs.update_obslog(obstype='science',
program='bright',
tileid=t)
elif flavor in ('gray', 'grey'):
obs.update_obslog(obstype='science', program='gray', tileid=t)
else:
obs.update_obslog(obstype='science', program='dark', tileid=t)
if pixprod is None:
if 'PIXPROD' in os.environ:
pixprod = os.environ['PIXPROD']
else:
raise ValueError('must provide pixprod or set $PIXPROD')
if desi_spectro_sim is None:
if 'DESI_SPECTRO_SIM' in os.environ:
desi_spectro_sim = os.environ['DESI_SPECTRO_SIM']
else:
raise ValueError(
'must provide desi_spectro_sim or set $DESI_SPECTRO_SIM')
log.info('output dir {}/{}/{}'.format(desi_spectro_sim, pixprod, night))
assert len(expids) == len(flavors)
cmd = "srun -n 1 -N 1 -c $nproc /usr/bin/time newexp-desi --night {night} --nspec {nspec} --flavor {flavor} --expid {expid} --tileid {tileid} --seed {seed}"
with open(batchfile, 'w') as fx:
fx.write("#!/bin/bash -l\n\n")
fx.write("#SBATCH --partition=debug\n")
fx.write("#SBATCH --account=desi\n")
fx.write("#SBATCH --nodes={}\n".format(nodes))
fx.write("#SBATCH --time={}\n".format(timestr))
fx.write("#SBATCH --job-name=newexp\n")
fx.write("#SBATCH --output={}\n".format(logfile))
fx.write("if [ ${NERSC_HOST} = edison ]; then\n")
fx.write(" nproc=24\n")
fx.write("else\n")
fx.write(" nproc=32\n")
fx.write("fi\n\n")
fx.write('export DESI_SPECTRO_SIM={}\n'.format(desi_spectro_sim))
fx.write('export PIXPROD={}\n'.format(pixprod))
fx.write('\n')
fx.write('echo Starting at `date`\n\n')
fx.write('mkdir -p $DESI_SPECTRO_SIM/$PIXPROD/etc\n')
fx.write('mkdir -p $DESI_SPECTRO_SIM/$PIXPROD/{}\n'.format(night))
fx.write('\n')
for expid, flavor, tileid, seed in zip(expids, flavors, tileids,
seeds):
fx.write(
cmd.format(nspec=nspec,
night=night,
expid=expid,
flavor=flavor,
tileid=tileid,
seed=seed) + ' &\n')
fx.write('\nwait\n')
fx.write('\necho Done at `date`\n')
return expids
def test_get_next_expid(self):
a = obs.get_next_expid()
b = obs.get_next_expid()
c = obs.get_next_expid()
self.assertNotEqual(a, b)
self.assertNotEqual(b, c)
def batch_pixsim(batchfile,
flavors,
nspec=5000,
night=None,
expids=None,
nodes=None,
pixprod=None,
desi_spectro_sim=None,
seed=None):
'''
Write a slurm batch script for run newexp-desi for the list of flavors
'''
nexp = len(flavors)
nspectrographs = (nspec - 1) // 500 + 1
ntasks = nexp * 3
timestr = '00:30:00'
logfile = '{}.%j.log'.format(batchfile)
np.random.seed(seed)
seeds = np.random.randint(2**32, size=nexp)
if night is None:
night = obs.get_night()
if expids is None:
expids = obs.get_next_expid(nexp)
assert len(expids) == len(flavors)
if pixprod is None:
if 'PIXPROD' in os.environ:
pixprod = os.environ['PIXPROD']
else:
raise ValueError('must provide pixprod or set $PIXPROD')
if desi_spectro_sim is None:
if 'DESI_SPECTRO_SIM' in os.environ:
desi_spectro_sim = os.environ['DESI_SPECTRO_SIM']
else:
raise ValueError(
'must provide desi_spectro_sim or set $DESI_SPECTRO_SIM')
outdir = os.path.normpath(os.path.join(desi_spectro_sim, pixprod, night))
log.info('output dir {}'.format(outdir))
if nodes is None:
# nodes = calc_nodes(ntasks, tasktime=5, maxtime=20)
nodes = nexp * nspectrographs
cmd = "srun -n {nspectrographs} -N {nspectrographs} -c $nproc /usr/bin/time pixsim-desi --mpi --verbose --cosmics --night {night} --expid {expid} --seed {seed}"
with open(batchfile, 'w') as fx:
fx.write("#!/bin/bash -l\n\n")
fx.write("#SBATCH --partition=debug\n")
fx.write("#SBATCH --account=desi\n")
fx.write("#SBATCH --nodes={}\n".format(nodes))
fx.write("#SBATCH --time={}\n".format(timestr))
fx.write("#SBATCH --job-name=pixsim\n")
fx.write("#SBATCH --output={}\n".format(logfile))
fx.write("if [ ${NERSC_HOST} = edison ]; then\n")
fx.write(" nproc=24\n")
fx.write("else\n")
fx.write(" nproc=32\n")
fx.write("fi\n\n")
fx.write('export DESI_SPECTRO_SIM={}\n'.format(desi_spectro_sim))
fx.write('export PIXPROD={}\n'.format(pixprod))
fx.write('\necho Starting at `date`\n')
for expid, flavor, seed in zip(expids, flavors, seeds):
fx.write('\n#- Exposure {} ({})\n'.format(expid, flavor))
cx = cmd.format(night=night,
expid=expid,
nspectrographs=nspectrographs,
seed=seed)
fx.write(cx + ' &\n')
fx.write('\nwait\n')
fx.write('\n#- Preprocess raw data\n')
cmd = 'srun -n 1 -N 1 desi_preproc --infile {infile} --outdir {outdir} --cameras {cameras}'
outdir = os.path.join(desi_spectro_sim, pixprod, night)
for expid in expids:
infile = desispec.io.findfile('raw',
night=night,
expid=expid,
outdir=outdir)
for i in range(nspectrographs):
cameras = 'b{},r{},z{}'.format(i, i, i)
cx = cmd.format(infile=infile, outdir=outdir, cameras=cameras)
fx.write(cx + ' &\n')
fx.write('\nwait\n')
fx.write('echo Done at `date`\n')
]
ntask = len(flavors) * len(nights)
comm = MPI.COMM_WORLD
if comm.size < ntask:
if comm.rank == 0:
print("Communicator size ({}) too small for {} tasks".format(comm.size, ntask), flush=True)
comm.Abort()
np.random.seed(123456)
seeds = np.random.randint(2**32, size=ntask)
expids = None
if comm.rank == 0:
expids = obs.get_next_expid(ntask)
expids = comm.bcast(expids, root=0)
tileids = list()
if comm.rank == 0:
for nt in nights:
for fl in flavors:
flavor = fl.lower()
t = obs.get_next_tileid(program=flavor)
tileids.append(t)
if flavor in ('arc', 'flat'):
obs.update_obslog(obstype=flavor, program='calib', tileid=t)
elif flavor in ('bright', 'bgs', 'mws'):
obs.update_obslog(obstype='science', program='bright', tileid=t)
elif flavor in ('gray', 'grey'):
obs.update_obslog(obstype='science', program='gray', tileid=t)
