Python xspectrum1d.XSpectrum1D类代码示例

def test_get_local_s2n():
    spec = XSpectrum1D.from_file(data_path('UM184_nF.fits'))
    wv0 = 4000 * u.AA
    s2n, sig_s2n = spec.get_local_s2n(wv0, 20, flux_th=0.9)
    np.testing.assert_allclose(s2n, 9.30119800567627, rtol=1e-5)
    np.testing.assert_allclose(sig_s2n, 1.0349911451339722, rtol=1e-5)
    # test with continuum
    spec.co = np.ones_like(spec.flux)
    s2n, sig_s2n = spec.get_local_s2n(wv0, 20, flux_th=0.9)
    np.testing.assert_allclose(s2n, 10.330545425415039, rtol=1e-5)
    np.testing.assert_allclose(sig_s2n, 0.4250050187110901, rtol=1e-5)
    # test errors
    # out of range
    with pytest.raises(IOError):
        spec.get_local_s2n(1215*u.AA, 20)
    # sig not defined
    spec = XSpectrum1D.from_tuple((spec.wavelength, spec.flux))
    with pytest.raises(ValueError):
        spec.get_local_s2n(wv0, 20)
    # bad shape for flux_th
    with pytest.raises(ValueError):
        spec.get_local_s2n(wv0, 20, flux_th=np.array([1,2,3,4,5]))
    # npix too big
    with pytest.raises(ValueError):
        spec.get_local_s2n(wv0, 1 + len(spec.wavelength))

def xspectrum1d_from_mpdaf_spec(sp, airvac='air'):
    """Gets a XSpectrum1D object in vacuum from an MPDAF Spectrum"""
    nomask = ~sp.mask
    fl = sp.data[nomask]
    er = np.sqrt(sp.var[nomask])
    wv = sp.wave.coord()[nomask]
    meta = dict(airvac=airvac)
    spec = XSpectrum1D.from_tuple((wv, fl, er), meta=meta)
    spec.airtovac()
    spec2 = XSpectrum1D.from_tuple((spec.wavelength, fl, er))
    return spec2

def test_wvmnx():
    npix = 1000
    # Without sig
    spec = XSpectrum1D.from_tuple((np.linspace(5000.,6000,npix), np.ones(npix)))
    assert spec.wvmin.value == 5000.
    assert spec.wvmax.value == 6000.
    # With sig
    spec = XSpectrum1D.from_tuple((np.linspace(5000.,6000,npix), np.ones(npix),
                                   np.ones(npix)*0.1))
    assert spec.wvmin.value == 5000.
    assert spec.wvmax.value == 6000.

def test_mask_edge():
    wave = 3000. + np.arange(1000)
    flux = np.ones_like(wave)
    sig = 0.1*np.ones_like(wave)
    sig[900:] = 0.
    sig[800:825] = 0.
    wave[900:] = 0.  # WARNING, the data are sorted first!
    #
    spec = XSpectrum1D.from_tuple((wave,flux,sig), masking='edges')
    assert len(spec.wavelength) == 900
    spec2 = XSpectrum1D.from_tuple((wave,flux,sig), masking='all')
    assert len(spec2.wavelength) == 875

def test_errors():

    # from_tuple
    try:
        spec = XSpectrum1D.from_tuple('this_is_not_a_tuple')
    except IOError:
        pass
    try:
        n_tuple = np.array([np.ones(5), np.ones(5)]), np.ones(5)
        spec = XSpectrum1D.from_tuple(n_tuple)
    except IOError:
        pass

    # wrong instances
    flux = [1,2,3]
    wave = [1,2,3]
    try:
        spec = XSpectrum1D(wave, flux)
    except IOError:
        pass

    #wrong shapes
    flux = np.ones(5)
    wave = np.array([1,2,3])
    try:
        spec = XSpectrum1D(wave, flux)
    except IOError:
        pass
    try:
        spec = XSpectrum1D(wave, np.ones(len(wave)), sig=np.ones(2))
    except IOError:
        pass
    try:
        spec = XSpectrum1D(wave, np.ones(len(wave)), co=np.ones(2), verbose = True) # test verbose here too
    except IOError:
        pass

    # wrong masking
    try:
        spec = XSpectrum1D(wave, np.ones(len(wave)), masking = 'wrong_masking')
    except IOError:
        pass

    #wrong units input
    try:
        spec = XSpectrum1D(wave, np.ones(len(wave)), units = 'not_a_dict')
    except IOError:
        pass
    try:
        spec = XSpectrum1D(wave, np.ones(len(wave)), units =dict(wrong_key=2))
    except IOError:
        pass

def test_from_tuple():
    idl = ascii.read(data_path('UM184.dat.gz'), names=['wave', 'flux', 'sig'])
    spec = XSpectrum1D.from_tuple((idl['wave'],idl['flux'],idl['sig']))
    #
    np.testing.assert_allclose(spec.dispersion.value, idl['wave'])
    np.testing.assert_allclose(spec.sig, idl['sig'], atol=2e-3, rtol=0)

    assert spec.dispersion.unit == u.Unit('AA')
    #
    spec = XSpectrum1D.from_tuple((idl['wave'],idl['flux']))
    np.testing.assert_allclose(spec.dispersion.value, idl['wave'])
    # continuum
    co = np.ones_like(idl['flux'])
    spec = XSpectrum1D.from_tuple((idl['wave'],idl['flux'],idl['sig'], co))
    np.testing.assert_allclose(spec.dispersion.value, idl['wave'])

def test_readwrite_meta_as_dicts(spec):
    sp = XSpectrum1D.from_tuple((np.array([5,6,7]), np.ones(3), np.ones(3)*0.1))
    sp.meta['headers'][0] = dict(a=1, b='abc')
    sp2 = XSpectrum1D.from_tuple((np.array([8,9,10]), np.ones(3), np.ones(3)*0.1))
    sp2.meta['headers'][0] = dict(c=2, d='efg')
    spec = ltsu.collate([sp,sp2])
    # Write
    spec.write_to_fits(data_path('tmp.fits'))
    spec.write_to_hdf5(data_path('tmp.hdf5'))
    # Read and test
    newspec = io.readspec(data_path('tmp.hdf5'))
    assert newspec.meta['headers'][0]['a'] == 1
    assert newspec.meta['headers'][0]['b'] == 'abc'
    newspec2 = io.readspec(data_path('tmp.fits'))
    assert 'METADATA' in newspec2.meta['headers'][0].keys()

def compare_s2n(pp,lrdx_sciobj,pypit_boxfile, iso):
    '''Compare boxcar S/N
    '''
    # Read/Load
    pypit_boxspec = lsio.readspec(pypit_boxfile)
    # Read LowRedux
    sig = np.sqrt(lrdx_sciobj['MASK_BOX']/(lrdx_sciobj['SIVAR_BOX'] + (lrdx_sciobj['MASK_BOX']==0)))
    lwrdx_boxspec = XSpectrum1D.from_tuple( (lrdx_sciobj['WAVE_BOX'], lrdx_sciobj['FLUX_BOX'], sig) )

    # Plot
    plt.clf()
    fig = plt.figure(figsize=(16,7))
    fig.suptitle("Instr={:s}, Setup={:s} :: Boxcar S/N for {:s} :: PYPIT ({:s})".format(iso[0], iso[1], iso[2], pypit.version), fontsize=18.)
    ax = plt.gca()
    ymax = np.median(pypit_boxspec.flux)*2.
    # PYPIT
    gdpy = pypit_boxspec.sig > 0.
    pys2n =  pypit_boxspec.flux[gdpy]/pypit_boxspec.sig[gdpy]
    ax.plot(pypit_boxspec.dispersion[gdpy],pys2n, 'k-', drawstyle='steps', label='PYPIT')
    # LowRedux
    gdlx = lwrdx_boxspec.sig > 0.
    ax.plot(lwrdx_boxspec.dispersion[gdlx], lwrdx_boxspec.flux[gdlx]/lwrdx_boxspec.sig[gdlx], 
            '-', color='blue', label='LowRedux')
    # Axes
    ax.set_xlim(np.min(pypit_boxspec.dispersion.value), np.max(pypit_boxspec.dispersion.value))
    ax.set_ylim(0.,np.median(pys2n)*2.)
    ax.set_xlabel('Wavelength',fontsize=17.)
    ax.set_ylabel('S/N per pixel',fontsize=17.)
    # Legend
    legend = plt.legend(loc='upper right', borderpad=0.3,
                handletextpad=0.3, fontsize='x-large')
    # Finish
    plt.tight_layout(pad=0.2,h_pad=0.,w_pad=0.1,rect=[0, 0.03, 1, 0.95])
    pp.savefig(bbox_inches='tight')
    plt.close()

def test_from_file():
    spec = XSpectrum1D.from_file(data_path('UM184_nF.fits'))
    idl = ascii.read(data_path('UM184.dat.gz'), names=['wave', 'flux', 'sig'])

    np.testing.assert_allclose(spec.dispersion.value, idl['wave'])
    np.testing.assert_allclose(spec.sig, idl['sig'], atol=2e-3, rtol=0)

    assert spec.dispersion.unit == u.Unit('AA')

def parse_linetools_spectrum_format(hdulist):
    """ Parse an old linetools-format spectrum from an hdulist

    Parameters
    ----------
    hdulist : FITS HDU list

    Returns
    -------
    xspec1d : XSpectrum1D
      Parsed spectrum

    """
    if 'WAVELENGTH' not in hdulist:
        pdb.set_trace()
        #spec1d = spec_read_fits.read_fits_spectrum1d(
        #    os.path.expanduser(datfil), dispersion_unit='AA')
        xspec1d = XSpectrum1D.from_spec1d(spec1d)
    else:
        wave = hdulist['WAVELENGTH'].data * u.AA
        fx = hdulist['FLUX'].data

    # Error array
    if 'ERROR' in hdulist:
        sig = hdulist['ERROR'].data
    else:
        sig = None

    if 'CONTINUUM' in hdulist:
        co = hdulist['CONTINUUM'].data
    else:
        co = None

    xspec1d = XSpectrum1D.from_tuple((wave, fx, sig, co))

    if 'METADATA' in hdulist[0].header:
        # import pdb; pdb.set_trace()
        # patch for reading continuum metadata; todo: should be fixed properly!!!
        if "contpoints" in hdulist[0].header['METADATA']:
            aux_s = hdulist[0].header['METADATA']
            if aux_s.endswith("}\' /"):
                aux_s = aux_s[:-3]  # delete these extra characters
                hdulist[0].header['METADATA'] = aux_s
        xspec1d.meta.update(json.loads(hdulist[0].header['METADATA']))

    return xspec1d

def wfc3_continuum(wfc3_indx=None, zqso=0., wave=None, smooth=3., NHI_max=17.5, rstate=None):
    '''Use the WFC3 data + models from O'Meara+13 to generate a continuum

    Parameters
    ----------
    wfc3_indx : int, optional
      Index of WFC3 data to use
    zqso : float, optional
      Redshift of the QSO
    wave : Quantity array, optional
      Wavelengths to rebin on
    smooth : float, optional
      Number of pixels to smooth on
    NHI_max : float, optional
      Maximum NHI for the sightline

    Returns
    -------
    wfc3_continuum : XSpectrum1D 
       of the continuum
    idx : int
      Index of the WFC3 spectrum used    
    '''
    # Random number
    if rstate is None:
        rstate = np.random.RandomState()
    # Open
    wfc3_models_hdu = fits.open(os.getenv('DROPBOX_DIR')+'XQ-100/LLS/wfc3_conti_models.fits')
    nwfc3 = len(wfc3_models_hdu)-1
    # Load up models
    wfc_models = []
    for ii in range(1,nwfc3-1):
        wfc_models.append( Table(wfc3_models_hdu[ii].data) )
    # Grab a random one
    if wfc3_indx is None:
        need_c = True
        while(need_c):
            idx = rstate.randint(0,nwfc3-1)
            if wfc_models[idx]['TOTNHI'] > NHI_max:
                continue
            if wfc_models[idx]['QSO'] in ['J122836.05+510746.2', 'J122015.50+460802.4']:
                continue # These QSOs are NG
            need_c=False
    else:
        idx = wfc3_indx

    # Generate spectrum
    wfc_spec = XSpectrum1D.from_tuple( (wfc_models[idx]['WREST'].flatten()*(1+zqso), 
        wfc_models[idx]['FLUX'].flatten()) )
    # Smooth
    wfc_smooth = wfc_spec.gauss_smooth(fwhm=smooth)

    # Rebin?
    if wave is not None:
        wfc_rebin = wfc_smooth.rebin(wave)
        return wfc_rebin, idx
    else:
        return wfc_smooth, idx

 def load_spec(self):
     '''Input the Spectrum
     '''
     from linetools.spectra.xspectrum1d import XSpectrum1D
     if self._specfil is None:
         self.get_specfil()
     #
     if self.verbose:
         print('SdssQso: Loading spectrum from {:s}'.format(self._specfil))
     self.spec = XSpectrum1D.from_file(self._specfil)

def main(args):
    from scipy.io.idl import readsav
    from linetools.spectra.xspectrum1d import XSpectrum1D

    # Read
    lrdx_sky = readsav(args.lowrdx_sky)
    # Generate
    xspec = XSpectrum1D.from_tuple((lrdx_sky['wave_calib'], lrdx_sky['sky_calib']))
    # Write
    xspec.write_to_fits(args.new_file)

def test_addmask():
    spec = XSpectrum1D.from_file(data_path('UM184_nF.fits'))
    assert not spec.data['flux'][0].mask[100]
    mask = spec.data['flux'][0].mask.copy()
    mask[100:110] = True
    spec.add_to_mask(mask)
    assert spec.data['flux'][0].mask[100]
    # Compressed
    badp = spec.flux < 0.1
    spec.add_to_mask(badp, compressed=True)
    assert np.sum(spec.data['flux'][0].mask) > 3000

def test_copy(spec):
    # From existing
    spec2 = spec.copy()
    assert spec.wavelength[0] == spec2.wavelength[0]
    assert spec.flux[-1] == spec2.flux[-1]
    #
    wave = np.arange(3000., 6500)
    npix = len(wave)
    spect = XSpectrum1D.from_tuple((wave*u.AA,np.ones(npix)))
    specf = spect.copy()
    assert specf.sig_is_set is False