diff --git a/pyorbital/astronomy.py b/pyorbital/astronomy.py index 3c212d4..b53bf12 100644 --- a/pyorbital/astronomy.py +++ b/pyorbital/astronomy.py @@ -1,30 +1,48 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- - +# # Copyright (c) 2011, 2013 - +# # Author(s): - +# # Martin Raspaud - +# # 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 . +"""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 @@ -32,26 +50,33 @@ 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,7 +132,7 @@ 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 @@ -115,8 +140,9 @@ def _local_hour_angle(utc_time, longitude, right_ascension): 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 diff --git a/pyorbital/tests/test_astronomy.py b/pyorbital/tests/test_astronomy.py index 17e8847..6eba33a 100644 --- a/pyorbital/tests/test_astronomy.py +++ b/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."""