Merge pull request #5 from bmorris3/moretests

Adding testing machinery for continuum normalization

aesop/spectra.py

@@ -45,7 +45,7 @@ def __init__(self, wavelength=None, flux=None, name=None, mask=None,
         """
         self.wavelength = wavelength
         self.wavelength_unit = wavelength.unit
-        self.flux = flux
+        self.flux = flux if hasattr(flux, 'unit') else u.Quantity(flux)
         self.name = name
         self.mask = mask
         self.wcs = wcs

aesop/tests/test_spectrum1d.py

@@ -6,8 +6,10 @@
 import astropy.units as u
 import pytest
 from astropy.units import UnitsError
+from astropy.tests.helper import remote_data
+from astropy.utils.data import download_file
 
-from ..spectra import Spectrum1D
+from ..spectra import Spectrum1D, EchelleSpectrum
 
 
 def test_constructor():
@@ -23,3 +25,63 @@ def test_constructor():
     with pytest.raises(UnitsError):
         example = Spectrum1D(wavelength=u.Quantity([0, 1, 2, 3], unit=u.day),
                              flux=f)
+
+
+@remote_data
+def test_read_fits():
+    url = ('http://staff.washington.edu/bmmorris/docs/'
+           'KIC8462852.0065.wfrmcpc.fits')
+    path = download_file(url, show_progress=False)
+
+    echelle_spectrum = EchelleSpectrum.from_fits(path)
+
+    assert hasattr(echelle_spectrum, 'header')
+    assert hasattr(echelle_spectrum, 'time')
+    assert hasattr(echelle_spectrum, 'name')
+    assert str(echelle_spectrum) == ('<EchelleSpectrum: 107 orders, '
+                                     '3506.8-10612.5 Angstrom>')
+
+    # There should be more flux in the 40th order (redder) than 0th (bluer)
+    assert echelle_spectrum[40].flux.mean() > echelle_spectrum[0].flux.mean()
+
+
+def generate_target_standard_pairs():
+    # realistic wavelength range, resolution for one order
+    wl = np.linspace(6000, 6117, 1650) * u.Angstrom
+    sig = 25 + 10 * (0.5 - np.random.rand())
+    amp = 500 + 200 * (0.5 - np.random.rand())
+    flux = amp * np.exp(-0.5 * (wl.value - wl.value.mean())**2 / sig**2)
+    flux += 10 * np.random.randn(flux.shape[0])
+
+    standard = Spectrum1D(wavelength=wl, flux=flux)
+
+    for i in range(4):
+        flux -= 100 * np.exp(-0.5 * (wl.value - wl.value.mean() -
+                                     20 * (i - 1.5))**2 / 0.7**2)
+
+    target = Spectrum1D(wavelength=wl, flux=flux)
+
+    return target, standard
+
+
+def test_continuum_norm():
+    target_orders = []
+    standard_orders = []
+    for i in range(10):
+        target, standard = generate_target_standard_pairs()
+        target_orders.append(target)
+        standard_orders.append(standard)
+
+    target_spectrum = EchelleSpectrum(target_orders)
+    standard_spectrum = EchelleSpectrum(standard_orders)
+
+    # Make sure that the continuum normalization is always good to 2%
+    polynomial_order = 8
+    target_spectrum.continuum_normalize(standard_spectrum, polynomial_order)
+
+    for order in target_spectrum.spectrum_list:
+        assert abs(np.median(order.flux) - 1) < 0.02
+
+    # Make sure spectral features are preserved in target spectra
+    for order in target_spectrum.spectrum_list:
+        assert np.mean(order.flux) < 1