Merge pull request #85 from crawfordsm/normalization

stitching

pyhrs/__init__.py

@@ -20,3 +20,4 @@
     from .hrstools import *
     from .extraction import *
     from .wavelengthsolution import *
+    from .velocitycorrection import *

pyhrs/extraction.py

@@ -13,13 +13,15 @@
 from astropy import stats
 
 from scipy import ndimage as nd
+from scipy import signal
 
 from .hrstools import *
 from .hrsorder import HRSOrder
 
 
-__all__=['extract_order', 'simple_extract_order', 'extract_science', 'normalize_order', 
-         'normalize_spectra', 'stitch_spectra', 'polyfitr', 'extract_normalize']
+__all__=['extract_order', 'simple_extract_order', 'extract_science',
+         'normalize_order', 'normalize_spectra', 'stitch_order',
+         'stitch_spectra', 'polyfitr', 'extract_normalize', 'resample']
 
 
 def extract_order(ccd, order_frame, n_order, ws, shift_dict, y1=3, y2=10, order=None, target=True, interp=False):
@@ -479,35 +481,214 @@ def normalize_spectra(spectra_dict, model=mod.models.Chebyshev1D(2),
 
     return spectra_dict
 
-def stitch_spectra(spectra_dict, n_min, n_max, normalize=False, model=None, fitter=None):
+
+def stitch_order(csp, sp, m_init, fitter=mod.fitting.LinearLSQFitter()):
+    """Stitch a single order onto the rest of a spectra
+
+    Parameters
+    ----------
+
+    csp: list
+        A list containing the wavelength, flux, flux_err, and summed flux
+        for the base spectra that will be stitched to
+
+    sp: list
+        A list containing the wavelength, flux, flux_err, and summed flux for 
+        the spectra that will be stitched
+
+    m_init: astropy.modeling.models
+        Model to use to match the orders
+
+    fitter: astropy.modeling.fitter
+        Fitter for model to use to match the orders
+
+    Returns
+    -------
+    csp: list
+        A list containing the wavelength, flux, flux_err, and summed flux
+        now with the csp and sp spectra stitched together
+
+    """
+
+
+    if sp[0].min() < csp[0].min():
+        cmask = (csp[0] < sp[0].max())
+    else:
+        cmask = (csp[0] > sp[0].min())
+
+
+    cw = csp[0][cmask]
+
+    #interpolate the flux of the spectra
+    sf = np.interp(cw, sp[0], sp[1])
+    m = fitter(m_init, csp[0][cmask], csp[1][cmask]/sf)
+    norm = m(sp[0])
+
+    #repeat for the summed version
+    sf = np.interp(cw, sp[0], sp[3])
+    sm = fitter(m_init, csp[0][cmask], csp[1][cmask]/sf)
+    sum_norm = sm(sp[0])
+
+    #set the rest of the values to the value for the 
+    #last overlapping pixel in the center
+    if sp[0].min() < csp[0].min():
+        smask = (sp[0] < csp[0].min())
+        norm[smask] = m(csp[0][cmask][0])
+        sum_norm[smask] = sm(csp[0][cmask][0])
+    else:
+        smask = (sp[0] > csp[0].max())
+        norm[smask] = m(csp[0][cmask][-1])
+        sum_norm[smask] = sm(csp[0][cmask][-1])
+
+    csp = [np.concatenate([csp[0], sp[0]]), np.concatenate([csp[1], norm*sp[1]]), np.concatenate([csp[2], norm*sp[2]]), np.concatenate([csp[3], sum_norm*sp[3]])]
+    return csp
+
+
+def stitch_spectra(sp_dict, center_order=103, trim=0, m_init = mod.models.Polynomial1D(1), fitter=mod.fitting.LinearLSQFitter()):
     """Give a spectra, stitch the spectra together
 
     Parameters
     ----------
     spectra_dict: dict
         Dictionary containing wavelenghts and fluxes
 
-    normalize_order: bool
-        Normalize the individual orders
+    center_order: int
+        Order of the central order
+
+    trim: int
+        Number of pixels to trim
+
+    m_init: astropy.modeling.models
+        Model to use to match the orders
+
+    fitter: astropy.modeling.fitter
+        Fitter for model to use to match the orders
+
+    Returns
+    -------
+    wavelength: numpy.ndarray
+        Wavelength for the spectra
+
+    flux: numpy.ndarray
+        Averaged flux for the spectra
+
+    flux_err: numpy.ndarray
+        Error for the flux for the spectra
+
+    sum: numpy.ndarray
+        Summed flux for the spectra
   
     """
-    warr = None
-    for o in range(n_min, n_max):
-          w,f,e,s = spectra_dict[o]
-          if np.all(np.isnan(f)): continue
-          f[np.isnan(f)] = 0
-          if normalize:
-              f = normalize_order(w, f, model=model, fitter=fitter)
-          if warr is None:
-             warr = 1.0 * w
-             farr = 1.0 * f
-             earr = 1.0 * e
-             sarr = 1.0 * s
-          else:
-             warr = np.concatenate([warr, w])
-             farr = np.concatenate([farr, f])
-             earr = np.concatenate([farr, e])
-             sarr = np.concatenate([farr, s])
-
-    i = warr.argsort()
-    return warr[i], farr[i], earr[i], sarr[i]
+
+
+    # trim the order and replace any bad values in the spectra
+    keys = np.array(sp_dict.keys())
+    if trim>0:
+       nsp_dict = {}
+       for i in keys:
+           nsp_dict[i] = [x[trim:-trim] for x in sp_dict[i]]
+           f = nsp_dict[i][1]
+           f[np.isnan(f)] = 0
+           nsp_dict[i][1] = f
+    else:
+       nsp_dict = sp_dict
+
+    # set up the initial spectra
+    csp = nsp_dict[center_order]
+
+    # for each order to the left and right of the order, normalize it 
+    # to the order and then add it into that order
+    for i in np.arange(center_order+1, keys.max()+1, 1):
+        csp = stitch_order(csp, nsp_dict[i], m_init, fitter=fitter)
+    for i in np.arange(center_order-1, keys.min()-1, -1):
+        csp = stitch_order(csp, nsp_dict[i], m_init, fitter=fitter)
+    j = csp[0].argsort()
+    csp = [csp[0][j], csp[1][j], csp[2][j], csp[3]]
+    return csp
+
+
+def resample(warr, farr, earr, R=15000, dr=2.0, median_clean=0):
+    """Resample an array and sum pixels falling in the same resolution element
+
+    Parameters
+    ----------
+    warr: numpy.ndarray
+        Wavelength for the spectra
+
+    farr: numpy.ndarray
+        Averaged flux for the spectra
+
+    earr: numpy.ndarray
+        Error for the flux for the spectra   
+
+    R: float
+        Resolution of the observations
+
+    dr: float
+        Amount to subsample resolution element 
+
+    median_clean: int 
+        Size of the filter to use for median cleaning
+
+    Returns  
+    -------
+    warr: numpy.ndarray
+        Resampled Wavelength for the spectra
+
+    farr: numpy.ndarray
+        Resampled Averaged flux for the spectra
+
+    earr: numpy.ndarray
+        Resampled error for the flux for the spectra 
+
+    """
+
+    mask = np.isnan(farr)
+    warr[mask] = 0
+    farr[mask] = 0
+    earr[mask] = 0
+
+    w1 = warr.min()
+    w2 = warr.max()
+
+    #create the resampled wavelength array
+    res_element = warr/R/dr
+    w = w1
+    warr_list = [w]
+
+    while w < w2:
+       dw = np.interp(w, warr, res_element)
+       w = w+dw
+       warr_list.append(w)
+    wave = np.array(warr_list)
+
+    #clean the array
+    if median_clean:
+        mfarr = signal.medfilt(farr, kernel_size=31)
+        mask = (mfarr/farr < 1.1)
+        mask = np.convolve(1.0*(~mask), np.ones(median_clean), mode='same')
+        mask = (mask == 0)
+
+        warr = warr[mask]
+        farr = farr[mask]
+        earr = earr[mask]
+
+    #sum the flux
+    flux = np.zeros_like(wave)
+    err = np.zeros_like(wave)
+    for i in range(len(wave)-1):
+        mask = (warr > wave[i]) * (warr < wave[i+1]) * (farr>0)
+        if mask.sum()==0: continue
+        w = 1.0/earr[mask]**2
+        f = (farr[mask] * w).sum()/w.sum()
+        flux[i] =  f
+        err[i] = (1.0/(earr[mask]**-2).sum())**0.5
+        wave[i] = 0.5*(wave[i]+wave[i+1])
+
+    wave = wave[:-1]
+    flux = flux[:-1]
+    err = err[:-1]
+
+    return wave, flux, err
+
+

pyhrs/hrstools.py

@@ -18,7 +18,8 @@
 __all__ = ['background', 'fit_order', 'normalize_image', 'xcross_fit', 'ncor',
            'iterfit1D', 'calc_weights', 'match_lines', 'zeropoint_shift', 
            'clean_flatimage', 'mode_setup_information', 'write_spdict',
-           'fit_wavelength_solution', 'create_linelists', 'collapse_array']
+           'fit_wavelength_solution', 'create_linelists', 'collapse_array', 
+           'read_spdict']
 
 def background(b_arr, niter=3):
     """Determine the background for an array
@@ -905,3 +906,27 @@ def write_spdict(outfile, sp_dict, header=None):
     thdulist = fits.HDUList([prihdu, tbhdu])
     thdulist.writeto(outfile, clobber=True)
 
+
+def read_spdict(infile):
+    """Read in a spectral dictionary
+
+    Parameters
+    ----------
+    infile: str
+       Name of infile with spectra
+
+    Returns
+    -------
+    sp_dict: dict
+       Dictionary containing wavelength, flux, and error as a function of order
+
+    """
+
+    hdu = fits.open(infile)
+    data = hdu[1].data
+
+    sp_dict = {}
+    for o in np.arange(data['ORDER'].min(), data['ORDER'].max()+1):
+        m = (data['ORDER'] == o)
+        sp_dict[o] = [data['WAVELENGTH'][m], data['FLUX'][m], data['ERROR'][m], data['SUM'][m]]
+    return sp_dict

pyhrs/velocitycorrection.py

@@ -0,0 +1,59 @@
+from astropy.time import Time
+from astropy.coordinates import SkyCoord, solar_system, EarthLocation, ICRS, UnitSphericalRepresentation, CartesianRepresentation
+from astropy import units as u
+from astropy import constants as c
+from astropy import coordinates
+from PySpectrograph.Spectra.Spectrum import air2vac
+
+__all__=['velcorr', 'convert_data']
+
+def velcorr(time, skycoord, location=None):
+  """Barycentric velocity correction.
+  
+  Uses the ephemeris set with  ``astropy.coordinates.solar_system_ephemeris.set`` for corrections. 
+  For more information see `~astropy.coordinates.solar_system_ephemeris`.
+  
+  Parameters
+  ----------
+  time : `~astropy.time.Time`
+    The time of observation.
+  skycoord: `~astropy.coordinates.SkyCoord`
+    The sky location to calculate the correction for.
+  location: `~astropy.coordinates.EarthLocation`, optional
+    The location of the observatory to calculate the correction for.
+    If no location is given, the ``location`` attribute of the Time
+    object is used
+    
+  Returns
+  -------
+  vel_corr : `~astropy.units.Quantity`
+    The velocity correction to convert to Barycentric velocities. Should be added to the original
+    velocity.
+  """
+
+  if location is None:
+    if time.location is None:
+        raise ValueError('An EarthLocation needs to be set or passed '
+                         'in to calculate bary- or heliocentric '
+                         'corrections')
+    location = time.location
+
+  # ensure sky location is ICRS compatible
+  if not skycoord.is_transformable_to(ICRS()):
+    raise ValueError("Given skycoord is not transformable to the ICRS")
+
+  ep, ev = solar_system.get_body_barycentric_posvel('earth', time) # ICRS position and velocity of Earth's geocenter
+  op, ov = location.get_gcrs_posvel(time) # GCRS position and velocity of observatory
+  # ICRS and GCRS are axes-aligned. Can add the velocities
+  velocity = ev + ov # relies on PR5434 being merged
+
+  # get unit ICRS vector in direction of SkyCoord
+  sc_cartesian = skycoord.icrs.represent_as(UnitSphericalRepresentation).represent_as(CartesianRepresentation)
+  return sc_cartesian.dot(velocity).to(u.km/u.s) # similarly requires PR5434
+
+def convert_data(wave, vslr):
+   """Convert wavelenght array for vacumm wavlength and to the v_slr 
+     
+   """
+   wave = air2vac(wave)
+   return wave * (1+vslr/c.c)