def test_scores(self):
#-
frame = self._get_frame()
scores, comments = compute_and_append_frame_scores(
frame, suffix="RAW", flux_per_angstrom=False)
scores, comments = compute_and_append_frame_scores(
frame, suffix="RAW", flux_per_angstrom=False)
scores, comments = compute_and_append_frame_scores(
frame, suffix="CALIB", flux_per_angstrom=True)
print(scores)
print(comments)
def main(args):
log = get_logger()
for filename in args.infile:
log.info("reading %s" % filename)
frame = io.read_frame(filename)
flux_per_angstrom = None
if args.flux_per_angstrom:
flux_per_angstrom = True
elif args.flux_per_pixel:
flux_per_angstrom = False
else:
flux_per_angstrom = None
scores, comments = compute_and_append_frame_scores(
frame,
suffix=args.suffix,
flux_per_angstrom=flux_per_angstrom,
overwrite=args.overwrite)
log.info("Adding or replacing SCORES extention with {} in {}".format(
scores.keys(), filename))
write_bintable(filename,
data=scores,
comments=comments,
extname="SCORES",
clobber=True)
def main(args) :
log = get_logger()
for filename in args.infile :
log.info("reading %s"%filename)
frame=io.read_frame(filename)
flux_per_angstrom=None
if args.flux_per_angstrom :
flux_per_angstrom=True
elif args.flux_per_pixel :
flux_per_angstrom=False
else :
flux_per_angstrom=None
scores,comments=compute_and_append_frame_scores(frame,suffix=args.suffix,flux_per_angstrom=flux_per_angstrom,overwrite=args.overwrite)
log.info("Adding or replacing SCORES extention with {} in {}".format(scores.keys(),filename))
write_bintable(filename,data=scores,comments=comments,extname="SCORES",clobber=True)
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 * 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.7
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=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
#- 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_mpi(args, comm=None, timing=None):
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 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)
mark_read_input = time.time()
# get spectral range
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 * 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.7
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.5))
psf_wavemax = np.min(psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y-0.5))
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()
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 {} 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=args.decorrelate_fibers,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True, nsubbundles=args.nsubbundles)
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)
#- 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
compute_and_append_frame_scores(frame,suffix="RAW")
mark_extraction = time.time()
#- 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)
log.info('extract: Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime()))
sys.stdout.flush()
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_mpi(args, comm=None, timing=None):
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 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)
mark_read_input = time.time()
# get spectral range
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 * psf.nspec + specmin
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)
#- 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))
# 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()
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 {} 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=args.decorrelate_fibers,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True, nsubbundles=args.nsubbundles)
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')
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)
#- 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
compute_and_append_frame_scores(frame,suffix="RAW")
mark_extraction = time.time()
#- 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)
log.info('extract: Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime()))
sys.stdout.flush()
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):
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 _extract_and_save(img, psf, bspecmin, bnspec, specmin, wave, raw_wave,
fibers, fibermap, outbundle, outmodel, bundlesize, args,
log):
'''
Performs the main extraction and saving of extracted frames found in the body of the
main loop. Refactored to be callable by both MPI and non-MPI versions of the code.
This should be viewed as a shorthand for the following commands.
'''
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin,
bnspec,
wave,
regularize=args.regularize,
ndecorr=args.decorrelate_fibers,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True,
nsubbundles=args.nsubbundles)
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 fibermap is not None:
bfibermap = fibermap[bspecmin - specmin:bspecmin + bnspec - specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin - specmin:bspecmin + bnspec - specmin]
#- Save the raw wavelength, not the corrected one (if corrected)
frame = Frame(raw_wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap,
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'] = 'electron/Angstrom'
#- Add scores to frame
if not args.no_scores:
compute_and_append_frame_scores(frame, suffix="RAW")
mark_extraction = time.time()
#- Write output
io.write_frame(outbundle, frame)
if args.model is not None:
fits.writeto(outmodel, results['modelimage'], header=frame.meta)
log.info('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(args.input), bspecmin, bspecmin + bnspec,
time.asctime()))
sys.stdout.flush()
return mark_extraction
def main(args=None):
log = get_logger()
if args is None:
args = parse()
log.info("reading frames ...")
frames = []
for filename in args.infile:
log.info("adding {}".format(filename))
frame = read_frame(filename)
if len(frames) == 0:
frames.append(frame)
continue
if "CAMERA" in frames[0].meta and (frame.meta["CAMERA"] !=
frames[0].meta["CAMERA"]):
log.error("ignore {} because not same camera".format(filename))
continue
if frame.wave.size != frames[0].wave.size or np.max(
np.abs(frame.wave - frames[0].wave)) > 0.01:
log.error("ignore {} because not same wave".format(filename))
continue
if frames[0].fibermap is not None:
if not np.all(frame.fibermap["TARGETID"] ==
frames[0].fibermap["TARGETID"]):
log.error(
"ignore {} because not same targets".format(filename))
continue
frames.append(frame)
ivar = frames[0].ivar * (frames[0].mask == 0)
ivarflux = ivar * frames[0].flux
ivarres = np.zeros(frames[0].resolution_data.shape)
ndiag = ivarres.shape[1]
for diag in range(ndiag):
ivarres[:, diag, :] = ivar * frames[0].resolution_data[:, diag, :]
mask = frames[0].mask
for frame in frames[1:]:
tmp_ivar = frame.ivar * (frame.mask == 0)
ivar += tmp_ivar
ivarflux += tmp_ivar * frame.flux
for diag in range(ndiag):
ivarres[:, diag, :] += tmp_ivar * frame.resolution_data[:, diag, :]
mask &= frame.mask # and mask
coadd = frames[0]
coadd.ivar = ivar
coadd.flux = ivarflux / (ivar + (ivar == 0))
coadd.mask = mask
for diag in range(ndiag):
coadd.resolution_data[:, diag, :] = ivarres[:, diag, :] / (ivar +
(ivar == 0))
if args.scores:
compute_and_append_frame_scores(coadd)
else:
coadd.scores = None
coadd.scores_comments = None
# remove info on chi2pix
coadd.chi2pix = None
log.info("writing {} ...".format(args.outfile))
write_frame(args.outfile, coadd)
log.info("done")
def main(args):
log = get_logger()
if (args.fiberflat is None) and (args.sky is None) and (args.calib is
None):
log.critical('no --fiberflat, --sky, or --calib; nothing to do ?!?')
sys.exit(12)
frame = read_frame(args.infile)
#- Raw scores already added in extraction, but just in case they weren't
#- it is harmless to rerun to make sure we have them.
compute_and_append_frame_scores(frame, suffix="RAW")
if args.cosmics_nsig > 0 and args.sky == None: # Reject cosmics (otherwise do it after sky subtraction)
log.info("cosmics ray 1D rejection")
reject_cosmic_rays_1d(frame, args.cosmics_nsig)
if args.fiberflat != None:
log.info("apply fiberflat")
# read fiberflat
fiberflat = read_fiberflat(args.fiberflat)
# apply fiberflat to all fibers
apply_fiberflat(frame, fiberflat)
compute_and_append_frame_scores(frame, suffix="FFLAT")
if args.sky != None:
# read sky
skymodel = read_sky(args.sky)
if args.cosmics_nsig > 0:
# first subtract sky without throughput correction
subtract_sky(frame, skymodel, throughput_correction=False)
# then find cosmics
log.info("cosmics ray 1D rejection after sky subtraction")
reject_cosmic_rays_1d(frame, args.cosmics_nsig)
if args.sky_throughput_correction:
# and (re-)subtract sky, but just the correction term
subtract_sky(frame,
skymodel,
throughput_correction=True,
default_throughput_correction=0.)
else:
# subtract sky
subtract_sky(frame,
skymodel,
throughput_correction=args.sky_throughput_correction)
compute_and_append_frame_scores(frame, suffix="SKYSUB")
if args.calib != None:
log.info("calibrate")
# read calibration
fluxcalib = read_flux_calibration(args.calib)
# apply calibration
apply_flux_calibration(frame, fluxcalib)
compute_and_append_frame_scores(frame, suffix="CALIB")
# save output
write_frame(args.outfile, frame, units='1e-17 erg/(s cm2 Angstrom)')
log.info("successfully wrote %s" % args.outfile)
def main(args):
log = get_logger()
if (args.fiberflat is None) and (args.sky is None) and (args.calib is None):
log.critical('no --fiberflat, --sky, or --calib; nothing to do ?!?')
sys.exit(12)
frame = read_frame(args.infile)
#- Raw scores already added in extraction, but just in case they weren't
#- it is harmless to rerun to make sure we have them.
compute_and_append_frame_scores(frame,suffix="RAW")
if args.cosmics_nsig>0 and args.sky==None : # Reject cosmics (otherwise do it after sky subtraction)
log.info("cosmics ray 1D rejection")
reject_cosmic_rays_1d(frame,args.cosmics_nsig)
if args.fiberflat!=None :
log.info("apply fiberflat")
# read fiberflat
fiberflat = read_fiberflat(args.fiberflat)
# apply fiberflat to all fibers
apply_fiberflat(frame, fiberflat)
compute_and_append_frame_scores(frame,suffix="FFLAT")
if args.sky!=None :
# read sky
skymodel=read_sky(args.sky)
if args.cosmics_nsig>0 :
# use a copy the frame (not elegant but robust)
copied_frame = copy.deepcopy(frame)
# first subtract sky without throughput correction
subtract_sky(copied_frame, skymodel, apply_throughput_correction = False)
# then find cosmics
log.info("cosmics ray 1D rejection after sky subtraction")
reject_cosmic_rays_1d(copied_frame,args.cosmics_nsig)
# copy mask
frame.mask = copied_frame.mask
# and (re-)subtract sky, but just the correction term
subtract_sky(frame, skymodel, apply_throughput_correction = (not args.no_sky_throughput_correction) )
else :
# subtract sky
subtract_sky(frame, skymodel, apply_throughput_correction = (not args.no_sky_throughput_correction) )
compute_and_append_frame_scores(frame,suffix="SKYSUB")
if args.calib!=None :
log.info("calibrate")
# read calibration
fluxcalib=read_flux_calibration(args.calib)
# apply calibration
apply_flux_calibration(frame, fluxcalib)
# Ensure that ivars are set to 0 for all values if any designated
# fibermask bit is set. Also flips a bits for each frame.mask value using specmask.BADFIBER
frame = get_fiberbitmasked_frame(frame,bitmask="flux",ivar_framemask=True)
compute_and_append_frame_scores(frame,suffix="CALIB")
# save output
write_frame(args.outfile, frame, units='10**-17 erg/(s cm2 Angstrom)')
log.info("successfully wrote %s"%args.outfile)
def main(args):
log = get_logger()
if (args.fiberflat is None) and (args.sky is None) and (args.calib is
None):
log.critical('no --fiberflat, --sky, or --calib; nothing to do ?!?')
sys.exit(12)
if (not args.no_tsnr) and (args.calib is None):
log.critical(
'need --fiberflat --sky and --calib to compute template SNR')
sys.exit(12)
frame = read_frame(args.infile)
if not args.no_tsnr:
# tsnr alpha calc. requires uncalibrated + no substraction rame.
uncalibrated_frame = copy.deepcopy(frame)
#- Raw scores already added in extraction, but just in case they weren't
#- it is harmless to rerun to make sure we have them.
compute_and_append_frame_scores(frame, suffix="RAW")
if args.cosmics_nsig > 0 and args.sky == None: # Reject cosmics (otherwise do it after sky subtraction)
log.info("cosmics ray 1D rejection")
reject_cosmic_rays_1d(frame, args.cosmics_nsig)
if args.fiberflat != None:
log.info("apply fiberflat")
# read fiberflat
fiberflat = read_fiberflat(args.fiberflat)
# apply fiberflat to all fibers
apply_fiberflat(frame, fiberflat)
compute_and_append_frame_scores(frame, suffix="FFLAT")
else:
fiberflat = None
if args.no_xtalk:
zero_ivar = (not args.no_zero_ivar)
else:
zero_ivar = False
if args.sky != None:
# read sky
skymodel = read_sky(args.sky)
if args.cosmics_nsig > 0:
# use a copy the frame (not elegant but robust)
copied_frame = copy.deepcopy(frame)
# first subtract sky without throughput correction
subtract_sky(copied_frame,
skymodel,
apply_throughput_correction=False,
zero_ivar=zero_ivar)
# then find cosmics
log.info("cosmics ray 1D rejection after sky subtraction")
reject_cosmic_rays_1d(copied_frame, args.cosmics_nsig)
# copy mask
frame.mask = copied_frame.mask
# and (re-)subtract sky, but just the correction term
subtract_sky(frame,
skymodel,
apply_throughput_correction=(
not args.no_sky_throughput_correction),
zero_ivar=zero_ivar)
else:
# subtract sky
subtract_sky(frame,
skymodel,
apply_throughput_correction=(
not args.no_sky_throughput_correction),
zero_ivar=zero_ivar)
compute_and_append_frame_scores(frame, suffix="SKYSUB")
if not args.no_xtalk:
log.info("fiber crosstalk correction")
correct_fiber_crosstalk(frame, fiberflat)
if not args.no_zero_ivar:
frame.ivar *= (frame.mask == 0)
if args.calib != None:
log.info("calibrate")
# read calibration
fluxcalib = read_flux_calibration(args.calib)
# apply calibration
apply_flux_calibration(frame, fluxcalib)
# Ensure that ivars are set to 0 for all values if any designated
# fibermask bit is set. Also flips a bits for each frame.mask value using specmask.BADFIBER
frame = get_fiberbitmasked_frame(
frame, bitmask="flux", ivar_framemask=(not args.no_zero_ivar))
compute_and_append_frame_scores(frame, suffix="CALIB")
if not args.no_tsnr:
log.info("calculating tsnr")
results, alpha = calc_tsnr2(uncalibrated_frame,
fiberflat=fiberflat,
skymodel=skymodel,
fluxcalib=fluxcalib,
alpha_only=args.alpha_only)
frame.meta['TSNRALPH'] = alpha
comments = {k: "from calc_frame_tsnr" for k in results.keys()}
append_frame_scores(frame, results, comments, overwrite=True)
# record inputs
frame.meta['IN_FRAME'] = shorten_filename(args.infile)
frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)
frame.meta['IN_SKY'] = shorten_filename(args.sky)
frame.meta['IN_CALIB'] = shorten_filename(args.calib)
# save output
write_frame(args.outfile, frame, units='10**-17 erg/(s cm2 Angstrom)')
log.info("successfully wrote %s" % args.outfile)