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)
@@ -267,15 +268,27 @@ def get_satellite_zenith_angle(data_arr: xr.DataArray) -> xr.DataArray:
     Note that this function can benefit from the ``satpy.config`` parameters
     :ref:`cache_lonlats <config_cache_lonlats_setting>` and
     :ref:`cache_sensor_angles <config_cache_sensor_angles_setting>`
-    being set to ``True``.
+    being set to ``True``. Values are in degrees.
 
     """
     satz = _get_sensor_angles(data_arr)[1]
     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``.
+
+    """
+    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]:
+def _get_valid_lonlats(area: PRGeometry, chunks: Union[int, str, tuple] = "auto") -> tuple[da.Array, da.Array]:
     with ignore_invalid_float_warnings():
         lons, lats = area.get_lonlats(chunks=chunks)
         lons = da.where(lons >= 1e30, np.nan, lons)
@@ -285,21 +298,36 @@ 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_sun_azimuth_ndarray, 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_ndarray,
+                            lons, lats, utc_time,
+                            meta=np.array((), dtype=lons.dtype),
+                            dtype=lons.dtype,
+                            chunks=lons.chunks)
+    return cos_sza
+
+
+def _cos_zen_ndarray(lons, lats, utc_time):
+    with ignore_invalid_float_warnings():
+        return pyob_cos_zen(utc_time, lons, lats)
+
+
+def _get_sun_azimuth_ndarray(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_sensor_angles(data_arr: xr.DataArray) -> tuple[xr.DataArray, xr.DataArray]:
@@ -316,13 +344,13 @@ def _get_sensor_angles(data_arr: xr.DataArray) -> tuple[xr.DataArray, xr.DataArr
 @cache_to_zarr_if("cache_sensor_angles", sanitize_args_func=_sanitize_observer_look_args)
 def _get_sensor_angles_from_sat_pos(sat_lon, sat_lat, sat_alt, start_time, area_def, chunks):
     lons, lats = _get_valid_lonlats(area_def, chunks)
-    res = da.map_blocks(_get_sensor_angles_wrapper, lons, lats, start_time, sat_lon, sat_lat, sat_alt,
+    res = da.map_blocks(_get_sensor_angles_ndarray, lons, lats, start_time, sat_lon, sat_lat, sat_alt,
                         dtype=lons.dtype, meta=np.array((), dtype=lons.dtype), new_axis=[0],
                         chunks=(2,) + lons.chunks)
     return res[0], res[1]
 
 
-def _get_sensor_angles_wrapper(lons, lats, start_time, sat_lon, sat_lat, sat_alt) -> np.ndarray:
+def _get_sensor_angles_ndarray(lons, lats, start_time, sat_lon, sat_lat, sat_alt) -> np.ndarray:
     with ignore_invalid_float_warnings():
         sata, satel = get_observer_look(
             sat_lon,
@@ -332,3 +360,51 @@ 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: da.Array,
+                    cos_zen: da.Array,
+                    limit: float = 88.,
+                    max_sza: Optional[float] = 95.) -> da.Array:
+    """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_ndarray,
+                         data, cos_zen, limit, max_sza,
+                         meta=np.array((), dtype=data.dtype),
+                         chunks=data.chunks)
+
+
+def _sunzen_corr_cos_ndarray(data: np.ndarray,
+                             cos_zen: np.ndarray,
+                             limit: float,
+                             max_sza: Optional[float]) -> np.ndarray:
+    # 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 = 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
+    return data * corr

satpy/modifiers/atmosphere.py

@@ -45,11 +45,9 @@ def __call__(self, projectables, optional_datasets=None, **info):
         if not optional_datasets or len(optional_datasets) != 4:
             vis, red = self.match_data_arrays(projectables)
             sata, satz, suna, sunz = get_angles(vis)
-            red.data = da.rechunk(red.data, vis.data.chunks)
         else:
             vis, red, sata, satz, suna, sunz = self.match_data_arrays(
                 projectables + optional_datasets)
-            sata, satz, suna, sunz = optional_datasets
 
         # get the dask array underneath
         sata = sata.data

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,10 @@ 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)
+            from .angles import get_cos_sza
+            coszen = get_cos_sza(vis)
             if self.max_sza is not None:
                 coszen = coszen.where(coszen >= self.max_sza_cos)
             self.coszen_cache[key] = coszen
@@ -130,7 +126,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/resample.py

@@ -193,7 +193,7 @@
 def hash_dict(the_dict, the_hash=None):
     """Calculate a hash for a dictionary."""
     if the_hash is None:
-        the_hash = hashlib.sha1()
+        the_hash = hashlib.sha1()  # nosec
     the_hash.update(json.dumps(the_dict, sort_keys=True).encode('utf-8'))
     return the_hash
 
@@ -991,7 +991,9 @@ def aggregate(d, y_size, x_size):
 
         new_chunks = (tuple(int(x / y_size) for x in d.chunks[0]),
                       tuple(int(x / x_size) for x in d.chunks[1]))
-        return da.core.map_blocks(_mean, d, y_size, x_size, dtype=d.dtype, chunks=new_chunks)
+        return da.core.map_blocks(_mean, d, y_size, x_size,
+                                  meta=np.array((), dtype=d.dtype),
+                                  dtype=d.dtype, chunks=new_chunks)
 
     @classmethod
     def expand_reduce(cls, d_arr, repeats):

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.