use units in ccd_offset function

eispac/core/eisfitresult.py

@@ -1,11 +1,12 @@
-__all__ = ['EISFitResult', 'create_fit_dict']
-
 import copy
 from datetime import datetime
+
+import astropy.units as u
 import numpy as np
 from scipy.ndimage import shift as shift_img
 from scipy.interpolate import interp1d
 import matplotlib.pyplot as plt
+
 from eispac import __version__ as eispac_version
 import eispac.core.fitting_functions as fit_fns
 from eispac.core.eisfittemplate import EISFitTemplate
@@ -14,6 +15,9 @@
 from eispac.instr.calc_velocity import calc_velocity
 from eispac.instr.ccd_offset import ccd_offset
 
+__all__ = ['EISFitResult', 'create_fit_dict']
+
+
 class EISFitResult:
     """Object containing the results from fitting one or more EIS window spectra
 
@@ -707,9 +711,10 @@ def get_aspect_ratio(self):
     def shift2wave(self, array, wave=195.119):
         """Shift an array from this fit to the desired wavelength
         """
+        # FIXME: This function should actually use units
         this_wave = self.fit['wave_range'].mean()
         disp = np.zeros(len(array.shape))
-        disp[0] = ccd_offset(wave) - ccd_offset(this_wave)
+        disp[0] = ccd_offset(wave*u.Angstrom) - ccd_offset(this_wave*u.Angstrom)
         array = shift_img(array, disp)
         return array
 

eispac/instr/ccd_offset.py

@@ -1,11 +1,14 @@
 """
 Functions for computing pixel offsets on CCD
 """
+import astropy.units as u
+import numpy as np
+
 __all__ = ['ccd_offset']
 
-import numpy as np
 
-def ccd_offset(wavelength):
+@u.quantity_input
+def ccd_offset(wavelength:u.Angstrom) ->u.pixel:
     """Calculate the spatial offset of a line relative to He II 256 Å
 
     Spatial offset of the specified wavelength relative
@@ -53,20 +56,23 @@ def ccd_offset(wavelength):
     .. [young09] Young, P.R., et al., 2009, A&A, 495, `587 <https://doi.org/10.1051/0004-6361:200810143>`_
     """
     # Calculate the pixel offset using an equation based on the wavelengths of
-    # specific reference lines (Fe VIII 185.21, Si 275.35, and He II 256.32)
+    # specific reference lines (Fe VIII 185.21, Si VII 275.35, and He II 256.32)
 
     # Note_1: The y-scale of EIS is always given as 1 [arcsec]/[pixel]. Since
     #         this offset is an instrumental effect specific to EIS, there is
     #         not much point in generalizing this function for use with other
     #         y-scale values.
     # Note_2: The value of 18.5 accounts for the base offset between the
     #         shortwave and longwave CCDs.
-    grating_tilt = -0.0792  # [pixels]/[angstrom]
+    grating_tilt = -0.0792*u.pixel/u.Angstrom
     wavelength = np.atleast_1d(wavelength)
     # Calculate the offset for all lines in the shortwave band
-    offset = grating_tilt * (wavelength-185.21) + 18.5 + grating_tilt * (275.35 - 256.32)
+    short_long_base_offset = 18.5 * u.pixel
+    offset_fe_8 = wavelength-185.21*u.Angstrom
+    offset_si_7_he_2 = 275.35*u.Angstrom - 256.32*u.Angstrom
+    offset = grating_tilt*offset_fe_8 + short_long_base_offset + grating_tilt*offset_si_7_he_2
     # Find and calculate the offset for all lines in the longwave band
-    i_ = np.where(wavelength > 230)
-    offset[i_] = grating_tilt * (wavelength[i_] - 256.32)
+    is_long_band = np.where(wavelength > 230*u.Angstrom)
+    offset[is_long_band] = grating_tilt * (wavelength[is_long_band] - 256.32*u.Angstrom)
 
     return offset

eispac/instr/tests/test_ccd_offset.py

@@ -1,18 +1,19 @@
 import pytest
 
+import astropy.units as u
 import numpy as np
 
 from eispac.instr import ccd_offset
 
 
 @pytest.mark.parametrize('wavelength, idl_offset',[
-    (185.21, 16.992825),
-    ([185.21], 16.992825),
-    (np.array([185.21]), 16.992825),
-    (np.array([185.21, 195.119, 256.317, 258.375, 284.160]),
-     [16.992825, 16.208033,  0.000239, -0.162755, -2.204928])
+    (185.21*u.Angstrom, 16.992825*u.pixel),
+    ([185.21]*u.Angstrom, 16.992825*u.pixel),
+    (np.array([185.21])*u.Angstrom, 16.992825*u.pixel),
+    (np.array([185.21, 195.119, 256.317, 258.375, 284.160])*u.Angstrom,
+     [16.992825, 16.208033,  0.000239, -0.162755, -2.204928]*u.pixel)
 ])
 def test_ccd_offset(wavelength, idl_offset):
     offset = ccd_offset(wavelength)
     idl_offset = np.atleast_1d(idl_offset)
-    assert np.allclose(offset, idl_offset, atol=1e-5, rtol=0)
+    assert u.allclose(offset, idl_offset, atol=1e-5*u.pixel, rtol=0)