def get_frame_data(nspec=10, objtype=None):
"""
Return basic test data for desispec.frame object:
"""
nwave = 100
wavemin, wavemax = 4000, 4100
wave, model_flux = get_models(nspec,
nwave,
wavemin=wavemin,
wavemax=wavemax)
resol_data = set_resolmatrix(nspec, nwave)
calib = np.sin((wave - wavemin) * np.pi / np.max(wave))
flux = np.zeros((nspec, nwave))
for i in range(nspec):
flux[i] = Resolution(resol_data[i]).dot(model_flux[i] * calib)
sigma = 0.01
# flux += np.random.normal(scale=sigma, size=flux.shape)
ivar = np.ones(flux.shape) / sigma**2
mask = np.zeros(flux.shape, dtype=int)
fibermap = empty_fibermap(nspec, 1500)
if objtype is None:
fibermap['OBJTYPE'] = 'QSO'
fibermap['OBJTYPE'][0:3] = 'STD' # For flux tests
else:
fibermap['OBJTYPE'] = objtype
frame = Frame(wave, flux, ivar, mask, resol_data, fibermap=fibermap)
frame.meta = {}
frame.meta['EXPTIME'] = 1. # For flux tests
return frame
def test_slice(self):
nspec = 5
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
mask = np.zeros(flux.shape, dtype=int)
rdata = np.ones((nspec, 5, nwave))
fibermap = desispec.io.fibermap.empty_fibermap(nspec)
frame = Frame(wave, flux, ivar, mask, rdata, spectrograph=0)
x = frame[1]
self.assertEqual(type(x), Spectrum)
x = frame[1:2]
self.assertEqual(type(x), Frame)
x = frame[[1, 2, 3]]
self.assertEqual(type(x), Frame)
x = frame[frame.fibers < 3]
self.assertEqual(type(x), Frame)
#- Slice fibermap too
frame = Frame(wave,
flux,
ivar,
mask,
rdata,
spectrograph=0,
fibermap=fibermap)
x = frame[frame.fibers < 3]
self.assertEqual(len(x.fibers), len(x.fibermap))
x = frame[[1, 2, 3]]
self.assertTrue(np.all(x.fibers == x.fibermap['FIBER']))
def _make_frame(self, camera='b0', flavor='dark', night=None, expid=None):
if night is None:
night = self.night
if expid is None:
expid = self.expid
nspec = 3
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
frame = Frame(wave, flux, ivar, spectrograph=0)
frame.meta = dict(CAMERA=camera, FLAVOR=flavor, NIGHT=night, EXPID=expid)
return frame
def compute_coadd_scores(coadd, update_coadd=True):
"""Compute scores for a coadded Spectra object
Args:
coadd: a Spectra object from a coadd
Options:
update_coadd: if True, update coadd.scores
Returns tuple of dictionaries (scores, comments); see compute_frame_scores
"""
scores = dict()
comments = dict()
if coadd.bands == ['brz']:
#- i.e. this is a coadd across cameras
fr = Frame(coadd.wave['brz'],
coadd.flux['brz'],
coadd.ivar['brz'],
fibermap=coadd.fibermap,
meta=coadd.meta,
resolution_data=coadd.resolution_data['brz'])
for band in ['b', 'r', 'z']:
bandscores, bandcomments = compute_frame_scores(
fr, band=band, suffix='COADD', flux_per_angstrom=True)
scores.update(bandscores)
comments.update(bandcomments)
else:
#- otherwise try individual bands, upper or lowercase
for band in ['b', 'r', 'z', 'B', 'R', 'Z']:
if band in coadd.bands:
fr = Frame(coadd.wave[band],
coadd.flux[band],
coadd.ivar[band],
fibermap=coadd.fibermap,
meta=coadd.meta,
resolution_data=coadd.resolution_data[band])
bandscores, bandcomments = compute_frame_scores(
fr, band=band, suffix='COADD', flux_per_angstrom=True)
scores.update(bandscores)
comments.update(bandcomments)
if update_coadd:
if hasattr(coadd, 'scores') and coadd.scores is not None:
for key in scores:
coadd.scores[key] = scores[key]
coadd.scores_comments[key] = comments[key]
else:
coadd.scores = scores
coadd.scores_comments = comments
return scores, comments
def test_frame_rw(self):
nspec, nwave, ndiag = 5, 10, 3
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.random.uniform(size=(nspec, nwave))
meta = dict(BLAT=0, FOO='abc', FIBERMIN=500)
mask_int = np.zeros((nspec, nwave), dtype=int)
mask_uint = np.zeros((nspec, nwave), dtype=np.uint32)
wave = np.arange(nwave)
R = np.random.uniform( size=(nspec, ndiag, nwave) )
for mask in (mask_int, mask_uint):
frx = Frame(wave, flux, ivar, mask, R, meta=meta)
desispec.io.write_frame(self.testfile, frx)
frame = desispec.io.read_frame(self.testfile)
flux2 = flux.astype('f4').astype('f8')
ivar2 = ivar.astype('f4').astype('f8')
wave2 = wave.astype('f4').astype('f8')
R2 = R.astype('f4').astype('f8')
self.assertTrue(frame.wave.dtype == np.float64)
self.assertTrue(frame.flux.dtype == np.float64)
self.assertTrue(frame.ivar.dtype == np.float64)
self.assertTrue(frame.resolution_data.dtype == np.float64)
self.assertTrue(np.all(flux2 == frame.flux))
self.assertTrue(np.all(ivar2 == frame.ivar))
self.assertTrue(np.all(wave2 == frame.wave))
self.assertTrue(np.all(mask == frame.mask))
self.assertTrue(np.all(R2 == frame.resolution_data))
self.assertTrue(frame.resolution_data.dtype.isnative)
self.assertEqual(frame.meta['BLAT'], meta['BLAT'])
self.assertEqual(frame.meta['FOO'], meta['FOO'])
#- Test float32 on disk vs. float64 in memory
for extname in ['FLUX', 'IVAR', 'WAVELENGTH', 'RESOLUTION']:
data = fits.getdata(self.testfile, extname)
self.assertEqual(data.dtype, np.dtype('>f4'), '{} not type >f4'.format(extname))
#- with and without fibermap
self.assertEqual(frame.fibermap, None)
fibermap = desispec.io.empty_fibermap(nspec)
fibermap['TARGETID'] = np.arange(nspec)*2
frx = Frame(wave, flux, ivar, mask, R, fibermap=fibermap)
desispec.io.write_frame(self.testfile, frx)
frame = desispec.io.read_frame(self.testfile)
for name in fibermap.dtype.names:
match = np.all(fibermap[name] == frame.fibermap[name])
self.assertTrue(match, 'Fibermap column {} mismatch'.format(name))
def test_resolution(self):
"""
Test that identical spectra convolved with different resolutions
results in identical fiberflats
"""
wave, flux, ivar, mask = _get_data()
nspec, nwave = flux.shape
#- Setup a Resolution matrix that varies with fiber and wavelength
#- Note: this is actually the transpose of the resolution matrix
#- I wish I was creating, but as long as we self-consistently
#- use it for convolving and solving, that shouldn't matter.
sigma = np.linspace(2, 10, nwave * nspec)
ndiag = 21
xx = np.linspace(-ndiag / 2.0, +ndiag / 2.0, ndiag)
Rdata = np.zeros((nspec, len(xx), nwave))
for i in range(nspec):
for j in range(nwave):
kernel = np.exp(-xx**2 / (2 * sigma[i * nwave + j]**2))
kernel /= sum(kernel)
Rdata[i, :, j] = kernel
#- Convolve the data with the resolution matrix
convflux = np.empty_like(flux)
for i in range(nspec):
convflux[i] = Resolution(Rdata[i]).dot(flux[i])
#- Run the code
frame = Frame(wave, convflux, ivar, mask, Rdata, spectrograph=0)
ff = compute_fiberflat(frame)
#- These fiber flats should all be ~1
self.assertTrue(np.all(np.abs(ff.fiberflat - 1) < 0.001))
def test_interface(self):
"""
Basic test that interface works and identical inputs result in
identical outputs
"""
wave, flux, ivar, mask = _get_data()
nspec, nwave = flux.shape
#- Setup data for a Resolution matrix
sigma = 4.0
ndiag = 11
xx = np.linspace(-(ndiag - 1) / 2.0, +(ndiag - 1) / 2.0, ndiag)
Rdata = np.zeros((nspec, ndiag, nwave))
kernel = np.exp(-xx**2 / (2 * sigma))
kernel /= sum(kernel)
for i in range(nspec):
for j in range(nwave):
Rdata[i, :, j] = kernel
#- Run the code
frame = Frame(wave, flux, ivar, mask, Rdata, spectrograph=0)
ff = compute_fiberflat(frame)
#- Check shape of outputs
self.assertEqual(ff.fiberflat.shape, flux.shape)
self.assertEqual(ff.ivar.shape, flux.shape)
self.assertEqual(ff.mask.shape, flux.shape)
self.assertEqual(len(ff.meanspec), nwave)
#- Identical inputs should result in identical ouputs
for i in range(1, nspec):
self.assertTrue(np.all(ff.fiberflat[i] == ff.fiberflat[0]))
self.assertTrue(np.all(ff.ivar[i] == ff.ivar[0]))
def _write_frame(self,
flavor='none',
camera='b',
expid=1,
night='20160607',
gaia_only=False):
"""Write a fake frame"""
flux = np.ones((self.nspec, self.nwave))
ivar = np.ones((self.nspec, self.nwave)) * 100 # S/N=10
mask = np.zeros((self.nspec, self.nwave), dtype=int)
Rdata = np.ones((self.nspec, 1, self.nwave))
fibermap = self._get_fibermap(gaia_only=gaia_only)
frame = Frame(self.wave,
flux,
ivar,
mask,
Rdata,
fibermap=fibermap,
meta=dict(FLAVOR=flavor,
CAMERA=camera,
EXPID=expid,
NIGHT=night,
EXPTIME=1000.,
DETECTOR='SIM'))
io.write_frame(self.framefile, frame)
def _make_frame(self, nspec=5, nwave=10, ndiag=3):
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.random.uniform(size=(nspec, nwave))
mask = np.zeros((nspec, nwave), dtype=int)
R = np.random.uniform(size=(nspec, ndiag, nwave))
return Frame(wave, flux, ivar, mask, R)
def _get_spectra(self):
#- Setup data for a Resolution matrix
sigma = 4.0
ndiag = 21
xx = np.linspace(-(ndiag - 1) / 2.0, +(ndiag - 1) / 2.0, ndiag)
Rdata = np.zeros((self.nspec, ndiag, self.nwave))
for i in range(self.nspec):
for j in range(self.nwave):
kernel = np.exp(-xx**2 / (2 * sigma))
kernel /= sum(kernel)
Rdata[i, :, j] = kernel
flux = np.zeros((self.nspec, self.nwave))
ivar = np.ones((self.nspec, self.nwave))
mask = np.zeros((self.nspec, self.nwave), dtype=int)
for i in range(self.nspec):
R = Resolution(Rdata[i])
flux[i] = R.dot(self.flux)
fibermap = desispec.io.empty_fibermap(self.nspec, 1500)
fibermap['OBJTYPE'][0::2] = 'SKY'
return Frame(self.wave,
flux,
ivar,
mask,
Rdata,
spectrograph=2,
fibermap=fibermap)
def _make_frame(self, camera='b0', flavor='dark', night=None, expid=None):
if night is None:
night = self.night
if expid is None:
expid = self.expid
nspec = 3
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
frame = Frame(wave, flux, ivar, spectrograph=0)
frame.meta = dict(CAMERA=camera,
FLAVOR=flavor,
NIGHT=night,
EXPID=expid)
return frame
def test_apply_fluxcalibration(self):
#get frame_data
wave = np.arange(5000, 6000)
nwave = len(wave)
nspec = 3
flux = np.random.uniform(0.9, 1.0, size=(nspec, nwave))
ivar = np.ones_like(flux)
origframe = Frame(wave, flux, ivar, spectrograph=0)
#define fluxcalib object
calib = np.ones_like(origframe.flux)
mask = np.zeros(origframe.flux.shape, dtype=np.uint32)
calib[0] *= 0.5
calib[1] *= 1.5
# fc with essentially no error
fcivar = 1e20 * np.ones_like(origframe.flux)
fc = FluxCalib(origframe.wave, calib, fcivar,mask)
frame = copy.deepcopy(origframe)
apply_flux_calibration(frame, fc)
self.assertTrue(np.allclose(frame.ivar, calib**2))
# origframe.flux=0 should result in frame.flux=0
fcivar = np.ones_like(origframe.flux)
calib = np.ones_like(origframe.flux)
fc = FluxCalib(origframe.wave, calib, fcivar, mask)
frame = copy.deepcopy(origframe)
frame.flux[0,0:10]=0.0
apply_flux_calibration(frame, fc)
self.assertTrue(np.all(frame.flux[0, 0:10] == 0.0))
#fcivar=0 should result in frame.ivar=0
fcivar=np.ones_like(origframe.flux)
calib=np.ones_like(origframe.flux)
fcivar[0,0:10]=0.0
fc=FluxCalib(origframe.wave,calib,fcivar,mask)
frame=copy.deepcopy(origframe)
apply_flux_calibration(frame,fc)
self.assertTrue(np.all(frame.ivar[0,0:10]==0.0))
# should also work even the calib =0 ??
#fcivar=np.ones_like(origframe.flux)
#calib=np.ones_like(origframe.flux)
#fcivar[0,0:10]=0.0
#calib[0,0:10]=0.0
#fc=FluxCalib(origframe.wave,calib,fcivar,mask)
#frame=copy.deepcopy(origframe)
#apply_flux_calibration(frame,fc)
#self.assertTrue(np.all(frame.ivar[0,0:10]==0.0))
# test different wavelength bins
frame=copy.deepcopy(origframe)
calib = np.ones_like(frame.flux)
fcivar=np.ones_like(frame.ivar)
mask=np.zeros(origframe.flux.shape, dtype=np.uint32)
fc=FluxCalib(origframe.wave+0.01,calib,fcivar,mask)
with self.assertRaises(SystemExit): #should be ValueError instead?
apply_flux_calibration(frame,fc)
def _write_frame(self):
"""Write a fake frame"""
wave = 5000+np.arange(self.nwave)
flux = np.ones((self.nspec, self.nwave))
ivar = np.ones((self.nspec, self.nwave))
mask = np.zeros((self.nspec, self.nwave), dtype=int)
Rdata = np.ones((self.nspec, 1, self.nwave))
frame = Frame(wave, flux, ivar, mask, Rdata)
io.write_frame(self.framefile, frame)
def test_io_qa_frame(self):
nspec = 3
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
frame = Frame(wave, flux, ivar, spectrograph=0)
frame.meta = dict(CAMERA='b0', FLAVOR='dark', NIGHT='20160607', EXPID=1)
#- Init
qaframe = QA_Frame(frame)
qaframe.init_skysub()
# Write
desio_qa.write_qa_frame(self.testyfile, qaframe)
# Read
xqaframe = desio_qa.read_qa_frame(self.testyfile)
# Check
self.assertTrue(qaframe.qa_data['SKYSUB']['PARAM']['PCHI_RESID'] == xqaframe.qa_data['SKYSUB']['PARAM']['PCHI_RESID'])
self.assertTrue(qaframe.flavor == xqaframe.flavor)
def _write_frame(self, flavor='none', camera='b', expid=1, night='20160607'):
"""Write a fake frame"""
wave = 5000+np.arange(self.nwave)
flux = np.ones((self.nspec, self.nwave))
ivar = np.ones((self.nspec, self.nwave))
mask = np.zeros((self.nspec, self.nwave), dtype=int)
Rdata = np.ones((self.nspec, 1, self.nwave))
fibermap = self._get_fibermap()
frame = Frame(wave, flux, ivar, mask, Rdata, fibermap=fibermap,
meta=dict(FLAVOR=flavor, CAMERA=camera, EXPID=expid, NIGHT=night))
io.write_frame(self.framefile, frame)
def test_throughput_resolution(self):
"""
Test that spectra with different throughputs and different resolutions
result in fiberflat variations that are only due to throughput.
"""
wave, flux, ivar, mask = _get_data()
nspec, nwave = flux.shape
#- Setup a Resolution matrix that varies with fiber and wavelength
#- Note: this is actually the transpose of the resolution matrix
#- I wish I was creating, but as long as we self-consistently
#- use it for convolving and solving, that shouldn't matter.
sigma = np.linspace(2, 10, nwave * nspec)
ndiag = 21
xx = np.linspace(-ndiag / 2.0, +ndiag / 2.0, ndiag)
Rdata = np.zeros((nspec, len(xx), nwave))
for i in range(nspec):
for j in range(nwave):
kernel = np.exp(-xx**2 / (2 * sigma[i * nwave + j]**2))
kernel /= sum(kernel)
Rdata[i, :, j] = kernel
#- Vary the input flux prior to calculating the fiber flat
flux[1] *= 1.1
flux[2] *= 1.2
flux[3] /= 1.1
flux[4] /= 1.2
#- Convolve the data with the varying resolution matrix
convflux = np.empty_like(flux)
for i in range(nspec):
convflux[i] = Resolution(Rdata[i]).dot(flux[i])
#- Run the code
frame = Frame(wave, convflux, ivar, mask, Rdata, spectrograph=0)
#- Set an accuracy for this
accuracy = 1.e-9
ff = compute_fiberflat(frame, accuracy=accuracy)
#- Compare variation with middle fiber
mid = ff.fiberflat.shape[0] // 2
diff = (ff.fiberflat[1] / 1.1 - ff.fiberflat[mid])
self.assertLess(np.max(np.abs(diff)), accuracy)
diff = (ff.fiberflat[2] / 1.2 - ff.fiberflat[mid])
self.assertLess(np.max(np.abs(diff)), accuracy)
diff = (ff.fiberflat[3] * 1.1 - ff.fiberflat[mid])
self.assertLess(np.max(np.abs(diff)), accuracy)
diff = (ff.fiberflat[4] * 1.2 - ff.fiberflat[mid])
self.assertLess(np.max(np.abs(diff)), accuracy)
def test_apply_fiberflat(self):
'''test apply_fiberflat interface and changes to flux and mask'''
wave = np.arange(5000, 5050)
nwave = len(wave)
nspec = 3
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones_like(flux)
frame = Frame(wave, flux, ivar, spectrograph=0)
fiberflat = np.ones_like(flux)
ffivar = 2 * np.ones_like(flux)
ffmask = np.zeros_like(flux)
fiberflat[0] *= 0.8
fiberflat[1] *= 1.2
fiberflat[2, 0:10] = 0 #- bad fiberflat
ffivar[2, 10:20] = 0 #- bad fiberflat
ffmask[2, 20:30] = 1 #- bad fiberflat
ff = FiberFlat(wave, fiberflat, ffivar)
origframe = copy.deepcopy(frame)
apply_fiberflat(frame, ff)
#- was fiberflat applied?
self.assertTrue(np.all(frame.flux[0] == origframe.flux[0] / 0.8))
self.assertTrue(np.all(frame.flux[1] == origframe.flux[1] / 1.2))
self.assertTrue(np.all(frame.flux[2] == origframe.flux[2]))
#- did mask get set?
ii = (ff.fiberflat == 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
ii = (ff.ivar == 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
ii = (ff.mask != 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
#- Should fail if frame and ff don't have a common wavelength grid
frame.wave = frame.wave + 0.1
with self.assertRaises(ValueError):
apply_fiberflat(frame, ff)
def test_apply_fiberflat(self):
'''test apply_fiberflat interface and changes to flux and mask'''
wave = np.arange(5000, 5050)
nwave = len(wave)
nspec = 3
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones_like(flux)
frame = Frame(wave, flux, ivar, spectrograph=0)
fiberflat = np.ones_like(flux)
ffivar = 2*np.ones_like(flux)
ffmask = np.zeros_like(flux)
fiberflat[0] *= 0.8
fiberflat[1] *= 1.2
fiberflat[2, 0:10] = 0 #- bad fiberflat
ffivar[2, 10:20] = 0 #- bad fiberflat
ffmask[2, 20:30] = 1 #- bad fiberflat
ff = FiberFlat(wave, fiberflat, ffivar)
origframe = copy.deepcopy(frame)
apply_fiberflat(frame, ff)
#- was fiberflat applied?
self.assertTrue(np.all(frame.flux[0] == origframe.flux[0]/0.8))
self.assertTrue(np.all(frame.flux[1] == origframe.flux[1]/1.2))
self.assertTrue(np.all(frame.flux[2] == origframe.flux[2]))
#- did mask get set?
ii = (ff.fiberflat == 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
ii = (ff.ivar == 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
ii = (ff.mask != 0)
self.assertTrue(np.all((frame.mask[ii] & specmask.BADFIBERFLAT) != 0))
#- Should fail if frame and ff don't have a common wavelength grid
frame.wave = frame.wave + 0.1
with self.assertRaises(ValueError):
apply_fiberflat(frame, ff)
def read_frame(filename, nspec=None):
"""Reads a frame fits file and returns its data.
Args:
filename: path to a file, or (night, expid, camera) tuple where
night = string YEARMMDD
expid = integer exposure ID
camera = b0, r1, .. z9
Returns:
desispec.Frame object with attributes wave, flux, ivar, etc.
"""
#- check if filename is (night, expid, camera) tuple instead
if not isinstance(filename, (str, unicode)):
night, expid, camera = filename
filename = findfile('frame', night, expid, camera)
if not os.path.isfile(filename):
raise IOError("cannot open" + filename)
fx = fits.open(filename, uint=True, memmap=False)
hdr = fx[0].header
flux = native_endian(fx['FLUX'].data.astype('f8'))
ivar = native_endian(fx['IVAR'].data.astype('f8'))
wave = native_endian(fx['WAVELENGTH'].data.astype('f8'))
if 'MASK' in fx:
mask = native_endian(fx['MASK'].data)
else:
mask = None #- let the Frame object create the default mask
resolution_data = native_endian(fx['RESOLUTION'].data.astype('f8'))
if 'FIBERMAP' in fx:
fibermap = fx['FIBERMAP'].data
else:
fibermap = None
fx.close()
if nspec is not None:
flux = flux[0:nspec]
ivar = ivar[0:nspec]
resolution_data = resolution_data[0:nspec]
# return flux,ivar,wave,resolution_data, hdr
return Frame(wave,
flux,
ivar,
mask,
resolution_data,
meta=hdr,
fibermap=fibermap)
def _get_spectra(self, with_gradient=False):
#- Setup data for a Resolution matrix
sigma2 = 4.0
ndiag = 21
xx = np.linspace(-(ndiag - 1) / 2.0, +(ndiag - 1) / 2.0, ndiag)
Rdata = np.zeros((self.nspec, ndiag, self.nwave))
for i in range(self.nspec):
kernel = np.exp(-(xx + float(i) / self.nspec * 0.3)**2 /
(2 * sigma2))
#kernel = np.exp(-xx**2/(2*sigma2))
kernel /= sum(kernel)
for j in range(self.nwave):
Rdata[i, :, j] = kernel
flux = np.zeros((self.nspec, self.nwave), dtype=float)
ivar = np.ones((self.nspec, self.nwave), dtype=float)
# Add a random component
for i in range(self.nspec):
ivar[i] += 0.4 * np.random.uniform(size=self.nwave)
mask = np.zeros((self.nspec, self.nwave), dtype=int)
fibermap = empty_fibermap(self.nspec, 1500)
fibermap['OBJTYPE'][0::2] = 'SKY'
x = fibermap["FIBERASSIGN_X"]
y = fibermap["FIBERASSIGN_Y"]
x = x - np.mean(x)
y = y - np.mean(y)
if np.std(x) > 0: x /= np.std(x)
if np.std(y) > 0: y /= np.std(y)
w = (self.wave - self.wave[0]) / (self.wave[-1] -
self.wave[0]) * 2. - 1
for i in range(self.nspec):
R = Resolution(Rdata[i])
if with_gradient:
scale = 1. + (0.1 * x[i] + 0.2 * y[i]) * (1 + 0.4 * w)
flux[i] = R.dot(scale * self.flux)
else:
flux[i] = R.dot(self.flux)
meta = {"camera": "r2"}
return Frame(self.wave,
flux,
ivar,
mask,
Rdata,
spectrograph=2,
fibermap=fibermap,
meta=meta)
def _get_frame(self):
nspec = 4
wave = np.linspace(3600, 6000, (6000 - 3600))
Rdata = np.ones((nspec, 1, wave.size))
flux = np.ones((nspec, wave.size))
ivar = np.ones((nspec, wave.size))
mask = np.zeros((nspec, wave.size), dtype=int)
fibermap = desispec.io.empty_fibermap(nspec)
return Frame(wave,
flux,
ivar,
mask,
Rdata,
spectrograph=0,
fibermap=fibermap)
def test_main(self):
"""
Test the main program.
"""
# generate the frame data
wave, flux, ivar, mask = _get_data()
nspec, nwave = flux.shape
#- Setup data for a Resolution matrix
sigma = 4.0
ndiag = 11
xx = np.linspace(-(ndiag - 1) / 2.0, +(ndiag - 1) / 2.0, ndiag)
Rdata = np.zeros((nspec, ndiag, nwave))
kernel = np.exp(-xx**2 / (2 * sigma))
kernel /= sum(kernel)
for i in range(nspec):
for j in range(nwave):
Rdata[i, :, j] = kernel
#- Convolve the data with the resolution matrix
convflux = np.empty_like(flux)
for i in range(nspec):
convflux[i] = Resolution(Rdata[i]).dot(flux[i])
# create a fake fibermap
fibermap = io.empty_fibermap(nspec, nwave)
for i in range(0, nspec):
fibermap['OBJTYPE'][i] = 'FAKE'
io.write_fibermap(self.testfibermap, fibermap)
#- write out the frame
frame = Frame(wave,
convflux,
ivar,
mask,
Rdata,
spectrograph=0,
fibermap=fibermap,
meta=dict(FLAVOR='flat'))
write_frame(self.testframe, frame, fibermap=fibermap)
# set program arguments
argstr = ['--infile', self.testframe, '--outfile', self.testflat]
# run it
args = ffscript.parse(options=argstr)
ffscript.main(args)
def test_slice(self):
nspec = 5
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
mask = np.zeros(flux.shape, dtype=int)
rdata = np.ones((nspec, 5, nwave))
frame = Frame(wave, flux, ivar, mask, rdata)
x = frame[1]
self.assertEqual(type(x), Spectrum)
x = frame[1:2]
self.assertEqual(type(x), Frame)
x = frame[[1,2,3]]
self.assertEqual(type(x), Frame)
x = frame[frame.fibers<3]
self.assertEqual(type(x), Frame)
def test_apply_fiberflat_ivar(self):
'''test error propagation in apply_fiberflat'''
wave = np.arange(5000, 5010)
nwave = len(wave)
nspec = 3
flux = np.random.uniform(0.9, 1.0, size=(nspec, nwave))
ivar = np.ones_like(flux)
origframe = Frame(wave, flux, ivar, spectrograph=0)
fiberflat = np.ones_like(flux)
ffmask = np.zeros_like(flux)
fiberflat[0] *= 0.5
fiberflat[1] *= 1.5
#- ff with essentially no error
ffivar = 1e20 * np.ones_like(flux)
ff = FiberFlat(wave, fiberflat, ffivar)
frame = copy.deepcopy(origframe)
apply_fiberflat(frame, ff)
self.assertTrue(np.allclose(frame.ivar, fiberflat**2))
#- ff with large error
ffivar = np.ones_like(flux)
ff = FiberFlat(wave, fiberflat, ffivar)
frame = copy.deepcopy(origframe)
apply_fiberflat(frame, ff)
#- c = a/b
#- (sigma_c/c)^2 = (sigma_a/a)^2 + (sigma_b/b)^2
var = frame.flux**2 * (1.0/(origframe.ivar * origframe.flux**2) + \
1.0/(ff.ivar * ff.fiberflat**2))
self.assertTrue(np.allclose(frame.ivar, 1 / var))
#- ff.ivar=0 should result in frame.ivar=0, even if ff.fiberflat=0 too
ffivar = np.ones_like(flux)
ffivar[0, 0:5] = 0.0
fiberflat[0, 0:5] = 0.0
ff = FiberFlat(wave, fiberflat, ffivar)
frame = copy.deepcopy(origframe)
apply_fiberflat(frame, ff)
self.assertTrue(np.all(frame.ivar[0, 0:5] == 0.0))
def test_frame_rw(self):
nspec, nwave, ndiag = 5, 10, 3
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.random.uniform(size=(nspec, nwave))
mask_int = np.zeros((nspec, nwave), dtype=int)
mask_uint = np.zeros((nspec, nwave), dtype=np.uint32)
wave = np.arange(nwave)
R = np.random.uniform( size=(nspec, ndiag, nwave) )
for mask in (mask_int, mask_uint):
frx = Frame(wave, flux, ivar, mask, R)
desispec.io.write_frame(self.testfile, frx)
frame = desispec.io.read_frame(self.testfile)
self.assertTrue(np.all(flux == frame.flux))
self.assertTrue(np.all(ivar == frame.ivar))
self.assertTrue(np.all(wave == frame.wave))
self.assertTrue(np.all(mask == frame.mask))
self.assertTrue(np.all(R == frame.resolution_data))
self.assertTrue(frame.resolution_data.dtype.isnative)
def get_frame_data(nspec=10, wavemin=4000, wavemax=4100, nwave=100, meta={}):
"""
Return basic test data for desispec.frame object:
"""
wave, model_flux = get_models(nspec,
nwave,
wavemin=wavemin,
wavemax=wavemax)
resol_data = set_resolmatrix(nspec, nwave)
calib = np.sin((wave - wavemin) * np.pi / np.max(wave))
flux = np.zeros((nspec, nwave))
for i in range(nspec):
flux[i] = Resolution(resol_data[i]).dot(model_flux[i] * calib)
sigma = 0.01
# flux += np.random.normal(scale=sigma, size=flux.shape)
ivar = np.ones(flux.shape) / sigma**2
mask = np.zeros(flux.shape, dtype=int)
fibermap = empty_fibermap(nspec, 1500)
fibermap['OBJTYPE'] = 'TGT'
fibermap['DESI_TARGET'] = desi_mask.QSO
fibermap['DESI_TARGET'][0:3] = desi_mask.STD_FAINT # For flux tests
fibermap['FIBER_X'] = np.arange(nspec) * 400. / nspec #mm
fibermap['FIBER_Y'] = np.arange(nspec) * 400. / nspec #mm
fibermap['DELTA_X'] = 0.005 * np.ones(nspec) #mm
fibermap['DELTA_Y'] = 0.003 * np.ones(nspec) #mm
if "EXPTIME" not in meta.keys():
meta['EXPTIME'] = 1.0
frame = Frame(wave,
flux,
ivar,
mask,
resol_data,
fibermap=fibermap,
meta=meta)
return frame
def test_init(self):
nspec = 3
nwave = 10
wave = np.arange(nwave)
flux = np.random.uniform(size=(nspec, nwave))
ivar = np.ones(flux.shape)
mask = np.zeros(flux.shape, dtype=int)
rdata = np.ones((nspec, 5, nwave))
frame = Frame(wave, flux, ivar, mask, rdata)
self.assertTrue(np.all(frame.wave == wave))
self.assertTrue(np.all(frame.flux == flux))
self.assertTrue(np.all(frame.ivar == ivar))
self.assertTrue(np.all(frame.resolution_data == rdata))
self.assertEqual(frame.nspec, nspec)
self.assertEqual(frame.nwave, nwave)
self.assertTrue(isinstance(frame.R[0], Resolution))
#- check dimensionality mismatches
self.assertRaises(AssertionError, lambda x: Frame(*x), (wave, wave, ivar, mask, rdata))
self.assertRaises(AssertionError, lambda x: Frame(*x), (wave, flux[0:2], ivar, mask, rdata))
#- Check constructing with defaults
frame = Frame(wave, flux, ivar)
self.assertEqual(frame.flux.shape, frame.mask.shape)
#- Check usage of fibers inputs
fibers = np.arange(nspec)
frame = Frame(wave, flux, ivar, fibers=fibers)
frame = Frame(wave, flux, ivar, fibers=fibers*2)
manyfibers = np.arange(2*nspec)
self.assertRaises(ValueError, lambda x: Frame(*x), (wave, flux, ivar, None, None, None, manyfibers))
#- Check usage of spectrograph input
for i in range(3):
frame = Frame(wave, flux, ivar, spectrograph=i)
self.assertEqual(len(frame.fibers), nspec)
self.assertEqual(frame.fibers[0], i*nspec)
def asframe(self, wavelength=None):
"""
Converts QFrame to a Frame
"""
if wavelength is None:
dwave = np.min(np.gradient(self.wave[self.nspec // 2]))
wmin = np.max(self.wave[:, 0])
wmax = np.min(self.wave[:, -1])
n = int((wmax - wmin) / dwave) + 1
wavelength = np.linspace(wmin, wmax, n)
rflux = np.zeros((self.nspec, wavelength.size))
rivar = np.zeros((self.nspec, wavelength.size))
if self.mask is None:
for i in range(self.nspec):
rflux[i], rivar[i] = resample_flux(wavelength,
self.wave[i],
self.flux[i],
self.ivar[i],
extrapolate=False)
else:
for i in range(self.nspec):
rflux[i], rivar[i] = resample_flux(wavelength,
self.wave[i],
self.flux[i],
self.ivar[i] *
(self.mask[i] == 0),
extrapolate=False)
return Frame(wave=wavelength,flux=rflux,ivar=rivar,mask=None,resolution_data=None,\
fibers=self.fibers, spectrograph=None, meta=self.meta, fibermap=self.fibermap,\
chi2pix=None,scores=None,scores_comments=None,\
wsigma=self.sigma,ndiag=1, suppress_res_warning=True)
def test_throughput(self):
"""
Test that spectra with different throughputs but the same resolution
produce a fiberflat mirroring the variations in throughput
"""
wave, flux, ivar, mask = _get_data()
nspec, nwave = flux.shape
#- Setup data for a Resolution matrix
sigma = 4.0
ndiag = 21
xx = np.linspace(-(ndiag - 1) / 2.0, +(ndiag - 1) / 2.0, ndiag)
Rdata = np.zeros((nspec, ndiag, nwave))
kernel = np.exp(-xx**2 / (2 * sigma))
kernel /= sum(kernel)
for i in range(nspec):
for j in range(nwave):
Rdata[i, :, j] = kernel
#- Vary the input flux prior to calculating the fiber flat
flux[1] *= 1.1
flux[2] *= 1.2
flux[3] *= 0.8
#- Convolve with the (common) resolution matrix
convflux = np.empty_like(flux)
for i in range(nspec):
convflux[i] = Resolution(Rdata[i]).dot(flux[i])
frame = Frame(wave, convflux, ivar, mask, Rdata, spectrograph=0)
ff = compute_fiberflat(frame)
#- flux[1] is brighter, so should fiberflat[1]. etc.
self.assertTrue(np.allclose(ff.fiberflat[0], ff.fiberflat[1] / 1.1))
self.assertTrue(np.allclose(ff.fiberflat[0], ff.fiberflat[2] / 1.2))
self.assertTrue(np.allclose(ff.fiberflat[0], ff.fiberflat[3] / 0.8))
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):
# Set up the logger
if args.verbose:
log = get_logger(DEBUG)
else:
log = get_logger()
# Make sure all necessary environment variables are set
DESI_SPECTRO_REDUX_DIR = "./quickGen"
if 'DESI_SPECTRO_REDUX' not in os.environ:
log.info('DESI_SPECTRO_REDUX environment is not set.')
else:
DESI_SPECTRO_REDUX_DIR = os.environ['DESI_SPECTRO_REDUX']
if os.path.exists(DESI_SPECTRO_REDUX_DIR):
if not os.path.isdir(DESI_SPECTRO_REDUX_DIR):
raise RuntimeError("Path %s Not a directory" %
DESI_SPECTRO_REDUX_DIR)
else:
try:
os.makedirs(DESI_SPECTRO_REDUX_DIR)
except:
raise
SPECPROD_DIR = 'specprod'
if 'SPECPROD' not in os.environ:
log.info('SPECPROD environment is not set.')
else:
SPECPROD_DIR = os.environ['SPECPROD']
prod_Dir = specprod_root()
if os.path.exists(prod_Dir):
if not os.path.isdir(prod_Dir):
raise RuntimeError("Path %s Not a directory" % prod_Dir)
else:
try:
os.makedirs(prod_Dir)
except:
raise
# Initialize random number generator to use.
np.random.seed(args.seed)
random_state = np.random.RandomState(args.seed)
# Derive spectrograph number from nstart if needed
if args.spectrograph is None:
args.spectrograph = args.nstart / 500
# Read fibermapfile to get object type, night and expid
if args.fibermap:
log.info("Reading fibermap file {}".format(args.fibermap))
fibermap = read_fibermap(args.fibermap)
objtype = get_source_types(fibermap)
stdindx = np.where(objtype == 'STD') # match STD with STAR
mwsindx = np.where(objtype == 'MWS_STAR') # match MWS_STAR with STAR
bgsindx = np.where(objtype == 'BGS') # match BGS with LRG
objtype[stdindx] = 'STAR'
objtype[mwsindx] = 'STAR'
objtype[bgsindx] = 'LRG'
NIGHT = fibermap.meta['NIGHT']
EXPID = fibermap.meta['EXPID']
else:
# Create a blank fake fibermap
fibermap = empty_fibermap(args.nspec)
targetids = random_state.randint(2**62, size=args.nspec)
fibermap['TARGETID'] = targetids
night = get_night()
expid = 0
log.info("Initializing SpecSim with config {}".format(args.config))
desiparams = load_desiparams()
qsim = get_simulator(args.config, num_fibers=1)
if args.simspec:
# Read the input file
log.info('Reading input file {}'.format(args.simspec))
simspec = desisim.io.read_simspec(args.simspec)
nspec = simspec.nspec
if simspec.flavor == 'arc':
log.warning("quickgen doesn't generate flavor=arc outputs")
return
else:
wavelengths = simspec.wave
spectra = simspec.flux
if nspec < args.nspec:
log.info("Only {} spectra in input file".format(nspec))
args.nspec = nspec
else:
# Initialize the output truth table.
spectra = []
wavelengths = qsim.source.wavelength_out.to(u.Angstrom).value
npix = len(wavelengths)
truth = dict()
meta = Table()
truth['OBJTYPE'] = np.zeros(args.nspec, dtype=(str, 10))
truth['FLUX'] = np.zeros((args.nspec, npix))
truth['WAVE'] = wavelengths
jj = list()
for thisobj in set(true_objtype):
ii = np.where(true_objtype == thisobj)[0]
nobj = len(ii)
truth['OBJTYPE'][ii] = thisobj
log.info('Generating {} template'.format(thisobj))
# Generate the templates
if thisobj == 'ELG':
elg = desisim.templates.ELG(wave=wavelengths,
add_SNeIa=args.add_SNeIa)
flux, tmpwave, meta1 = elg.make_templates(
nmodel=nobj,
seed=args.seed,
zrange=args.zrange_elg,
sne_rfluxratiorange=args.sne_rfluxratiorange)
elif thisobj == 'LRG':
lrg = desisim.templates.LRG(wave=wavelengths,
add_SNeIa=args.add_SNeIa)
flux, tmpwave, meta1 = lrg.make_templates(
nmodel=nobj,
seed=args.seed,
zrange=args.zrange_lrg,
sne_rfluxratiorange=args.sne_rfluxratiorange)
elif thisobj == 'QSO':
qso = desisim.templates.QSO(wave=wavelengths)
flux, tmpwave, meta1 = qso.make_templates(
nmodel=nobj, seed=args.seed, zrange=args.zrange_qso)
elif thisobj == 'BGS':
bgs = desisim.templates.BGS(wave=wavelengths,
add_SNeIa=args.add_SNeIa)
flux, tmpwave, meta1 = bgs.make_templates(
nmodel=nobj,
seed=args.seed,
zrange=args.zrange_bgs,
rmagrange=args.rmagrange_bgs,
sne_rfluxratiorange=args.sne_rfluxratiorange)
elif thisobj == 'STD':
std = desisim.templates.STD(wave=wavelengths)
flux, tmpwave, meta1 = std.make_templates(nmodel=nobj,
seed=args.seed)
elif thisobj == 'QSO_BAD': # use STAR template no color cuts
star = desisim.templates.STAR(wave=wavelengths)
flux, tmpwave, meta1 = star.make_templates(nmodel=nobj,
seed=args.seed)
elif thisobj == 'MWS_STAR' or thisobj == 'MWS':
mwsstar = desisim.templates.MWS_STAR(wave=wavelengths)
flux, tmpwave, meta1 = mwsstar.make_templates(nmodel=nobj,
seed=args.seed)
elif thisobj == 'WD':
wd = desisim.templates.WD(wave=wavelengths)
flux, tmpwave, meta1 = wd.make_templates(nmodel=nobj,
seed=args.seed)
elif thisobj == 'SKY':
flux = np.zeros((nobj, npix))
meta1 = Table(dict(REDSHIFT=np.zeros(nobj, dtype=np.float32)))
elif thisobj == 'TEST':
flux = np.zeros((args.nspec, npix))
indx = np.where(wave > 5800.0 - 1E-6)[0][0]
ref_integrated_flux = 1E-10
ref_cst_flux_density = 1E-17
single_line = (np.arange(args.nspec) % 2 == 0).astype(
np.float32)
continuum = (np.arange(args.nspec) % 2 == 1).astype(np.float32)
for spec in range(args.nspec):
flux[spec, indx] = single_line[
spec] * ref_integrated_flux / np.gradient(wavelengths)[
indx] # single line
flux[spec] += continuum[
spec] * ref_cst_flux_density # flat continuum
meta1 = Table(
dict(REDSHIFT=np.zeros(args.nspec, dtype=np.float32),
LINE=wave[indx] *
np.ones(args.nspec, dtype=np.float32),
LINEFLUX=single_line * ref_integrated_flux,
CONSTFLUXDENSITY=continuum * ref_cst_flux_density))
else:
log.fatal('Unknown object type {}'.format(thisobj))
sys.exit(1)
# Pack it in.
truth['FLUX'][ii] = flux
meta = vstack([meta, meta1])
jj.append(ii.tolist())
# Sanity check on units; templates currently return ergs, not 1e-17 ergs...
# assert (thisobj == 'SKY') or (np.max(truth['FLUX']) < 1e-6)
# Sort the metadata table.
jj = sum(jj, [])
meta_new = Table()
for k in range(args.nspec):
index = int(np.where(np.array(jj) == k)[0])
meta_new = vstack([meta_new, meta[index]])
meta = meta_new
# Add TARGETID and the true OBJTYPE to the metadata table.
meta.add_column(
Column(true_objtype, dtype=(str, 10), name='TRUE_OBJTYPE'))
meta.add_column(Column(targetids, name='TARGETID'))
# Rename REDSHIFT -> TRUEZ anticipating later table joins with zbest.Z
meta.rename_column('REDSHIFT', 'TRUEZ')
# explicitly set location on focal plane if needed to support airmass
# variations when using specsim v0.5
if qsim.source.focal_xy is None:
qsim.source.focal_xy = (u.Quantity(0, 'mm'), u.Quantity(100, 'mm'))
# Set simulation parameters from the simspec header or desiparams
bright_objects = ['bgs', 'mws', 'bright', 'BGS', 'MWS', 'BRIGHT_MIX']
gray_objects = ['gray', 'grey']
if args.simspec is None:
object_type = objtype
flavor = None
elif simspec.flavor == 'science':
object_type = None
flavor = simspec.header['PROGRAM']
else:
object_type = None
flavor = simspec.flavor
log.warning(
'Maybe using an outdated simspec file with flavor={}'.format(
flavor))
# Set airmass
if args.airmass is not None:
qsim.atmosphere.airmass = args.airmass
elif args.simspec and 'AIRMASS' in simspec.header:
qsim.atmosphere.airmass = simspec.header['AIRMASS']
else:
qsim.atmosphere.airmass = 1.25 # Science Req. Doc L3.3.2
# Set exptime
if args.exptime is not None:
qsim.observation.exposure_time = args.exptime * u.s
elif args.simspec and 'EXPTIME' in simspec.header:
qsim.observation.exposure_time = simspec.header['EXPTIME'] * u.s
elif objtype in bright_objects:
qsim.observation.exposure_time = desiparams['exptime_bright'] * u.s
else:
qsim.observation.exposure_time = desiparams['exptime_dark'] * u.s
# Set Moon Phase
if args.moon_phase is not None:
qsim.atmosphere.moon.moon_phase = args.moon_phase
elif args.simspec and 'MOONFRAC' in simspec.header:
qsim.atmosphere.moon.moon_phase = simspec.header['MOONFRAC']
elif flavor in bright_objects or object_type in bright_objects:
qsim.atmosphere.moon.moon_phase = 0.7
elif flavor in gray_objects:
qsim.atmosphere.moon.moon_phase = 0.1
else:
qsim.atmosphere.moon.moon_phase = 0.5
# Set Moon Zenith
if args.moon_zenith is not None:
qsim.atmosphere.moon.moon_zenith = args.moon_zenith * u.deg
elif args.simspec and 'MOONALT' in simspec.header:
qsim.atmosphere.moon.moon_zenith = simspec.header['MOONALT'] * u.deg
elif flavor in bright_objects or object_type in bright_objects:
qsim.atmosphere.moon.moon_zenith = 30 * u.deg
elif flavor in gray_objects:
qsim.atmosphere.moon.moon_zenith = 80 * u.deg
else:
qsim.atmosphere.moon.moon_zenith = 100 * u.deg
# Set Moon - Object Angle
if args.moon_angle is not None:
qsim.atmosphere.moon.separation_angle = args.moon_angle * u.deg
elif args.simspec and 'MOONSEP' in simspec.header:
qsim.atmosphere.moon.separation_angle = simspec.header[
'MOONSEP'] * u.deg
elif flavor in bright_objects or object_type in bright_objects:
qsim.atmosphere.moon.separation_angle = 50 * u.deg
elif flavor in gray_objects:
qsim.atmosphere.moon.separation_angle = 60 * u.deg
else:
qsim.atmosphere.moon.separation_angle = 60 * u.deg
# Initialize per-camera output arrays that will be saved
waves, trueflux, noisyflux, obsivar, resolution, sflux = {}, {}, {}, {}, {}, {}
maxbin = 0
nmax = args.nspec
for camera in qsim.instrument.cameras:
# Lookup this camera's resolution matrix and convert to the sparse
# format used in desispec.
R = Resolution(camera.get_output_resolution_matrix())
resolution[camera.name] = np.tile(R.to_fits_array(),
[args.nspec, 1, 1])
waves[camera.name] = (camera.output_wavelength.to(
u.Angstrom).value.astype(np.float32))
nwave = len(waves[camera.name])
maxbin = max(maxbin, len(waves[camera.name]))
nobj = np.zeros((nmax, 3, maxbin)) # object photons
nsky = np.zeros((nmax, 3, maxbin)) # sky photons
nivar = np.zeros((nmax, 3, maxbin)) # inverse variance (object+sky)
cframe_observedflux = np.zeros(
(nmax, 3, maxbin)) # calibrated object flux
cframe_ivar = np.zeros(
(nmax, 3, maxbin)) # inverse variance of calibrated object flux
cframe_rand_noise = np.zeros(
(nmax, 3, maxbin)) # random Gaussian noise to calibrated flux
sky_ivar = np.zeros((nmax, 3, maxbin)) # inverse variance of sky
sky_rand_noise = np.zeros(
(nmax, 3, maxbin)) # random Gaussian noise to sky only
frame_rand_noise = np.zeros(
(nmax, 3, maxbin)) # random Gaussian noise to nobj+nsky
trueflux[camera.name] = np.empty(
(args.nspec, nwave)) # calibrated flux
noisyflux[camera.name] = np.empty(
(args.nspec, nwave)) # observed flux with noise
obsivar[camera.name] = np.empty(
(args.nspec, nwave)) # inverse variance of flux
if args.simspec:
for i in range(10):
cn = camera.name + str(i)
if cn in simspec.cameras:
dw = np.gradient(simspec.cameras[cn].wave)
break
else:
raise RuntimeError(
'Unable to find a {} camera in input simspec'.format(
camera))
else:
sflux = np.empty((args.nspec, npix))
#- Check if input simspec is for a continuum flat lamp instead of science
#- This does not convolve to per-fiber resolution
if args.simspec:
if simspec.flavor == 'flat':
log.info("Simulating flat lamp exposure")
for i, camera in enumerate(qsim.instrument.cameras):
channel = camera.name #- from simspec, b/r/z not b0/r1/z9
assert camera.output_wavelength.unit == u.Angstrom
num_pixels = len(waves[channel])
phot = list()
for j in range(10):
cn = camera.name + str(j)
if cn in simspec.cameras:
camwave = simspec.cameras[cn].wave
dw = np.gradient(camwave)
phot.append(simspec.cameras[cn].phot)
if len(phot) == 0:
raise RuntimeError(
'Unable to find a {} camera in input simspec'.format(
camera))
else:
phot = np.vstack(phot)
meanspec = resample_flux(waves[channel], camwave,
np.average(phot / dw, axis=0))
fiberflat = random_state.normal(loc=1.0,
scale=1.0 / np.sqrt(meanspec),
size=(nspec, num_pixels))
ivar = np.tile(meanspec, [nspec, 1])
mask = np.zeros((simspec.nspec, num_pixels), dtype=np.uint32)
for kk in range((args.nspec + args.nstart - 1) // 500 + 1):
camera = channel + str(kk)
outfile = desispec.io.findfile('fiberflat', NIGHT, EXPID,
camera)
start = max(500 * kk, args.nstart)
end = min(500 * (kk + 1), nmax)
if (args.spectrograph <= kk):
log.info(
"Writing files for channel:{}, spectrograph:{}, spectra:{} to {}"
.format(channel, kk, start, end))
ff = FiberFlat(waves[channel],
fiberflat[start:end, :],
ivar[start:end, :],
mask[start:end, :],
meanspec,
header=dict(CAMERA=camera))
write_fiberflat(outfile, ff)
filePath = desispec.io.findfile("fiberflat", NIGHT, EXPID,
camera)
log.info("Wrote file {}".format(filePath))
sys.exit(0)
# Repeat the simulation for all spectra
fluxunits = 1e-17 * u.erg / (u.s * u.cm**2 * u.Angstrom)
for j in range(args.nspec):
thisobjtype = objtype[j]
sys.stdout.flush()
if flavor == 'arc':
qsim.source.update_in('Quickgen source {0}'.format, 'perfect',
wavelengths * u.Angstrom,
spectra * fluxunits)
else:
qsim.source.update_in('Quickgen source {0}'.format(j),
thisobjtype.lower(),
wavelengths * u.Angstrom,
spectra[j, :] * fluxunits)
qsim.source.update_out()
qsim.simulate()
qsim.generate_random_noise(random_state)
for i, output in enumerate(qsim.camera_output):
assert output['observed_flux'].unit == 1e17 * fluxunits
# Extract the simulation results needed to create our uncalibrated
# frame output file.
num_pixels = len(output)
nobj[j, i, :num_pixels] = output['num_source_electrons'][:, 0]
nsky[j, i, :num_pixels] = output['num_sky_electrons'][:, 0]
nivar[j, i, :num_pixels] = 1.0 / output['variance_electrons'][:, 0]
# Get results for our flux-calibrated output file.
cframe_observedflux[
j, i, :num_pixels] = 1e17 * output['observed_flux'][:, 0]
cframe_ivar[
j,
i, :num_pixels] = 1e-34 * output['flux_inverse_variance'][:, 0]
# Fill brick arrays from the results.
camera = output.meta['name']
trueflux[camera][j][:] = 1e17 * output['observed_flux'][:, 0]
noisyflux[camera][j][:] = 1e17 * (
output['observed_flux'][:, 0] +
output['flux_calibration'][:, 0] *
output['random_noise_electrons'][:, 0])
obsivar[camera][j][:] = 1e-34 * output['flux_inverse_variance'][:,
0]
# Use the same noise realization in the cframe and frame, without any
# additional noise from sky subtraction for now.
frame_rand_noise[
j, i, :num_pixels] = output['random_noise_electrons'][:, 0]
cframe_rand_noise[j, i, :num_pixels] = 1e17 * (
output['flux_calibration'][:, 0] *
output['random_noise_electrons'][:, 0])
# The sky output file represents a model fit to ~40 sky fibers.
# We reduce the variance by a factor of 25 to account for this and
# give the sky an independent (Gaussian) noise realization.
sky_ivar[
j,
i, :num_pixels] = 25.0 / (output['variance_electrons'][:, 0] -
output['num_source_electrons'][:, 0])
sky_rand_noise[j, i, :num_pixels] = random_state.normal(
scale=1.0 / np.sqrt(sky_ivar[j, i, :num_pixels]),
size=num_pixels)
armName = {"b": 0, "r": 1, "z": 2}
for channel in 'brz':
#Before writing, convert from counts/bin to counts/A (as in Pixsim output)
#Quicksim Default:
#FLUX - input spectrum resampled to this binning; no noise added [1e-17 erg/s/cm2/s/Ang]
#COUNTS_OBJ - object counts in 0.5 Ang bin
#COUNTS_SKY - sky counts in 0.5 Ang bin
num_pixels = len(waves[channel])
dwave = np.gradient(waves[channel])
nobj[:, armName[channel], :num_pixels] /= dwave
frame_rand_noise[:, armName[channel], :num_pixels] /= dwave
nivar[:, armName[channel], :num_pixels] *= dwave**2
nsky[:, armName[channel], :num_pixels] /= dwave
sky_rand_noise[:, armName[channel], :num_pixels] /= dwave
sky_ivar[:, armName[channel], :num_pixels] /= dwave**2
# Now write the outputs in DESI standard file system. None of the output file can have more than 500 spectra
# Looping over spectrograph
for ii in range((args.nspec + args.nstart - 1) // 500 + 1):
start = max(500 * ii,
args.nstart) # first spectrum for a given spectrograph
end = min(500 * (ii + 1),
nmax) # last spectrum for the spectrograph
if (args.spectrograph <= ii):
camera = "{}{}".format(channel, ii)
log.info(
"Writing files for channel:{}, spectrograph:{}, spectra:{} to {}"
.format(channel, ii, start, end))
num_pixels = len(waves[channel])
# Write frame file
framefileName = desispec.io.findfile("frame", NIGHT, EXPID,
camera)
frame_flux=nobj[start:end,armName[channel],:num_pixels]+ \
nsky[start:end,armName[channel],:num_pixels] + \
frame_rand_noise[start:end,armName[channel],:num_pixels]
frame_ivar = nivar[start:end, armName[channel], :num_pixels]
sh1 = frame_flux.shape[
0] # required for slicing the resolution metric, resolusion matrix has (nspec,ndiag,wave)
# for example if nstart =400, nspec=150: two spectrographs:
# 400-499=> 0 spectrograph, 500-549 => 1
if (args.nstart == start):
resol = resolution[channel][:sh1, :, :]
else:
resol = resolution[channel][-sh1:, :, :]
# must create desispec.Frame object
frame=Frame(waves[channel], frame_flux, frame_ivar,\
resolution_data=resol, spectrograph=ii, \
fibermap=fibermap[start:end], \
meta=dict(CAMERA=camera, FLAVOR=simspec.flavor) )
desispec.io.write_frame(framefileName, frame)
framefilePath = desispec.io.findfile("frame", NIGHT, EXPID,
camera)
log.info("Wrote file {}".format(framefilePath))
if args.frameonly or simspec.flavor == 'arc':
continue
# Write cframe file
cframeFileName = desispec.io.findfile("cframe", NIGHT, EXPID,
camera)
cframeFlux = cframe_observedflux[
start:end,
armName[channel], :num_pixels] + cframe_rand_noise[
start:end, armName[channel], :num_pixels]
cframeIvar = cframe_ivar[start:end,
armName[channel], :num_pixels]
# must create desispec.Frame object
cframe = Frame(waves[channel], cframeFlux, cframeIvar, \
resolution_data=resol, spectrograph=ii,
fibermap=fibermap[start:end],
meta=dict(CAMERA=camera, FLAVOR=simspec.flavor) )
desispec.io.frame.write_frame(cframeFileName, cframe)
cframefilePath = desispec.io.findfile("cframe", NIGHT, EXPID,
camera)
log.info("Wrote file {}".format(cframefilePath))
# Write sky file
skyfileName = desispec.io.findfile("sky", NIGHT, EXPID, camera)
skyflux=nsky[start:end,armName[channel],:num_pixels] + \
sky_rand_noise[start:end,armName[channel],:num_pixels]
skyivar = sky_ivar[start:end, armName[channel], :num_pixels]
skymask = np.zeros(skyflux.shape, dtype=np.uint32)
# must create desispec.Sky object
skymodel = SkyModel(waves[channel],
skyflux,
skyivar,
skymask,
header=dict(CAMERA=camera))
desispec.io.sky.write_sky(skyfileName, skymodel)
skyfilePath = desispec.io.findfile("sky", NIGHT, EXPID, camera)
log.info("Wrote file {}".format(skyfilePath))
# Write calib file
calibVectorFile = desispec.io.findfile("calib", NIGHT, EXPID,
camera)
flux = cframe_observedflux[start:end,
armName[channel], :num_pixels]
phot = nobj[start:end, armName[channel], :num_pixels]
calibration = np.zeros_like(phot)
jj = (flux > 0)
calibration[jj] = phot[jj] / flux[jj]
#- TODO: what should calibivar be?
#- For now, model it as the noise of combining ~10 spectra
calibivar = 10 / cframe_ivar[start:end,
armName[channel], :num_pixels]
#mask=(1/calibivar>0).astype(int)??
mask = np.zeros(calibration.shape, dtype=np.uint32)
# write flux calibration
fluxcalib = FluxCalib(waves[channel], calibration, calibivar,
mask)
write_flux_calibration(calibVectorFile, fluxcalib)
calibfilePath = desispec.io.findfile("calib", NIGHT, EXPID,
camera)
log.info("Wrote file {}".format(calibfilePath))
def main(args=None):
'''
Converts simspec -> frame files; see fastframe --help for usage options
'''
#- TODO: use desiutil.log
if isinstance(args, (list, tuple, type(None))):
args = parse(args)
print('Reading files')
simspec = desisim.io.read_simspec(args.simspec)
if simspec.flavor == 'arc':
print('arc exposure; no frames to output')
return
fibermap = simspec.fibermap
obs = simspec.obs
night = simspec.header['NIGHT']
expid = simspec.header['EXPID']
firstspec = args.firstspec
nspec = min(args.nspec, len(fibermap) - firstspec)
print('Simulating spectra {}-{}'.format(firstspec, firstspec + nspec))
wave = simspec.wave['brz']
flux = simspec.flux
ii = slice(firstspec, firstspec + nspec)
if simspec.flavor == 'science':
sim = desisim.simexp.simulate_spectra(wave,
flux[ii],
fibermap=fibermap[ii],
obsconditions=obs,
dwave_out=1.0)
elif simspec.flavor in ['arc', 'flat', 'calib']:
x = fibermap['X_TARGET']
y = fibermap['Y_TARGET']
fiber_area = desisim.simexp.fiber_area_arcsec2(fibermap['X_TARGET'],
fibermap['Y_TARGET'])
surface_brightness = (flux.T / fiber_area).T
config = desisim.simexp._specsim_config_for_wave(wave, dwave_out=1.0)
# sim = specsim.simulator.Simulator(config, num_fibers=nspec)
sim = desisim.specsim.get_simulator(config, num_fibers=nspec)
sim.observation.exposure_time = simspec.header['EXPTIME'] * u.s
sbunit = 1e-17 * u.erg / (u.Angstrom * u.s * u.cm**2 * u.arcsec**2)
xy = np.vstack([x, y]).T * u.mm
sim.simulate(calibration_surface_brightness=surface_brightness[ii] *
sbunit,
focal_positions=xy[ii])
else:
raise ValueError('Unknown simspec flavor {}'.format(simspec.flavor))
sim.generate_random_noise()
for i, results in enumerate(sim.camera_output):
results = sim.camera_output[i]
wave = results['wavelength']
phot = (results['num_source_electrons'] + \
results['num_sky_electrons'] + \
results['num_dark_electrons'] + \
results['random_noise_electrons']).T
ivar = 1.0 / results['variance_electrons'].T
R = Resolution(
sim.instrument.cameras[i].get_output_resolution_matrix())
Rdata = np.tile(R.data.T, nspec).T.reshape(nspec, R.data.shape[0],
R.data.shape[1])
assert np.all(Rdata[0] == R.data)
assert phot.shape == (nspec, len(wave))
for spectro in range(10):
imin = max(firstspec, spectro * 500) - firstspec
imax = min(firstspec + nspec, (spectro + 1) * 500) - firstspec
if imax <= imin:
continue
xphot = phot[imin:imax]
xivar = ivar[imin:imax]
xfibermap = fibermap[ii][imin:imax]
camera = '{}{}'.format(sim.camera_names[i], spectro)
meta = simspec.header.copy()
meta['CAMERA'] = camera
frame = Frame(wave,
xphot,
xivar,
resolution_data=Rdata[0:imax - imin],
spectrograph=spectro,
fibermap=xfibermap,
meta=meta)
outfile = desispec.io.findfile('frame',
night,
expid,
camera,
outdir=args.outdir)
print('writing {}'.format(outfile))
desispec.io.write_frame(outfile, frame)
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