def test_io(self):
np.random.seed(123456789)
test_dir = test_subdir_create("assign_test_io")
input_mtl = os.path.join(test_dir, "mtl.fits")
input_std = os.path.join(test_dir, "standards.fits")
input_sky = os.path.join(test_dir, "sky.fits")
input_suppsky = os.path.join(test_dir, "suppsky.fits")
tgoff = 0
nscience = sim_targets(input_mtl,
TARGET_TYPE_SCIENCE,
tgoff,
density=self.density_science)
tgoff += nscience
nstd = sim_targets(input_std,
TARGET_TYPE_STANDARD,
tgoff,
density=self.density_standards)
tgoff += nstd
nsky = sim_targets(input_sky,
TARGET_TYPE_SKY,
tgoff,
density=self.density_sky)
tgoff += nsky
nsuppsky = sim_targets(input_suppsky,
TARGET_TYPE_SUPPSKY,
tgoff,
density=self.density_suppsky)
tgs = Targets()
load_target_file(tgs, input_mtl)
load_target_file(tgs, input_std)
load_target_file(tgs, input_sky)
load_target_file(tgs, input_suppsky)
# Create a hierarchical triangle mesh lookup of the targets positions
tree = TargetTree(tgs, 0.01)
# Compute the targets available to each fiber for each tile.
fp, exclude, state = sim_focalplane(rundate=test_assign_date)
hw = load_hardware(focalplane=(fp, exclude, state))
tfile = os.path.join(test_dir, "footprint.fits")
sim_tiles(tfile)
tiles = load_tiles(tiles_file=tfile)
tgsavail = TargetsAvailable(hw, tgs, tiles, tree)
# Free the tree
del tree
# Compute the fibers on all tiles available for each target
favail = LocationsAvailable(tgsavail)
# Pass empty map of STUCK positioners that land on good sky
stucksky = {}
# First pass assignment
asgn = Assignment(tgs, tgsavail, favail, stucksky)
asgn.assign_unused(TARGET_TYPE_SCIENCE)
# Write out, merge, read back in and verify
write_assignment_ascii(tiles,
asgn,
out_dir=test_dir,
out_prefix="test_io_ascii_")
write_assignment_fits(tiles,
asgn,
out_dir=test_dir,
out_prefix="basic_",
all_targets=False)
write_assignment_fits(tiles,
asgn,
out_dir=test_dir,
out_prefix="full_",
all_targets=True)
plotpetals = [0]
# plotpetals = None
plot_tiles(hw,
tiles,
result_dir=test_dir,
result_prefix="basic_",
plot_dir=test_dir,
plot_prefix="basic_",
result_split_dir=False,
petals=plotpetals,
serial=True)
plot_tiles(hw,
tiles,
result_dir=test_dir,
result_prefix="full_",
plot_dir=test_dir,
plot_prefix="full_",
result_split_dir=False,
petals=plotpetals,
serial=True)
target_files = [input_mtl, input_sky, input_std]
tile_ids = list(tiles.id)
merge_results(target_files,
list(),
tile_ids,
result_dir=test_dir,
result_prefix="basic_",
out_dir=test_dir,
out_prefix="basic_tile-",
copy_fba=False)
merge_results(target_files,
list(),
tile_ids,
result_dir=test_dir,
result_prefix="full_",
out_dir=test_dir,
out_prefix="full_tile-",
copy_fba=False)
# Here we test reading with the standard reading function
for tid in tile_ids:
tdata = asgn.tile_location_target(tid)
avail = tgsavail.tile_data(tid)
# Check basic format
infile = os.path.join(test_dir,
"basic_tile-{:06d}.fits".format(tid))
inhead, fiber_data, targets_data, avail_data, gfa_targets = \
read_assignment_fits_tile((tid, infile))
for lid, tgid, tgra, tgdec in zip(fiber_data["LOCATION"],
fiber_data["TARGETID"],
fiber_data["TARGET_RA"],
fiber_data["TARGET_DEC"]):
if tgid >= 0:
self.assertEqual(tgid, tdata[lid])
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
# Check full format
infile = os.path.join(test_dir,
"full_tile-{:06d}.fits".format(tid))
inhead, fiber_data, targets_data, avail_data, gfa_targets = \
read_assignment_fits_tile((tid, infile))
for lid, tgid, tgra, tgdec in zip(fiber_data["LOCATION"],
fiber_data["TARGETID"],
fiber_data["TARGET_RA"],
fiber_data["TARGET_DEC"]):
if tgid >= 0:
self.assertEqual(tgid, tdata[lid])
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
# Now read the files directly with fitsio and verify against the input
# target data.
for tid in tile_ids:
tdata = asgn.tile_location_target(tid)
avail = tgsavail.tile_data(tid)
# Check basic format
infile = os.path.join(test_dir,
"basic_tile-{:06d}.fits".format(tid))
fdata = fitsio.FITS(infile, "r")
fassign = fdata["FIBERASSIGN"].read()
ftargets = fdata["TARGETS"].read()
for lid, tgid, tgra, tgdec, tgsub, tgprior, tgobs in zip(
fassign["LOCATION"], fassign["TARGETID"],
fassign["TARGET_RA"], fassign["TARGET_DEC"],
fassign["SUBPRIORITY"], fassign["PRIORITY"],
fassign["OBSCONDITIONS"]):
if tgid >= 0:
self.assertEqual(tgid, tdata[lid])
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
self.assertEqual(tgsub, props.subpriority)
self.assertEqual(tgprior, props.priority)
self.assertEqual(tgobs, props.obscond)
for tgid, tgra, tgdec, tgsub, tgprior, tgobs in zip(
ftargets["TARGETID"], ftargets["RA"], ftargets["DEC"],
ftargets["SUBPRIORITY"], ftargets["PRIORITY"],
ftargets["OBSCONDITIONS"]):
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
self.assertEqual(tgsub, props.subpriority)
self.assertEqual(tgprior, props.priority)
self.assertEqual(tgobs, props.obscond)
# Check full format
infile = os.path.join(test_dir,
"full_tile-{:06d}.fits".format(tid))
fdata = fitsio.FITS(infile, "r")
fassign = fdata["FIBERASSIGN"].read()
ftargets = fdata["TARGETS"].read()
for lid, tgid, tgra, tgdec, tgsub, tgprior, tgobs in zip(
fassign["LOCATION"], fassign["TARGETID"],
fassign["TARGET_RA"], fassign["TARGET_DEC"],
fassign["SUBPRIORITY"], fassign["PRIORITY"],
fassign["OBSCONDITIONS"]):
if tgid >= 0:
self.assertEqual(tgid, tdata[lid])
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
self.assertEqual(tgsub, props.subpriority)
self.assertEqual(tgprior, props.priority)
self.assertEqual(tgobs, props.obscond)
for tgid, tgra, tgdec, tgsub, tgprior, tgobs in zip(
ftargets["TARGETID"], ftargets["RA"], ftargets["DEC"],
ftargets["SUBPRIORITY"], ftargets["PRIORITY"],
ftargets["OBSCONDITIONS"]):
props = tgs.get(tgid)
self.assertEqual(tgra, props.ra)
self.assertEqual(tgdec, props.dec)
self.assertEqual(tgsub, props.subpriority)
self.assertEqual(tgprior, props.priority)
self.assertEqual(tgobs, props.obscond)
plot_tiles(hw,
tiles,
result_dir=test_dir,
result_prefix="basic_tile-",
plot_dir=test_dir,
plot_prefix="basic_tile-",
result_split_dir=False,
petals=plotpetals,
serial=True)
plot_tiles(hw,
tiles,
result_dir=test_dir,
result_prefix="full_tile-",
plot_dir=test_dir,
plot_prefix="full_tile-",
result_split_dir=False,
petals=plotpetals,
serial=True)
return
def main():
log = Logger.get()
mpi_procs = MPI.COMM_WORLD.size
mpi_rank = MPI.COMM_WORLD.rank
parser = argparse.ArgumentParser()
parser.add_argument(
"--survey_log",
type=str,
required=False,
help="Eventually we would pass in a file containing the log"
" of when each fiber assignment was run and for which tiles, "
"along with the options that were used.")
parser.add_argument(
"--sky",
type=str,
required=False,
help="Input file with sky or supp_sky targets. "
"These target files are assumed to be constant and not "
"tracked by the MTL ledger.")
parser.add_argument(
"--mtl",
type=str,
required=True,
help="The MTL ledger. This is still a work in progress and"
" I am not sure what the interface will be, but given the "
"fiber assignment dates in the survey log, we should be able"
" to get the MTL state at that time. For now, this option"
" is just one or more target files.")
parser.add_argument("--footprint",
type=str,
required=False,
default=None,
help="Optional FITS file defining the footprint. If"
" not specified, the default footprint from desimodel"
" is used.")
parser.add_argument("--tiles",
type=str,
required=False,
default=None,
help="Optional text file containing a subset of the"
" tile IDs to use in the footprint, one ID per line."
" Default uses all tiles in the footprint.")
parser.add_argument("--out",
type=str,
required=False,
default=None,
help="Output directory.")
parser.add_argument("--realizations",
type=int,
required=False,
default=10,
help="Number of realizations.")
args = parser.parse_args()
if args.sky is None:
args.sky = list()
# Set output directory
if args.out is None:
args.out = "."
# Read tiles we are using
tileselect = None
if args.tiles is not None:
tileselect = list()
with open(args.tiles, "r") as f:
for line in f:
# Try to convert the first column to an integer.
try:
tileselect.append(int(line.split()[0]))
except ValueError:
pass
tiles = load_tiles(
tiles_file=args.footprint,
select=tileselect,
)
# Create empty target list
tgs = Targets()
# Append each input target file. These target files must all be of the
# same survey type, and will set the Targets object to be of that survey.
print(args.mtl)
print(args.sky)
#for tgfile in args.targets:
# load_target_file(tgs, tgfile)
load_target_file(tgs, args.mtl)
# Just the science target IDs
tg_science = tgs.ids()
tg_science2indx = {y: x for x, y in enumerate(tg_science)}
# Now load the sky target files.
survey = tgs.survey()
#for tgfile in args.sky:
# load_target_file(tgs, tgfile)
load_target_file(tgs, args.sky)
# Divide up realizations among the processes.
n_realization = args.realizations
realizations = np.arange(n_realization, dtype=np.int32)
my_realizations = np.array_split(realizations, mpi_procs)[mpi_rank]
# Bitarray for all targets and realizations
#tgarray = bitarray(len(tg_science) * n_realization)
#tgarray.setall(False)
tgarray = np.zeros(len(tg_science) * n_realization, dtype='bool')
# Target tree
tree = TargetTree(tgs)
hw = load_hardware()
for realization in my_realizations:
# Set the seed based on the realization, so that the result is reproducible
# regardless of which process is working on the realization.
np.random.seed(realization)
# Comment out the next block to avoid randomizing subpriority
# ----
# Randomize science target subpriority for this realization
new_subpriority = np.random.random_sample(size=len(tg_science))
for indx, tgid in enumerate(tg_science):
tg = tgs.get(tgid)
tg.subpriority = new_subpriority[indx]
# Comment out the next block to avoid dithering tiles
# ----
# Dither tiles centers by the same
# Compute available targets / locations
tgsavail = TargetsAvailable(hw, tgs, tiles, tree)
favail = LocationsAvailable(tgsavail)
asgn = Assignment(tgs, tgsavail, favail)
# Replay the survey log for each time fiber assignment was run. For now, this
# is just doing the tiles all at once.
for assign_event in range(1):
# In the future, load MTL updates to the obs remaining for each target here
# Read hardware properties- in the future, pass in the assignment run date
# to this function.
hw = load_hardware()
# Run assignment for this event.
run(asgn)
# Update bit arrays for assigned science targets
for tile_id in tiles.id: #():
adata = asgn.tile_location_target(tile_id)
for loc, tgid in adata.items():
try:
idx = tg_science2indx[tgid]
tgarray[idx * n_realization + realization] = True
except KeyError:
# Not a science target
pass
# Reduce bitarrays to root process. The bitarray type conforms to the
# buffer protocol.
tgall = None
#if mpi_rank == 0:
# tgall = bitarray(tgarray)
# tgall.setall(False)
MPI.COMM_WORLD.Reduce(tgarray, tgall, op=MPI.BOR, root=0)
# Write it out
if mpi_rank == 0:
#pass
print(len(tgall))