Implement 1D wavelength calibration classes (#162)

Author: rosteen

docs/conf.py

@@ -174,7 +174,8 @@
     {
         'astropy': ('https://docs.astropy.org/en/stable/', None),
         'ccdproc': ('https://ccdproc.readthedocs.io/en/stable/', None),
-        'specutils': ('https://specutils.readthedocs.io/en/stable/', None)
+        'specutils': ('https://specutils.readthedocs.io/en/stable/', None),
+        'gwcs': ('https://gwcs.readthedocs.io/en/stable/', None)
     }
 )
 #

docs/index.rst

@@ -50,6 +50,7 @@ Calibration
 .. toctree::
     :maxdepth: 1
 
+    wavelength_calibration.rst
     extinction.rst
     specphot_standards.rst
 

docs/wavelength_calibration.rst

@@ -0,0 +1,53 @@
+.. _wavelength_calibration:
+
+Wavelength Calibration
+======================
+
+Wavelength calibration is currently supported for 1D spectra. Given a list of spectral
+lines with known wavelengths and estimated pixel positions on an input calibration
+spectrum, you can currently use ``specreduce`` to:
+
+#. Fit an ``astropy`` model to the wavelength/pixel pairs to generate a spectral WCS
+   solution for the dispersion.
+#. Apply the generated spectral WCS to other `~specutils.Spectrum1D` objects.
+
+1D Wavelength Calibration
+-------------------------
+
+The `~specreduce.wavelength_calibration.WavelengthCalibration1D` class can be used
+to fit a dispersion model to a list of line positions and wavelengths. Future development
+will implement catalogs of known lamp spectra for use in matching observed lines. In the
+example below, the line positions (``pixel_centers``) have already been extracted from
+``lamp_spectrum``::
+
+    import astropy.units as u
+	 from specreduce import WavelengthCalibration1D
+	 pixel_centers = [10, 22, 31, 43]
+	 wavelengths = [5340, 5410, 5476, 5543]*u.AA
+	 test_cal = WavelengthCalibration1D(lamp_spectrum, line_pixels=pixel_centers,
+										line_wavelengths=wavelengths)
+    calibrated_spectrum = test_cal.apply_to_spectrum(science_spectrum)
+
+The example above uses the default model (`~astropy.modeling.functional_models.Linear1D`)
+to fit the input spectral lines, and then applies the calculated WCS solution to a second
+spectrum (``science_spectrum``). Any other 1D ``astropy`` model can be provided as the
+input ``model`` parameter to the `~specreduce.wavelength_calibration.WavelengthCalibration1D`.
+In the above example, the model fit and WCS construction is all done as part of the
+``apply_to_spectrum()`` call, but you could also access the `~gwcs.wcs.WCS` object itself
+by calling::
+
+    test_cal.wcs
+
+The calculated WCS is a cached property that will be cleared if the ``line_list``, ``model``,
+or ``input_spectrum`` properties are updated, since these will alter the calculated dispersion
+fit.
+
+You can also provide the input pixel locations and wavelengths of the lines as an
+`~astropy.table.QTable` with (at minimum) columns ``pixel_center`` and ``wavelength``,
+using the ``matched_line_list`` input argument::
+
+    from astropy.table import QTable
+    pixels = [10, 20, 30, 40]*u.pix
+    wavelength = [5340, 5410, 5476, 5543]*u.AA
+    line_list = QTable([pixels, wavelength], names=["pixel_center", "wavelength"])
+    test_cal = WavelengthCalibration1D(lamp_spectrum, matched_line_list=line_list)

specreduce/__init__.py

@@ -7,3 +7,4 @@
 # ----------------------------------------------------------------------------
 
 from specreduce.core import *  # noqa
+from specreduce.wavelength_calibration import * # noqa

specreduce/conftest.py

@@ -6,6 +6,10 @@
 import os
 
 from astropy.version import version as astropy_version
+import astropy.units as u
+import numpy as np
+import pytest
+from specutils import Spectrum1D
 
 # For Astropy 3.0 and later, we can use the standalone pytest plugin
 if astropy_version < '3.0':
@@ -20,6 +24,37 @@
         ASTROPY_HEADER = False
 
 
+@pytest.fixture
+def spec1d():
+    np.random.seed(7)
+    flux = np.random.random(50)*u.Jy
+    sa = np.arange(0, 50)*u.pix
+    spec = Spectrum1D(flux, spectral_axis=sa)
+    return spec
+
+
+@pytest.fixture
+def spec1d_with_emission_line():
+    np.random.seed(7)
+    sa = np.arange(0, 200)*u.pix
+    flux = (np.random.randn(200) +
+            10*np.exp(-0.01*((sa.value-130)**2)) +
+            sa.value/100) * u.Jy
+    spec = Spectrum1D(flux, spectral_axis=sa)
+    return spec
+
+
+@pytest.fixture
+def spec1d_with_absorption_line():
+    np.random.seed(7)
+    sa = np.arange(0, 200)*u.pix
+    flux = (np.random.randn(200) -
+            10*np.exp(-0.01*((sa.value-130)**2)) +
+            sa.value/100) * u.Jy
+    spec = Spectrum1D(flux, spectral_axis=sa)
+    return spec
+
+
 def pytest_configure(config):
 
     if ASTROPY_HEADER:

specreduce/tests/test_wavelength_calibration.py

@@ -0,0 +1,85 @@
+from numpy.testing import assert_allclose
+import pytest
+
+from astropy.table import QTable
+import astropy.units as u
+from astropy.modeling.models import Polynomial1D
+from astropy.modeling.fitting import LinearLSQFitter
+from astropy.tests.helper import assert_quantity_allclose
+
+from specreduce import WavelengthCalibration1D
+
+
+def test_linear_from_list(spec1d):
+    centers = [0, 10, 20, 30]
+    w = [5000, 5100, 5198, 5305]*u.AA
+    test = WavelengthCalibration1D(spec1d, line_pixels=centers, line_wavelengths=w)
+    spec2 = test.apply_to_spectrum(spec1d)
+
+    assert_quantity_allclose(spec2.spectral_axis[0], 4998.8*u.AA)
+    assert_quantity_allclose(spec2.spectral_axis[-1], 5495.169999*u.AA)
+
+
+def test_wavelength_from_table(spec1d):
+    centers = [0, 10, 20, 30]
+    w = [5000, 5100, 5198, 5305]*u.AA
+    table = QTable([w], names=["wavelength"])
+    WavelengthCalibration1D(spec1d, line_pixels=centers, line_wavelengths=table)
+
+
+def test_linear_from_table(spec1d):
+    centers = [0, 10, 20, 30]
+    w = [5000, 5100, 5198, 5305]*u.AA
+    table = QTable([centers, w], names=["pixel_center", "wavelength"])
+    test = WavelengthCalibration1D(spec1d, matched_line_list=table)
+    spec2 = test.apply_to_spectrum(spec1d)
+
+    assert_quantity_allclose(spec2.spectral_axis[0], 4998.8*u.AA)
+    assert_quantity_allclose(spec2.spectral_axis[-1], 5495.169999*u.AA)
+
+
+def test_poly_from_table(spec1d):
+    # This test is mostly to prove that you can use other models
+    centers = [0, 10, 20, 30, 40]
+    w = [5005, 5110, 5214, 5330, 5438]*u.AA
+    table = QTable([centers, w], names=["pixel_center", "wavelength"])
+
+    test = WavelengthCalibration1D(spec1d, matched_line_list=table,
+                                   model=Polynomial1D(2), fitter=LinearLSQFitter())
+    test.apply_to_spectrum(spec1d)
+
+    assert_allclose(test.model.parameters, [5.00477143e+03, 1.03457143e+01, 1.28571429e-02])
+
+
+def test_replace_spectrum(spec1d, spec1d_with_emission_line):
+    centers = [0, 10, 20, 30]*u.pix
+    w = [5000, 5100, 5198, 5305]*u.AA
+    test = WavelengthCalibration1D(spec1d, line_pixels=centers, line_wavelengths=w)
+    # Accessing this property causes fits the model and caches the resulting WCS
+    test.wcs
+    assert "wcs" in test.__dict__
+
+    # Replace the input spectrum, which should clear the cached properties
+    test.input_spectrum = spec1d_with_emission_line
+    assert "wcs" not in test.__dict__
+
+
+def test_expected_errors(spec1d):
+    centers = [0, 10, 20, 30, 40]
+    w = [5005, 5110, 5214, 5330, 5438]*u.AA
+    table = QTable([centers, w], names=["pixel_center", "wavelength"])
+
+    with pytest.raises(ValueError, match="Cannot specify line_wavelengths separately"):
+        WavelengthCalibration1D(spec1d, matched_line_list=table, line_wavelengths=w)
+
+    with pytest.raises(ValueError, match="must have the same length"):
+        w2 = [5005, 5110, 5214, 5330, 5438, 5500]*u.AA
+        WavelengthCalibration1D(spec1d, line_pixels=centers, line_wavelengths=w2)
+
+    with pytest.raises(ValueError, match="astropy.units.Quantity array or"
+                                         " as an astropy.table.QTable"):
+        w2 = [5005, 5110, 5214, 5330, 5438]
+        WavelengthCalibration1D(spec1d, line_pixels=centers, line_wavelengths=w2)
+
+    with pytest.raises(ValueError, match="specify at least one"):
+        WavelengthCalibration1D(spec1d, line_pixels=centers)