def setUpClass(cls):
cls.psf = load_psf(resource_filename("specter.test", "t/psf-spot.fits"))
np.random.seed(0)
nspec = 10
wmin = min(cls.psf.wavelength(0, y=0), cls.psf.wavelength(nspec-1, y=0))
### ww = psf.wavelength(0, y=np.arange(10,60))
wmin, wmax = cls.psf.wavelength(0, y=(10,50))
ww = np.arange(wmin, wmax, 0.5)
nwave = len(ww)
phot_shape = (nspec, nwave)
phot = np.random.uniform(1, 1000, size=phot_shape)
image_orig = cls.psf.project(ww, phot, verbose=False)
var = 1.0 + image_orig
image = image_orig + np.random.normal(scale=np.sqrt(var))
cls.phot = phot
cls.image_orig = image_orig
cls.image = image
cls.ivar = 1.0 / var
cls.ww = ww
cls.nspec = nspec
# construct a symmetric test matrix
cls.dim = 100
cls.a1 = np.random.uniform(low=0.01, high=100.0, size=(cls.dim, cls.dim))
cls.a2 = np.random.uniform(low=0.01, high=100.0, size=(cls.dim, cls.dim))
cls.sym = np.dot(np.transpose(cls.a1), cls.a1)
cls.sym += np.dot(np.transpose(cls.a2), cls.a2)
def setUpClass(self):
np.random.seed(0)
psf = load_psf(test_data_dir() + "/psf-spot.fits")
nspec = 10
wmin = min(psf.wavelength(0, y=0), psf.wavelength(nspec-1, y=0))
### ww = psf.wavelength(0, y=np.arange(10,60))
wmin, wmax = psf.wavelength(0, y=(10,50))
ww = np.arange(wmin, wmax, 0.5)
nwave = len(ww)
phot_shape = (nspec, nwave)
phot = np.random.uniform(1, 1000, size=phot_shape)
image_orig = psf.project(ww, phot, verbose=False)
var = 1.0 + image_orig
image = image_orig + np.random.normal(scale=np.sqrt(var))
self.phot = phot
self.image_orig = image_orig
self.image = image
self.ivar = 1.0 / var
self.psf = psf
self.ww = ww
self.nspec = nspec
# construct a symmetric test matrix
self.dim = 100
self.a1 = np.random.uniform(low=0.01, high=100.0, size=(self.dim, self.dim))
self.a2 = np.random.uniform(low=0.01, high=100.0, size=(self.dim, self.dim))
self.sym = np.dot(np.transpose(self.a1), self.a1)
self.sym += np.dot(np.transpose(self.a2), self.a2)
def setUpClass(cls):
cls.psf = load_psf(resource_filename("specter.test",
"t/psf-spot.fits"))
np.random.seed(0)
nspec = 10
### wmin = min(cls.psf.wavelength(0, y=0), cls.psf.wavelength(nspec-1, y=0))
### ww = psf.wavelength(0, y=np.arange(10,60))
wmin, wmax = cls.psf.wavelength(0, y=(10, 90))
ww = np.arange(wmin, wmax, 1.0)
nwave = len(ww)
phot_shape = (nspec, nwave)
phot = np.random.uniform(1, 1000, size=phot_shape)
image_orig = cls.psf.project(ww, phot, verbose=False)
var = 1.0 + image_orig
image = image_orig + np.random.normal(scale=np.sqrt(var))
cls.phot = phot
cls.image_orig = image_orig
cls.image = image
cls.ivar = 1.0 / var
cls.ww = ww
cls.nspec = nspec
# construct a symmetric test matrix
cls.dim = 100
cls.a1 = np.random.uniform(low=0.01,
high=100.0,
size=(cls.dim, cls.dim))
cls.a2 = np.random.uniform(low=0.01,
high=100.0,
size=(cls.dim, cls.dim))
cls.sym = np.dot(np.transpose(cls.a1), cls.a1)
cls.sym += np.dot(np.transpose(cls.a2), cls.a2)
def test_psfbias(self):
psf = load_psf(resource_filename('specter.test', "t/psf-pix.fits"))
wmid = 0.5*(psf.wmin+psf.wmax)
ww = np.linspace(wmid-10, wmid+10)
phot = np.ones(len(ww))
phot[10] = 20
bias = util.psfbias(psf, psf, ww, phot)
absbias, R = util.psfabsbias(psf, psf, ww, phot)
def test_psfbias(self):
psf = load_psf(resource_filename('specter.test', "t/psf-pix.fits"))
wmid = 0.5 * (psf.wmin + psf.wmax)
ww = np.linspace(wmid - 10, wmid + 10)
phot = np.ones(len(ww))
phot[10] = 20
bias = util.psfbias(psf, psf, ww, phot)
absbias, R = util.psfabsbias(psf, psf, ww, phot)
def test_psfbias(self):
psf = load_psf(test_data_dir() + "/psf-pix.fits")
wmid = 0.5*(psf.wmin+psf.wmax)
ww = np.linspace(wmid-10, wmid+10)
phot = np.ones(len(ww))
phot[10] = 20
bias = util.psfbias(psf, psf, ww, phot)
absbias, R = util.psfabsbias(psf, psf, ww, phot)
def test_boxcar(self):
from lvmspec.boxcar import do_boxcar
psf = load_psf(self.psffile)
pix = np.random.normal(0, 3.0, size=(psf.npix_y, psf.npix_x))
ivar = np.ones_like(pix) / 3.0**2
mask = np.zeros(pix.shape, dtype=np.uint32)
img = lvmspec.image.Image(pix, ivar, mask, camera='z0')
outwave = np.arange(7500, 7600)
nwave = len(outwave)
nspec = 5
flux, ivar, resolution = do_boxcar(img,
psf,
outwave,
boxwidth=2.5,
nspec=nspec)
self.assertEqual(flux.shape, (nspec, nwave))
self.assertEqual(ivar.shape, (nspec, nwave))
self.assertEqual(resolution.shape[0], nspec)
# resolution.shape[1] is number of diagonals; picked by algorithm
self.assertEqual(resolution.shape[2], nwave)
def main_mpi(args, comm=None):
log = get_logger()
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(
psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError(
'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
format(wstart, psf_wavemin))
if psf_wavemax < wstop:
raise ValueError(
'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
format(wstop, psf_wavemax))
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(
np.floor_divide(np.arange(specmin, specmax),
bundlesize * np.ones(nspec)).astype(int))
bundles = sorted(checkbundles)
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b + 1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle // nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
(rank - leftover))
if rank == 0:
#- Print parameters
log.info("extract: input = {}".format(input_file))
log.info("extract: psf = {}".format(psf_file))
log.info("extract: specmin = {}".format(specmin))
log.info("extract: nspec = {}".format(nspec))
log.info("extract: wavelength = {},{},{}".format(wstart, wstop, dw))
log.info("extract: nwavestep = {}".format(args.nwavestep))
log.info("extract: regularize = {}".format(args.regularize))
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError(
"extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.normpath(os.path.dirname(outroot))
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle + mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
log.info('extract: Rank {} starting {} spectra {}:{} at {}'.format(
rank,
os.path.basename(input_file),
bspecmin[b],
bspecmin[b] + bnspec[b],
time.asctime(),
))
sys.stdout.flush()
#- The actual extraction
try:
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin[b],
bnspec[b],
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX
mask[chi2pix > 100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
if fibermap is not None:
bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] +
bnspec[b] - specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] -
specmin]
frame = Frame(wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(outbundle, frame, units='photon/bin')
if args.model is not None:
from astropy.io import fits
fits.writeto(outmodel,
results['modelimage'],
header=frame.meta)
log.info('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b],
bspecmin[b] + bnspec[b], time.asctime()))
sys.stdout.flush()
except:
# Log the error and increment the number of failures
log.error(
"extract: FAILED bundle {}, spectrum range {}:{}".format(
b, bspecmin[b], bspecmin[b] + bnspec[b]))
exc_type, exc_value, exc_traceback = sys.exc_info()
lines = traceback.format_exception(exc_type, exc_value,
exc_traceback)
log.error(''.join(lines))
failcount += 1
sys.stdout.flush()
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
mergeopts = ['--output', args.output, '--force', '--delete']
mergeopts.extend(
["{}_{:02d}.fits".format(outroot, b) for b in bundles])
mergeargs = mergebundles.parse(mergeopts)
mergebundles.main(mergeargs)
if args.model is not None:
model = None
for b in bundles:
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
if model is None:
model = fits.getdata(outmodel)
else:
#- TODO: test and warn if models overlap for pixels with
#- non-zero values
model += fits.getdata(outmodel)
os.remove(outmodel)
fits.writeto(args.model, model)
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), specmin, specmin + nspec,
time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(
psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError(
'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
format(wstart, psf_wavemin))
if psf_wavemax < wstop:
raise ValueError(
'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
format(wstop, psf_wavemax))
#- Print parameters
print("""\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file,
psf=psf_file,
output=args.output,
wstart=wstart,
wstop=wstop,
dw=dw,
specmin=specmin,
nspec=nspec,
regularize=args.regularize))
#- The actual extraction
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
specmin,
nspec,
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX
mask[chi2pix > 100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=fibers,
meta=img.meta,
fibermap=fibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(args.output, frame, units='photon/bin')
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model,
results['modelimage'],
header=frame.meta,
clobber=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin + nspec,
time.asctime()))
def main(args):
if args.mpi:
from mpi4py import MPI
comm = MPI.COMM_WORLD
return main_mpi(args, comm)
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
if args.fibermap_index is not None :
fibermin = args.fibermap_index
else :
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * 500 + specmin
print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
else:
try:
fibermap = io.read_fibermap(args.input)
except (AttributeError, IOError, KeyError):
fibermap = None
#- Trim fibermap to matching fiber range and create fibers array
if fibermap:
ii = np.in1d(fibermap['FIBER'], np.arange(fibermin, fibermin+nspec))
fibermap = fibermap[ii]
fibers = fibermap['FIBER']
else:
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.7
if args.heliocentric_correction :
heliocentric_correction_factor = heliocentric_correction_multiplicative_factor(img.meta)
wstart /= heliocentric_correction_factor
wstop /= heliocentric_correction_factor
dw /= heliocentric_correction_factor
else :
heliocentric_correction_factor = 1.
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y-1))
if psf_wavemin-5 > wstart:
raise ValueError('Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin))
if psf_wavemax+5 < wstop:
raise ValueError('Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax))
#- Print parameters
print("""\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file, psf=psf_file, output=args.output,
wstart=wstart, wstop=wstop, dw=dw,
specmin=specmin, nspec=nspec,
regularize=args.regularize))
#- The actual extraction
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave,
regularize=args.regularize, ndecorr=args.decorrelate_fibers,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True, nsubbundles=args.nsubbundles,psferr=args.psferr)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
if heliocentric_correction_factor != 1 :
#- Apply heliocentric correction factor to the wavelength
#- without touching the spectra, that is the whole point
wave *= heliocentric_correction_factor
wstart *= heliocentric_correction_factor
wstop *= heliocentric_correction_factor
dw *= heliocentric_correction_factor
img.meta['HELIOCOR'] = heliocentric_correction_factor
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=fibers, meta=img.meta, fibermap=fibermap,
chi2pix=chi2pix)
#- Add unit
# In specter.extract.ex2d one has flux /= dwave
# to convert the measured total number of electrons per
# wavelength node to an electron 'density'
frame.meta['BUNIT'] = 'count/Angstrom'
#- Add scores to frame
if not args.no_scores :
compute_and_append_frame_scores(frame,suffix="RAW")
#- Write output
io.write_frame(args.output, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model, results['modelimage'], header=frame.meta, overwrite=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.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
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)
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
group_cameras = np.array_split(np.arange(ncamera, dtype=np.int32), ngroup)[group]
# Compute which spectrographs our group needs to store based on which
# cameras we are processing.
group_spectro = np.unique(np.array([int(args.cameras[c][1]) for c in group_cameras], dtype=np.int32))
# 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()
psf = None
if args.psf is not None:
from specter.psf import load_psf
psf = load_psf(args.psf)
# Read and distribute the simspec data
simspec = None
if comm is not None:
simspec = io.read_simspec_mpi(args.simspec, comm, spectrographs=group_spectro)
else:
simspec = io.read_simspec(args.simspec)
# Read the fibermap
fibers = None
if rank == 0:
if args.fibermap is not None:
fibermap = lvmspec.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)
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)
image = {}
rawpix = {}
truepix = {}
lastchannel = None
for c in group_cameras:
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:
if channel != lastchannel:
psf = lvmmodel.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
lastchannel = channel
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
# Each process in the group is responsible for a subset of the
# fibers.
if group_rank == 0:
group_size = 1
if comm_group is not None:
group_size = comm_group.size
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)
if args.psf is None:
del psf
# Write the cameras in order. Only the rank zero process in each
# group has the data. If we are appending new cameras to an existing
# raw file, then copy the original to the temporary output first.
# 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)
# Wait for all processes to finish their cameras
if comm is not None:
comm.barrier()
for c in np.arange(ncamera, dtype=np.int32):
camera = args.cameras[c]
if c in group_cameras:
if group_rank == 0:
lvmspec.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')
if len(group_cameras) > 0:
if group_rank == 0:
for c in group_cameras:
camera = args.cameras[c]
pixfile = lvmspec.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 main_mpi(args, comm=None):
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin+nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax)
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(np.floor_divide(np.arange(specmin, specmax), bundlesize*np.ones(nspec)).astype(int))
bundles = sorted(list(checkbundles))
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b+1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle // nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover))
if rank == 0:
#- Print parameters
print "extract: input = {}".format(input_file)
print "extract: psf = {}".format(psf_file)
print "extract: specmin = {}".format(specmin)
print "extract: nspec = {}".format(nspec)
print "extract: wavelength = {},{},{}".format(wstart, wstop, dw)
print "extract: nwavestep = {}".format(args.nwavestep)
print "extract: regularize = {}".format(args.regularize)
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError("extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.normpath(os.path.dirname(outroot))
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle+mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
print('extract: Rank {} starting {} spectra {}:{} at {}'.format(
rank, os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime(),
) )
#- The actual extraction
try:
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, bspecmin[b],
bnspec[b], wave, regularize=args.regularize, ndecorr=True,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
if fibermap is not None:
bfibermap = fibermap[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin]
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=bfibers, meta=img.meta, fibermap=bfibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(outbundle, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(outmodel, results['modelimage'], header=frame.meta)
print('extract: Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime()))
except:
failcount += 1
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
mergeopts = [
'--output', args.output,
'--force',
'--delete'
]
mergeopts.extend([ "{}_{:02d}.fits".format(outroot, b) for b in bundles ])
mergeargs = mergebundles.parse(mergeopts)
mergebundles.main(mergeargs)
if args.model is not None:
model = None
for b in bundles:
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
if model is None:
model = fits.getdata(outmodel)
else:
#- TODO: test and warn if models overlap for pixels with
#- non-zero values
model += fits.getdata(outmodel)
os.remove(outmodel)
fits.writeto(args.model, model)
def setUpClass(cls):
cls.psf = load_psf(resource_filename("specter.test", "t/psf-monospot.fits"))
def run(self, *args, **kwargs):
if len(args) == 0:
raise qlexceptions.ParameterException("Missing input parameter")
if not self.is_compatible(type(args[0])):
raise qlexceptions.ParameterException(
"Incompatible input. Was expecting %s got %s" %
(type(self.__inpType__), type(args[0])))
if "PSFFile_sp" not in kwargs:
raise qlexceptions.ParameterException("Need PSF File")
from specter.psf import load_psf
input_image = args[0]
psffile = kwargs["PSFFile_sp"]
psf = load_psf(psffile)
if "Wavelength" not in kwargs:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
else:
wavelength = kwargs["Wavelength"]
if kwargs[
"Wavelength"] is not None: #- should be in wstart,wstop,dw format
wstart, wstop, dw = map(float, wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
if "Specmin" not in kwargs:
specmin = 0
else:
specmin = kwargs["Specmin"]
if kwargs["Specmin"] is None:
specmin = 0
if "Nspec" not in kwargs:
nspec = psf.nspec
else:
nspec = kwargs["Nspec"]
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = input_image.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * 500 + specmin
if "FiberMap" not in kwargs:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
else:
fibermap = kwargs["FiberMap"]
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
if "Regularize" in kwargs:
regularize = kwargs["Regularize"]
else:
regularize = False
if "ndecorr" in kwargs:
ndecorr = ndecorr
else:
ndecorr = True
bundlesize = 25 #- hard coded
if "Outfile" in kwargs:
outfile = kwargs["Outfile"]
else:
outfile = None
if "Nwavestep" in kwargs:
wavesize = kwargs["Nwavestep"]
else:
wavesize = 50
return self.run_pa(input_image,
psf,
specmin,
nspec,
wave,
regularize=regularize,
ndecorr=ndecorr,
bundlesize=bundlesize,
wavesize=wavesize,
outfile=outfile,
fibers=fibers,
fibermap=fibermap)
def run(self,*args,**kwargs):
if len(args) == 0 :
raise qlexceptions.ParameterException("Missing input parameter")
if not self.is_compatible(type(args[0])):
raise qlexceptions.ParameterException("Incompatible input. Was expecting %s got %s"%(type(self.__inpType__),type(args[0])))
if "PSFFile_sp" not in kwargs:
raise qlexceptions.ParameterException("Need PSF File")
from specter.psf import load_psf
input_image=args[0]
psffile=kwargs["PSFFile_sp"]
psf=load_psf(psffile)
if "Wavelength" not in kwargs:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
else:
wavelength=kwargs["Wavelength"]
if kwargs["Wavelength"] is not None: #- should be in wstart,wstop,dw format
wstart, wstop, dw = map(float, wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop+dw/2.0, dw)
if "Specmin" not in kwargs:
specmin=0
else:
specmin=kwargs["Specmin"]
if kwargs["Specmin"] is None:
specmin=0
if "Nspec" not in kwargs:
nspec = psf.nspec
else:
nspec=kwargs["Nspec"]
if nspec is None:
nspec=psf.nspec
specmax = specmin + nspec
camera = input_image.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph*500 + specmin
if "FiberMap" not in kwargs:
fibermap = None
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
else:
fibermap=kwargs["FiberMap"]
fibermap = fibermap[fibermin:fibermin+nspec]
fibers = fibermap['FIBER']
if "Regularize" in kwargs:
regularize=kwargs["Regularize"]
else:
regularize=False
if "ndecorr" in kwargs:
ndecorr=ndecorr
else:
ndecorr=True
bundlesize=25 #- hard coded
if "Outfile" in kwargs:
outfile=kwargs["Outfile"]
else:
outfile=None
if "Nwavestep" in kwargs:
wavesize=kwargs["Nwavestep"]
else:
wavesize=50
return self.run_pa(input_image,psf,specmin,nspec,wave,regularize=regularize,ndecorr=ndecorr, bundlesize=bundlesize, wavesize=wavesize,outfile=outfile,fibers=fibers,fibermap=fibermap)
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 main_mpi(args, comm=None):
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(
psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(
wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(
wstop, psf_wavemax)
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(
np.floor_divide(np.arange(specmin, specmax),
bundlesize * np.ones(nspec)).astype(int))
bundles = sorted(list(checkbundles))
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b + 1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle / nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
(rank - leftover))
if rank == 0:
#- Print parameters
print "extract: input = {}".format(input_file)
print "extract: psf = {}".format(psf_file)
print "extract: specmin = {}".format(specmin)
print "extract: nspec = {}".format(nspec)
print "extract: wavelength = {},{},{}".format(wstart, wstop, dw)
print "extract: nwavestep = {}".format(args.nwavestep)
print "extract: regularize = {}".format(args.regularize)
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError(
"extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.dirname(outroot)
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle + mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
print('extract: Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b],
time.asctime()))
#- The actual extraction
try:
flux, ivar, Rdata = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin[b],
bnspec[b],
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose)
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] +
bnspec[b] - specmin]
bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] -
specmin]
frame = Frame(wave,
flux,
ivar,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap)
#- Write output
io.write_frame(outbundle, frame)
print('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b],
bspecmin[b] + bnspec[b], time.asctime()))
except:
failcount += 1
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
opts = ['--output', args.output, '--force', '--delete']
opts.extend(["{}_{:02d}.fits".format(outroot, b) for b in bundles])
args = merge.parse(opts)
merge.main(args)
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), specmin, specmin + nspec,
time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(
psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(
wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(
wstop, psf_wavemax)
#- Print parameters
print """\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file,
psf=psf_file,
output=args.output,
wstart=wstart,
wstop=wstop,
dw=dw,
specmin=specmin,
nspec=nspec,
regularize=args.regularize)
#- The actual extraction
flux, ivar, Rdata = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
specmin,
nspec,
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose)
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave,
flux,
ivar,
resolution_data=Rdata,
fibers=fibers,
meta=img.meta,
fibermap=fibermap)
#- Write output
io.write_frame(args.output, frame)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin + nspec,
time.asctime()))
def setUp(self):
self.rawfile = 'test-raw-abcd.fits'
self.pixfile = 'test-pix-abcd.fits'
self.xwfile = 'test-xw-abcd.fits'
self.framefile = 'test-frame-abcd.fits'
self.fibermapfile = 'test-fibermap-abcd.fits'
self.skyfile = 'test-sky-abcd.fits'
self.qafile = 'test_qa.yaml'
self.qafig = 'test_qa.png'
#- use specter psf for this test
self.psffile = resource_filename('specter', 'test/t/psf-monospot.fits')
#self.psffile=os.environ['DESIMODEL']+'/data/specpsf/psf-b.fits'
self.config = {}
#- rawimage
hdr = dict()
hdr['CAMERA'] = 'z1'
hdr['DATE-OBS'] = '2018-09-23T08:17:03.988'
#- Dimensions per amp
ny = self.ny = 500
nx = self.nx = 400
noverscan = nover = 50
hdr['BIASSEC1'] = xy2hdr(np.s_[0:ny, nx:nx + nover])
hdr['DATASEC1'] = xy2hdr(np.s_[0:ny, 0:nx])
hdr['CCDSEC1'] = xy2hdr(np.s_[0:ny, 0:nx])
hdr['BIASSEC2'] = xy2hdr(np.s_[0:ny, nx + nover:nx + 2 * nover])
hdr['DATASEC2'] = xy2hdr(np.s_[0:ny,
nx + 2 * nover:nx + 2 * nover + nx])
hdr['CCDSEC2'] = xy2hdr(np.s_[0:ny, nx:nx + nx])
hdr['BIASSEC3'] = xy2hdr(np.s_[ny:ny + ny, nx:nx + nover])
hdr['DATASEC3'] = xy2hdr(np.s_[ny:ny + ny, 0:nx])
hdr['CCDSEC3'] = xy2hdr(np.s_[ny:ny + ny, 0:nx])
hdr['BIASSEC4'] = xy2hdr(np.s_[ny:ny + ny, nx + nover:nx + 2 * nover])
hdr['DATASEC4'] = xy2hdr(np.s_[ny:ny + ny,
nx + 2 * nover:nx + 2 * nover + nx])
hdr['CCDSEC4'] = xy2hdr(np.s_[ny:ny + ny, nx:nx + nx])
hdr['NIGHT'] = '20180923'
hdr['EXPID'] = 1
hdr['PROGRAM'] = 'dark'
hdr['FLAVOR'] = 'science'
hdr['EXPTIME'] = 100.0
rawimage = np.zeros((2 * ny, 2 * nx + 2 * noverscan))
offset = {'1': 100.0, '2': 100.5, '3': 50.3, '4': 200.4}
gain = {'1': 1.0, '2': 1.5, '3': 0.8, '4': 1.2}
rdnoise = {'1': 2.0, '2': 2.2, '3': 2.4, '4': 2.6}
quad = {
'1': np.s_[0:ny, 0:nx],
'2': np.s_[0:ny, nx:nx + nx],
'3': np.s_[ny:ny + ny, 0:nx],
'4': np.s_[ny:ny + ny, nx:nx + nx],
}
for amp in ('1', '2', '3', '4'):
hdr['GAIN' + amp] = gain[amp]
hdr['RDNOISE' + amp] = rdnoise[amp]
xy = _parse_sec_keyword(hdr['BIASSEC' + amp])
shape = [xy[0].stop - xy[0].start, xy[1].stop - xy[1].start]
rawimage[xy] += offset[amp]
rawimage[xy] += np.random.normal(scale=rdnoise[amp],
size=shape) / gain[amp]
xy = _parse_sec_keyword(hdr['DATASEC' + amp])
shape = [xy[0].stop - xy[0].start, xy[1].stop - xy[1].start]
rawimage[xy] += offset[amp]
rawimage[xy] += np.random.normal(scale=rdnoise[amp],
size=shape) / gain[amp]
#- set CCD parameters
self.ccdsec1 = hdr["CCDSEC1"]
self.ccdsec2 = hdr["CCDSEC2"]
self.ccdsec3 = hdr["CCDSEC3"]
self.ccdsec4 = hdr["CCDSEC4"]
#- raw data are integers, not floats
rawimg = rawimage.astype(np.int32)
self.expid = hdr["EXPID"]
self.camera = hdr["CAMERA"]
#- Confirm that all regions were correctly offset
assert not np.any(rawimage == 0.0)
#- write to the rawfile and read it in QA test
hdr['DOSVER'] = 'SIM'
hdr['FEEVER'] = 'SIM'
hdr['DETECTOR'] = 'SIM'
desispec.io.write_raw(self.rawfile, rawimg, hdr, camera=self.camera)
self.rawimage = fits.open(self.rawfile)
#- read psf, should use specter.PSF.load_psf instead of desispec.PSF(), otherwise need to create a psfboot somewhere.
psf = self.psf = load_psf(self.psffile)
#- make the test pixfile, fibermap file
img_pix = rawimg
img_ivar = np.ones_like(img_pix) / 3.0**2
img_mask = np.zeros(img_pix.shape, dtype=np.uint32)
img_mask[200] = 1
self.image = desispec.image.Image(img_pix,
img_ivar,
img_mask,
camera='z1',
meta=hdr)
desispec.io.write_image(self.pixfile, self.image)
self.fibermap = desispec.io.empty_fibermap(30)
self.fibermap['OBJTYPE'][::2] = 'ELG'
self.fibermap['OBJTYPE'][::3] = 'STD'
self.fibermap['OBJTYPE'][::5] = 'QSO'
self.fibermap['OBJTYPE'][::9] = 'LRG'
self.fibermap['OBJTYPE'][::7] = 'SKY'
#- add a filter and arbitrary magnitude
self.fibermap['MAG'][:29] = np.tile(np.random.uniform(18, 20, 29),
5).reshape(29,
5) #- Last fiber left
self.fibermap['FILTER'][:29] = np.tile(
['DECAM_R', '..', '..', '..', '..'], (29, 1)) #- last fiber left
desispec.io.write_fibermap(self.fibermapfile, self.fibermap)
#- make a test frame file
self.night = hdr['NIGHT']
self.nspec = nspec = 30
wave = np.arange(7600.0, 9800.0, 1.0) #- z channel
nwave = self.nwave = len(wave)
flux = np.random.uniform(size=(nspec, nwave)) + 100.
ivar = np.ones_like(flux)
resolution_data = np.ones((nspec, 13, nwave))
self.frame = desispec.frame.Frame(wave,
flux,
ivar,
resolution_data=resolution_data,
fibermap=self.fibermap)
self.frame.meta = dict(CAMERA=self.camera,
PROGRAM='dark',
FLAVOR='science',
NIGHT=self.night,
EXPID=self.expid,
EXPTIME=100,
CCDSEC1=self.ccdsec1,
CCDSEC2=self.ccdsec2,
CCDSEC3=self.ccdsec3,
CCDSEC4=self.ccdsec4)
desispec.io.write_frame(self.framefile, self.frame)
#- make a skymodel
sky = np.ones_like(self.frame.flux) * 0.5
skyivar = np.ones_like(sky)
self.mask = np.zeros(sky.shape, dtype=np.uint32)
self.skymodel = desispec.sky.SkyModel(wave, sky, skyivar, self.mask)
self.skyfile = desispec.io.write_sky(self.skyfile, self.skymodel)
#- Make a dummy boundary map for wavelength-flux in pixel space
self.map2pix = {}
self.map2pix["LEFT_MAX_FIBER"] = 14
self.map2pix["RIGHT_MIN_FIBER"] = 17
self.map2pix["BOTTOM_MAX_WAVE_INDEX"] = 900
self.map2pix["TOP_MIN_WAVE_INDEX"] = 1100
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 setUpClass(cls):
cls.psf = load_psf(resource_filename("specter.test", "t/psf-gausshermite2.fits"))
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 setUp(self):
#- use specter psf for this test
self.psffile=resource_filename('specter', 'test/t/psf-monospot.fits')
#self.psffile=os.environ['DESIMODEL']+'/data/specpsf/psf-b.fits'
self.config={"kwargs":{
"refKey":None,
"param":{},
"qso_resid":None
}}
#- rawimage
hdr = dict()
hdr['CAMERA'] = 'z1'
hdr['DATE-OBS'] = '2018-09-23T08:17:03.988'
hdr['PROGRAM'] = 'dark'
hdr['EXPTIME'] = 100
#- Dimensions per amp
ny = self.ny = 500
nx = self.nx = 400
noverscan = nover = 50
hdr['BIASSEC1'] = xy2hdr(np.s_[0:ny, nx:nx+nover])
hdr['DATASEC1'] = xy2hdr(np.s_[0:ny, 0:nx])
hdr['CCDSEC1'] = xy2hdr(np.s_[0:ny, 0:nx])
hdr['BIASSEC2'] = xy2hdr(np.s_[0:ny, nx+nover:nx+2*nover])
hdr['DATASEC2'] = xy2hdr(np.s_[0:ny, nx+2*nover:nx+2*nover+nx])
hdr['CCDSEC2'] = xy2hdr(np.s_[0:ny, nx:nx+nx])
hdr['BIASSEC3'] = xy2hdr(np.s_[ny:ny+ny, nx:nx+nover])
hdr['DATASEC3'] = xy2hdr(np.s_[ny:ny+ny, 0:nx])
hdr['CCDSEC3'] = xy2hdr(np.s_[ny:ny+ny, 0:nx])
hdr['BIASSEC4'] = xy2hdr(np.s_[ny:ny+ny, nx+nover:nx+2*nover])
hdr['DATASEC4'] = xy2hdr(np.s_[ny:ny+ny, nx+2*nover:nx+2*nover+nx])
hdr['CCDSEC4'] = xy2hdr(np.s_[ny:ny+ny, nx:nx+nx])
hdr['NIGHT'] = '20180923'
hdr['EXPID'] = 1
hdr['PROGRAM'] = 'dark'
hdr['FLAVOR'] = 'science'
hdr['EXPTIME'] = 100.0
rawimage = np.zeros((2*ny, 2*nx+2*noverscan))
offset = {'1':100.0, '2':100.5, '3':50.3, '4':200.4}
gain = {'1':1.0, '2':1.5, '3':0.8, '4':1.2}
rdnoise = {'1':2.0, '2':2.2, '3':2.4, '4':2.6}
obsrdn = {'1':3.4, '2':3.3, '3':3.6, '4':3.3}
quad = {
'1': np.s_[0:ny, 0:nx], '2': np.s_[0:ny, nx:nx+nx],
'3': np.s_[ny:ny+ny, 0:nx], '4': np.s_[ny:ny+ny, nx:nx+nx],
}
for amp in ('1', '2', '3', '4'):
hdr['GAIN'+amp] = gain[amp]
hdr['RDNOISE'+amp] = rdnoise[amp]
hdr['OBSRDN'+amp] = obsrdn[amp]
xy = _parse_sec_keyword(hdr['BIASSEC'+amp])
shape = [xy[0].stop-xy[0].start, xy[1].stop-xy[1].start]
rawimage[xy] += offset[amp]
rawimage[xy] += np.random.normal(scale=rdnoise[amp], size=shape)/gain[amp]
xy = _parse_sec_keyword(hdr['DATASEC'+amp])
shape = [xy[0].stop-xy[0].start, xy[1].stop-xy[1].start]
rawimage[xy] += offset[amp]
rawimage[xy] += np.random.normal(scale=rdnoise[amp], size=shape)/gain[amp]
#- set CCD parameters
self.ccdsec1=hdr["CCDSEC1"]
self.ccdsec2=hdr["CCDSEC2"]
self.ccdsec3=hdr["CCDSEC3"]
self.ccdsec4=hdr["CCDSEC4"]
#- raw data are integers, not floats
rawimg = rawimage.astype(np.int32)
self.expid=hdr["EXPID"]
self.camera=hdr["CAMERA"]
#- Confirm that all regions were correctly offset
assert not np.any(rawimage == 0.0)
#- write to the rawfile and read it in QA test
hdr['DOSVER'] = 'SIM'
hdr['FEEVER'] = 'SIM'
hdr['DETECTOR'] = 'SIM'
desispec.io.write_raw(self.rawfile,rawimg,hdr,camera=self.camera)
self.rawimage=fits.open(self.rawfile)
#- read psf, should use specter.PSF.load_psf instead of desispec.PSF(), otherwise need to create a psfboot somewhere.
self.psf = load_psf(self.psffile)
#- make the test pixfile, fibermap file
img_pix = rawimg
img_ivar = np.ones_like(img_pix) / 3.0**2
img_mask = np.zeros(img_pix.shape, dtype=np.uint32)
img_mask[200] = 1
self.image = desispec.image.Image(img_pix, img_ivar, img_mask, camera='z1',meta=hdr)
desispec.io.write_image(self.pixfile, self.image)
#- Create a fibermap with purposefully overlapping targeting bits
n = 30
self.fibermap = desispec.io.empty_fibermap(n)
self.fibermap['OBJTYPE'][:] = 'TGT'
self.fibermap['DESI_TARGET'][::2] |= desi_mask.ELG
self.fibermap['DESI_TARGET'][::5] |= desi_mask.QSO
self.fibermap['DESI_TARGET'][::7] |= desi_mask.LRG
#- add some arbitrary fluxes
for key in ['FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2']:
self.fibermap[key] = 10**((22.5 - np.random.uniform(18, 21, size=n))/2.5)
#- Make some standards; these still have OBJTYPE = 'TGT'
ii = [6,18,29]
self.fibermap['DESI_TARGET'][ii] = desi_mask.STD_FAINT
#- set some targets to SKY
ii = self.skyfibers = [5,10,21]
self.fibermap['OBJTYPE'][ii] = 'SKY'
self.fibermap['DESI_TARGET'][ii] = desi_mask.SKY
for key in ['FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2']:
self.fibermap[key][ii] = np.random.normal(scale=100, size=len(ii))
desispec.io.write_fibermap(self.fibermapfile, self.fibermap)
#- make a test frame file
self.night=hdr['NIGHT']
self.nspec = nspec = 30
wave=np.arange(7600.0,9800.0,1.0) #- z channel
nwave = self.nwave = len(wave)
flux=np.random.uniform(size=(nspec,nwave))+100.
ivar=np.ones_like(flux)
resolution_data=np.ones((nspec,13,nwave))
self.frame=desispec.frame.Frame(wave,flux,ivar,resolution_data=resolution_data,fibermap=self.fibermap)
self.frame.meta = hdr
self.frame.meta['WAVESTEP']=0.5
desispec.io.write_frame(self.framefile, self.frame)
#- make a skymodel
sky=np.ones_like(self.frame.flux)*0.5
skyivar=np.ones_like(sky)
self.mask=np.zeros(sky.shape,dtype=np.uint32)
self.skymodel=desispec.sky.SkyModel(wave,sky,skyivar,self.mask)
self.skyfile=desispec.io.write_sky(self.skyfile,self.skymodel)
#- Make a dummy boundary map for wavelength-flux in pixel space
self.map2pix={}
self.map2pix["LEFT_MAX_FIBER"] = 14
self.map2pix["RIGHT_MIN_FIBER"] = 17
self.map2pix["BOTTOM_MAX_WAVE_INDEX"] = 900
self.map2pix["TOP_MIN_WAVE_INDEX"] = 1100
def main_mpi(args, comm=None, timing=None):
freeze_iers()
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mark_start = time.time()
log = get_logger()
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if rank == 0:
img = io.read_image(input_file)
if comm is not None:
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
mark_read_input = time.time()
# get spectral range
if nspec is None:
nspec = psf.nspec
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
else:
try:
fibermap = io.read_fibermap(args.input)
except (AttributeError, IOError, KeyError):
fibermap = None
if fibermap is not None:
fibermap = fibermap[specmin:specmin + nspec]
if nspec > len(fibermap):
log.warning(
"nspec {} > len(fibermap) {}; reducing nspec to {}".format(
nspec, len(fibermap), len(fibermap)))
nspec = len(fibermap)
fibers = fibermap['FIBER']
else:
fibers = np.arange(specmin, specmin + nspec)
specmax = specmin + nspec
#- Get wavelength grid from options
if args.wavelength is not None:
raw_wstart, raw_wstop, raw_dw = [
float(tmp) for tmp in args.wavelength.split(',')
]
else:
raw_wstart = np.ceil(psf.wmin_all)
raw_wstop = np.floor(psf.wmax_all)
raw_dw = 0.7
raw_wave = np.arange(raw_wstart, raw_wstop + raw_dw / 2.0, raw_dw)
nwave = len(raw_wave)
bundlesize = args.bundlesize
if args.barycentric_correction:
if ('RA' in img.meta) or ('TARGTRA' in img.meta):
barycentric_correction_factor = \
barycentric_correction_multiplicative_factor(img.meta)
#- Early commissioning has RA/TARGTRA in fibermap but not HDU 0
elif fibermap is not None and \
(('RA' in fibermap.meta) or ('TARGTRA' in fibermap.meta)):
barycentric_correction_factor = \
barycentric_correction_multiplicative_factor(fibermap.meta)
else:
msg = 'Barycentric corr requires (TARGT)RA in HDU 0 or fibermap'
log.critical(msg)
raise KeyError(msg)
else:
barycentric_correction_factor = 1.
# Explictly define the correct wavelength values to avoid confusion of reference frame
# If correction applied, otherwise divide by 1 and use the same raw values
wstart = raw_wstart / barycentric_correction_factor
wstop = raw_wstop / barycentric_correction_factor
dw = raw_dw / barycentric_correction_factor
wave = raw_wave / barycentric_correction_factor
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=-0.5))
psf_wavemax = np.min(
psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 0.5))
if psf_wavemin - 5 > wstart:
raise ValueError(
'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
format(wstart, psf_wavemin))
if psf_wavemax + 5 < wstop:
raise ValueError(
'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
format(wstop, psf_wavemax))
if rank == 0:
#- Print parameters
log.info("extract: input = {}".format(input_file))
log.info("extract: psf = {}".format(psf_file))
log.info("extract: specmin = {}".format(specmin))
log.info("extract: nspec = {}".format(nspec))
log.info("extract: wavelength = {},{},{}".format(wstart, wstop, dw))
log.info("extract: nwavestep = {}".format(args.nwavestep))
log.info("extract: regularize = {}".format(args.regularize))
if barycentric_correction_factor != 1.:
img.meta['HELIOCOR'] = barycentric_correction_factor
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (raw_wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (raw_wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (raw_dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (io.shorten_filename(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = io.shorten_filename(input_file)
depend.add_dependencies(img.meta)
#- Check if input PSF was itself a traceshifted version of another PSF
orig_psf = None
if rank == 0:
try:
psfhdr = fits.getheader(psf_file, 'PSF')
orig_psf = psfhdr['IN_PSF']
except KeyError:
#- could happen due to PSF format not having "PSF" extension,
#- or due to PSF header not having 'IN_PSF' keyword. Either is OK
pass
if comm is not None:
orig_psf = comm.bcast(orig_psf, root=0)
if orig_psf is not None:
img.meta['ORIG_PSF'] = orig_psf
#- If not using MPI, use a single call to each of these and then end this function call
# Otherwise, continue on to splitting things up for the different ranks
if comm is None:
_extract_and_save(img, psf, specmin, nspec, specmin, wave, raw_wave,
fibers, fibermap, args.output, args.model,
bundlesize, args, log)
#- This is it if we aren't running MPI, so return
return
#else:
# # Continue to the MPI section, which could go under this else statment
# # But to save on indentation we'll just pass on to the rest of the function
# # since the alternative has already returned
# pass
# Now we divide our spectra into bundles
checkbundles = set()
checkbundles.update(
np.floor_divide(np.arange(specmin, specmax),
bundlesize * np.ones(nspec)).astype(int))
bundles = sorted(checkbundles)
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b + 1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
mynbundle = int(nbundle // nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
(rank - leftover))
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError(
"extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.normpath(os.path.dirname(outroot))
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
mark_preparation = time.time()
time_total_extraction = 0.0
time_total_write_output = 0.0
failcount = 0
for b in range(myfirstbundle, myfirstbundle + mynbundle):
mark_iteration_start = time.time()
outbundle = "{}_{:02d}.fits".format(outroot, b)
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
log.info('extract: Rank {} extracting {} spectra {}:{} at {}'.format(
rank,
os.path.basename(input_file),
bspecmin[b],
bspecmin[b] + bnspec[b],
time.asctime(),
))
sys.stdout.flush()
#- The actual extraction
try:
mark_extraction = _extract_and_save(img, psf, bspecmin[b],
bnspec[b], specmin, wave,
raw_wave, fibers, fibermap,
outbundle, outmodel,
bundlesize, args, log)
mark_write_output = time.time()
time_total_extraction += mark_extraction - mark_iteration_start
time_total_write_output += mark_write_output - mark_extraction
except:
# Log the error and increment the number of failures
log.error(
"extract: FAILED bundle {}, spectrum range {}:{}".format(
b, bspecmin[b], bspecmin[b] + bnspec[b]))
exc_type, exc_value, exc_traceback = sys.exc_info()
lines = traceback.format_exception(exc_type, exc_value,
exc_traceback)
log.error(''.join(lines))
failcount += 1
sys.stdout.flush()
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
time_merge = None
if rank == 0:
mark_merge_start = time.time()
mergeopts = ['--output', args.output, '--force', '--delete']
mergeopts.extend(
["{}_{:02d}.fits".format(outroot, b) for b in bundles])
mergeargs = mergebundles.parse(mergeopts)
mergebundles.main(mergeargs)
if args.model is not None:
model = None
for b in bundles:
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
if model is None:
model = fits.getdata(outmodel)
else:
#- TODO: test and warn if models overlap for pixels with
#- non-zero values
model += fits.getdata(outmodel)
os.remove(outmodel)
fits.writeto(args.model, model)
mark_merge_end = time.time()
time_merge = mark_merge_end - mark_merge_start
# Resolve difference timer data
if type(timing) is dict:
timing["read_input"] = mark_read_input - mark_start
timing["preparation"] = mark_preparation - mark_read_input
timing["total_extraction"] = time_total_extraction
timing["total_write_output"] = time_total_write_output
timing["merge"] = time_merge
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin+nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax)
#- Print parameters
print """\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file, psf=psf_file, output=args.output,
wstart=wstart, wstop=wstop, dw=dw,
specmin=specmin, nspec=nspec,
regularize=args.regularize)
#- The actual extraction
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave,
regularize=args.regularize, ndecorr=True,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=fibers, meta=img.meta, fibermap=fibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(args.output, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model, results['modelimage'], header=frame.meta, clobber=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
parser = argparse.ArgumentParser(usage="{prog} [options]")
# parser.add_argument("-i", "--input", type=str, help="input data")
parser.add_argument("-n",
"--nproc",
type=int,
help="number of parallel processes to use")
opts = parser.parse_args()
if opts.nproc is None:
opts.nproc = multiprocessing.cpu_count() // 2
print('Using nproc={}'.format(opts.nproc))
#- Load input data
t0 = time.time()
print('{} Reading input data'.format(time.asctime()))
psf = load_psf('psfmodel.fits')
image = fits.getdata('image.fits', 'IMAGE')
image_ivar = fits.getdata('image.fits', 'IVAR')
#- Split data into sub-regions to extract
t1 = time.time()
print('{} Splitting data into subregions'.format(time.asctime()))
wavelengths = np.arange(psf.wmin_all, psf.wmax_all)
nspec = 25
nwave = 200
args = list()
for specmin in range(0, 500, nspec):
for i in range(0, len(wavelengths), nwave):
ww = wavelengths[i:i + nwave]
xyrange = psf.xyrange((specmin, specmin + nspec), ww)
xmin, xmax, ymin, ymax = xyrange
