add unit support to airtovac

pydl/goddard/astro.py

@@ -25,22 +25,40 @@ def airtovac(air):
     * Values of wavelength below 2000 Å are not converted.
     """
     from numpy import zeros
+    from astropy.units import Angstrom
     try:
-        vacuum = zeros(air.shape, dtype=air.dtype) + air
-        g = vacuum < 2000.0
+        u = air.unit
     except AttributeError:
-        # Most likely, vacuum is simply a float.
+        u = None
+    try:
+        t = air.dtype
+    except AttributeError:
+        # Most likely, air is simply a float.
+        t = None
+    if t is None:
+        if air < 2000.0:
+            return air
         vacuum = air
+        a = air
         g = None
-        if air < 2000.0:
+    else:
+        try:
+            a = air.to(Angstrom).value
+        except AttributeError:
+            a = air
+        g = a < 2000.0
+        if g.all():
             return air
+        vacuum = zeros(air.shape, dtype=t) + a
     for k in range(2):
         sigma2 = (1.0e4/vacuum)**2
         fact = (1.0 + 5.792105e-2/(238.0185 - sigma2) +
                 1.67917e-3/(57.362 - sigma2))
-        vacuum = air * fact
+        vacuum = a * fact
     if g is not None:
-        vacuum[g] = air[g]
+        vacuum[g] = a[g]
+    if u is not None:
+        vacuum = (vacuum * Angstrom).to(u)
     return vacuum
 
 
@@ -155,19 +173,37 @@ def vactoair(vacuum):
     * Values of wavelength below 2000 Å are not converted.
     """
     from numpy import zeros
+    from astropy.units import Angstrom
     try:
-        air = zeros(vacuum.shape, dtype=vacuum.dtype)
-        g = vacuum < 2000.0
+        u = vacuum.unit
+    except AttributeError:
+        u = None
+    try:
+        t = vacuum.dtype
     except AttributeError:
         # Most likely, vacuum is simply a float.
+        t = None
+    if t is None:
+        if vacuum < 2000.0:
+            return vacuum
         air = vacuum
+        v = vacuum
         g = None
-        if vacuum < 2000.0:
-            return air
-    sigma2 = (1.0e4/vacuum)**2
+    else:
+        try:
+            v = vacuum.to(Angstrom).value
+        except AttributeError:
+            v = vacuum
+        g = v < 2000.0
+        if g.all():
+            return vacuum
+        air = zeros(vacuum.shape, dtype=t) + v
+    sigma2 = (1.0e4/v)**2
     fact = (1.0 + 5.792105e-2/(238.0185 - sigma2) +
             1.67917e-3/(57.362 - sigma2))
-    air = vacuum/fact
+    air = v / fact
     if g is not None:
-        air[g] = vacuum[g]
+        air[g] = v[g]
+    if u is not None:
+        air = (air * Angstrom).to(u)
     return air

pydl/goddard/tests/test_goddard.py

@@ -1,6 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 # -*- coding: utf-8 -*-
 import numpy as np
+import astropy.units as u
 from astropy.tests.helper import raises
 from ..astro import airtovac, gcirc, get_juldate, vactoair
 from ..math import flegendre
@@ -29,6 +30,8 @@ def test_airtovac(self):
                                      2100.6664, 2200.6868, 2300.7083]))
         vacuum = airtovac(6056.125)
         assert np.allclose(vacuum, 6057.8019)
+        vacuum = airtovac(np.array([1800.0, 1850.0, 1900.0]))
+        assert np.allclose(vacuum, np.array([1800.0, 1850.0, 1900.0]))
         #
         # Regression test for #8.
         #
@@ -37,6 +40,23 @@ def test_airtovac(self):
         assert np.allclose(vacuum,
                            np.array([[1800.0, 1900.0, 2000.6475],
                                      [2100.6664, 2200.6868, 2300.7083]]))
+        #
+        # Test with units.
+        #
+        air = air * u.Angstrom
+        vacuum = airtovac(air)
+        assert np.allclose(vacuum.value,
+                           np.array([1800.0, 1900.0, 2000.6475,
+                                     2100.6664, 2200.6868, 2300.7083]))
+        assert vacuum.unit is u.Angstrom
+        vacuum = airtovac(air.to(u.nm))
+        #
+        # Due to numeric funkiness, 2000 -> 1999.999999999 < 2000.
+        #
+        assert np.allclose(vacuum.value,
+                           np.array([180.0, 190.0, 200.0,
+                                     210.06664, 220.06868, 230.07083]))
+        assert vacuum.unit is u.nm
 
     def test_cirrange(self):
         ra1 = np.linspace(-4.0*np.pi, 4.0*np.pi, 100)
@@ -143,6 +163,8 @@ def test_vactoair(self):
         air = vactoair(vacuum)
         assert np.allclose(air, np.array([1800.0, 1900.0, 1999.3526,
                                           2099.3337, 2199.3133, 2299.2918]))
+        air = vactoair(np.array([1800.0, 1850.0, 1900.0]))
+        assert np.allclose(air, np.array([1800.0, 1850.0, 1900.0]))
         #
         # Regression test for #8.
         #
@@ -151,3 +173,20 @@ def test_vactoair(self):
         assert np.allclose(air,
                            np.array([[1800.0, 1900.0, 1999.3526],
                                      [2099.3337, 2199.3133, 2299.2918]]))
+        #
+        # Test with units.
+        #
+        vacuum = vacuum * u.Angstrom
+        air = vactoair(vacuum)
+        assert np.allclose(air.value,
+                           np.array([1800.0, 1900.0, 1999.3526,
+                                     2099.3337, 2199.3133, 2299.2918]))
+        assert air.unit is u.Angstrom
+        air = vactoair(vacuum.to(u.nm))
+        #
+        # Due to numeric funkiness, 2000 -> 1999.999999999 < 2000.
+        #
+        assert np.allclose(air.value,
+                           np.array([180.0, 190.0, 200.0,
+                                     209.93337, 219.93133, 229.92918]))
+        assert air.unit is u.nm