Refactor SZA and cos(SZA) generation to reduce duplicate computations

satpy/modifiers/angles.py

@@ -27,10 +27,11 @@
 from typing import Any, Callable, Optional, Union
 
 import dask
-import dask.array as da
 import numpy as np
 import xarray as xr
-from pyorbital.astronomy import get_alt_az, sun_zenith_angle
+from dask import array as da
+from pyorbital.astronomy import cos_zen as pyob_cos_zen
+from pyorbital.astronomy import get_alt_az
 from pyorbital.orbital import get_observer_look
 from pyresample.geometry import AreaDefinition, SwathDefinition
 
@@ -253,7 +254,7 @@ def get_angles(data_arr: xr.DataArray) -> tuple[xr.DataArray, xr.DataArray, xr.D
 
     Returns:
         Four DataArrays representing sensor azimuth angle, sensor zenith angle,
-        solar azimuth angle, and solar zenith angle.
+        solar azimuth angle, and solar zenith angle. All values are in degrees.
 
     """
     sata, satz = _get_sensor_angles(data_arr)
@@ -274,6 +275,18 @@ def get_satellite_zenith_angle(data_arr: xr.DataArray) -> xr.DataArray:
     return satz
 
 
+def get_cos_sza(data_arr: xr.DataArray) -> xr.DataArray:
+    """Generate the cosine of the solar zenith angle for the provided data.
+
+    Returns:
+        DataArray with the same shape as ``data_arr``. Units are radians.
+
+    """
+    lons, lats = _get_valid_lonlats(data_arr.attrs["area"], data_arr.chunks)
+    cos_sza = _get_cos_sza(data_arr.attrs["start_time"], lons, lats)
+    return _geo_dask_to_data_array(cos_sza)
+
+
 @cache_to_zarr_if("cache_lonlats")
 def _get_valid_lonlats(area: PRGeometry, chunks: Union[int, str] = "auto") -> tuple[da.Array, da.Array]:
     with ignore_invalid_float_warnings():
@@ -285,21 +298,43 @@ def _get_valid_lonlats(area: PRGeometry, chunks: Union[int, str] = "auto") -> tu
 
 def _get_sun_angles(data_arr: xr.DataArray) -> tuple[xr.DataArray, xr.DataArray]:
     lons, lats = _get_valid_lonlats(data_arr.attrs["area"], data_arr.data.chunks)
-    res = da.map_blocks(_get_sun_angles_wrapper, lons, lats,
-                        data_arr.attrs["start_time"],
-                        dtype=lons.dtype, meta=np.array((), dtype=lons.dtype),
-                        new_axis=[0], chunks=(2,) + lons.chunks)
-    suna = _geo_dask_to_data_array(res[0])
-    sunz = _geo_dask_to_data_array(res[1])
+    suna = da.map_blocks(_get_alt_az_wrapper, lons, lats,
+                         data_arr.attrs["start_time"],
+                         dtype=lons.dtype, meta=np.array((), dtype=lons.dtype),
+                         chunks=lons.chunks)
+    cos_sza = _get_cos_sza(data_arr.attrs["start_time"], lons, lats)
+    sunz = np.rad2deg(np.arccos(cos_sza))
+    suna = _geo_dask_to_data_array(suna)
+    sunz = _geo_dask_to_data_array(sunz)
     return suna, sunz
 
 
-def _get_sun_angles_wrapper(lons: da.Array, lats: da.Array, start_time: datetime) -> np.ndarray:
+def _get_cos_sza(utc_time, lons, lats):
+    cos_sza = da.map_blocks(_cos_zen_wrapper,
+                            lons, lats, utc_time,
+                            meta=np.array((), dtype=lons.dtype),
+                            dtype=lons.dtype,
+                            chunks=lons.chunks)
+    return cos_sza
+
+
+def _cos_zen_wrapper(lons, lats, utc_time):
+    with ignore_invalid_float_warnings():
+        return pyob_cos_zen(utc_time, lons, lats)
+
+
+def _get_alt_az_wrapper(lons: np.ndarray, lats: np.ndarray, start_time: datetime) -> np.ndarray:
     with ignore_invalid_float_warnings():
         suna = get_alt_az(start_time, lons, lats)[1]
         suna = np.rad2deg(suna)
-        sunz = sun_zenith_angle(start_time, lons, lats)
-        return np.stack([suna, sunz])
+    return suna
+
+
+# def _get_sza_wrapper(lons: np.ndarray, lats: np.ndarray, start_time: datetime) -> np.ndarray:
+#     with ignore_invalid_float_warnings():
+#         # return np.rad2deg(np.arccos(cos_zen(utc_time, lon, lat)))
+#         sunz = sun_zenith_angle(start_time, lons, lats)
+#     return sunz
 
 
 def _get_sensor_angles(data_arr: xr.DataArray) -> tuple[xr.DataArray, xr.DataArray]:
@@ -332,3 +367,49 @@ def _get_sensor_angles_wrapper(lons, lats, start_time, sat_lon, sat_lat, sat_alt
             lons, lats, 0)
         satz = 90 - satel
         return np.stack([sata, satz])
+
+
+def sunzen_corr_cos(data, cos_zen, limit=88., max_sza=95.):
+    """Perform Sun zenith angle correction.
+
+    The correction is based on the provided cosine of the zenith
+    angle (``cos_zen``).  The correction is limited
+    to ``limit`` degrees (default: 88.0 degrees).  For larger zenith
+    angles, the correction is the same as at the ``limit`` if ``max_sza``
+    is `None`. The default behavior is to gradually reduce the correction
+    past ``limit`` degrees up to ``max_sza`` where the correction becomes
+    0. Both ``data`` and ``cos_zen`` should be 2D arrays of the same shape.
+
+    """
+    return da.map_blocks(_sunzen_corr_cos_wrapper,
+                         data, cos_zen, limit, max_sza,
+                         meta=np.array((), dtype=data.dtype),
+                         chunks=data.chunks)
+
+
+def _sunzen_corr_cos_wrapper(data, cos_zen, limit, max_sza):
+    # Convert the zenith angle limit to cosine of zenith angle
+    limit_rad = np.deg2rad(limit)
+    limit_cos = np.cos(limit_rad)
+    max_sza_rad = np.deg2rad(max_sza) if max_sza is not None else max_sza
+
+    # Cosine correction
+    corr = 1. / cos_zen
+    if max_sza is not None:
+        # gradually fall off for larger zenith angle
+        grad_factor = (np.arccos(cos_zen) - limit_rad) / (max_sza_rad - limit_rad)
+        # invert the factor so maximum correction is done at `limit` and falls off later
+        grad_factor = 1. - np.log(grad_factor + 1) / np.log(2)
+        # make sure we don't make anything negative
+        grad_factor = grad_factor.clip(0.)
+    else:
+        # Use constant value (the limit) for larger zenith angles
+        grad_factor = 1.
+    # corr[cos_zen <= limit_cos] = grad_factor / limit_cos
+    # corr = corr.where(cos_zen > limit_cos, grad_factor / limit_cos)
+    corr = np.where(cos_zen > limit_cos, corr, grad_factor / limit_cos)
+    # Force "night" pixels to 0 (where SZA is invalid)
+    corr[np.isnan(cos_zen)] = 0
+    # corr = corr.where(cos_zen.notnull(), 0)
+
+    return data * corr

satpy/modifiers/geometry.py

@@ -17,6 +17,8 @@
 # satpy.  If not, see <http://www.gnu.org/licenses/>.
 """Modifier classes for corrections based on sun and other angles."""
 
+from __future__ import annotations
+
 import logging
 import time
 from datetime import datetime
@@ -27,7 +29,8 @@
 import xarray as xr
 
 from satpy.modifiers import ModifierBase
-from satpy.utils import atmospheric_path_length_correction, sunzen_corr_cos
+from satpy.modifiers.angles import sunzen_corr_cos
+from satpy.utils import atmospheric_path_length_correction
 
 logger = logging.getLogger(__name__)
 
@@ -63,17 +66,21 @@ def __call__(self, projectables, **info):
         logger.debug("Applying sun zen correction")
         coszen = self.coszen_cache.get(key)
         if coszen is None and not info.get('optional_datasets'):
-            # we were not given SZA, generate SZA then calculate cos(SZA)
-            from pyorbital.astronomy import cos_zen
+            # we were not given SZA, generate cos(SZA)
             logger.debug("Computing sun zenith angles.")
-            lons, lats = vis.attrs["area"].get_lonlats(chunks=vis.data.chunks)
-
-            coords = {}
-            if 'y' in vis.coords and 'x' in vis.coords:
-                coords['y'] = vis['y']
-                coords['x'] = vis['x']
-            coszen = xr.DataArray(cos_zen(vis.attrs["start_time"], lons, lats),
-                                  dims=['y', 'x'], coords=coords)
+            # coords = {}
+            # if 'y' in vis.coords and 'x' in vis.coords:
+            #     coords['y'] = vis['y']
+            #     coords['x'] = vis['x']
+
+            from .angles import get_cos_sza
+            coszen = get_cos_sza(vis)
+            # lons, lats = vis.attrs["area"].get_lonlats(chunks=vis.data.chunks)
+            # coszen = xr.DataArray(cos_zen(vis.attrs["start_time"], lons, lats),
+            #                       dims=['y', 'x'], coords=coords)
+            # import dask.array as da
+            # coszen = xr.DataArray(da.ones_like(lons),
+            #                       dims=['y', 'x'], coords=coords)
             if self.max_sza is not None:
                 coszen = coszen.where(coszen >= self.max_sza_cos)
             self.coszen_cache[key] = coszen
@@ -130,7 +137,9 @@ def __init__(self, correction_limit=88., **kwargs):
 
     def _apply_correction(self, proj, coszen):
         logger.debug("Apply the standard sun-zenith correction [1/cos(sunz)]")
-        return sunzen_corr_cos(proj, coszen, limit=self.correction_limit, max_sza=self.max_sza)
+        res = proj.copy()
+        res.data = sunzen_corr_cos(proj.data, coszen.data, limit=self.correction_limit, max_sza=self.max_sza)
+        return res
 
 
 class EffectiveSolarPathLengthCorrector(SunZenithCorrectorBase):

satpy/utils.py

@@ -240,42 +240,6 @@ def _get_sunz_corr_li_and_shibata(cos_zen):
     return 24.35 / (2. * cos_zen + np.sqrt(498.5225 * cos_zen**2 + 1))
 
 
-def sunzen_corr_cos(data, cos_zen, limit=88., max_sza=95.):
-    """Perform Sun zenith angle correction.
-
-    The correction is based on the provided cosine of the zenith
-    angle (``cos_zen``).  The correction is limited
-    to ``limit`` degrees (default: 88.0 degrees).  For larger zenith
-    angles, the correction is the same as at the ``limit`` if ``max_sza``
-    is `None`. The default behavior is to gradually reduce the correction
-    past ``limit`` degrees up to ``max_sza`` where the correction becomes
-    0. Both ``data`` and ``cos_zen`` should be 2D arrays of the same shape.
-
-    """
-    # Convert the zenith angle limit to cosine of zenith angle
-    limit_rad = np.deg2rad(limit)
-    limit_cos = np.cos(limit_rad)
-    max_sza_rad = np.deg2rad(max_sza) if max_sza is not None else max_sza
-
-    # Cosine correction
-    corr = 1. / cos_zen
-    if max_sza is not None:
-        # gradually fall off for larger zenith angle
-        grad_factor = (np.arccos(cos_zen) - limit_rad) / (max_sza_rad - limit_rad)
-        # invert the factor so maximum correction is done at `limit` and falls off later
-        grad_factor = 1. - np.log(grad_factor + 1) / np.log(2)
-        # make sure we don't make anything negative
-        grad_factor = grad_factor.clip(0.)
-    else:
-        # Use constant value (the limit) for larger zenith angles
-        grad_factor = 1.
-    corr = corr.where(cos_zen > limit_cos, grad_factor / limit_cos)
-    # Force "night" pixels to 0 (where SZA is invalid)
-    corr = corr.where(cos_zen.notnull(), 0)
-
-    return data * corr
-
-
 def atmospheric_path_length_correction(data, cos_zen, limit=88., max_sza=95.):
     """Perform Sun zenith angle correction.