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 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