Nowcast refactoring (#291)

* Remove unnecessary encoding lines

* Change variable names to improve PEP 8 compatibility

* Use 4 spaces in docstring indentations

* Implement a common utility function for computing the S-PROG mask

* Fix incorrect imports

* Use f-strings instead of old-style printing

* Implement utility method for computing nowcasts with non-integer time steps

* Use state model for the nowcast loop function

* Fix typos

* Replace R_thr keyword argument with precip_thr

* Remove extra index increment leading to possible IndexError

* Use np.abs instead of abs

* Minor docstring revisions

* Fix typo

* Implement supplying additional parameters to the state functions

* Argument refactoring

* Fix incorrect indices

* Implement S-PROG by using utils.nowcast_main_loop

* Add comment lines

* Fix measurement of main loop time

* Use PEP 8-compatible variable names

* Remove unnecessary code lines

* Use PEP 8-compatible variable names

* Use PEP 8-compatible variable names in _ckeck_inputs

* Move worker definition outside the nowcast loop

* Supply parameters directly instead of dictionary

* Use lowercase variable name for phi

* Remove unnecessary code lines

* Simplify code

* Use a single update/forecast function for the nowcast loop

* Use state and params dictionaries with the update method

* First version of nowcast_main_loop with support for ensembles and velocity perturbations

* Multiple fixes

* Use dask for parallelization

* Remove unnecessary encoding line

* Use more descriptive variable name for rain rate

* Refactored version of ANVILthat uses utils.nowcast_main_loop

* Use proper check for deterministic/ensemble nowcast

* Add return_output option to nowcast_main_loop

* Docstring polishing

* Implement support for callback functions in nowcast_main_loop

* Docstring update

* Bug fixes

* Fix incorrect order of code lines

* Fix incorrect initialization of starttime

* Refactored STEPS module using utils.nowcast_main_loop

* Remove unnecessary encoding line

* Add checks for ensemble size and number of workers

* Revise variable names to be compatible with the other nowcasting methods

* Ensure that return_displacement is set to True

* Remove unnecessary lines

* First version of refactored LINDA using utils.nowcast_main_loop

* Update docstring of the callback option

* Implement callback and return_output options

* Change order of input arguments

* Fix parallelization bug

* Add ensemble size check

* Change num_ens_members to n_ens_members for compatibility with steps

* Add test for callback function

* Update black to newest version

* Run black

* Fix possible division by zero

* Set allow_nonfinite_values automatically based on the input data

* Replace None default argument with dict

* Fix typo

* Fix typo

* Use float mask as in steps

* Implement common utility method for incremental mask

* Explicitly check that precip_thr is given

* Explicitly check that precip_thr is specified when it is required

* Remove redundant code

* Use different name to avoid redefining built-in 'filter'

* Remove unused variable

* Replace filter with bp_filter

* Use None default argument value instead of dict()

* Change worker function names to avoid redefinition

* Fix typo

Co-authored-by: Ruben Imhoff <31476760+RubenImhoff@users.noreply.github.com>

* Increase test coverage for anvil

* Remove unnecessary encoding line

* Improve variable naming in stack_cascades

* Replace R with precip

* Use consistent variable naming for forecasts

* Run black

* Remove unnecessary loop indices

* Minor refactoring of printing functions

* Use more descriptive variable name for velocity perturbators

* Minor refactoring of seed generation

* Use lowercase variable names for masks

* Remove TODO comment

* Use more decriptive variable names

* Use lowercase variable names for gamma and psi

* Use longer variable name for velocity perturbation generator

* Replace precip_c and precip_d with more descriptive names

* Replace precip_cascade with precip_cascades

* Remove extra underscore from vp_

* Use longer names for lambda functions

* Replace P with pert_gen

* Replace pp with pert_gen

* List all methods in the autosummary

* Docstring addition

* Fix incorrect initialization of precip_forecast_prev

* Explicitly require nowcast type and ensemble size in nowcast_main_loop

* First version of tests for nowcasts.utils

* Add docstring

* Remove unused metadata

* Add future warnings for new argument names; support old names till 1.8.0

* Remoe generic filter for warnings

* Pin micromamba to 0.24 since 0.25 fails with windows

Co-authored-by: Daniele Nerini <daniele.nerini@gmail.com>
Co-authored-by: Ruben Imhoff <31476760+RubenImhoff@users.noreply.github.com>
Co-authored-by: ned <daniele.nerini@meteoswiss.ch>

.github/workflows/test_pysteps.yml

@@ -36,6 +36,7 @@ jobs:
       - name: Install mamba and create environment
         uses: mamba-org/provision-with-micromamba@main
         with:
+          micromamba-version: 0.24.0
           environment-file: ci/ci_test_env.yml
           environment-name: test_environment
           extra-specs: python=${{ matrix.python-version }}

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 22.6.0
     hooks:
     - id: black
       language_version: python3

pysteps/decorators.py

@@ -16,6 +16,7 @@
 """
 import inspect
 import uuid
+import warnings
 from collections import defaultdict
 from functools import wraps
 
@@ -282,3 +283,35 @@ def _func_with_cache(*args, **kwargs):
         return _func_with_cache
 
     return _memoize
+
+
+def deprecate_args(old_new_args, deprecation_release):
+    """
+    Support deprecated argument names while issuing deprecation warnings.
+
+    Parameters
+    ----------
+    old_new_args: dict[str, str]
+        Mapping from old to new argument names.
+    deprecation_release: str
+        Specify which future release will convert this warning into an error.
+    """
+
+    def _deprecate(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            kwargs_names = list(kwargs.keys())
+            for key_old in kwargs_names:
+                if key_old in old_new_args:
+                    key_new = old_new_args[key_old]
+                    kwargs[key_new] = kwargs.pop(key_old)
+                    warnings.warn(
+                        f"Argument '{key_old}' has been renamed to '{key_new}'. "
+                        f"This will raise a TypeError in pysteps {deprecation_release}.",
+                        FutureWarning,
+                    )
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return _deprecate

pysteps/nowcasts/__init__.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """Implementations of deterministic and ensemble nowcasting methods."""
 
 from pysteps.nowcasts.interface import get_method

pysteps/nowcasts/anvil.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 pysteps.nowcasts.anvil
 ======================
@@ -23,7 +22,7 @@
 import numpy as np
 from scipy.ndimage import gaussian_filter
 from pysteps import cascade, extrapolation
-from pysteps.nowcasts import utils as nowcast_utils
+from pysteps.nowcasts.utils import nowcast_main_loop
 from pysteps.timeseries import autoregression
 from pysteps import utils
 
@@ -149,36 +148,36 @@ def forecast(
     if filter_kwargs is None:
         filter_kwargs = dict()
 
-    print("Computing ANVIL nowcast:")
-    print("------------------------")
+    print("Computing ANVIL nowcast")
+    print("-----------------------")
     print("")
 
-    print("Inputs:")
-    print("-------")
-    print("input dimensions: %dx%d" % (vil.shape[1], vil.shape[2]))
+    print("Inputs")
+    print("------")
+    print(f"input dimensions: {vil.shape[1]}x{vil.shape[2]}")
     print("")
 
-    print("Methods:")
-    print("--------")
-    print("extrapolation:   %s" % extrap_method)
-    print("FFT:             %s" % fft_method)
+    print("Methods")
+    print("-------")
+    print(f"extrapolation:   {extrap_method}")
+    print(f"FFT:             {fft_method}")
     print("")
 
-    print("Parameters:")
-    print("-----------")
+    print("Parameters")
+    print("----------")
     if isinstance(timesteps, int):
-        print("number of time steps:        %d" % timesteps)
+        print(f"number of time steps:        {timesteps}")
     else:
-        print("time steps:                  %s" % timesteps)
-    print("parallel threads:            %d" % num_workers)
-    print("number of cascade levels:    %d" % n_cascade_levels)
-    print("order of the ARI(p,1) model: %d" % ar_order)
+        print(f"time steps:                  {timesteps}")
+    print(f"parallel threads:            {num_workers}")
+    print(f"number of cascade levels:    {n_cascade_levels}")
+    print(f"order of the ARI(p,1) model: {ar_order}")
     if type(ar_window_radius) == int:
-        print("ARI(p,1) window radius:      %d" % ar_window_radius)
+        print(f"ARI(p,1) window radius:      {ar_window_radius}")
     else:
         print("ARI(p,1) window radius:      none")
 
-    print("R(VIL) window radius:        %d" % r_vil_window_radius)
+    print(f"R(VIL) window radius:        {r_vil_window_radius}")
 
     if measure_time:
         starttime_init = time.time()
@@ -188,11 +187,15 @@ def forecast(
 
     if rainrate is None and apply_rainrate_mask:
         rainrate_mask = vil[-1, :] < 0.1
+    else:
+        rainrate_mask = None
 
     if rainrate is not None:
         # determine the coefficients fields of the relation R=a*VIL+b by
         # localized linear regression
         r_vil_a, r_vil_b = _r_vil_regression(vil[-1, :], rainrate, r_vil_window_radius)
+    else:
+        r_vil_a, r_vil_b = None, None
 
     # transform the input fields to Lagrangian coordinates by extrapolation
     extrapolator = extrapolation.get_method(extrap_method)
@@ -287,129 +290,41 @@ def worker(vil, i):
 
     print("Starting nowcast computation.")
 
-    if measure_time:
-        starttime_mainloop = time.time()
-
-    r_f = []
-
-    if isinstance(timesteps, int):
-        timesteps = range(timesteps + 1)
-        timestep_type = "int"
-    else:
-        original_timesteps = [0] + list(timesteps)
-        timesteps = nowcast_utils.binned_timesteps(original_timesteps)
-        timestep_type = "list"
+    rainrate_f = []
 
-    if rainrate is not None:
-        r_f_prev = r_vil_a * vil[-1, :] + r_vil_b
-    else:
-        r_f_prev = vil[-1, :]
     extrap_kwargs["return_displacement"] = True
 
-    dp = None
-    t_prev = 0.0
-
-    for t, subtimestep_idx in enumerate(timesteps):
-        if timestep_type == "list":
-            subtimesteps = [original_timesteps[t_] for t_ in subtimestep_idx]
-        else:
-            subtimesteps = [t]
-
-        if (timestep_type == "list" and subtimesteps) or (
-            timestep_type == "int" and t > 0
-        ):
-            is_nowcast_time_step = True
-        else:
-            is_nowcast_time_step = False
-
-        if is_nowcast_time_step:
-            print(
-                "Computing nowcast for time step %d... " % t,
-                end="",
-                flush=True,
-            )
-
-        if measure_time:
-            starttime = time.time()
-
-        # iterate the ARI models for each cascade level
-        for i in range(n_cascade_levels):
-            vil_dec[i, :] = autoregression.iterate_ar_model(vil_dec[i, :], phi[i])
-
-        # recompose the cascade to obtain the forecast field
-        vil_dec_dict = {}
-        vil_dec_dict["cascade_levels"] = vil_dec[:, -1, :]
-        vil_dec_dict["domain"] = "spatial"
-        vil_dec_dict["normalized"] = False
-        vil_f = recomp_method(vil_dec_dict)
-        vil_f[~mask] = np.nan
-
-        if rainrate is not None:
-            # convert VIL to rain rate
-            r_f_new = r_vil_a * vil_f + r_vil_b
-        else:
-            r_f_new = vil_f
-            if apply_rainrate_mask:
-                r_f_new[rainrate_mask] = 0.0
-
-        r_f_new[r_f_new < 0.0] = 0.0
-
-        # advect the recomposed field to obtain the forecast for the current
-        # time step (or subtimesteps if non-integer time steps are given)
-        for t_sub in subtimesteps:
-            if t_sub > 0:
-                t_diff_prev_int = t_sub - int(t_sub)
-                if t_diff_prev_int > 0.0:
-                    r_f_ip = (
-                        1.0 - t_diff_prev_int
-                    ) * r_f_prev + t_diff_prev_int * r_f_new
-                else:
-                    r_f_ip = r_f_prev
-
-                t_diff_prev = t_sub - t_prev
-
-                extrap_kwargs["displacement_prev"] = dp
-                extrap_kwargs["allow_nonfinite_values"] = (
-                    True if np.any(~np.isfinite(r_f_ip)) else False
-                )
-
-                r_f_ep, dp = extrapolator(
-                    r_f_ip,
-                    velocity,
-                    [t_diff_prev],
-                    **extrap_kwargs,
-                )
-                r_f.append(r_f_ep[0])
-                t_prev = t_sub
-
-        # advect the forecast field by one time step if no subtimesteps in the
-        # current interval were found
-        if not subtimesteps:
-            t_diff_prev = t + 1 - t_prev
-            extrap_kwargs["displacement_prev"] = dp
-            _, dp = extrapolator(
-                None,
-                velocity,
-                [t_diff_prev],
-                **extrap_kwargs,
-            )
-            t_prev = t + 1
-
-        r_f_prev = r_f_new
-
-        if is_nowcast_time_step:
-            if measure_time:
-                print("%.2f seconds." % (time.time() - starttime))
-            else:
-                print("done.")
-
+    state = {"vil_dec": vil_dec}
+    params = {
+        "apply_rainrate_mask": apply_rainrate_mask,
+        "mask": mask,
+        "n_cascade_levels": n_cascade_levels,
+        "phi": phi,
+        "rainrate": rainrate,
+        "rainrate_mask": rainrate_mask,
+        "recomp_method": recomp_method,
+        "r_vil_a": r_vil_a,
+        "r_vil_b": r_vil_b,
+    }
+
+    rainrate_f = nowcast_main_loop(
+        vil[-1, :],
+        velocity,
+        state,
+        timesteps,
+        extrap_method,
+        _update,
+        extrap_kwargs=extrap_kwargs,
+        params=params,
+        measure_time=measure_time,
+    )
     if measure_time:
-        mainloop_time = time.time() - starttime_mainloop
+        rainrate_f, mainloop_time = rainrate_f
 
     if measure_time:
-        return np.stack(r_f), init_time, mainloop_time
+        return np.stack(rainrate_f), init_time, mainloop_time
     else:
-        return np.stack(r_f)
+        return np.stack(rainrate_f)
 
 
 def _check_inputs(vil, rainrate, velocity, timesteps, ar_order):
@@ -559,3 +474,31 @@ def _r_vil_regression(vil, r, window_radius):
     b[~mask_vil] = 0.0
 
     return a, b
+
+
+def _update(state, params):
+    # iterate the ARI models for each cascade level
+    for i in range(params["n_cascade_levels"]):
+        state["vil_dec"][i, :] = autoregression.iterate_ar_model(
+            state["vil_dec"][i, :], params["phi"][i]
+        )
+
+    # recompose the cascade to obtain the forecast field
+    vil_dec_dict = {}
+    vil_dec_dict["cascade_levels"] = state["vil_dec"][:, -1, :]
+    vil_dec_dict["domain"] = "spatial"
+    vil_dec_dict["normalized"] = False
+    vil_f = params["recomp_method"](vil_dec_dict)
+    vil_f[~params["mask"]] = np.nan
+
+    if params["rainrate"] is not None:
+        # convert VIL to rain rate
+        rainrate_f_new = params["r_vil_a"] * vil_f + params["r_vil_b"]
+    else:
+        rainrate_f_new = vil_f
+        if params["apply_rainrate_mask"]:
+            rainrate_f_new[params["rainrate_mask"]] = 0.0
+
+    rainrate_f_new[rainrate_f_new < 0.0] = 0.0
+
+    return rainrate_f_new, state

pysteps/nowcasts/extrapolation.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 pysteps.nowcasts.extrapolation
 ==============================
@@ -34,24 +33,24 @@ def forecast(
     Parameters
     ----------
     precip: array-like
-      Two-dimensional array of shape (m,n) containing the input precipitation
-      field.
+        Two-dimensional array of shape (m,n) containing the input precipitation
+        field.
     velocity: array-like
-      Array of shape (2,m,n) containing the x- and y-components of the
-      advection field. The velocities are assumed to represent one time step
-      between the inputs.
+        Array of shape (2,m,n) containing the x- and y-components of the
+        advection field. The velocities are assumed to represent one time step
+        between the inputs.
     timesteps: int or list of floats
-      Number of time steps to forecast or a list of time steps for which the
-      forecasts are computed (relative to the input time step). The elements of
-      the list are required to be in ascending order.
+        Number of time steps to forecast or a list of time steps for which the
+        forecasts are computed (relative to the input time step). The elements
+        of the list are required to be in ascending order.
     extrap_method: str, optional
-      Name of the extrapolation method to use. See the documentation of
-      pysteps.extrapolation.interface.
+        Name of the extrapolation method to use. See the documentation of
+        pysteps.extrapolation.interface.
     extrap_kwargs: dict, optional
-      Optional dictionary that is expanded into keyword arguments for the
-      extrapolation method.
+        Optional dictionary that is expanded into keyword arguments for the
+        extrapolation method.
     measure_time: bool, optional
-      If True, measure, print, and return the computation time.
+        If True, measure, print, and return the computation time.
 
     Returns
     -------

pysteps/nowcasts/interface.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 pysteps.nowcasts.interface
 ==========================

pysteps/nowcasts/lagrangian_probability.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 pysteps.nowcasts.lagrangian_probability
 =======================================