def boxcar_extraction_from_filenames(image_filename,
psf_filename,
fibers=None,
width=7):
"""
Fast boxcar extraction of spectra from a preprocessed image and a trace set
Args:
image_filename : input preprocessed fits filename
psf_filename : input PSF fits filename
Optional:
fibers : 1D np.array of int (default is all fibers, the first fiber is always = 0)
width : extraction boxcar width, default is 7
Returns:
flux : 2D np.array of shape (nfibers,n0=image.shape[0]), sum of pixel values per row of length=width per fiber
ivar : 2D np.array of shape (nfibers,n0), ivar[f,j] = 1/( sum_[j,b:e] (1/image.ivar) ), ivar=0 if at least 1 pixel in the row has image.ivar=0 or image.mask!=0
wave : 2D np.array of shape (nfibers,n0), determined from the traces
"""
tset = read_xytraceset(psf_filename)
image = read_image(image_filename)
qframe = qproc_boxcar_extraction(xytraceset,
image,
fibers=fibers,
width=width)
return qframe.flux, qframe.ivar, qframe.wave
def test_preproc_script(self):
io.write_raw(self.rawfile, self.rawimage, self.header, camera='b0')
io.write_raw(self.rawfile, self.rawimage, self.header, camera='b1')
args = ['--infile', self.rawfile, '--cameras', 'b1',
'--pixfile', self.pixfile]
if os.path.exists(self.pixfile):
os.remove(self.pixfile)
desispec.scripts.preproc.main(args)
img = io.read_image(self.pixfile)
self.assertEqual(img.pix.shape, (2*self.ny, 2*self.nx))
def test_preproc_script(self):
io.write_raw(self.rawfile, self.rawimage, self.header, primary_header = self.primary_header, camera='b0')
io.write_raw(self.rawfile, self.rawimage, self.header, primary_header = self.primary_header, camera='b1')
args = ['--infile', self.rawfile, '--cameras', 'b0',
'--outfile', self.pixfile]
if os.path.exists(self.pixfile):
os.remove(self.pixfile)
desispec.scripts.preproc.main(args)
img = io.read_image(self.pixfile)
self.assertEqual(img.pix.shape, (2*self.ny, 2*self.nx))
def main(args) :
log= get_logger()
log.info("degxx={} degxy={} degyx={} degyy={}".format(args.degxx,args.degxy,args.degyx,args.degyy))
# read preprocessed image
image=read_image(args.image)
log.info("read image {}".format(args.image))
if image.mask is not None :
image.ivar *= (image.mask==0)
tset = fit_trace_shifts(image=image,args=args)
if args.outpsf is not None :
write_traces_in_psf(args.psf,args.outpsf,tset)
log.info("wrote modified PSF in %s"%args.outpsf)
def boxcar_extraction_from_filenames(image_filename,psf_filename,fibers=None, width=7) :
"""
Fast boxcar extraction of spectra from a preprocessed image and a trace set
Args:
image_filename : input preprocessed fits filename
psf_filename : input PSF fits filename
Optional:
fibers : 1D np.array of int (default is all fibers, the first fiber is always = 0)
width : extraction boxcar width, default is 7
Returns:
flux : 2D np.array of shape (nfibers,n0=image.shape[0]), sum of pixel values per row of length=width per fiber
ivar : 2D np.array of shape (nfibers,n0), ivar[f,j] = 1/( sum_[j,b:e] (1/image.ivar) ), ivar=0 if at least 1 pixel in the row has image.ivar=0 or image.mask!=0
wave : 2D np.array of shape (nfibers,n0), determined from the traces
"""
tset = read_xytraceset(psf_filename)
image = read_image(image_filename)
qframe = qproc_boxcar_extraction(xytraceset,image,fibers=fibers,width=width)
return qframe.flux, qframe.ivar, qframe.wave
def main(args):
global psfs
log = get_logger()
log.info("starting")
# read preprocessed image
image = read_image(args.image)
log.info("read image {}".format(args.image))
if image.mask is not None:
image.ivar *= (image.mask == 0)
xcoef = None
ycoef = None
psf = None
wavemin = None
wavemax = None
nfibers = None
lines = None
psf, xcoef, ycoef, wavemin, wavemax = read_psf_and_traces(args.psf)
nfibers = xcoef.shape[0]
log.info(
"read PSF trace with xcoef.shape = {} , ycoef.shape = {} , and wavelength range {}:{}"
.format(xcoef.shape, ycoef.shape, int(wavemin), int(wavemax)))
if args.lines is not None:
log.info("We will fit the image using the psf model and lines")
# read lines
lines = np.loadtxt(args.lines, usecols=[0])
ok = (lines > wavemin) & (lines < wavemax)
log.info(
"read {} lines in {}, with {} of them in traces wavelength range".
format(len(lines), args.lines, np.sum(ok)))
lines = lines[ok]
else:
log.info(
"We will do an internal calibration of trace coordinates without using the psf shape in a first step"
)
internal_wavelength_calib = (not args.continuum)
external_wavelength_calib = args.sky | (args.spectrum is not None)
if args.auto:
log.debug("read flavor of input image {}".format(args.image))
hdus = pyfits.open(args.image)
if "FLAVOR" not in hdus[0].header:
log.error(
"no FLAVOR keyword in image header, cannot run with --auto option"
)
raise KeyError(
"no FLAVOR keyword in image header, cannot run with --auto option"
)
flavor = hdus[0].header["FLAVOR"].strip().lower()
hdus.close()
log.info("Input is a '{}' image".format(flavor))
if flavor == "flat":
internal_wavelength_calib = False
external_wavelength_calib = False
elif flavor == "arc":
internal_wavelength_calib = True
external_wavelength_calib = False
else:
internal_wavelength_calib = True
external_wavelength_calib = True
log.info("wavelength calib, internal={}, external={}".format(
internal_wavelength_calib, external_wavelength_calib))
spectrum_filename = args.spectrum
if external_wavelength_calib and spectrum_filename is None:
srch_file = "data/spec-sky.dat"
if not resource_exists('desispec', srch_file):
log.error("Cannot find sky spectrum file {:s}".format(srch_file))
raise RuntimeError(
"Cannot find sky spectrum file {:s}".format(srch_file))
else:
spectrum_filename = resource_filename('desispec', srch_file)
log.info("Use external calibration from cross-correlation with {}".
format(spectrum_filename))
if args.nfibers is not None:
nfibers = args.nfibers # FOR DEBUGGING
fibers = np.arange(nfibers)
if lines is not None:
# use a forward modeling of the image
# it's slower and works only for individual lines
# it's in principle more accurate
# but gives systematic residuals for complex spectra like the sky
x, y, dx, ex, dy, ey, fiber_xy, wave_xy = compute_dx_dy_using_psf(
psf, image, fibers, lines)
x_for_dx = x
y_for_dx = y
fiber_for_dx = fiber_xy
wave_for_dx = wave_xy
x_for_dy = x
y_for_dy = y
fiber_for_dy = fiber_xy
wave_for_dy = wave_xy
else:
# internal calibration method that does not use the psf
# nor a prior set of lines. this method is much faster
# measure x shifts
x_for_dx, y_for_dx, dx, ex, fiber_for_dx, wave_for_dx = compute_dx_from_cross_dispersion_profiles(
xcoef,
ycoef,
wavemin,
wavemax,
image=image,
fibers=fibers,
width=7,
deg=args.degxy)
if internal_wavelength_calib:
# measure y shifts
x_for_dy, y_for_dy, dy, ey, fiber_for_dy, wave_for_dy = compute_dy_using_boxcar_extraction(
xcoef,
ycoef,
wavemin,
wavemax,
image=image,
fibers=fibers,
width=7)
else:
# duplicate dx results with zero shift to avoid write special case code below
x_for_dy = x_for_dx.copy()
y_for_dy = y_for_dx.copy()
dy = np.zeros(dx.shape)
ey = 1.e-6 * np.ones(ex.shape)
fiber_for_dy = fiber_for_dx.copy()
wave_for_dy = wave_for_dx.copy()
degxx = args.degxx
degxy = args.degxy
degyx = args.degyx
degyy = args.degyy
while (True): # loop because polynomial degrees could be reduced
log.info(
"polynomial fit of measured offsets with degx=(%d,%d) degy=(%d,%d)"
% (degxx, degxy, degyx, degyy))
try:
dx_coeff, dx_coeff_covariance, dx_errorfloor, dx_mod, dx_mask = polynomial_fit(
z=dx, ez=ex, xx=x_for_dx, yy=y_for_dx, degx=degxx, degy=degxy)
dy_coeff, dy_coeff_covariance, dy_errorfloor, dy_mod, dy_mask = polynomial_fit(
z=dy, ez=ey, xx=x_for_dy, yy=y_for_dy, degx=degyx, degy=degyy)
log.info("dx dy error floor = %4.3f %4.3f pixels" %
(dx_errorfloor, dy_errorfloor))
log.info("check fit uncertainties are ok on edge of CCD")
merr = 0.
for fiber in [0, nfibers - 1]:
for rw in [-1, 1]:
tx = legval(rw, xcoef[fiber])
ty = legval(rw, ycoef[fiber])
m = monomials(tx, ty, degxx, degxy)
tdx = np.inner(dx_coeff, m)
tsx = np.sqrt(np.inner(m, dx_coeff_covariance.dot(m)))
m = monomials(tx, ty, degyx, degyy)
tdy = np.inner(dy_coeff, m)
tsy = np.sqrt(np.inner(m, dy_coeff_covariance.dot(m)))
merr = max(merr, tdx)
merr = max(merr, tdy)
#log.info("fiber=%d wave=%dA x=%d y=%d dx=%4.3f+-%4.3f dy=%4.3f+-%4.3f"%(fiber,int(wave),int(x),int(y),dx,sx,dy,sy))
log.info("max edge shift error = %4.3f pixels" % merr)
if degxx == 0 and degxy == 0 and degyx == 0 and degyy == 0:
break
except (LinAlgError, ValueError):
log.warning(
"polynomial fit failed with degx=(%d,%d) degy=(%d,%d)" %
(degxx, degxy, degyx, degyy))
if degxx == 0 and degxy == 0 and degyx == 0 and degyy == 0:
log.error(
"polynomial degrees are already 0. we can fit the offsets")
raise RuntimeError(
"polynomial degrees are already 0. we can fit the offsets")
merr = 100000. # this will lower the pol. degree.
if merr > args.max_error:
if merr != 100000.:
log.warning(
"max edge shift error = %4.3f pixels is too large, reducing degrees"
% merr)
if degxy > 0 and degyy > 0 and degxy > degxx and degyy > degyx: # first along wavelength
if degxy > 0: degxy -= 1
if degyy > 0: degyy -= 1
else: # then along fiber
if degxx > 0: degxx -= 1
if degyx > 0: degyx -= 1
else:
# error is ok, so we quit the loop
break
# write this for debugging
if args.outoffsets:
file = open(args.outoffsets, "w")
file.write(
"# axis wave fiber x y delta error polval (axis 0=y axis1=x)\n")
for e in range(dy.size):
file.write("0 %f %d %f %f %f %f %f\n" %
(wave_for_dy[e], fiber_for_dy[e], x_for_dy[e],
y_for_dy[e], dy[e], ey[e], dy_mod[e]))
for e in range(dx.size):
file.write("1 %f %d %f %f %f %f %f\n" %
(wave_for_dx[e], fiber_for_dx[e], x_for_dx[e],
y_for_dx[e], dx[e], ex[e], dx_mod[e]))
file.close()
log.info("wrote offsets in ASCII file %s" % args.outoffsets)
# print central shift
mx = np.median(x_for_dx)
my = np.median(y_for_dx)
m = monomials(mx, my, degxx, degxy)
mdx = np.inner(dx_coeff, m)
mex = np.sqrt(np.inner(m, dx_coeff_covariance.dot(m)))
mx = np.median(x_for_dy)
my = np.median(y_for_dy)
m = monomials(mx, my, degyx, degyy)
mdy = np.inner(dy_coeff, m)
mey = np.sqrt(np.inner(m, dy_coeff_covariance.dot(m)))
log.info("central shifts dx = %4.3f +- %4.3f dy = %4.3f +- %4.3f " %
(mdx, mex, mdy, mey))
# for each fiber, apply offsets and recompute legendre polynomial
log.info("for each fiber, apply offsets and recompute legendre polynomial")
xcoef, ycoef = recompute_legendre_coefficients(xcoef=xcoef,
ycoef=ycoef,
wavemin=wavemin,
wavemax=wavemax,
degxx=degxx,
degxy=degxy,
degyx=degyx,
degyy=degyy,
dx_coeff=dx_coeff,
dy_coeff=dy_coeff)
# use an input spectrum as an external calibration of wavelength
if spectrum_filename:
log.info(
"write and reread PSF to be sure predetermined shifts were propagated"
)
write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin,
wavemax)
psf, xcoef, ycoef, wavemin, wavemax = read_psf_and_traces(args.outpsf)
ycoef = shift_ycoef_using_external_spectrum(
psf=psf,
xcoef=xcoef,
ycoef=ycoef,
wavemin=wavemin,
wavemax=wavemax,
image=image,
fibers=fibers,
spectrum_filename=spectrum_filename,
degyy=args.degyy,
width=7)
write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin,
wavemax)
log.info("wrote modified PSF in %s" % args.outpsf)
else:
write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin,
wavemax)
log.info("wrote modified PSF in %s" % args.outpsf)
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_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) :
global psfs
log=get_logger()
log.info("starting")
# read preprocessed image
image=read_image(args.image)
log.info("read image {}".format(args.image))
if image.mask is not None :
image.ivar *= (image.mask==0)
xytraceset = read_xytraceset(args.psf)
wavemin = xytraceset.wavemin
wavemax = xytraceset.wavemax
xcoef = xytraceset.x_vs_wave_traceset._coeff
ycoef = xytraceset.y_vs_wave_traceset._coeff
nfibers=xcoef.shape[0]
log.info("read PSF trace with xcoef.shape = {} , ycoef.shape = {} , and wavelength range {}:{}".format(xcoef.shape,ycoef.shape,int(wavemin),int(wavemax)))
lines=None
if args.lines is not None :
log.info("We will fit the image using the psf model and lines")
# read lines
lines=np.loadtxt(args.lines,usecols=[0])
ok=(lines>wavemin)&(lines<wavemax)
log.info("read {} lines in {}, with {} of them in traces wavelength range".format(len(lines),args.lines,np.sum(ok)))
lines=lines[ok]
else :
log.info("We will do an internal calibration of trace coordinates without using the psf shape in a first step")
internal_wavelength_calib = (not args.continuum)
external_wavelength_calib = args.sky | ( args.spectrum is not None )
if args.auto :
log.debug("read flavor of input image {}".format(args.image))
hdus = pyfits.open(args.image)
if "FLAVOR" not in hdus[0].header :
log.error("no FLAVOR keyword in image header, cannot run with --auto option")
raise KeyError("no FLAVOR keyword in image header, cannot run with --auto option")
flavor = hdus[0].header["FLAVOR"].strip().lower()
hdus.close()
log.info("Input is a '{}' image".format(flavor))
if flavor == "flat" :
internal_wavelength_calib = False
external_wavelength_calib = False
elif flavor == "arc" :
internal_wavelength_calib = True
external_wavelength_calib = False
else :
internal_wavelength_calib = True
external_wavelength_calib = True
log.info("wavelength calib, internal={}, external={}".format(internal_wavelength_calib,external_wavelength_calib))
spectrum_filename = args.spectrum
if external_wavelength_calib and spectrum_filename is None :
srch_file = "data/spec-sky.dat"
if not resource_exists('desispec', srch_file):
log.error("Cannot find sky spectrum file {:s}".format(srch_file))
raise RuntimeError("Cannot find sky spectrum file {:s}".format(srch_file))
else :
spectrum_filename=resource_filename('desispec', srch_file)
log.info("Use external calibration from cross-correlation with {}".format(spectrum_filename))
if args.nfibers is not None :
nfibers = args.nfibers # FOR DEBUGGING
fibers=np.arange(nfibers)
if lines is not None :
# use a forward modeling of the image
# it's slower and works only for individual lines
# it's in principle more accurate
# but gives systematic residuals for complex spectra like the sky
psf = read_specter_psf(args.psf)
x,y,dx,ex,dy,ey,fiber_xy,wave_xy=compute_dx_dy_using_psf(psf,image,fibers,lines)
x_for_dx=x
y_for_dx=y
fiber_for_dx=fiber_xy
wave_for_dx=wave_xy
x_for_dy=x
y_for_dy=y
fiber_for_dy=fiber_xy
wave_for_dy=wave_xy
else :
# internal calibration method that does not use the psf
# nor a prior set of lines. this method is much faster
# measure x shifts
x_for_dx,y_for_dx,dx,ex,fiber_for_dx,wave_for_dx = compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image=image, fibers=fibers, width=args.width, deg=args.degxy,image_rebin=args.ccd_rows_rebin)
if internal_wavelength_calib :
# measure y shifts
x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(xytraceset, image=image, fibers=fibers, width=args.width)
mdy = np.median(dy)
log.info("Subtract median(dy)={}".format(mdy))
dy -= mdy # remove median, because this is an internal calibration
else :
# duplicate dx results with zero shift to avoid write special case code below
x_for_dy = x_for_dx.copy()
y_for_dy = y_for_dx.copy()
dy = np.zeros(dx.shape)
ey = 1.e-6*np.ones(ex.shape)
fiber_for_dy = fiber_for_dx.copy()
wave_for_dy = wave_for_dx.copy()
degxx=args.degxx
degxy=args.degxy
degyx=args.degyx
degyy=args.degyy
while(True) : # loop because polynomial degrees could be reduced
log.info("polynomial fit of measured offsets with degx=(%d,%d) degy=(%d,%d)"%(degxx,degxy,degyx,degyy))
try :
dx_coeff,dx_coeff_covariance,dx_errorfloor,dx_mod,dx_mask=polynomial_fit(z=dx,ez=ex,xx=x_for_dx,yy=y_for_dx,degx=degxx,degy=degxy)
dy_coeff,dy_coeff_covariance,dy_errorfloor,dy_mod,dy_mask=polynomial_fit(z=dy,ez=ey,xx=x_for_dy,yy=y_for_dy,degx=degyx,degy=degyy)
log.info("dx dy error floor = %4.3f %4.3f pixels"%(dx_errorfloor,dy_errorfloor))
log.info("check fit uncertainties are ok on edge of CCD")
merr=0.
for fiber in [0,nfibers-1] :
for rw in [-1,1] :
tx = legval(rw,xcoef[fiber])
ty = legval(rw,ycoef[fiber])
m=monomials(tx,ty,degxx,degxy)
tdx=np.inner(dx_coeff,m)
tsx=np.sqrt(np.inner(m,dx_coeff_covariance.dot(m)))
m=monomials(tx,ty,degyx,degyy)
tdy=np.inner(dy_coeff,m)
tsy=np.sqrt(np.inner(m,dy_coeff_covariance.dot(m)))
merr=max(merr,tsx)
merr=max(merr,tsy)
log.info("max edge shift error = %4.3f pixels"%merr)
if degxx==0 and degxy==0 and degyx==0 and degyy==0 :
break
except ( LinAlgError , ValueError ) :
log.warning("polynomial fit failed with degx=(%d,%d) degy=(%d,%d)"%(degxx,degxy,degyx,degyy))
if degxx==0 and degxy==0 and degyx==0 and degyy==0 :
log.error("polynomial degrees are already 0. we can fit the offsets")
raise RuntimeError("polynomial degrees are already 0. we can fit the offsets")
merr = 100000. # this will lower the pol. degree.
if merr > args.max_error :
if merr != 100000. :
log.warning("max edge shift error = %4.3f pixels is too large, reducing degrees"%merr)
if degxy>0 and degyy>0 and degxy>degxx and degyy>degyx : # first along wavelength
if degxy>0 : degxy-=1
if degyy>0 : degyy-=1
else : # then along fiber
if degxx>0 : degxx-=1
if degyx>0 : degyx-=1
else :
# error is ok, so we quit the loop
break
# write this for debugging
if args.outoffsets :
file=open(args.outoffsets,"w")
file.write("# axis wave fiber x y delta error polval (axis 0=y axis1=x)\n")
for e in range(dy.size) :
file.write("0 %f %d %f %f %f %f %f\n"%(wave_for_dy[e],fiber_for_dy[e],x_for_dy[e],y_for_dy[e],dy[e],ey[e],dy_mod[e]))
for e in range(dx.size) :
file.write("1 %f %d %f %f %f %f %f\n"%(wave_for_dx[e],fiber_for_dx[e],x_for_dx[e],y_for_dx[e],dx[e],ex[e],dx_mod[e]))
file.close()
log.info("wrote offsets in ASCII file %s"%args.outoffsets)
# print central shift
mx=np.median(x_for_dx)
my=np.median(y_for_dx)
m=monomials(mx,my,degxx,degxy)
mdx=np.inner(dx_coeff,m)
mex=np.sqrt(np.inner(m,dx_coeff_covariance.dot(m)))
mx=np.median(x_for_dy)
my=np.median(y_for_dy)
m=monomials(mx,my,degyx,degyy)
mdy=np.inner(dy_coeff,m)
mey=np.sqrt(np.inner(m,dy_coeff_covariance.dot(m)))
log.info("central shifts dx = %4.3f +- %4.3f dy = %4.3f +- %4.3f "%(mdx,mex,mdy,mey))
# for each fiber, apply offsets and recompute legendre polynomial
log.info("for each fiber, apply offsets and recompute legendre polynomial")
# compute x y to record max deviations
u = np.linspace(-1,1,5)
x0 = np.zeros((xcoef.shape[0],u.size))
y0 = np.zeros((ycoef.shape[0],u.size))
for f in range(xcoef.shape[0]) :
x0[f]=legval(u,xcoef[f])
y0[f]=legval(u,ycoef[f])
xcoef,ycoef = recompute_legendre_coefficients(xcoef=xcoef,ycoef=ycoef,wavemin=wavemin,wavemax=wavemax,degxx=degxx,degxy=degxy,degyx=degyx,degyy=degyy,dx_coeff=dx_coeff,dy_coeff=dy_coeff)
# use an input spectrum as an external calibration of wavelength
if spectrum_filename :
log.info("write and reread PSF to be sure predetermined shifts were propagated")
write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax)
#psf,xcoef,ycoef,wavemin,wavemax = read_psf_and_traces(args.outpsf)
xytraceset = read_xytraceset(args.outpsf)
wavemin = xytraceset.wavemin
wavemax = xytraceset.wavemax
xcoef = xytraceset.x_vs_wave_traceset._coeff
ycoef = xytraceset.y_vs_wave_traceset._coeff
psf = read_specter_psf(args.outpsf)
ycoef=shift_ycoef_using_external_spectrum(psf=psf,xytraceset=xytraceset,
image=image,fibers=fibers,spectrum_filename=spectrum_filename,degyy=args.degyy,width=7)
x = np.zeros((xcoef.shape[0],u.size))
y = np.zeros((ycoef.shape[0],u.size))
for f in range(xcoef.shape[0]) :
x[f]=legval(u,xcoef[f])
y[f]=legval(u,ycoef[f])
dx = x-x0
dy = y-y0
header_keywords = {}
header_keywords["MEANDX"]=np.mean(dx)
header_keywords["MINDX"]=np.min(dx)
header_keywords["MAXDX"]=np.max(dx)
header_keywords["MEANDY"]=np.mean(dy)
header_keywords["MINDY"]=np.min(dy)
header_keywords["MAXDY"]=np.max(dy)
write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax,header_keywords=header_keywords)
log.info("wrote modified PSF in %s"%args.outpsf)
else :
x = np.zeros((xcoef.shape[0],u.size))
y = np.zeros((ycoef.shape[0],u.size))
for f in range(xcoef.shape[0]) :
x[f]=legval(u,xcoef[f])
y[f]=legval(u,ycoef[f])
dx = x-x0
dy = y-y0
header_keywords = {}
header_keywords["MEANDX"]=np.mean(dx)
header_keywords["MINDX"]=np.min(dx)
header_keywords["MAXDX"]=np.max(dx)
header_keywords["MEANDY"]=np.mean(dy)
header_keywords["MINDY"]=np.min(dy)
header_keywords["MAXDY"]=np.max(dy)
write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax,header_keywords=header_keywords)
log.info("wrote modified PSF in %s"%args.outpsf)
def main(args=None):
if args is None:
args = parse()
elif isinstance(args, (list, tuple)):
args = parse(args)
t0 = time.time()
log = get_logger()
# guess if it is a preprocessed or a raw image
hdulist = fits.open(args.image)
is_input_preprocessed = ("IMAGE" in hdulist) & ("IVAR" in hdulist)
primary_header = hdulist[0].header
hdulist.close()
if is_input_preprocessed:
image = read_image(args.image)
else:
if args.camera is None:
print(
"ERROR: Need to specify camera to open a raw fits image (with all cameras in different fits HDUs)"
)
print(
"Try adding the option '--camera xx', with xx in {brz}{0-9}, like r7, or type 'desi_qproc --help' for more options"
)
sys.exit(12)
image = read_raw(args.image, args.camera, fill_header=[
1,
])
if args.auto:
log.debug("AUTOMATIC MODE")
try:
night = image.meta['NIGHT']
if not 'EXPID' in image.meta:
if 'EXPNUM' in image.meta:
log.warning('using EXPNUM {} for EXPID'.format(
image.meta['EXPNUM']))
image.meta['EXPID'] = image.meta['EXPNUM']
expid = image.meta['EXPID']
except KeyError as e:
log.error(
"Need at least NIGHT and EXPID (or EXPNUM) to run in auto mode. Retry without the --auto option."
)
log.error(str(e))
sys.exit(12)
indir = os.path.dirname(args.image)
if args.fibermap is None:
filename = '{}/fibermap-{:08d}.fits'.format(indir, expid)
if os.path.isfile(filename):
log.debug("auto-mode: found a fibermap, {}, using it!".format(
filename))
args.fibermap = filename
if args.output_preproc is None:
if not is_input_preprocessed:
args.output_preproc = '{}/preproc-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera.lower(), expid)
log.debug("auto-mode: will write preproc in " +
args.output_preproc)
else:
log.debug(
"auto-mode: will not write preproc because input is a preprocessed image"
)
if args.auto_output_dir != '.':
if not os.path.isdir(args.auto_output_dir):
log.debug("auto-mode: creating directory " +
args.auto_output_dir)
os.makedirs(args.auto_output_dir)
if args.output_preproc is not None:
write_image(args.output_preproc, image)
cfinder = None
if args.psf is None:
if cfinder is None:
cfinder = CalibFinder([image.meta, primary_header])
args.psf = cfinder.findfile("PSF")
log.info(" Using PSF {}".format(args.psf))
tset = read_xytraceset(args.psf)
# add fibermap
if args.fibermap:
if os.path.isfile(args.fibermap):
fibermap = read_fibermap(args.fibermap)
else:
log.error("no fibermap file {}".format(args.fibermap))
fibermap = None
else:
fibermap = None
if "OBSTYPE" in image.meta:
obstype = image.meta["OBSTYPE"].upper()
image.meta["OBSTYPE"] = obstype # make sure it's upper case
qframe = None
else:
log.warning("No OBSTYPE keyword, trying to guess ...")
qframe = qproc_boxcar_extraction(tset,
image,
width=args.width,
fibermap=fibermap)
obstype = check_qframe_flavor(
qframe, input_flavor=image.meta["FLAVOR"]).upper()
image.meta["OBSTYPE"] = obstype
log.info("OBSTYPE = '{}'".format(obstype))
if args.auto:
# now set the things to do
if obstype == "SKY" or obstype == "TWILIGHT" or obstype == "SCIENCE":
args.shift_psf = True
args.output_psf = '{}/psf-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.apply_fiberflat = True
args.skysub = True
args.output_skyframe = '{}/qsky-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.fluxcalib = True
args.outframe = '{}/qcframe-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
elif obstype == "ARC" or obstype == "TESTARC":
args.shift_psf = True
args.output_psf = '{}/psf-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.compute_lsf_sigma = True
elif obstype == "FLAT" or obstype == "TESTFLAT":
args.shift_psf = True
args.output_psf = '{}/psf-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
args.compute_fiberflat = '{}/qfiberflat-{}-{:08d}.fits'.format(
args.auto_output_dir, args.camera, expid)
if args.shift_psf:
# using the trace shift script
if args.auto:
options = option_list({
"psf":
args.psf,
"image":
"dummy",
"outpsf":
"dummy",
"continuum": ((obstype == "FLAT") | (obstype == "TESTFLAT")),
"sky": ((obstype == "SCIENCE") | (obstype == "SKY"))
})
else:
options = option_list({
"psf": args.psf,
"image": "dummy",
"outpsf": "dummy"
})
tmp_args = trace_shifts_script.parse(options=options)
tset = trace_shifts_script.fit_trace_shifts(image=image, args=tmp_args)
qframe = qproc_boxcar_extraction(tset,
image,
width=args.width,
fibermap=fibermap)
if tset.meta is not None:
# add traceshift info in the qframe, this will be saved in the qframe header
if qframe.meta is None:
qframe.meta = dict()
for k in tset.meta.keys():
qframe.meta[k] = tset.meta[k]
if args.output_rawframe is not None:
write_qframe(args.output_rawframe, qframe)
log.info("wrote raw extracted frame in {}".format(
args.output_rawframe))
if args.compute_lsf_sigma:
tset = process_arc(qframe, tset, linelist=None, npoly=2, nbins=2)
if args.output_psf is not None:
for k in qframe.meta:
if k not in tset.meta:
tset.meta[k] = qframe.meta[k]
write_xytraceset(args.output_psf, tset)
if args.compute_fiberflat is not None:
fiberflat = qproc_compute_fiberflat(qframe)
#write_qframe(args.compute_fiberflat,qflat)
write_fiberflat(args.compute_fiberflat, fiberflat, header=qframe.meta)
log.info("wrote fiberflat in {}".format(args.compute_fiberflat))
if args.apply_fiberflat or args.input_fiberflat:
if args.input_fiberflat is None:
if cfinder is None:
cfinder = CalibFinder([image.meta, primary_header])
try:
args.input_fiberflat = cfinder.findfile("FIBERFLAT")
except KeyError as e:
log.error("no FIBERFLAT for this spectro config")
sys.exit(12)
log.info("applying fiber flat {}".format(args.input_fiberflat))
flat = read_fiberflat(args.input_fiberflat)
qproc_apply_fiberflat(qframe, flat)
if args.skysub:
log.info("sky subtraction")
if args.output_skyframe is not None:
skyflux = qproc_sky_subtraction(qframe, return_skymodel=True)
sqframe = QFrame(qframe.wave, skyflux, np.ones(skyflux.shape))
write_qframe(args.output_skyframe, sqframe)
log.info("wrote sky model in {}".format(args.output_skyframe))
else:
qproc_sky_subtraction(qframe)
if args.fluxcalib:
if cfinder is None:
cfinder = CalibFinder([image.meta, primary_header])
# check for flux calib
if cfinder.haskey("FLUXCALIB"):
fluxcalib_filename = cfinder.findfile("FLUXCALIB")
fluxcalib = read_average_flux_calibration(fluxcalib_filename)
log.info("read average calib in {}".format(fluxcalib_filename))
seeing = qframe.meta["SEEING"]
airmass = qframe.meta["AIRMASS"]
exptime = qframe.meta["EXPTIME"]
exposure_calib = fluxcalib.value(seeing=seeing, airmass=airmass)
for q in range(qframe.nspec):
fiber_calib = np.interp(qframe.wave[q], fluxcalib.wave,
exposure_calib) * exptime
inv_calib = (fiber_calib > 0) / (fiber_calib +
(fiber_calib == 0))
qframe.flux[q] *= inv_calib
qframe.ivar[q] *= fiber_calib**2 * (fiber_calib > 0)
# add keyword in header giving the calibration factor applied at a reference wavelength
band = qframe.meta["CAMERA"].upper()[0]
if band == "B":
refwave = 4500
elif band == "R":
refwave = 6500
else:
refwave = 8500
calvalue = np.interp(refwave, fluxcalib.wave,
exposure_calib) * exptime
qframe.meta["CALWAVE"] = refwave
qframe.meta["CALVALUE"] = calvalue
else:
log.error(
"Cannot calibrate fluxes because no FLUXCALIB keywork in calibration files"
)
fibers = parse_fibers(args.fibers)
if fibers is None:
fibers = qframe.flux.shape[0]
else:
ii = np.arange(qframe.fibers.size)[np.in1d(qframe.fibers, fibers)]
if ii.size == 0:
log.error("no such fibers in frame,")
log.error("fibers are in range [{}:{}]".format(
qframe.fibers[0], qframe.fibers[-1] + 1))
sys.exit(12)
qframe = qframe[ii]
if args.outframe is not None:
write_qframe(args.outframe, qframe)
log.info("wrote {}".format(args.outframe))
t1 = time.time()
log.info("all done in {:3.1f} sec".format(t1 - t0))
if args.plot:
log.info("plotting {} spectra".format(qframe.wave.shape[0]))
import matplotlib.pyplot as plt
fig = plt.figure()
for i in range(qframe.wave.shape[0]):
j = (qframe.ivar[i] > 0)
plt.plot(qframe.wave[i, j], qframe.flux[i, j])
plt.grid()
plt.xlabel("wavelength")
plt.ylabel("flux")
plt.show()
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 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 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()))
def main(args=None):
log = get_logger()
if args is None:
args = parse()
first_preproc = None
images = []
ivars = []
masks = []
log.info("read inputs ...")
for filename in args.infile:
log.info(" read {}".format(filename))
tmp = read_image(filename)
if first_preproc is None:
first_preproc = tmp
images.append(tmp.pix.ravel())
# make sure we don't include data with ivar=0
mask = tmp.mask + (tmp.ivar == 0).astype(int)
masks.append(mask.ravel())
ivars.append((tmp.ivar * (tmp.mask == 0)).ravel())
images = np.array(images)
masks = np.array(masks)
ivars = np.array(ivars)
smask = np.sum(masks, axis=0)
log.info("compute masked median image ...")
medimage = masked_median(images, masks)
log.info("use masked median image to discard outlier ...")
good = ((images - medimage)**2 * (ivars > 0) /
((medimage > 0) * (medimage * args.fracerr)**2 + 1. /
(ivars + (ivars == 0)))) < args.nsig**3
ivars *= good.astype(float)
if args.weighted:
log.info("compute weighted mean ...")
sw = np.sum(ivars, axis=0)
swf = np.sum(ivars * images, axis=0)
meanimage = swf / (sw + (sw == 0))
meanivar = sw
else:
log.info("compute unweighted mean ...")
s1 = np.sum((ivars > 0).astype(int), axis=0)
sf = np.sum((ivars > 0) * images, axis=0)
meanimage = sf / (s1 + (s1 == 0))
meanvar = np.sum(
(ivars > 0) / (ivars + (ivars == 0))) / (s1 + (s1 == 0))**2
meanivar = (meanvar > 0) / (meanvar + (meanvar == 0))
log.info("write nimages.fits ...")
fitsio.write("nimages.fits",
s1.reshape(first_preproc.pix.shape),
clobber=True)
log.info("compute mask ...")
meanmask = masks[0]
for mask in masks[1:]:
meanmask &= mask
log.info("write image ...")
preproc = first_preproc
shape = preproc.pix.shape
preproc.pix = meanimage.reshape(shape)
preproc.ivar = meanivar.reshape(shape)
preproc.mask = meanmask.reshape(shape)
write_image(args.outfile, preproc)
log.info("wrote {}".format(args.outfile))
