def simulate_frame(night, expid, camera, ccdshape=None, **kwargs):
"""
Simulate a single frame, including I/O
Args:
night: YEARMMDD string
expid: integer exposure ID
camera: b0, r1, .. z9
Options:
ccdshape = (npix_y, npix_x) primarily used to limit memory while testing
Additional keyword args are passed to pixsim.simulate()
Reads:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
Writes:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
$DESI_SPECTRO_SIM/$PIXPROD/{night}/desi-{expid}.fits
$DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
For a lower-level pixel simulation interface that doesn't perform I/O,
see pixsim.simulate()
"""
#- night, expid, camera -> input file names
simspecfile = io.findfile('simspec', night=night, expid=expid)
#- Read inputs
psf = desimodel.io.load_psf(camera[0])
simspec = io.read_simspec(simspecfile)
#- Trim effective CCD size; mainly to limit memory for testing
if ccdshape is not None:
psf.npix_y, psf.npix_x = ccdshape
if 'cosmics' in kwargs:
shape = (psf.npix_y, psf.npix_x)
kwargs['cosmics'] = io.read_cosmics(kwargs['cosmics'],
expid,
shape=shape)
image, rawpix, truepix = simulate(camera, simspec, psf, **kwargs)
#- Outputs; force "real" data files into simulation directory
simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
io.write_simpix(simpixfile, truepix, camera=camera, meta=image.meta)
simdir = io.simdir(night=night)
rawfile = desispec.io.findfile('desi', night=night, expid=expid)
rawfile = os.path.join(simdir, os.path.basename(rawfile))
desispec.io.write_raw(rawfile, rawpix, image.meta, camera=camera)
pixfile = desispec.io.findfile('pix',
night=night,
expid=expid,
camera=camera)
pixfile = os.path.join(simdir, os.path.basename(pixfile))
desispec.io.write_image(pixfile, image)
def simulate_exposure(simspecfile, rawfile, cameras=None,
ccdshape=None, simpixfile=None, addcosmics=None, comm=None,
**kwargs):
"""
Simulate frames from an exposure, including I/O
Args:
simspecfile: input simspec format file with spectra
rawfile: output raw data file to write
Options:
cameras: str or list of str, e.g. b0, r1, .. z9
ccdshape: (npix_y, npix_x) primarily used to limit memory while testing
simpixfile: output file for noiseless truth pixels
addcosmics: if True (must be specified via command input), add cosmics from real data
comm: MPI communicator object
Additional keyword args are passed to pixsim.simulate()
For a lower-level pixel simulation interface that doesn't perform I/O,
see pixsim.simulate()
Note: call desi_preproc or desispec.preproc.preproc to pre-process the
output desi*.fits file for overscan subtraction, noise estimation, etc.
"""
#- Split communicator by nodes; each node processes N frames
#- Assumes / requires equal number of ranks per node
if comm is not None:
rank, size = comm.rank, comm.size
num_nodes = mpi_count_nodes(comm)
comm_node, node_index, num_nodes = mpi_split_by_node(comm, 1)
node_rank = comm_node.rank
node_size = comm_node.size
else:
log.debug('Not using MPI')
rank, size = 0, 1
comm_node = None
node_index = 0
num_nodes = 1
node_rank = 0
node_size = 1
if rank == 0:
log.debug('Starting simulate_exposure at {}'.format(asctime()))
if cameras is None:
if rank == 0:
from astropy.io import fits
fibermap = fits.getdata(simspecfile, 'FIBERMAP')
cameras = io.fibers2cameras(fibermap['FIBER'])
log.debug('Found cameras {} in input simspec file'.format(cameras))
if len(cameras) % num_nodes != 0:
raise ValueError('Number of cameras {} should be evenly divisible by number of nodes {}'.format(
len(cameras), num_nodes))
if comm is not None:
cameras = comm.bcast(cameras, root=0)
#- Fail early if camera alreaady in output file
if rank == 0 and os.path.exists(rawfile):
from astropy.io import fits
err = False
fx = fits.open(rawfile)
for camera in cameras:
if camera in fx:
log.error('Camera {} already in {}'.format(camera, rawfile))
err = True
if err:
raise ValueError('Some cameras already in output file')
#- Read simspec input; I/O layer handles MPI broadcasting
if rank == 0:
log.debug('Reading simspec at {}'.format(asctime()))
mycameras = cameras[node_index::num_nodes]
if node_rank == 0:
log.info("Assigning cameras {} to comm_exp node {}".format(mycameras, node_index))
simspec = io.read_simspec(simspecfile, cameras=mycameras,
readflux=False, comm=comm)
night = simspec.header['NIGHT']
expid = simspec.header['EXPID']
if rank == 0:
log.debug('Reading PSFs at {}'.format(asctime()))
psfs = dict()
#need to initialize previous channel
previous_channel = 'a'
for camera in mycameras:
#- Note: current PSF object can't be pickled and thus every
#- rank must read it instead of rank 0 read + bcast
channel = camera[0]
if channel not in psfs:
log.info('Reading {} PSF at {}'.format(channel, asctime()))
psfs[channel] = desimodel.io.load_psf(channel)
#- Trim effective CCD size; mainly to limit memory for testing
if ccdshape is not None:
psfs[channel].npix_y, psfs[channel].npix_x = ccdshape
psf = psfs[channel]
cosmics=None
#avoid re-broadcasting cosmics if we can
if previous_channel != channel:
if (addcosmics is True) and (node_rank == 0):
cosmics_file = io.find_cosmics(camera, simspec.header['EXPTIME'])
log.info('Reading cosmics templates {} at {}'.format(
cosmics_file, asctime()))
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file, expid, shape=shape)
if (addcosmics is True) and (comm_node is not None):
if node_rank == 0:
log.info('Broadcasting cosmics at {}'.format(asctime()))
cosmics = comm_node.bcast(cosmics, root=0)
else:
log.debug("Cosmics not requested")
if node_rank == 0:
log.info("Starting simulate for camera {} on node {}".format(camera,node_index))
image, rawpix, truepix = simulate(camera, simspec, psf, comm=comm_node, preproc=False, cosmics=cosmics, **kwargs)
#- Use input communicator as barrier since multiple sub-communicators
#- will write to the same output file
if rank == 0:
log.debug('Writing outputs at {}'.format(asctime()))
tmprawfile = rawfile + '.tmp'
if comm is not None:
for i in range(comm.size):
if (i == comm.rank) and (comm_node.rank == 0):
desispec.io.write_raw(tmprawfile, rawpix, image.meta,
camera=camera)
if simpixfile is not None:
io.write_simpix(simpixfile, truepix, camera=camera,
meta=image.meta)
comm.barrier()
else:
desispec.io.write_raw(tmprawfile, rawpix, image.meta, camera=camera)
if simpixfile is not None:
io.write_simpix(simpixfile, truepix, camera=camera,
meta=image.meta)
if rank == 0:
log.info('Wrote {}'.format(rawfile))
log.debug('done at {}'.format(asctime()))
previous_channel = channel
#- All done; rename temporary raw file to final location
if comm is None or comm.rank == 0:
os.rename(tmprawfile, rawfile)
def simulate_exposure(simspecfile, rawfile, cameras=None,
ccdshape=None, simpixfile=None, addcosmics=None, comm=None,
**kwargs):
"""
Simulate frames from an exposure, including I/O
Args:
simspecfile: input simspec format file with spectra
rawfile: output raw data file to write
Options:
cameras: str or list of str, e.g. b0, r1, .. z9
ccdshape: (npix_y, npix_x) primarily used to limit memory while testing
simpixfile: output file for noiseless truth pixels
addcosmics: if True (must be specified via command input), add cosmics from real data
comm: MPI communicator object
Additional keyword args are passed to pixsim.simulate()
For a lower-level pixel simulation interface that doesn't perform I/O,
see pixsim.simulate()
Note: call desi_preproc or desispec.preproc.preproc to pre-process the
output desi*.fits file for overscan subtraction, noise estimation, etc.
"""
#- Split communicator by nodes; each node processes N frames
#- Assumes / requires equal number of ranks per node
if comm is not None:
rank, size = comm.rank, comm.size
num_nodes = mpi_count_nodes(comm)
comm_node, node_index, num_nodes = mpi_split_by_node(comm, 1)
node_rank = comm_node.rank
node_size = comm_node.size
else:
log.debug('Not using MPI')
rank, size = 0, 1
comm_node = None
node_index = 0
num_nodes = 1
node_rank = 0
node_size = 1
if rank == 0:
log.debug('Starting simulate_exposure at {}'.format(asctime()))
if cameras is None:
if rank == 0:
from astropy.io import fits
fibermap = fits.getdata(simspecfile, 'FIBERMAP')
cameras = io.fibers2cameras(fibermap['FIBER'])
log.debug('Found cameras {} in input simspec file'.format(cameras))
if len(cameras) % num_nodes != 0:
raise ValueError('Number of cameras {} should be evenly divisible by number of nodes {}'.format(
len(cameras), num_nodes))
if comm is not None:
cameras = comm.bcast(cameras, root=0)
#- Fail early if camera alreaady in output file
if rank == 0 and os.path.exists(rawfile):
from astropy.io import fits
err = False
fx = fits.open(rawfile)
for camera in cameras:
if camera in fx:
log.error('Camera {} already in {}'.format(camera, rawfile))
err = True
if err:
raise ValueError('Some cameras already in output file')
#- Read simspec input; I/O layer handles MPI broadcasting
if rank == 0:
log.debug('Reading simspec at {}'.format(asctime()))
mycameras = cameras[node_index::num_nodes]
if node_rank == 0:
log.info("Assigning cameras {} to comm_exp node {}".format(mycameras, node_index))
simspec = io.read_simspec(simspecfile, cameras=mycameras,
readflux=False, comm=comm)
night = simspec.header['NIGHT']
expid = simspec.header['EXPID']
if rank == 0:
log.debug('Reading PSFs at {}'.format(asctime()))
psfs = dict()
#need to initialize previous channel
previous_channel = 'a'
for camera in mycameras:
#- Note: current PSF object can't be pickled and thus every
#- rank must read it instead of rank 0 read + bcast
channel = camera[0]
if channel not in psfs:
psfs[channel] = desimodel.io.load_psf(channel)
#- Trim effective CCD size; mainly to limit memory for testing
if ccdshape is not None:
psfs[channel].npix_y, psfs[channel].npix_x = ccdshape
psf = psfs[channel]
cosmics=None
#avoid re-broadcasting cosmics if we can
if previous_channel != channel:
if (addcosmics is True) and (node_rank == 0):
cosmics_file = io.find_cosmics(camera, simspec.header['EXPTIME'])
log.info('cosmics templates {}'.format(cosmics_file))
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file, expid, shape=shape)
if (addcosmics is True) and (comm_node is not None):
cosmics = comm_node.bcast(cosmics, root=0)
else:
log.debug("Cosmics not requested")
if node_rank == 0:
log.info("Starting simulate for camera {} on node {}".format(camera,node_index))
image, rawpix, truepix = simulate(camera, simspec, psf, comm=comm_node, preproc=False, cosmics=cosmics, **kwargs)
#- Use input communicator as barrier since multiple sub-communicators
#- will write to the same output file
if rank == 0:
log.debug('Writing outputs at {}'.format(asctime()))
tmprawfile = rawfile + '.tmp'
if comm is not None:
for i in range(comm.size):
if (i == comm.rank) and (comm_node.rank == 0):
desispec.io.write_raw(tmprawfile, rawpix, image.meta,
camera=camera)
if simpixfile is not None:
io.write_simpix(simpixfile, truepix, camera=camera,
meta=image.meta)
comm.barrier()
else:
desispec.io.write_raw(tmprawfile, rawpix, image.meta, camera=camera)
if simpixfile is not None:
io.write_simpix(simpixfile, truepix, camera=camera,
meta=image.meta)
if rank == 0:
log.info('Wrote {}'.format(rawfile))
log.debug('done at {}'.format(asctime()))
previous_channel = channel
#- All done; rename temporary raw file to final location
if comm is None or comm.rank == 0:
os.rename(tmprawfile, rawfile)
def main(args=None):
log.info('Starting pixsim at {}'.format(asctime()))
if isinstance(args, (list, tuple, type(None))):
args = parse(args)
if args.verbose:
import logging
log.setLevel(logging.DEBUG)
if args.mpi:
from mpi4py import MPI
mpicomm = MPI.COMM_WORLD
log.debug('Using mpi4py MPI communicator')
else:
from desisim.util import _FakeMPIComm
log.debug('Using fake MPI communicator')
mpicomm = _FakeMPIComm()
log.info('Starting pixsim rank {} at {}'.format(mpicomm.rank, asctime()))
log.debug('MPI rank {} size {} / {} {}'.format(mpicomm.rank, mpicomm.size,
mpicomm.Get_rank(),
mpicomm.Get_size()))
#- Pre-flight check that these cameras haven't been done yet
if (mpicomm.rank == 0) and (not args.overwrite) and os.path.exists(
args.rawfile):
log.debug('Checking if cameras are already in output file')
from astropy.io import fits
fx = fits.open(args.rawfile)
oops = False
for camera in args.cameras:
if camera.upper() in fx:
log.error('Camera {} already in {}'.format(
camera, args.rawfile))
oops = True
fx.close()
if oops:
log.fatal('Exiting due to repeat cameras already in output file')
mpicomm.Abort()
ncameras = len(args.cameras)
if ncameras % mpicomm.size != 0:
log.fatal('Processing cameras {}'.format(args.cameras))
log.fatal(
'Number of cameras {} must be evenly divisible by MPI size {}'.
format(ncameras, mpicomm.size))
mpicomm.Abort()
#- Use original seed to generate different random seeds for each MPI rank
np.random.seed(args.seed)
while True:
seeds = np.random.randint(0, 2**32 - 1, size=mpicomm.size)
if np.unique(seeds).size == mpicomm.size:
random.seed(seeds[mpicomm.rank])
np.random.seed(seeds[mpicomm.rank])
break
if args.psf is not None:
from specter.psf import load_psf
psf = load_psf(args.psf)
simspec = io.read_simspec(args.simspec)
if args.fibermap:
fibermap = desispec.io.read_fibermap(args.fibermap)
fibers = fibermap['FIBER']
if args.nspec is not None:
fibers = fibers[0:args.nspec]
else:
fibers = None
if args.overwrite and os.path.exists(args.rawfile):
log.debug('removing {}'.format(args.rawfile))
os.remove(args.rawfile)
if args.overwrite and os.path.exists(args.simpixfile):
log.debug('removing {}'.format(args.simpixfile))
os.remove(args.simpixfile)
for i in range(mpicomm.rank, ncameras, mpicomm.size):
camera = args.cameras[i]
log.debug('Rank {} processing camera {}'.format(mpicomm.rank, camera))
channel = camera[0].lower()
assert channel in (
'b', 'r',
'z'), "Unknown camera {} doesn't start with b,r,z".format(camera)
#- Read inputs for this camera
if args.psf is None:
psf = desimodel.io.load_psf(channel)
if args.ccd_npix_x is not None:
psf.npix_x = args.ccd_npix_x
if args.ccd_npix_y is not None:
psf.npix_y = args.ccd_npix_y
if args.cosmics:
if args.cosmics_file is None:
cosmics_file = io.find_cosmics(camera,
simspec.header['EXPTIME'],
cosmics_dir=args.cosmics_dir)
log.info('cosmics templates {}'.format(cosmics_file))
else:
cosmics_file = args.cosmics_file
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file, args.expid, shape=shape)
else:
cosmics = None
#- Do the actual simulation
image, rawpix, truepix = desisim.pixsim.simulate(camera,
simspec,
psf,
fibers=fibers,
nspec=args.nspec,
ncpu=args.ncpu,
cosmics=cosmics,
wavemin=args.wavemin,
wavemax=args.wavemax,
preproc=False)
#- Synchronize with other MPI threads before continuing with output
mpicomm.Barrier()
#- Loop over MPI ranks, letting each one take its turn writing output
for rank in range(mpicomm.size):
if rank == mpicomm.rank:
desispec.io.write_raw(args.rawfile,
rawpix,
camera=camera,
header=image.meta,
primary_header=simspec.header)
log.info('Wrote {} image to {}'.format(camera, args.rawfile))
io.write_simpix(args.simpixfile,
truepix,
camera=camera,
meta=simspec.header)
log.info('Wrote {} image to {}'.format(camera,
args.simpixfile))
mpicomm.Barrier()
else:
mpicomm.Barrier()
#- Call preproc separately from pixsim.simulate to ensure that it is
#- done identically to real data.
#- Note: This does lose the advantage of parallel preprocessing; we could
#- change this if that is problematic
if args.preproc and mpicomm.rank == 0:
log.info('Preprocessing raw -> pix files')
from desispec.scripts import preproc
preproc_opts = ['--infile', args.rawfile, '--outdir', args.preproc_dir]
preproc_opts += ['--cameras', ','.join(args.cameras)]
preproc.main(preproc.parse(preproc_opts))
if mpicomm.rank == 0:
log.info('Finished pixsim {} expid {} at {}'.format(
args.night, args.expid, asctime()))
def main(args, comm=None):
if args.verbose:
import logging
log.setLevel(logging.DEBUG)
rank = 0
nproc = 1
if comm is not None:
rank = comm.rank
nproc = comm.size
if rank == 0:
log.info('Starting pixsim at {}'.format(asctime()))
#- Pre-flight check that these cameras haven't been done yet
if (rank == 0) and (not args.overwrite) and os.path.exists(args.rawfile):
log.debug('Checking if cameras are already in output file')
from astropy.io import fits
fx = fits.open(args.rawfile)
oops = False
for camera in args.cameras:
if camera.upper() in fx:
log.error('Camera {} already in {}'.format(camera,
args.rawfile))
oops = True
fx.close()
if oops:
log.fatal('Exiting due to repeat cameras already in output file')
if comm is not None:
comm.Abort()
else:
sys.exit(1)
ncamera = len(args.cameras)
comm_group = comm
comm_rank = None
group = 0
ngroup = 1
group_rank = 0
if comm is not None:
if args.mpi_camera > 1:
ngroup = int(comm.size / args.mpi_camera)
group = int(comm.rank / args.mpi_camera)
group_rank = comm.rank % args.mpi_camera
comm_group = comm.Split(color=group, key=group_rank)
comm_rank = comm.Split(color=group_rank, key=group)
else:
group = comm.rank
ngroup = comm.size
comm_group = MPI.COMM_SELF
comm_rank = comm
mycameras = np.array_split(np.arange(ncamera, dtype=np.int32),
ngroup)[group]
rawtemp = "{}.tmp".format(args.rawfile)
simpixtemp = "{}.tmp".format(args.simpixfile)
if rank == 0:
if args.overwrite and os.path.exists(args.rawfile):
log.debug('removing {}'.format(args.rawfile))
os.remove(args.rawfile)
if args.overwrite and os.path.exists(args.simpixfile):
log.debug('removing {}'.format(args.simpixfile))
os.remove(args.simpixfile)
# cleanup stale temp files
if os.path.isfile(rawtemp):
os.remove(rawtemp)
if comm is not None:
comm.barrier()
psf = None
if args.psf is not None:
from specter.psf import load_psf
psf = load_psf(args.psf)
simspec = None
if rank == 0:
simspec = io.read_simspec(args.simspec)
if comm is not None:
# Broadcast one array at a time, since this is a
# very large object.
flv = None
wv = None
pht = None
if simspec is not None:
flv = simspec.flavor
wv = simspec.wave
pht = simspec.phot
flv = comm.bcast(flv, root=0)
wv = comm.bcast(wv, root=0)
pht = comm.bcast(pht, root=0)
if simspec is None:
simspec = SimSpec(flv, wv, pht)
simspec.flux = comm.bcast(simspec.flux, root=0)
simspec.skyflux = comm.bcast(simspec.skyflux, root=0)
simspec.skyphot = comm.bcast(simspec.skyphot, root=0)
simspec.metadata = comm.bcast(simspec.metadata, root=0)
simspec.fibermap = comm.bcast(simspec.fibermap, root=0)
simspec.obs = comm.bcast(simspec.obs, root=0)
simspec.header = comm.bcast(simspec.header, root=0)
fibers = None
if args.fibermap:
if rank == 0:
fibermap = desispec.io.read_fibermap(args.fibermap)
fibers = fibermap['FIBER']
if args.nspec is not None:
fibers = fibers[0:args.nspec]
if comm is not None:
fibers = comm.bcast(fibers, root=0)
# Use original seed to generate different random seeds for each camera
np.random.seed(args.seed)
seeds = np.random.randint(0, 2**32-1, size=ncamera)
image = {}
rawpix = {}
truepix = {}
for c in mycameras:
camera = args.cameras[c]
if group_rank == 0:
log.debug('Processing camera {}'.format(camera))
channel = camera[0].lower()
# Set the seed for this camera (regardless of which process is
# performing the simulation).
np.random.seed(seeds[c])
# Get the random cosmic expids. The actual values will be
# remapped internally with the modulus operator.
cosexpid = np.random.randint(0, 100, size=1)[0]
# Read inputs for this camera. Unfortunately psf
# objects are not serializable, so we read it on all
# processes.
if args.psf is None:
psf = desimodel.io.load_psf(channel)
if args.ccd_npix_x is not None:
psf.npix_x = args.ccd_npix_x
if args.ccd_npix_y is not None:
psf.npix_y = args.ccd_npix_y
cosmics = None
if args.cosmics:
if group_rank == 0:
if args.cosmics_file is None:
cosmics_file = io.find_cosmics(camera,
simspec.header['EXPTIME'],
cosmics_dir=args.cosmics_dir)
log.info('cosmics templates {}'.format(cosmics_file))
else:
cosmics_file = args.cosmics_file
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file, cosexpid,
shape=shape)
if comm_group is not None:
cosmics = comm_group.bcast(cosmics, root=0)
#- Do the actual simulation
image[camera], rawpix[camera], truepix[camera] = \
desisim.pixsim.simulate(camera, simspec, psf, fibers=fibers,
nspec=args.nspec, ncpu=args.ncpu, cosmics=cosmics,
wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
comm=comm_group)
if args.psf is None:
del psf
# Wait for all processes to finish their cameras
if comm is not None:
comm.barrier()
# Write the cameras in order. Only the rank zero process in each
# group has the data.
for c in np.arange(ncamera, dtype=np.int32):
camera = args.cameras[c]
if c in mycameras:
if group_rank == 0:
desispec.io.write_raw(rawtemp, rawpix[camera],
camera=camera, header=image[camera].meta,
primary_header=simspec.header)
log.info('Wrote {} image to {}'.format(camera, args.rawfile))
io.write_simpix(simpixtemp, truepix[camera],
camera=camera, meta=simspec.header)
log.info('Wrote {} image to {}'.format(camera,
args.simpixfile))
if comm is not None:
comm.barrier()
# Move temp files into place
if rank == 0:
os.rename(simpixtemp, args.simpixfile)
os.rename(rawtemp, args.rawfile)
if comm is not None:
comm.barrier()
# Apply preprocessing
if args.preproc:
if rank == 0:
log.info('Preprocessing raw -> pix files')
from desispec.scripts import preproc
if len(mycameras) > 0:
if group_rank == 0:
for c in mycameras:
camera = args.cameras[c]
pixfile = desispec.io.findfile('pix', night=args.night,
expid=args.expid, camera=camera)
preproc_opts = ['--infile', args.rawfile, '--outdir',
args.preproc_dir, '--pixfile', pixfile]
preproc_opts += ['--cameras', camera]
preproc.main(preproc.parse(preproc_opts))
if comm is not None:
comm.barrier()
# Python is terrible with garbage collection, but at least
# encourage it...
del image
del rawpix
del truepix
if rank == 0:
log.info('Finished pixsim {} expid {} at {}'.format(args.night, args.expid, asctime()))
def simulate(night, expid, camera, nspec=None, verbose=False, ncpu=None,
trimxy=False, cosmics=None):
"""
Run pixel-level simulation of input spectra
Args:
night (string) : YEARMMDD
expid (integer) : exposure id
camera (str) : e.g. b0, r1, z9
nspec (int) : number of spectra to simulate
verbose (boolean) : if True, print status messages
ncpu (int) : number of CPU cores to use in parallel
trimxy (boolean) : trim image to just pixels with input signal
cosmics (str) : filename with dark images with cosmics to add
Reads:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
Writes:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
$DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
"""
if verbose:
print "Reading input files"
channel = camera[0].lower()
ispec = int(camera[1])
assert channel in 'brz'
assert 0 <= ispec < 10
#- Load DESI parameters
params = desimodel.io.load_desiparams()
nfibers = params['spectro']['nfibers']
#- Load simspec file
simfile = io.findfile('simspec', night=night, expid=expid)
simspec = io.read_simspec(simfile)
wave = simspec.wave[channel]
if simspec.skyphot is not None:
phot = simspec.phot[channel] + simspec.skyphot[channel]
else:
phot = simspec.phot[channel]
if ispec*nfibers >= simspec.nspec:
print "ERROR: camera {} not in the {} spectra in {}/{}".format(
camera, simspec.nspec, night, os.path.basename(simfile))
return
#- Load PSF
psf = desimodel.io.load_psf(channel)
#- Trim to just the spectra for this spectrograph
if nspec is None:
ii = slice(nfibers*ispec, nfibers*(ispec+1))
else:
ii = slice(nfibers*ispec, nfibers*ispec + nspec)
phot = phot[ii]
#- check if simulation has less than 500 input spectra
if phot.shape[0] < nspec:
nspec = phot.shape[0]
#- Project to image and append that to file
if verbose:
print "Projecting photons onto CCD"
img = parallel_project(psf, wave, phot, ncpu=ncpu)
if trimxy:
xmin, xmax, ymin, ymax = psf.xyrange((0,nspec), wave)
img = img[0:ymax, 0:xmax]
# img = img[ymin:ymax, xmin:xmax]
# hdr['CRVAL1'] = xmin+1
# hdr['CRVAL2'] = ymin+1
#- Prepare header
hdr = simspec.header
tmp = '/'.join(simfile.split('/')[-3:]) #- last 3 elements of path
hdr['SIMFILE'] = (tmp, 'Input simulation file')
#- Strip unnecessary keywords
for key in ('EXTNAME', 'LOGLAM', 'AIRORVAC', 'CRVAL1', 'CDELT1'):
if key in hdr:
del hdr[key]
#- Write noiseless output
simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
io.write_simpix(simpixfile, img, meta=hdr)
#- Add cosmics from library of dark images
#- in this case, don't add readnoise since the dark image already has it
if cosmics is not None:
cosmics = io.read_cosmics(cosmics, expid, shape=img.shape)
pix = np.random.poisson(img) + cosmics.pix
readnoise = cosmics.meta['RDNOISE']
ivar = 1.0/(pix.clip(0) + readnoise**2)
mask = cosmics.mask
#- Or just add noise
else:
params = desimodel.io.load_desiparams()
channel = camera[0].lower()
readnoise = params['ccd'][channel]['readnoise']
pix = np.random.poisson(img) + np.random.normal(scale=readnoise, size=img.shape)
ivar = 1.0/(pix.clip(0) + readnoise**2)
mask = np.zeros(img.shape, dtype=np.uint16)
#- Metadata to be included in pix file header is in the fibermap header
#- TODO: this is fragile; consider updating fibermap to use astropy Table
#- that includes the header rather than directly assuming FITS as the
#- underlying format.
fibermapfile = desispec.io.findfile('fibermap', night=night, expid=expid)
fm, fmhdr = desispec.io.read_fibermap(fibermapfile, header=True)
meta = dict()
try:
meta['TELRA'] = simspec.header['TELRA']
meta['TELDEC'] = simspec.header['TELDEC']
except KeyError: #- temporary backwards compatibilty
meta['TELRA'] = fmhdr['TELERA']
meta['TELDEC'] = fmhdr['TELEDEC']
meta['TILEID'] = simspec.header['TILEID']
meta['DATE-OBS'] = simspec.header['DATE-OBS']
meta['FLAVOR'] = simspec.header['FLAVOR']
meta['EXPTIME'] = simspec.header['EXPTIME']
meta['AIRMASS'] = simspec.header['AIRMASS']
image = Image(pix, ivar, mask, readnoise=readnoise, camera=camera, meta=meta)
pixfile = desispec.io.findfile('pix', night=night, camera=camera, expid=expid)
desispec.io.write_image(pixfile, image)
if verbose:
print "Wrote "+pixfile
def main(args, comm=None):
if args.verbose:
import logging
log.setLevel(logging.DEBUG)
#we do this so we can use operator.itemgetter
import operator
#we do this so our print statements can have timestamps
import time
rank = 0
nproc = 1
if comm is not None:
import mpi4py
rank = comm.rank
nproc = comm.size
else:
if args.ncpu == 0:
import multiprocessing as mp
args.ncpu = mp.cpu_count() // 2
if rank == 0:
log.info('Starting pixsim at {}'.format(asctime()))
#- Pre-flight check that these cameras haven't been done yet
if (rank == 0) and (not args.overwrite) and os.path.exists(args.rawfile):
log.debug('Checking if cameras are already in output file')
from astropy.io import fits
fx = fits.open(args.rawfile)
oops = False
for camera in args.cameras:
if camera.upper() in fx:
log.error('Camera {} already in {}'.format(
camera, args.rawfile))
oops = True
fx.close()
if oops:
log.fatal('Exiting due to repeat cameras already in output file')
if comm is not None:
comm.Abort()
else:
sys.exit(1)
#assign communicators, groups, etc, used for both mpi and multiprocessing
ncamera = len(args.cameras)
comm_group = comm
comm_rank = None
group = 0
ngroup = 1
group_rank = 0
if comm is not None:
if args.mpi_camera > 1:
ngroup = int(comm.size / args.mpi_camera)
group = int(comm.rank / args.mpi_camera)
group_rank = comm.rank % args.mpi_camera
comm_group = comm.Split(color=group, key=group_rank)
comm_rank = comm.Split(color=group_rank, key=group)
else:
group = comm.rank
ngroup = comm.size
comm_group = MPI.COMM_SELF
comm_rank = comm
#check to make sure that the number of frames is evenly divisible by nubmer of nodes
#if not, abort and provide a helpful error message
#otherwise due to barrier logic the program will hang
if comm is not None:
if comm.rank == 0:
if ncamera % ngroup != 0:
msg = 'Number of frames (ncamera) must be evenly divisible by number of nodes (N)'
log.error(msg)
raise ValueError(msg)
comm.Abort()
# Remove outputs and or temp files
rawtemp = "{}.tmp".format(args.rawfile)
simpixtemp = "{}.tmp".format(args.simpixfile)
if rank == 0:
if args.overwrite and os.path.exists(args.rawfile):
log.debug('removing {}'.format(args.rawfile))
os.remove(args.rawfile)
if args.overwrite and os.path.exists(args.simpixfile):
log.debug('removing {}'.format(args.simpixfile))
os.remove(args.simpixfile)
# cleanup stale temp files
if os.path.isfile(rawtemp):
os.remove(rawtemp)
if os.path.isfile(simpixtemp):
os.remove(simpixtemp)
if comm is not None:
comm.barrier()
#load psf for both mpi and multiprocessing
psf = None
if args.psf is not None:
from specter.psf import load_psf
psf = load_psf(args.psf)
# create list of tuples
camera_channel_args = []
camera_channel_list = []
if comm is not None:
#need to parse these together so we preserve the order
for item in args.cameras:
#0th element is camera (b,r,z)
#1st element is channel (0-9)
camera_channel_args.append((item[0], item[1]))
#divide cameras among nodes
#first sort so zs get distributed first
#reverse the order so z comes first
#sort in reverse order so we get z first
camera_channel_list = np.asarray(
sorted(camera_channel_args,
key=operator.itemgetter(0),
reverse=True))
#also handle multiprocessing case
if comm is None:
camera_channel_list = []
for item in args.cameras:
camera_channel_list.append((item[0], item[1]))
#if mpi and N>1, divide cameras between nodes
if comm is not None:
if ngroup > 1:
node_cameras = None
for i in range(ngroup):
if i == group:
node_cameras = camera_channel_list[i::ngroup]
if ngroup == 1:
node_cameras = camera_channel_list
#also handle multiprocessing case
if comm is None:
node_cameras = camera_channel_list
group_spectro = np.array([int(c[1]) for c in node_cameras], dtype=np.int32)
#preallocate things needed for multiprocessing (we handle mpi later)
if comm is None:
simspec = None
cosmics = None
image = {}
rawpix = {}
truepix = {}
lastcamera = None
# Read the fibermap
#this is cheap, do here both mpi non mpi
fibers = None
if rank == 0:
if args.fibermap is not None:
fibermap = desispec.io.read_fibermap(args.fibermap)
fibers = np.array(fibermap['FIBER'], dtype=np.int32)
else:
# Get the fiber list from the simspec file
from astropy.io import fits
from astropy.table import Table
fx = fits.open(args.simspec, memmap=True)
if 'FIBERMAP' in fx:
fibermap = Table(fx['FIBERMAP'].data)
fibers = np.array(fibermap['FIBER'], dtype=np.int32)
else:
# Get the number of fibers from one of the photon HDUs
fibers = np.arange(fx['PHOT_B'].header['NAXIS2'],
dtype=np.int32)
fx.close()
if args.nspec > 0:
fibers = fibers[0:args.nspec]
if comm is not None:
fibers = comm.bcast(fibers, root=0)
#do multiprocessing fibers now, handle mpi inside the loop
if comm is None:
fs = np.in1d(fibers // 500, group_spectro)
group_fibers = fibers[fs]
# Use original seed to generate different random seeds for each camera
np.random.seed(args.seed)
seeds = np.random.randint(0, 2**32 - 1, size=ncamera)
seed_counter = 0 #need to initialize counter
#this loop handles the mpi case
if comm is not None:
#now that each communicator (group) knows which cameras it
#is supposed to process, we can get started
previous_camera = 'a' #need to initialize
for i in range(len(node_cameras)): #may be different in each group
if i > 1:
#keep track in case we can avoid re-broadcasting stuff
previous_camera = node_cameras[i - 1]
current_camera = node_cameras[i]
camera = current_camera[0] + current_camera[1]
channel = current_camera[0]
#since we clear, we need to pre-allocate every time
simspec = None
image = {}
rawpix = {}
truepix = {}
#since we handle only one spectra at a time, just load the one we need
simspec = io.read_simspec_mpi(args.simspec,
comm_group,
channel,
spectrographs=group_spectro[i])
if group_rank == 0:
log.debug('Processing camera {}'.format(camera))
# Set the seed for this camera (regardless of which process is
# performing the simulation).
np.random.seed(seeds[(seed_counter)])
# Get the random cosmic expids. The actual values will be
# remapped internally with the modulus operator.
cosexpid = np.random.randint(0, 100, size=1)[0]
# Read inputs for this camera. Unfortunately psf
# objects are not serializable, so we read it on all
# processes.
if args.psf is None:
psf = desimodel.io.load_psf(channel)
if args.ccd_npix_x is not None:
psf.npix_x = args.ccd_npix_x
if args.ccd_npix_y is not None:
psf.npix_y = args.ccd_npix_y
#if we have already loaded the right cosmics camera, don't bcast again
cosmics = None
if args.cosmics:
if current_camera[0] != previous_camera[0]:
cosmics = None
if group_rank == 0:
if args.cosmics_file is None:
cosmics_file = io.find_cosmics(
camera,
simspec.header['EXPTIME'],
cosmics_dir=args.cosmics_dir)
log.info(
'cosmics templates {}'.format(cosmics_file))
else:
cosmics_file = args.cosmics_file
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file,
cosexpid,
shape=shape)
cosmics = comm_group.bcast(cosmics, root=0)
if group_rank == 0:
group_size = comm_group.size
log.info("Group {} ({} processes) simulating camera "
"{}".format(group, group_size, camera))
#calc group_fibers inside the loop
fs = np.in1d(fibers // 500, group_spectro[i])
group_fibers = fibers[fs]
image[camera], rawpix[camera], truepix[camera] = \
simulate(camera, simspec, psf, fibers=group_fibers,
ncpu=args.ncpu, nspec=args.nspec, cosmics=cosmics,
wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
comm=comm_group)
#to avoid overflowing the memory let's write after each instance of simulate
#the data are already at rank 0 in each communicator
#then have each communicator open and write to the file one at a time
#this is the part that will fail if the number of nodes divided by number of frames
#is not evenly divisible (hence the check above)
for i in range(ngroup): #number of total communicators
if group == i: #only one group at a time should open/write/close
if group_rank == 0: #write only from rank 0 where the data are
#write the raw file using desispec write_raw in raw.py
desispec.io.write_raw(rawtemp,
rawpix[camera],
camera=camera,
header=image[camera].meta,
primary_header=simspec.header)
log.info('Wrote {} image to {} at time {}'.format(
camera, args.rawfile, time.asctime()))
#write the simpix file using desisim write_simpix in io.py
io.write_simpix(simpixtemp,
truepix[camera],
camera=camera,
meta=simspec.header)
log.info('Wrote {} image to {} at time {}'.format(
camera, args.simpixfile, time.asctime()))
#delete for each group after we have written the output files
del rawpix, image, truepix, simspec
#to avoid bugs this barrier statement must align with if group==i!!
comm.Barrier() #all ranks should be done writing
if args.psf is None:
del psf
#iterate random number counter
seed_counter = seed_counter + 1
#initalize random number seeds
seed_counter = 0 #need to initialize counter
if comm is None:
simspec = io.read_simspec(args.simspec)
previous_camera = 'a'
for i in range(len(node_cameras)):
current_camera = node_cameras[i]
camera = current_camera[0] + current_camera[1]
channel = current_camera[0]
if group_rank == 0:
log.debug('Processing camera {}'.format(camera))
# Set the seed for this camera (regardless of which process is
# performing the simulation).
np.random.seed(seeds[(seed_counter)])
# Get the random cosmic expids. The actual values will be
# remapped internally with the modulus operator.
cosexpid = np.random.randint(0, 100, size=1)[0]
# Read inputs for this camera. Unfortunately psf
# objects are not serializable, so we read it on all
# processes.
if args.psf is None:
#add this so that it won't fail if we enter two repeated channels
if camera != previous_camera:
psf = desimodel.io.load_psf(channel)
if args.ccd_npix_x is not None:
psf.npix_x = args.ccd_npix_x
if args.ccd_npix_y is not None:
psf.npix_y = args.ccd_npix_y
previous_camera = camera
if args.cosmics:
if group_rank == 0:
if args.cosmics_file is None:
cosmics_file = io.find_cosmics(
camera,
simspec.header['EXPTIME'],
cosmics_dir=args.cosmics_dir)
log.info('cosmics templates {}'.format(cosmics_file))
else:
cosmics_file = args.cosmics_file
shape = (psf.npix_y, psf.npix_x)
cosmics = io.read_cosmics(cosmics_file,
cosexpid,
shape=shape)
#- Do the actual simulation
# Each process in the group is responsible for a subset of the
# fibers.
if group_rank == 0:
group_size = 1
log.info("Group {} ({} processes) simulating camera "
"{}".format(group, group_size, camera))
image[camera], rawpix[camera], truepix[camera] = \
simulate(camera, simspec, psf, fibers=group_fibers,
ncpu=args.ncpu, nspec=args.nspec, cosmics=cosmics,
wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
comm=comm_group)
#like we did in mpi, move the write in here so we don't have to
#keep accumulating all the data
if group_rank == 0:
desispec.io.write_raw(rawtemp,
rawpix[camera],
camera=camera,
header=image[camera].meta,
primary_header=simspec.header)
log.info('Wrote {} image to {}'.format(camera, args.rawfile))
io.write_simpix(simpixtemp,
truepix[camera],
camera=camera,
meta=simspec.header)
log.info('Wrote {} image to {}'.format(camera,
args.simpixfile))
if args.psf is None:
del psf
#iterate random number counter
seed_counter = seed_counter + 1
#done with both mpi and multiprocessing
# Move temp files into place
if rank == 0:
# Copy the original files into place if we are appending
if not args.overwrite and os.path.exists(args.rawfile):
shutil.copy2(args.rawfile, rawtemp)
if not args.overwrite and os.path.exists(args.simpixfile):
shutil.copy2(args.simpixfile, simpixtemp)
if rank == 0:
os.rename(simpixtemp, args.simpixfile)
os.rename(rawtemp, args.rawfile)
if comm is not None:
comm.barrier()
# Apply preprocessing
if args.preproc:
if rank == 0:
log.info('Preprocessing raw -> pix files')
from desispec.scripts import preproc
if len(node_cameras) > 0:
if group_rank == 0:
for i in range(len(node_cameras)):
current_camera = node_cameras[i]
camera = current_camera[0] + current_camera[1]
#file should be preproc instead of pix
pixfile = desispec.io.findfile('preproc',
night=args.night,
expid=args.expid,
camera=camera)
if args.preproc_dir:
pixfile = os.path.join(args.preproc_dir,
os.path.basename(pixfile))
pixdir = os.path.dirname(pixfile)
if not os.path.isdir(pixdir):
os.makedirs(pixdir)
preproc_opts = [
'--infile', args.rawfile, '--outfile', pixfile,
'--cameras', camera
]
preproc.main(preproc.parse(preproc_opts))
# Python is terrible with garbage collection, but at least
# encourage it...
if comm is None:
del image
del rawpix
del truepix
#have already deleted these for the mpi case
if rank == 0:
log.info('Finished pixsim {} expid {} at {}'.format(
args.night, args.expid, asctime()))
def simulate(night,
expid,
camera,
nspec=None,
verbose=False,
ncpu=None,
trimxy=False,
cosmics=None):
"""
Run pixel-level simulation of input spectra
Args:
night (string) : YEARMMDD
expid (integer) : exposure id
camera (str) : e.g. b0, r1, z9
nspec (int) : number of spectra to simulate
verbose (boolean) : if True, print status messages
ncpu (int) : number of CPU cores to use in parallel
trimxy (boolean) : trim image to just pixels with input signal
cosmics (str) : filename with dark images with cosmics to add
Reads:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
Writes:
$DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
$DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
"""
if verbose:
print "Reading input files"
channel = camera[0].lower()
ispec = int(camera[1])
assert channel in 'brz'
assert 0 <= ispec < 10
#- Load DESI parameters
params = desimodel.io.load_desiparams()
nfibers = params['spectro']['nfibers']
#- Load simspec file
simfile = io.findfile('simspec', night=night, expid=expid)
simspec = io.read_simspec(simfile)
wave = simspec.wave[channel]
if simspec.skyphot is not None:
phot = simspec.phot[channel] + simspec.skyphot[channel]
else:
phot = simspec.phot[channel]
if ispec * nfibers >= simspec.nspec:
print "ERROR: camera {} not in the {} spectra in {}/{}".format(
camera, simspec.nspec, night, os.path.basename(simfile))
return
#- Load PSF
psf = desimodel.io.load_psf(channel)
#- Trim to just the spectra for this spectrograph
if nspec is None:
ii = slice(nfibers * ispec, nfibers * (ispec + 1))
else:
ii = slice(nfibers * ispec, nfibers * ispec + nspec)
phot = phot[ii]
#- check if simulation has less than 500 input spectra
if phot.shape[0] < nspec:
nspec = phot.shape[0]
#- Project to image and append that to file
if verbose:
print "Projecting photons onto CCD"
img = parallel_project(psf, wave, phot, ncpu=ncpu)
if trimxy:
xmin, xmax, ymin, ymax = psf.xyrange((0, nspec), wave)
img = img[0:ymax, 0:xmax]
# img = img[ymin:ymax, xmin:xmax]
# hdr['CRVAL1'] = xmin+1
# hdr['CRVAL2'] = ymin+1
#- Prepare header
hdr = simspec.header
tmp = '/'.join(simfile.split('/')[-3:]) #- last 3 elements of path
hdr['SIMFILE'] = (tmp, 'Input simulation file')
#- Strip unnecessary keywords
for key in ('EXTNAME', 'LOGLAM', 'AIRORVAC', 'CRVAL1', 'CDELT1'):
if key in hdr:
del hdr[key]
#- Write noiseless output
simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
io.write_simpix(simpixfile, img, meta=hdr)
#- Add cosmics from library of dark images
#- in this case, don't add readnoise since the dark image already has it
if cosmics is not None:
cosmics = io.read_cosmics(cosmics, expid, shape=img.shape)
pix = np.random.poisson(img) + cosmics.pix
readnoise = cosmics.meta['RDNOISE']
ivar = 1.0 / (pix.clip(0) + readnoise**2)
mask = cosmics.mask
#- Or just add noise
else:
params = desimodel.io.load_desiparams()
channel = camera[0].lower()
readnoise = params['ccd'][channel]['readnoise']
pix = np.random.poisson(img) + np.random.normal(scale=readnoise,
size=img.shape)
ivar = 1.0 / (pix.clip(0) + readnoise**2)
mask = np.zeros(img.shape, dtype=np.uint16)
#- Metadata to be included in pix file header is in the fibermap header
#- TODO: this is fragile; consider updating fibermap to use astropy Table
#- that includes the header rather than directly assuming FITS as the
#- underlying format.
fibermapfile = desispec.io.findfile('fibermap', night=night, expid=expid)
fm, fmhdr = desispec.io.read_fibermap(fibermapfile, header=True)
meta = dict()
try:
meta['TELRA'] = simspec.header['TELRA']
meta['TELDEC'] = simspec.header['TELDEC']
except KeyError: #- temporary backwards compatibilty
meta['TELRA'] = fmhdr['TELERA']
meta['TELDEC'] = fmhdr['TELEDEC']
meta['TILEID'] = simspec.header['TILEID']
meta['DATE-OBS'] = simspec.header['DATE-OBS']
meta['FLAVOR'] = simspec.header['FLAVOR']
meta['EXPTIME'] = simspec.header['EXPTIME']
meta['AIRMASS'] = simspec.header['AIRMASS']
image = Image(pix,
ivar,
mask,
readnoise=readnoise,
camera=camera,
meta=meta)
pixfile = desispec.io.findfile('pix',
night=night,
camera=camera,
expid=expid)
desispec.io.write_image(pixfile, image)
if verbose:
print "Wrote " + pixfile