Add type annotations and dtype consistency to astronomy functions

pyorbital/astronomy.py

@@ -1,57 +1,82 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-
+#
 # Copyright (c) 2011, 2013
-
+#
 # Author(s):
-
+#
 #   Martin Raspaud <[email protected]>
-
+#
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
 # the Free Software Foundation, either version 3 of the License, or
 # (at your option) any later version.
-
+#
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
-
+#
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""Angle and time-based astronomy functions.
 
-"""Astronomy module.
 Parts taken from http://www.geoastro.de/elevaz/basics/index.htm
+
+Note on argument types
+----------------------
+
+Many of these functions accept Python datetime objects,
+numpy datetime64 objects, or anything that can be turned
+into a numpy array of datetime64 objects. These objects are inherently
+64-bit so if other arguments (ex. longitude and latitude arrays) are
+32-bit floats internal operations will be automatically promoted to
+64-bit floating point numbers. Where possible these are then converted
+back to 32-bit before being returned. In general scalar inputs will also
+produce scalar outputs.
+
 """
+from __future__ import annotations
+
+import datetime
+from typing import TypeAlias
 
 import numpy as np
+import numpy.typing as npt
 
 from pyorbital import dt2np
 
 F = 1 / 298.257223563  # Earth flattening WGS-84
 A = 6378.137  # WGS84 Equatorial radius
 MFACTOR = 7.292115E-5
 
+ArrayOrFloat: TypeAlias = npt.ArrayLike | float
+# numpy datetime or python datetime
+DatetimeArrayLike: TypeAlias = npt.ArrayLike | np.datetime64 | datetime.datetime
+TimedeltaArrayLike: TypeAlias = npt.ArrayLike | np.timedelta64
 
-def jdays2000(utc_time):
+
+def jdays2000(utc_time: DatetimeArrayLike) -> np.ArrayLike[np.timedelta64]:
     """Get the days since year 2000.
     """
     return _days(dt2np(utc_time) - np.datetime64('2000-01-01T12:00'))
 
 
-def jdays(utc_time):
+def jdays(utc_time: DatetimeArrayLike) -> float:
     """Get the julian day of *utc_time*.
     """
-    return jdays2000(utc_time) + 2451545
+    return jdays2000(utc_time) + 2451545.0
 
 
-def _days(dt):
+def _days(dt: TimedeltaArrayLike) -> np.ArrayLike[np.float64]:
     """Get the days (floating point) from *d_t*.
     """
+    if hasattr(dt, "shape"):
+        dt = np.asanyarray(dt, dtype=np.timedelta64)
     return dt / np.timedelta64(1, 'D')
 
 
-def gmst(utc_time):
+def gmst(utc_time: DatetimeArrayLike) -> npt.ArrayLike[np.float64]:
     """Greenwich mean sidereal utc_time, in radians.
 
     As defined in the AIAA 2006 implementation:
@@ -63,14 +88,14 @@ def gmst(utc_time):
     return np.deg2rad(theta / 240.0) % (2 * np.pi)
 
 
-def _lmst(utc_time, longitude):
+def _lmst(utc_time: DatetimeArrayLike, longitude: ArrayOrFloat) -> npt.ArrayLike[np.float64]:
     """Local mean sidereal time, computed from *utc_time* and *longitude*.
     In radians.
     """
     return gmst(utc_time) + longitude
 
 
-def sun_ecliptic_longitude(utc_time):
+def sun_ecliptic_longitude(utc_time: datetime.datetime) -> npt.ArrayLike[np.float64]:
     """Ecliptic longitude of the sun at *utc_time*.
     """
     jdate = jdays2000(utc_time) / 36525.0
@@ -88,7 +113,7 @@ def sun_ecliptic_longitude(utc_time):
     return np.deg2rad(l__)
 
 
-def sun_ra_dec(utc_time):
+def sun_ra_dec(utc_time: datetime.datetime) -> tuple[npt.ArrayLike[np.float64], npt.ArrayLike[np.float64]]:
     """Right ascension and declination of the sun at *utc_time*.
     """
     jdate = jdays2000(utc_time) / 36525.0
@@ -107,16 +132,17 @@ def sun_ra_dec(utc_time):
     return right_ascension, declination
 
 
-def _local_hour_angle(utc_time, longitude, right_ascension):
+def _local_hour_angle(utc_time: DatetimeArrayLike, longitude: ArrayOrFloat, right_ascension: npt.ArrayLike[np.float64]) -> ArrayOrFloat:
     """Hour angle at *utc_time* for the given *longitude* and
     *right_ascension*
     longitude in radians
     """
     return _lmst(utc_time, longitude) - right_ascension
 
 
-def get_alt_az(utc_time, lon, lat):
+def get_alt_az(utc_time: DatetimeArrayLike, lon: ArrayOrFloat, lat: ArrayOrFloat) -> tuple[ArrayOrFloat, ArrayOrFloat]:
     """Return sun altitude and azimuth from *utc_time*, *lon*, and *lat*.
+
     lon,lat in degrees
     The returned angles are given in radians.
     """
@@ -125,13 +151,16 @@ def get_alt_az(utc_time, lon, lat):
 
     ra_, dec = sun_ra_dec(utc_time)
     h__ = _local_hour_angle(utc_time, lon, ra_)
-    return (np.arcsin(np.sin(lat) * np.sin(dec) +
-                      np.cos(lat) * np.cos(dec) * np.cos(h__)),
-            np.arctan2(-np.sin(h__), (np.cos(lat) * np.tan(dec) -
-                                      np.sin(lat) * np.cos(h__))))
+    alt_az = (np.arcsin(np.sin(lat) * np.sin(dec) +
+                        np.cos(lat) * np.cos(dec) * np.cos(h__)),
+              np.arctan2(-np.sin(h__), (np.cos(lat) * np.tan(dec) -
+                                        np.sin(lat) * np.cos(h__))))
+    if not isinstance(lon, float):
+        alt_az = (alt_az[0].astype(lon.dtype), alt_az[1].astype(lon.dtype))
+    return alt_az
 
 
-def cos_zen(utc_time, lon, lat):
+def cos_zen(utc_time: DatetimeArrayLike, lon: ArrayOrFloat, lat: ArrayOrFloat) -> ArrayOrFloat:
     """Cosine of the sun-zenith angle for *lon*, *lat* at *utc_time*.
     utc_time: datetime.datetime instance of the UTC time
     lon and lat in degrees.
@@ -141,21 +170,26 @@ def cos_zen(utc_time, lon, lat):
 
     r_a, dec = sun_ra_dec(utc_time)
     h__ = _local_hour_angle(utc_time, lon, r_a)
-    return (np.sin(lat) * np.sin(dec) + np.cos(lat) * np.cos(dec) * np.cos(h__))
+    csza = (np.sin(lat) * np.sin(dec) + np.cos(lat) * np.cos(dec) * np.cos(h__))
+    if not isinstance(lon, float):
+        csza = csza.astype(lon.dtype)
+    return csza
 
 
-def sun_zenith_angle(utc_time, lon, lat):
+def sun_zenith_angle(utc_time: DatetimeArrayLike, lon: ArrayOrFloat, lat: ArrayOrFloat) -> ArrayOrFloat:
     """Sun-zenith angle for *lon*, *lat* at *utc_time*.
     lon,lat in degrees.
     The angle returned is given in degrees
     """
-    return np.rad2deg(np.arccos(cos_zen(utc_time, lon, lat)))
+    sza = np.rad2deg(np.arccos(cos_zen(utc_time, lon, lat)))
+    if not isinstance(lon, float):
+        sza = sza.astype(lon.dtype)
+    return sza
 
 
-def sun_earth_distance_correction(utc_time):
+def sun_earth_distance_correction(utc_time: DatetimeArrayLike) -> ArrayOrFloat:
     """Calculate the sun earth distance correction, relative to 1 AU.
     """
-
     # Computation according to
     #  https://web.archive.org/web/20150117190838/http://curious.astro.cornell.edu/question.php?number=582
     # with
@@ -175,11 +209,12 @@ def sun_earth_distance_correction(utc_time):
     #       "=" 1 - 0.0167 * np.cos(theta)
 
     corr = 1 - 0.0167 * np.cos(2 * np.pi * (jdays2000(utc_time) - 3) / 365.25636)
-
     return corr
 
 
-def observer_position(time, lon, lat, alt):
+def observer_position(
+        utc_time: DatetimeArrayLike, lon: ArrayOrFloat, lat: ArrayOrFloat, alt: ArrayOrFloat
+) -> tuple[tuple[ArrayOrFloat, ArrayOrFloat, ArrayOrFloat], tuple[ArrayOrFloat, ArrayOrFloat, ArrayOrFloat]]:
     """Calculate observer ECI position.
 
     http://celestrak.com/columns/v02n03/
@@ -188,7 +223,7 @@ def observer_position(time, lon, lat, alt):
     lon = np.deg2rad(lon)
     lat = np.deg2rad(lat)
 
-    theta = (gmst(time) + lon) % (2 * np.pi)
+    theta = (gmst(utc_time) + lon) % (2 * np.pi)
     c = 1 / np.sqrt(1 + F * (F - 2) * np.sin(lat)**2)
     sq = c * (1 - F)**2
 
@@ -199,6 +234,35 @@ def observer_position(time, lon, lat, alt):
 
     vx = -MFACTOR * y  # kilometers/second
     vy = MFACTOR * x
-    vz = 0
-
+    vz = _float_to_sibling_result(0.0, vx)
+
+    if not isinstance(lon, float):
+        x = x.astype(lon.dtype, copy=False)
+        y = y.astype(lon.dtype, copy=False)
+        z = z.astype(lon.dtype, copy=False)
+        vx = vx.astype(lon.dtype, copy=False)
+        vy = vy.astype(lon.dtype, copy=False)
+        vz = vz.astype(lon.dtype, copy=False)  # type: ignore[union-attr]
     return (x, y, z), (vx, vy, vz)
+
+
+def _float_to_sibling_result(
+        result_to_convert: float,
+        template_result: ArrayOrFloat,
+) -> ArrayOrFloat:
+    """Convert a scalar to the same type as another return type.
+
+    This is mostly used to make a static value consistent with the types of
+    other returned values.
+
+    """
+    if isinstance(template_result, float):
+        return result_to_convert
+    # get any array like object that might be wrapped by our template (ex. xarray DataArray)
+    array_like = template_result if hasattr(template_result, "__array_function__") else template_result.data
+    array_convert = np.asarray(result_to_convert, like=array_like)
+    if not hasattr(template_result, "__array_function__"):
+        # the template result has some wrapper class (likely xarray DataArray)
+        # recreate the wrapper object
+        array_convert = template_result.__class__(array_convert)
+    return array_convert

pyorbital/tests/test_astronomy.py

@@ -22,11 +22,33 @@
 
 from datetime import datetime
 
+import dask.array as da
 import numpy as np
+import numpy.typing as npt
 import pytest
 
 import pyorbital.astronomy as astr
 
+try:
+    from xarray import DataArray
+except ImportError:
+    DataArray = None
+
+
+def _create_dask_array(input_list: list, dtype: npt.DTypeLike) -> da.Array:
+    np_arr = np.array(input_list, dtype=dtype)
+    return da.from_array(np_arr)
+
+
+def _create_xarray_numpy(input_list: list, dtype: npt.DTypeLike) -> DataArray:
+    np_arr = np.array(input_list, dtype=dtype)
+    return DataArray(np_arr)
+
+
+def _create_xarray_dask(input_list: list, dtype: npt.DTypeLike) -> DataArray:
+    dask_arr = _create_dask_array(input_list, dtype)
+    return DataArray(dask_arr)
+
 
 class TestAstronomy:
 
@@ -50,14 +72,30 @@ def test_jdays(self, dt, exp_jdays, exp_j2000):
             (0.0, 0.0, 1.8751916863323426),
         ]
     )
-    @pytest.mark.parametrize("dtype", [None, np.float32, np.float64])
-    def test_sunangles(self, lon, lat, exp_theta, dtype):
+    @pytest.mark.parametrize(
+        ("dtype", "array_construct"),
+        [
+            (None, None),
+            (np.float32, np.array),
+            (np.float64, np.array),
+            (np.float32, _create_dask_array),
+            (np.float64, _create_dask_array),
+            (np.float32, _create_xarray_numpy),
+            (np.float64, _create_xarray_numpy),
+            (np.float32, _create_xarray_dask),
+            (np.float64, _create_xarray_dask),
+        ]
+    )
+    def test_sunangles(self, lon, lat, exp_theta, dtype, array_construct):
         """Test the sun-angle calculations."""
+        if array_construct is None and dtype is not None:
+            pytest.skip(reason="Xarray dependency unavailable")
+
         time_slot = datetime(2011, 9, 23, 12, 0)
         abs_tolerance = 1e-8
         if dtype is not None:
-            lon = np.array([lon], dtype=dtype)
-            lat = np.array([lat], dtype=dtype)
+            lon = array_construct([lon], dtype=dtype)
+            lat = array_construct([lat], dtype=dtype)
             if np.dtype(dtype).itemsize < 8:
                 abs_tolerance = 1e-4
 
@@ -68,6 +106,7 @@ def test_sunangles(self, lon, lat, exp_theta, dtype):
         else:
             assert sun_theta.dtype == dtype
             np.testing.assert_allclose(sun_theta, exp_theta, atol=abs_tolerance)
+            assert isinstance(sun_theta, type(lon))
 
     def test_sun_earth_distance_correction(self):
         """Test the sun-earth distance correction."""