Add better dtype preservation in astronomy functions with numpy 2

.github/workflows/ci.yaml

@@ -13,10 +13,10 @@ jobs:
       fail-fast: true
       matrix:
         os: ["windows-latest", "ubuntu-latest", "macos-latest"]
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
         experimental: [false]
         include:
-          - python-version: "3.11"
+          - python-version: "3.12"
             os: "ubuntu-latest"
             experimental: true
 

.pre-commit-config.yaml

@@ -6,3 +6,15 @@ repos:
     hooks:
     - id: flake8
       additional_dependencies: [flake8-docstrings, flake8-debugger, flake8-bugbear]
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: 'v1.10.0'  # Use the sha / tag you want to point at
+  hooks:
+    - id: mypy
+      additional_dependencies:
+        - types-docutils
+        - types-pkg-resources
+        - types-PyYAML
+        - types-requests
+        - types-mock
+        - types-python-dateutil
+      args: ["--python-version", "3.10", "--ignore-missing-imports"]

continuous_integration/environment.yaml

@@ -19,7 +19,6 @@ dependencies:
   - coverage
   - codecov
   - behave
-  - mock
   - zarr
   - geoviews
   - pytest

doc/source/conf.py

@@ -11,6 +11,7 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 """Configurations for sphinx based documentation."""
+from __future__ import annotations
 
 import sys
 import os
@@ -69,7 +70,7 @@
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
-exclude_patterns = []
+exclude_patterns: list[str] = []
 
 # The reST default role (used for this markup: `text`) to use for all documents.
 # #default_role = None

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 @@
     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 @@
     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 @@
     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 @@
 
     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 @@
 
     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 @@
     #       "=" 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 @@
     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 @@
 
     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/__init__.py

@@ -20,25 +20,3 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """The tests package."""
-
-from pyorbital.tests import (test_aiaa, test_tlefile, test_orbital,
-                             test_astronomy, test_geoloc)
-import unittest
-
-
-def suite():
-    """The global test suite."""
-    mysuite = unittest.TestSuite()
-    # Test the documentation strings
-    # mysuite.addTests(doctest.DocTestSuite(image))
-    # Use the unittests also
-    mysuite.addTests(test_aiaa.suite())
-    mysuite.addTests(test_tlefile.suite())
-    mysuite.addTests(test_orbital.suite())
-    mysuite.addTests(test_astronomy.suite())
-    mysuite.addTests(test_geoloc.suite())
-    return mysuite
-
-
-if __name__ == '__main__':
-    unittest.TextTestRunner(verbosity=2).run(suite())