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))