diff --git a/pydl/goddard/astro.py b/pydl/goddard/astro.py index 9bedaf58..a94dd483 100644 --- a/pydl/goddard/astro.py +++ b/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 diff --git a/pydl/goddard/tests/test_goddard.py b/pydl/goddard/tests/test_goddard.py index b7cc0c88..b6d8f46a 100644 --- a/pydl/goddard/tests/test_goddard.py +++ b/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