def assign_quickcat_mtl_loop(destination_dir=None, targets_file=None, truth_file=None, zcat_file=None):
from astropy.table import Table, Column
truth = Table.read(truth_file)
targets = Table.read(targets_file)
# this part tests another PR https://github.com/desihub/desitarget/pull/32
if zcat_file is None:
mtl = desitarget.mtl.make_mtl(targets)
mtl.write('mtl.fits', overwrite=True)
else:
zcat = Table.read(zcat_file, format='fits')
mtl = desitarget.mtl.make_mtl(targets, zcat)
mtl.write('mtl.fits', overwrite=True)
import glob
import shutil
# for f in glob.glob("*.fits"):
shutil.copy('mtl.fits', "/project/projectdirs/desi/users/forero/mtl")
# this now tests another PR https://github.com/desihub/fiberassign/pull/24
# I had to go first into fiberassign to run make && make install
import subprocess
p = subprocess.call(["/global/homes/f/forero/fiberassign/bin/./fiberassign", "fa_features.txt"], stdout=subprocess.PIPE)
#find the tiles
from glob import glob
tiledir = '/project/projectdirs/desi/users/forero/fiberassign'
tilefiles = sorted(glob(tiledir+'/tile*.fits'))
len(tilefiles)
import imp
# this ugly hack avoids import from desi stack
desisim = imp.load_source("desisim.quickcat", "/global/homes/f/forero/desisim/py/desisim/quickcat.py")
# import os
# os.chdir('/global/homes/f/forero/desisim/py')
#from desisim.quickcat import quickcat
if zcat_file is None:
zcat = desisim.quickcat(tilefiles, targets, truth, zcat=None, perfect=False)
zcat.write('zcat.fits', overwrite=True)
else:
zcat = Table.read(zcat_file, format='fits')
newzcat = desisim.quickcat(tilefiles, targets, truth, zcat=zcat, perfect=False)
newzcat.write('zcat.fits', format='fits', overwrite=True)
shutil.copy('zcat.fits', "/project/projectdirs/desi/users/forero/zcat")
shutil.copy('zcat.fits', destination_dir)
#clean files
import os
for tilefile in tilefiles:
os.remove(tilefile)
if (not os.path.exists(mtlfile)) or (not os.path.exists(brightfile)) or (
not os.path.exists(darkfile)):
targetdata = fitsio.read(targetfile, 'TARGETS', columns=columns)
print('Done reading target data')
#compute MTL
if not os.path.exists(mtlfile):
import desitarget.mtl
mtl = desitarget.mtl.make_mtl(targetdata)
mtl.meta['EXTNAME'] = 'MTL'
# rewrite NUMOBS for BGS targets
ii = mtl['BGS_TARGET'] != 0
mtl['NUMOBS_MORE'][ii] = 4
mtl.write(mtlfile)
#print some stats
print('MWS_TARGETS: {}'.format(
np.count_nonzero(targetdata['MWS_TARGET'] != 0)))
print('BGS_TARGETS: {}'.format(
np.count_nonzero(targetdata['BGS_TARGET'] != 0)))
print('DESI_TARGETS: {}'.format(
np.count_nonzero(targetdata['DESI_TARGET'] != 0)))
#standards
if not os.path.exists(darkfile):
darkstd = (targetdata['DESI_TARGET']
& desi_mask.mask('STD_FSTAR|STD_WD')) != 0
darkdata = targetdata[darkstd]
obscond = np.ones(len(darkdata), dtype=np.int64)
from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("-s", "--src", help="Tractor/sweeps file or root directory with tractor/sweeps files")
ap.add_argument("-d", "--dest", help="Output target selection file")
ns = ap.parse_args()
iter_truth = dtio.iter_files(ns.src, 'truth')
iter_target = dtio.iter_files(ns.src, 'targets')
alltruth = []
alltargets = []
for truth_file, target_file in zip(iter_truth, iter_target):
alltruth.append(Table.read(truth_file))
alltargets.append(Table.read(target_file))
print(truth_file, target_file)
targets = vstack(alltargets)
truth = vstack(alltruth)
mtl = mtl.make_mtl(targets)
out_targets = os.path.join(ns.dest,'targets.fits')
out_truth = os.path.join(ns.dest,'truth.fits')
out_mtl = os.path.join(ns.dest,'mtl.fits')
targets.write(out_targets, overwrite=True)
truth.write(out_truth, overwrite=True)
mtl.write(out_mtl, overwrite=True)
not os.path.exists(out_mtl):
import desitarget
targets = fitsio.read(out_targets, 'TARGETS')
if not os.path.exists(out_stddark):
print('Generating ' + out_stddark)
isSTD = (targets['DESI_TARGET']
& desitarget.desi_mask['STD_FSTAR']) != 0
tmpout = out_stddark + '.tmp'
fitsio.write(tmpout, targets[isSTD], extname='STD')
os.rename(tmpout, out_stddark)
if not os.path.exists(out_stdbright):
print('Generating ' + out_stdbright)
isSTD = (targets['DESI_TARGET']
& desitarget.desi_mask['STD_BRIGHT']) != 0
tmpout = out_stdbright + '.tmp'
fitsio.write(tmpout, targets[isSTD], extname='STD')
os.rename(tmpout, out_stdbright)
if not os.path.exists(out_mtl):
print('Generating ' + out_mtl)
import desitarget.mtl
mtl = desitarget.mtl.make_mtl(targets)
tmpout = out_mtl + '.tmp'
mtl.meta['EXTNAME'] = 'MTL'
mtl.write(tmpout, format='fits')
os.rename(tmpout, out_mtl)
MPI.Finalize()
#print(tiles_pass)
truth = Table.read("truth.fits")
for i_pass in passes:
pass_tilefile = 'pass_{:02d}_{}'.format(i_pass, tilefile)
input_targets_file = 'mtl_targets_{:02d}.fits'.format(i_pass)
output_targets_file = 'mtl_targets_{:02d}.fits'.format(i_pass+1)
fba_output_dir = 'fiberassign_pass_{:02d}'.format(i_pass)
old_zcat_file = 'zcat_{:02}.fits'.format(i_pass-1)
new_zcat_file = 'zcat_{:02}.fits'.format(i_pass)
p = subprocess.call(["fba_run", "--targets", input_targets_file,'--sky', 'sky.fits',
'--footprint', pass_tilefile,
'--dir', fba_output_dir, '--overwrite'])
fba_files = glob.glob(os.path.join(fba_output_dir,"fiberassign*.fits"))
print(i_pass, fba_files)
targets = Table.read(input_targets_file)
if i_pass==0:
zcat = quickcat(fba_files, targets, truth)
else:
old_zcat = Table.read(old_zcat_file)
zcat = quickcat(fba_files, targets, truth, zcat=old_zcat, perfect=True)
zcat.write(new_zcat_file, overwrite=True)
mtl = desitarget.mtl.make_mtl(targets, zcat)
mtl.write(output_targets_file, overwrite=True)
print 'targets.colnames = ', targets.colnames
print 'truth.colnames = ', truth.colnames
iisky = (targets['DESI_TARGET'] & desi_mask.SKY) != 0
iistd = (targets['DESI_TARGET'] & desi_mask.STD_FSTAR) != 0
iiobj = (~iisky & ~iistd)
print np.count_nonzero(iisky), np.count_nonzero(iistd), np.count_nonzero(iiobj)
targets[iiobj].write('targets.fits', overwrite=True)
targets[iisky].write('sky.fits', overwrite=True)
targets[iistd].write('stdstars.fits', overwrite=True)
truth[iiobj].write('truth.fits', overwrite=True)
# this part tests another PR https://github.com/desihub/desitarget/pull/32
mtl = desitarget.mtl.make_mtl(targets)
mtl.write('mtl.fits', overwrite=True)
import glob
import shutil
for f in glob.glob("*.fits"):
shutil.copy(f, "/project/projectdirs/desi/users/forero/mtl")
# this now tests another PR https://github.com/desihub/fiberassign/pull/24
# I had to go first into fiberassign to run make && make install
import subprocess
p = subprocess.call(["./fiberassign/bin/./fiberassign", "fa_features.txt"],
stdout=subprocess.PIPE)
#find the tiles
from glob import glob
tiledir = '/project/projectdirs/desi/users/forero/fiberassign'
print 'targets.colnames = ', targets.colnames
print 'truth.colnames = ', truth.colnames
iisky = (targets['DESI_TARGET'] & desi_mask.SKY) != 0
iistd = (targets['DESI_TARGET'] & desi_mask.STD_FSTAR) != 0
iiobj = (~iisky & ~iistd)
print np.count_nonzero(iisky), np.count_nonzero(iistd), np.count_nonzero(iiobj)
targets[iiobj].write('targets.fits', overwrite=True)
targets[iisky].write('sky.fits', overwrite=True)
targets[iistd].write('stdstars.fits', overwrite=True)
truth[iiobj].write('truth.fits', overwrite=True)
# this part tests another PR https://github.com/desihub/desitarget/pull/32
mtl = desitarget.mtl.make_mtl(targets)
mtl.write('mtl.fits', overwrite=True)
! cp *.fits /project/projectdirs/desi/users/forero/mtl
# this now tests another PR https://github.com/desihub/fiberassign/pull/24
# I had to go first into fiberassign to run make && make install
! ./fiberassign/bin/./fiberassign fa_features.txt
#find the tiles
from glob import glob
tiledir = '/project/projectdirs/desi/users/forero/fiberassign'
tilefiles = sorted(glob(tiledir+'/tile*.fits'))
len(tilefiles)
# this ugly hack avoids import from desi stack
import os
def simulate_epoch(self, epoch, truth, targets, perfect=False, zcat=None):
"""Core routine simulating a DESI epoch,
Args:
epoch (int): epoch to simulate
perfect (boolean): Default: False. Selects whether how redshifts are taken from the truth file.
True: redshifts are taken without any error from the truth file.
False: redshifts include uncertainties.
truth (Table): Truth data
targets (Table): Targets data
zcat (Table): Redshift Catalog Data
Notes:
This routine simulates three steps:
* Merged target list creation
* Fiber allocation
* Redshift catalogue construction
"""
# create the MTL file
print("{} Starting MTL".format(asctime()))
self.mtl_file = os.path.join(self.tmp_output_path, 'mtl.fits')
if zcat is None:
mtl = desitarget.mtl.make_mtl(targets)
else:
mtl = desitarget.mtl.make_mtl(targets, zcat)
mtl.write(self.mtl_file, overwrite=True)
del mtl
gc.collect()
print("{} Finished MTL".format(asctime()))
# clean files and prepare fiberasign inputs
tilefiles = sorted(glob.glob(self.tmp_fiber_path+'/tile*.fits'))
if tilefiles:
for tilefile in tilefiles:
os.remove(tilefile)
# setup the tileids for the current observation epoch
self.create_surveyfile(epoch)
# launch fiberassign
print("{} Launching fiberassign".format(asctime()))
f = open('fiberassign.log','a')
p = subprocess.call([self.fiberassign,
'--mtl', os.path.join(self.tmp_output_path, 'mtl.fits'),
'--stdstar', self.stdfile,
'--sky', self.skyfile,
'--surveytiles', self.surveyfile,
'--outdir',os.path.join(self.tmp_output_path, 'fiberassign'),
'--fibstatusfile', self.fibstatusfile],
stdout=f)
print("{} Finished fiberassign".format(asctime()))
f.close()
# create a list of fiberassign tiles to read and update zcat
self.update_observed_tiles(epoch)
#update obsconditions from progress_data
# progress_data = Table.read(self.progress_files[epoch + 1])
# ii = np.in1d(progress_data['TILEID'], self.observed_tiles)
# obsconditions = progress_data[ii]
obsconditions = None
print('tilefiles', len(self.tilefiles))
# write the zcat, it uses the tilesfiles constructed in the last step
self.zcat_file = os.path.join(self.tmp_output_path, 'zcat.fits')
print("{} starting quickcat".format(asctime()))
newzcat = quickcat(self.tilefiles, targets, truth, zcat=zcat,
obsconditions=obsconditions, perfect=perfect)
print("{} writing zcat".format(asctime()))
newzcat.write(self.zcat_file, format='fits', overwrite=True)
print("{} Finished zcat".format(asctime()))
del newzcat
gc.collect()
return
def simulate_epoch(self,
epoch,
truth,
targets,
perfect=False,
mtl=None,
zcat=None):
"""Core routine simulating a DESI epoch,
Args:
epoch (int): epoch to simulate
perfect (boolean): Default: False. Selects whether how redshifts are taken from the truth file.
True: redshifts are taken without any error from the truth file.
False: redshifts include uncertainties.
truth (Table): Truth data
targets (Table): Targets data
mtl (Table): Merged Targets List data
zcat (Table): Redshift Catalog Data
Notes:
This routine simulates three steps:
* Merged target list creation
* Fiber allocation
* Redshift catalogue construction
"""
# create the MTL file
print("{} Starting MTL".format(asctime()))
self.mtl_file = os.path.join(self.tmp_output_path, 'mtl.fits')
mtl = desitarget.mtl.make_mtl(targets, zcat)
mtl.write(self.mtl_file, overwrite=True)
print("{} Finished MTL".format(asctime()))
# clean files and prepare fiberasign inputs
tilefiles = sorted(glob.glob(self.tmp_fiber_path + '/tile*.fits'))
if tilefiles:
for tilefile in tilefiles:
os.remove(tilefile)
# setup the tileids for the current observation epoch
self.create_surveyfile(epoch)
self.create_fiberassign_input()
# launch fiberassign
print("{} Launching fiberassign".format(asctime()))
f = open('fiberassign.log', 'a')
p = subprocess.call([
self.fiberassign_exec,
os.path.join(self.tmp_output_path, 'fa_features.txt')
],
stdout=f) # stdout=subprocess.PIPE)
print("{} Finished fiberassign".format(asctime()))
f.close()
# create a list of fiberassign tiles to read and update zcat
self.update_observed_tiles(epoch)
#update obsconditions from progress_data
# progress_data = Table.read(self.progress_files[epoch + 1])
# ii = np.in1d(progress_data['TILEID'], self.observed_tiles)
# obsconditions = progress_data[ii]
obsconditions = None
# write the zcat, it uses the tilesfiles constructed in the last step
self.zcat_file = os.path.join(self.tmp_output_path, 'zcat.fits')
print("{} starting quickcat".format(asctime()))
newzcat = quickcat(self.tilefiles,
targets,
truth,
zcat=zcat,
obsconditions=obsconditions,
perfect=perfect)
print("{} writing zcat".format(asctime()))
newzcat.write(self.zcat_file, format='fits', overwrite=True)
print("{} Finished zcat".format(asctime()))
return mtl, newzcat
def run_strategy(initial_mtl_file,
truth_file,
sky_file,
output_path="./",
batch_path="./",
program='dark',
sbatch=0,
mxbatch=100):
os.makedirs(output_path, exist_ok=True)
targets_path = '{}/targets'.format(output_path)
zcat_path = '{}/zcat'.format(output_path)
os.makedirs(targets_path, exist_ok=True)
os.makedirs(zcat_path, exist_ok=True)
os.makedirs('{}/fiberassign'.format(output_path), exist_ok=True)
batch_files = glob.glob(batch_path + "/batch_*_" + program + ".fits")
batch_files.sort()
# Read targets and truth
print('reading truth file')
truth = Table.read(truth_file)
print('done readying truth file')
#obsconditions
obsconditions = None
if program == 'dark':
obsconditions = 'DARK|GRAY'
if program == 'bright':
obsconditions = 'BRIGHT'
n_batch = len(batch_files)
if mxbatch < n_batch:
n_batch = mxbatch
for i_batch in range(sbatch, n_batch):
print()
print("Batch {}".format(i_batch))
fiberassign_path = '{}/fiberassign/{:04d}'.format(output_path, i_batch)
os.makedirs(fiberassign_path, exist_ok=True)
footprint = batch_files[i_batch]
mtl_filename = os.path.join(targets_path,
'{:04d}_mtl.fits'.format(i_batch))
new_mtl_filename = os.path.join(targets_path,
'{:04d}_mtl.fits'.format(i_batch + 1))
zcat_filename = os.path.join(zcat_path,
'{:04d}_zcat.fits'.format(i_batch))
old_zcat_filename = os.path.join(
zcat_path, '{:04d}_zcat.fits'.format(i_batch - 1))
if i_batch == 0:
shutil.copyfile(initial_mtl_file, mtl_filename)
print(footprint)
fba_run = 'fba_run --targets {} --sky {} --footprint {} --dir {} --rundate 2020-01-01T00:00:00 --overwrite'.format(
mtl_filename, sky_file, footprint, fiberassign_path)
print(fba_run)
os.system(fba_run)
# Gather fiberassign files
fba_files = np.sort(
glob.glob(os.path.join(fiberassign_path, "fba-*.fits")))
# read the current mtl file
targets = Table.read(mtl_filename)
# Compute zcat
if i_batch == 0:
zcat = desisim.quickcat.quickcat(fba_files,
targets,
truth,
fassignhdu='FASSIGN',
perfect=True)
else:
old_zcat = Table.read(old_zcat_filename)
zcat = desisim.quickcat.quickcat(fba_files,
targets,
truth,
fassignhdu='FASSIGN',
zcat=old_zcat,
perfect=True)
del old_zcat
zcat.write(zcat_filename, overwrite=True)
mtl = desitarget.mtl.make_mtl(targets, obsconditions, zcat=zcat)
del targets
del zcat
#gc.collect()
print('writing mtl ')
mtl.write(new_mtl_filename, overwrite=True)
print('wrote mtl')
del mtl
gc.collect()
return True
import numpy as np
import desitarget.mtl
from astropy.table import Table, Column
import os.path
import os
epoch_path = "/project/projectdirs/desi/mocks/preliminary/mtl/lowfat/"
targets_path = "/project/projectdirs/desi/mocks/preliminary/mtl/lowfat/"
output_path = "/project/projectdirs/desi/users/forero/lowfat_perfect/"
print(os.path.join(output_path, '1/zcat.fits'))
zcat = Table.read(os.path.join(output_path, '1/zcat.fits'), format='fits')
targets = Table.read(os.path.join(targets_path,'targets.fits'))
mtl = desitarget.mtl.make_mtl(targets, zcat, trim=True)
mtl.write('a.fits', overwrite=True)
mtl_test = Table.read('a.fits')
iielg = mtl_test['DESI_TARGET'] == 2
print(mtl_test['NUMOBS_MORE'][iielg])
mtl = desitarget.mtl.make_mtl(targets, zcat, trim=False)
mtl.write('b.fits', overwrite=True)
mtl_test = Table.read('b.fits')
iielg = mtl_test['DESI_TARGET'] == 2
print(mtl_test['NUMOBS_MORE'][iielg])
os.remove('a.fits')
os.remove('b.fits')
def run_strategy(footprint_names, pass_names, obsconditions, strategy, initial_mtl_file, initial_sky_file, initial_std_file,
fiberassign_script='fiberassign_legacy', legacy=None):
for i_pass in range(len(footprint_names)-1):
footprint_name = footprint_names[i_pass]
old_pass_name = pass_names[i_pass-1]
pass_name = pass_names[i_pass]
new_pass_name = pass_names[i_pass+1]
os.makedirs('{}/fiberassign_{}'.format(strategy, pass_name), exist_ok=True)
os.makedirs('{}/targets'.format(strategy), exist_ok=True)
os.makedirs('{}/zcat'.format(strategy), exist_ok=True)
assign_footprint_filename = 'footprint/subset_{}.fits'.format(footprint_name)
zcat_footprint_filename = 'footprint/subset_{}.fits'.format(pass_name)
fiberassign_dir = '{}/fiberassign_{}/'.format(strategy, pass_name)
mtl_filename = '{}/targets/{}_subset_dr8_mtl_dark_gray_NGC.fits'.format(strategy, pass_name)
new_mtl_filename = '{}/targets/{}_subset_dr8_mtl_dark_gray_NGC.fits'.format(strategy, new_pass_name)
old_zcat_filename = '{}/zcat/{}_zcat.fits'.format(strategy, old_pass_name)
zcat_filename = '{}/zcat/{}_zcat.fits'.format(strategy, pass_name)
if i_pass == 0:
shutil.copyfile(initial_mtl_file, mtl_filename)
# Run fiberassign
if legacy==True:
cmd = '{} --mtl {} --sky {} --std {}'.format(fiberassign_script, mtl_filename, initial_sky_file, initial_std_file)
cmd += ' --footprint {} --outdir {} --overwrite '.format(assign_footprint_filename, fiberassign_dir)
cmd += ' --fibstatusfile fiberstatus.ecsv --starmask 60129542144'
if legacy==False:
cmd = 'fiberassign --mtl {} --sky targets/subset_dr8_sky.fits '.format(mtl_filename)
cmd +=' --footprint {} --outdir {} --overwrite'.format(assign_footprint_filename, fiberassign_dir)
print(cmd)
os.system(cmd)
# Gather fiberassign files
fba_files = np.sort(glob.glob(os.path.join(fiberassign_dir,"fiberassign*.fits")))
# remove tilefiles that are not in the list of tiles to build zcat
footprint = Table.read(zcat_footprint_filename)
to_keep = []
for i_file, fba_file in enumerate(fba_files):
fibassign, header = fits.getdata(fba_file, header=True)
tileid = int(header['TILEID'])
if tileid in footprint['TILEID']:
print(tileid, 'in list', zcat_footprint_filename)
print('keeping {}'.format(fba_file))
to_keep.append(i_file)
else:
print('renaming {}'.format(fba_file))
fiberassign_file = fba_file.replace('fiberassign-', 'fba_')
renamed_file = fiberassign_file.replace('.fits', '_unused.fits')
print('renaming', fba_file, renamed_file)
os.rename(fba_file, renamed_file)
fba_files = fba_files[to_keep]
print('Files to keep', len(fba_files))
# Read targets and truth
targets = Table.read(mtl_filename)
truth = Table.read(initial_truth_file)
# Compute zcat
if i_pass==0:
zcat = desisim.quickcat.quickcat(fba_files, targets, truth, perfect=True)
else:
old_zcat = Table.read(old_zcat_filename)
zcat = desisim.quickcat.quickcat(fba_files, targets, truth, zcat=old_zcat, perfect=True)
zcat.write(zcat_filename, overwrite=True)
mtl = desitarget.mtl.make_mtl(targets, obsconditions[i_pass], zcat=zcat)
mtl.write(new_mtl_filename, overwrite=True)
from astropy.table import Table
import numpy as np
import matplotlib.pyplot as plt
import glob
import os
import fitsio
import desimodel.io
import desitarget.mtl
from desitarget.targetmask import desi_mask, obsconditions
data = fitsio.FITS('targets.fits', 'r')
target_data = data[1].read(columns=[
'TARGETID', 'DESI_TARGET', 'MWS_TARGET', 'BGS_TARGET', 'SUBPRIORITY',
'NUMOBS_INIT', 'PRIORITY_INIT', 'RA', 'DEC'
])
data.close()
mtl = desitarget.mtl.make_mtl(target_data)
mtl_file = 'mtl_targets.fits'
mtl.write(mtl_file, overwrite=True)