From bbc3f6a337561b517f441fea6eadf0f7e22c916c Mon Sep 17 00:00:00 2001 From: John Moustakas Date: Fri, 24 Feb 2017 17:05:32 -0500 Subject: [PATCH 1/2] refactor lya_spectra to support optional inputs --- py/desisim/lya_spectra.py | 157 ++++++++++++++++++++------------------ 1 file changed, 84 insertions(+), 73 deletions(-) diff --git a/py/desisim/lya_spectra.py b/py/desisim/lya_spectra.py index dc83a4fe8..e4e524bd5 100644 --- a/py/desisim/lya_spectra.py +++ b/py/desisim/lya_spectra.py @@ -1,93 +1,104 @@ -from desisim.templates import QSO -from desisim.io import empty_metatable -import scipy as sp -from scipy import interpolate -from speclite import filters -import numpy as np -import astropy.table - -def get_spectra(infile, first=0, nqso=None, seed=None): - - ''' - read input lyman-alpha absorption files, - return normalized spectra including forest absorption. +""" +desisim.lya_spectra +=================== - Args: - infile: name of the lyman-alpha spectra file +Function to simulate a QSO spectrum including Lyman-alpha absorption. +""" - Options: - first: first spectrum to read - nqso: number of spectra to read - seed: random seed for template generation +from __future__ import division, print_function - returns (flux, wave, meta) - spectra with forest absorption normalized according to the magnitude - - Note: - meta is metadata table from QSO continuum template generation - plus RA and DEC columns - ''' +def get_spectra(lyafile, templateid=None, wave=None, normfilter='sdss2010-g', + rand=None, qso=None): + '''Generate a QSO spectrum which includes Lyman-alpha absorption. - import fitsio - h = fitsio.FITS(infile) - if nqso is None: - nqso = len(h)-1 + Args: + lyafile (str): name of the Lyman-alpha spectrum file to read. + templateid (int numpy.ndarray, optional): indices of the spectra + (0-indexed) to read from LYAFILE (default is to read everything). + wave (numpy.ndarray, optional): desired output wavelength vector. + normfilter (str, optional): normalization filter + rand (numpy.RandomState, optional): RandomState object used for the + random number generation. + qso (desisim.templates.QSO, optional): object with which to generate + individual spectra/templates. + + Returns: + flux (numpy.ndarray): Array [nmodel, npix] of observed-frame spectra + (erg/s/cm2/A). + wave (numpy.ndarray): Observed-frame [npix] wavelength array (Angstrom). + meta (astropy.Table): Table of meta-data [nmodel] for each output spectrum + with columns defined in desisim.io.empty_metatable *plus* RA, DEC. - if first<0: - raise ValueError("first must be >= 0") + ''' + import numpy as np + from scipy.interpolate import interp1d - heads = [head.read_header() for head in h[first+1:first+1+nqso]] + import fitsio - zqso = [head["ZQSO"] for head in heads] - ra = [head["RA"] for head in heads] - dec = [head["DEC"] for head in heads] - mag_g = [head["MAG_G"] for head in heads] + from speclite.filters import load_filters + from desisim.templates import QSO + from desisim.io import empty_metatable - assert(len(zqso) == nqso) + h = fitsio.FITS(lyafile) + if templateid is None: + nqso = len(h)-1 + else: + nqso = len(templateid) - rand = sp.random.RandomState(seed) + if rand is None: + rand = np.random.RandomState() seed = rand.randint(2**32, size=nqso) - input_meta = empty_metatable(objtype='QSO', nmodel=1) - - filter_name = 'sdss2010-g' - - normfilt = filters.load_filters(filter_name) - qso = QSO(normfilter=filter_name) - - flux = np.zeros([nqso, len(qso.wave)], dtype='f4') - meta = None - for i,head in enumerate(h[first+1:first+1+nqso]): - f, wave, meta_qso = qso.make_templates(nmodel=1, - redshift=np.array([zqso[i]]), mag=np.array([mag_g[i]]), seed=seed[i]) - - meta_qso['TEMPLATEID'] = first + i + 1 - if meta is None: - meta = meta_qso.copy() - else: - meta = astropy.table.vstack([meta, meta_qso]) + #heads = [head.read_header() for head in h[templateid + 1]] + heads = [] + for indx in templateid: + heads.append(h[indx + 1].read_header()) + + zqso = np.array([head['ZQSO'] for head in heads]) + ra = np.array([head['RA'] for head in heads]) + dec = np.array([head['DEC'] for head in heads]) + mag_g = np.array([head['MAG_G'] for head in heads]) + + # Hard-coded filtername! + normfilt = load_filters(normfilter) + + if qso is None: + qso = QSO(normfilter=normfilter, wave=wave) + + wave = qso.wave + flux = np.zeros([nqso, len(wave)], dtype='f4') + + meta = empty_metatable(objtype='QSO', nmodel=nqso) + meta['TEMPLATEID'] = templateid + meta['REDSHIFT'] = zqso + meta['MAG'] = mag_g + meta['SEED'] = seed + meta['RA'] = ra + meta['DEC'] = dec + + for ii, indx in enumerate(templateid): + flux1, _, meta1 = qso.make_templates(nmodel=1, redshift=np.array([zqso[ii]]), + mag=np.array([mag_g[ii]]), seed=seed[ii]) # read lambda and forest transmission - la = head["LAMBDA"][:] - tr = head["FLUX"][:] + data = h[indx + 1].read() + la = data['LAMBDA'][:] + tr = data['FLUX'][:] + if len(tr): - # will interpolate the transmission at the spectral wavelengths, - # if outside the forest, the transmission is 1 - itr = interpolate.interp1d(la,tr,bounds_error=False,fill_value=1) - f *= itr(wave) + # Interpolate the transmission at the spectral wavelengths, if + # outside the forest, the transmission is 1. + itr = interp1d(la, tr, bounds_error=False, fill_value=1.0) + flux1 *= itr(wave) - padflux, padwave = normfilt.pad_spectrum(f, wave, method='edge') - normmaggies = sp.array(normfilt.get_ab_maggies(padflux, padwave, - mask_invalid=True)[filter_name]) - f *= 10**(-0.4 * mag_g[i]) / normmaggies - flux[i,:] = f[:] + padflux, padwave = normfilt.pad_spectrum(flux1, wave, method='edge') + normmaggies = np.array(normfilt.get_ab_maggies(padflux, padwave, + mask_invalid=True)[normfilter]) + flux1 *= 10**(-0.4 * mag_g[ii]) / normmaggies + flux[ii, :] = flux1[:] h.close() - # Add RA,DEC to output meta - meta['RA'] = ra - meta['DEC'] = dec - - return flux,wave,meta + return flux, wave, meta From 339bbf5b91491a59225c6f6268604b66650e4ad4 Mon Sep 17 00:00:00 2001 From: John Moustakas Date: Sun, 26 Feb 2017 08:58:41 -0500 Subject: [PATCH 2/2] update tests --- py/desisim/lya_spectra.py | 41 +++++++++++++++++++++++-------------- py/desisim/test/test_lya.py | 30 +++++++++++++++++---------- 2 files changed, 45 insertions(+), 26 deletions(-) diff --git a/py/desisim/lya_spectra.py b/py/desisim/lya_spectra.py index e4e524bd5..151ae133d 100644 --- a/py/desisim/lya_spectra.py +++ b/py/desisim/lya_spectra.py @@ -7,27 +7,33 @@ from __future__ import division, print_function -def get_spectra(lyafile, templateid=None, wave=None, normfilter='sdss2010-g', - rand=None, qso=None): +def get_spectra(lyafile, nqso=None, wave=None, templateid=None, normfilter='sdss2010-g', + seed=None, rand=None, qso=None): '''Generate a QSO spectrum which includes Lyman-alpha absorption. Args: lyafile (str): name of the Lyman-alpha spectrum file to read. - templateid (int numpy.ndarray, optional): indices of the spectra - (0-indexed) to read from LYAFILE (default is to read everything). + nqso (int, optional): number of spectra to generate (starting from the + first spectrum; if more flexibility is needed use TEMPLATEID). wave (numpy.ndarray, optional): desired output wavelength vector. + templateid (int numpy.ndarray, optional): indices of the spectra + (0-indexed) to read from LYAFILE (default is to read everything). If + provided together with NQSO, TEMPLATEID wins. normfilter (str, optional): normalization filter + seed (int, optional): Seed for random number generator. rand (numpy.RandomState, optional): RandomState object used for the - random number generation. + random number generation. If provided together with SEED, this + optional input superseeds the numpy.RandomState object instantiated by + SEED. qso (desisim.templates.QSO, optional): object with which to generate individual spectra/templates. Returns: - flux (numpy.ndarray): Array [nmodel, npix] of observed-frame spectra - (erg/s/cm2/A). - wave (numpy.ndarray): Observed-frame [npix] wavelength array (Angstrom). - meta (astropy.Table): Table of meta-data [nmodel] for each output spectrum - with columns defined in desisim.io.empty_metatable *plus* RA, DEC. + flux (numpy.ndarray): Array [nmodel, npix] of observed-frame spectra + (erg/s/cm2/A). + wave (numpy.ndarray): Observed-frame [npix] wavelength array (Angstrom). + meta (astropy.Table): Table of meta-data [nmodel] for each output spectrum + with columns defined in desisim.io.empty_metatable *plus* RA, DEC. ''' import numpy as np @@ -41,17 +47,22 @@ def get_spectra(lyafile, templateid=None, wave=None, normfilter='sdss2010-g', h = fitsio.FITS(lyafile) if templateid is None: - nqso = len(h)-1 + if nqso is None: + nqso = len(h)-1 + templateid = np.arange(nqso) else: + templateid = np.array(templateid) nqso = len(templateid) + print(templateid) if rand is None: - rand = np.random.RandomState() - seed = rand.randint(2**32, size=nqso) + rand = np.random.RandomState(seed) + templateseed = rand.randint(2**32, size=nqso) #heads = [head.read_header() for head in h[templateid + 1]] heads = [] for indx in templateid: + print(indx + 1) heads.append(h[indx + 1].read_header()) zqso = np.array([head['ZQSO'] for head in heads]) @@ -72,13 +83,13 @@ def get_spectra(lyafile, templateid=None, wave=None, normfilter='sdss2010-g', meta['TEMPLATEID'] = templateid meta['REDSHIFT'] = zqso meta['MAG'] = mag_g - meta['SEED'] = seed + meta['SEED'] = templateseed meta['RA'] = ra meta['DEC'] = dec for ii, indx in enumerate(templateid): flux1, _, meta1 = qso.make_templates(nmodel=1, redshift=np.array([zqso[ii]]), - mag=np.array([mag_g[ii]]), seed=seed[ii]) + mag=np.array([mag_g[ii]]), seed=templateseed[ii]) # read lambda and forest transmission data = h[indx + 1].read() diff --git a/py/desisim/test/test_lya.py b/py/desisim/test/test_lya.py index b30042bae..df08f614d 100644 --- a/py/desisim/test/test_lya.py +++ b/py/desisim/test/test_lya.py @@ -19,30 +19,38 @@ def setUpClass(cls): fx = fitsio.FITS(cls.infile) cls.nspec = len(fx) - 1 fx.close() - + cls.wavemin = 5000 + cls.wavemax = 8000 + cls.dwave = 2.0 + cls.wave = np.arange(cls.wavemin, cls.wavemax+cls.dwave/2, cls.dwave) + cls.nspec = 5 + cls.templateid = [3, 10, 500] + cls.seed = np.random.randint(2**32) + cls.rand = np.random.RandomState(cls.seed) + @unittest.skipIf(missing_fitsio, 'fitsio not installed; skipping lya_spectra tests') def test_read_lya(self): - flux, wave, meta = lya_spectra.get_spectra(self.infile) + flux, wave, meta = lya_spectra.get_spectra(self.infile, wave=self.wave, seed=self.seed) self.assertEqual(flux.shape[0], self.nspec) self.assertEqual(wave.shape[0], flux.shape[1]) self.assertEqual(len(meta), self.nspec) - nqso = 3 - flux, wave, meta = lya_spectra.get_spectra(self.infile, nqso=nqso) + flux, wave, meta = lya_spectra.get_spectra(self.infile, templateid=templateid, + wave=self.wave, seed=self.seed) self.assertEqual(flux.shape[0], nqso) self.assertEqual(wave.shape[0], flux.shape[1]) self.assertEqual(len(meta), nqso) - flux, wave, meta = lya_spectra.get_spectra(self.infile, nqso=nqso, first=2) - self.assertEqual(flux.shape[0], nqso) - self.assertEqual(wave.shape[0], flux.shape[1]) - self.assertEqual(len(meta), nqso) + #flux, wave, meta = lya_spectra.get_spectra(self.infile, nqso=nqso, first=2) + #self.assertEqual(flux.shape[0], nqso) + #self.assertEqual(wave.shape[0], flux.shape[1]) + #self.assertEqual(len(meta), nqso) @unittest.skipIf(missing_fitsio, 'fitsio not installed; skipping lya_spectra tests') def test_read_lya_seed(self): - flux1a, wave1a, meta1a = lya_spectra.get_spectra(self.infile, nqso=3, seed=1) - flux1b, wave1b, meta1b = lya_spectra.get_spectra(self.infile, nqso=3, seed=1) - flux2, wave2, meta2 = lya_spectra.get_spectra(self.infile, nqso=3, seed=2) + flux1a, wave1a, meta1a = lya_spectra.get_spectra(self.infile, wave=self.wave, nqso=3, seed=1) + flux1b, wave1b, meta1b = lya_spectra.get_spectra(self.infile, wave=self.wave, nqso=3, seed=1) + flux2, wave2, meta2 = lya_spectra.get_spectra(self.infile, wave=self.wave, nqso=3, seed=2) self.assertTrue(np.all(flux1a == flux1b)) self.assertTrue(np.any(flux1a != flux2))