Merge pull request #212 from desihub/lyaread

update the API of lya_spectra to work with mock-->spectra connection

py/desisim/lya_spectra.py

@@ -1,93 +1,115 @@
-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.
-
-    Args:
-        infile: name of the lyman-alpha spectra file
-
-    Options:
-        first: first spectrum to read
-        nqso: number of spectra to read
-        seed: random seed for template generation
-
-    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
-    '''
-
-    import fitsio
-    h = fitsio.FITS(infile)
-    if nqso is None:
-        nqso = len(h)-1
-
-    if first<0:
-        raise ValueError("first must be >= 0")
-
-    heads = [head.read_header() for head in h[first+1:first+1+nqso]]
+"""
+desisim.lya_spectra
+===================
 
-    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]
+Function to simulate a QSO spectrum including Lyman-alpha absorption.
+"""
 
-    assert(len(zqso) == nqso)
+from __future__ import division, print_function
 
-    rand = sp.random.RandomState(seed)
-    seed = rand.randint(2**32, size=nqso)
+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.
 
-    input_meta = empty_metatable(objtype='QSO', nmodel=1)
-
-    filter_name = 'sdss2010-g'
+    Args:
+        lyafile (str): name of the Lyman-alpha spectrum file to read.
+        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.  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.
 
-    normfilt = filters.load_filters(filter_name)
-    qso = QSO(normfilter=filter_name)
+    '''
+    import numpy as np
+    from scipy.interpolate import interp1d
 
-    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])
+    import fitsio
 
-        meta_qso['TEMPLATEID'] = first + i + 1
-        if meta is None:
-            meta = meta_qso.copy()
-        else:
-            meta = astropy.table.vstack([meta, meta_qso])
+    from speclite.filters import load_filters
+    from desisim.templates import QSO
+    from desisim.io import empty_metatable
+
+    h = fitsio.FITS(lyafile)
+    if templateid is None:
+        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)
+    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])
+    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'] = 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=templateseed[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
 
 

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