From 7b52ab543dab000276286559f3ef5d0c14b51e9c Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 11:29:55 +0200
Subject: [PATCH 01/18] allow for significantly less memory usage in steps
 blending

---
 pysteps/blending/steps.py | 203 +++++++++++++++++---------------------
 1 file changed, 92 insertions(+), 111 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 7e4535dd..1a92c27e 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -46,10 +46,13 @@
 import math
 import time
 from copy import deepcopy
+from functools import partial
+from multiprocessing import Pool
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
+from scipy.ndimage import (binary_dilation, generate_binary_structure,
+                           iterate_structure)
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils
@@ -557,22 +560,19 @@ def forecast(
         fft,
     )
 
-    # 2.2 If necessary, decompose (NWP) model forecasts and stack cascades
-    (
-        precip_models_cascade,
-        mu_models,
-        sigma_models,
-        precip_models_pm,
-    ) = _compute_cascade_decomposition_nwp(
-        precip_models, bp_filter, decompositor, recompositor, fft, domain
-    )
+    # 2.2 If necessary, recompose (NWP) model forecasts
+    precip_models_cascade = None
+
+    if precip_models.ndim != 4:
+        precip_models_cascade = precip_models
+        precip_models = _compute_cascade_recomposition_nwp(precip_models, recompositor)
 
     # 2.3 Check for zero input fields in the radar and NWP data.
     zero_precip_radar = blending.utils.check_norain(precip, precip_thr, norain_thr)
     # The norain fraction threshold used for nwp is the default value of 0.0,
     # since nwp does not suffer from clutter.
     zero_model_fields = blending.utils.check_norain(
-        precip_models_pm, precip_thr, norain_thr
+        precip_models, precip_thr, norain_thr
     )
 
     if isinstance(timesteps, int):
@@ -633,37 +633,14 @@ def forecast(
         # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain
         # only nans or infs. If so, fill it with the zero value.
         if zero_precip_radar:
-            precip_cascade = np.nan_to_num(
-                precip_cascade,
-                copy=True,
-                nan=np.nanmin(precip_models_cascade),
-                posinf=np.nanmin(precip_models_cascade),
-                neginf=np.nanmin(precip_models_cascade),
-            )
-
-        # 2.3.4 Check if the NWP fields contain nans or infinite numbers. If so,
-        # fill these with the minimum value present in precip (corresponding to
-        # zero rainfall in the radar observations)
-        (
-            precip_models_cascade,
-            precip_models_pm,
-            mu_models,
-            sigma_models,
-        ) = _fill_nans_infs_nwp_cascade(
-            precip_models_cascade,
-            precip_models_pm,
-            precip_cascade,
-            precip,
-            mu_models,
-            sigma_models,
-        )
+            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(precip)
 
         # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP
         # forecast. I.e., velocity (radar) equals velocity_model at the first time
         # step.
         if zero_precip_radar:
             # Use the velocity from velocity_models at time step 0
-            velocity = velocity_models[:, 0, :, :, :]
+            velocity = velocity_models[:, 0, :, :, :].astype(np.float64, copy=False)
             # Take the average over the first axis, which corresponds to n_models
             # (hence, the model average)
             velocity = np.mean(velocity, axis=0)
@@ -675,7 +652,7 @@ def forecast(
         # rainfall data
         if zero_precip_radar:
             precip_noise_input = _determine_max_nr_rainy_cells_nwp(
-                precip_models_pm, precip_thr, precip_models_pm.shape[0], timesteps
+                precip_models, precip_thr, precip_models.shape[0], timesteps
             )
             # Make sure precip_noise_input is three dimensional
             if len(precip_noise_input.shape) != 3:
@@ -809,6 +786,63 @@ def forecast(
             if measure_time:
                 starttime = time.time()
 
+            if precip_models_cascade is not None:
+                decomp_precip_models = list(precip_models_cascade[:,t])
+
+            if precip_models.shape[0] == 1:
+                decomp_precip_models = [
+                    decompositor(
+                        precip_models[0, t, :, :],
+                        bp_filter=bp_filter,
+                        fft_method=fft,
+                        output_domain=domain,
+                        normalize=True,
+                        compute_stats=True,
+                        compact_output=True,
+                    )
+                ]
+            else:
+                with Pool(num_workers) as pool:
+                    decomp_precip_models = pool.map(
+                        partial(
+                            decompositor,
+                            bp_filter=bp_filter,
+                            fft_method=fft,
+                            output_domain=domain,
+                            normalize=True,
+                            compute_stats=True,
+                            compact_output=True,
+                        ),
+                        list(precip_models[:, t, :, :]),
+                    )
+
+            precip_models_cascade_temp = np.array(
+                [decomp["cascade_levels"] for decomp in decomp_precip_models]
+            )
+            mu_models_temp = np.array(
+                [decomp["means"] for decomp in decomp_precip_models]
+            )
+            sigma_models_temp = np.array(
+                [decomp["stds"] for decomp in decomp_precip_models]
+            )
+
+            # 2.3.4 Check if the NWP fields contain nans or infinite numbers. If so,
+            # fill these with the minimum value present in precip (corresponding to
+            # zero rainfall in the radar observations)
+            (
+                precip_models_cascade_temp,
+                precip_models_temp,
+                mu_models_temp,
+                sigma_models_temp,
+            ) = _fill_nans_infs_nwp_cascade(
+                precip_models_cascade_temp,
+                precip_models[:, t, :, :].astype(np.float64, copy=False),
+                precip_cascade,
+                precip,
+                mu_models_temp,
+                sigma_models_temp,
+            )
+
             # 8.1.1 Before calling the worker for the forecast loop, determine which (NWP)
             # models will be combined with which nowcast ensemble members. With the
             # way it is implemented at this moment: n_ens_members of the output equals
@@ -821,11 +855,11 @@ def forecast(
                 sigma_models_temp,
                 n_model_indices,
             ) = _find_nwp_combination(
-                precip_models_cascade[:, t, :, :, :],
-                precip_models_pm[:, t, :, :],
-                velocity_models[:, t, :, :, :],
-                mu_models[:, t, :],
-                sigma_models[:, t, :],
+                precip_models_cascade_temp,
+                precip_models_temp,
+                velocity_models[:, t, :, :, :].astype(np.float64, copy=False),
+                mu_models_temp,
+                sigma_models_temp,
                 n_ens_members,
                 ar_order,
                 n_cascade_levels,
@@ -2043,78 +2077,25 @@ def _compute_cascade_decomposition_radar(
     return R_c, mu_extrapolation, sigma_extrapolation
 
 
-def _compute_cascade_decomposition_nwp(
-    precip_models, bp_filter, decompositor, recompositor, fft, domain
-):
-    """If necessary, decompose (NWP) model forecasts and stack cascades."""
-    if precip_models.ndim == 4:
-        # Keep the model fields for the probability matching later on
-        precip_models_pm = precip_models.copy()
-        # Decompose the (NWP) model forecasts
-        precip_models_cascade = []
-        # Loop through the n_models
-        for i in range(precip_models.shape[0]):
-            precip_model_cascade = []
-            # Loop through the time steps
-            for j in range(precip_models.shape[1]):
-                precip_model_cascade.append(
-                    decompositor(
-                        field=precip_models[i, j, :, :],
-                        bp_filter=bp_filter,
-                        fft_method=fft,
-                        output_domain=domain,
-                        normalize=True,
-                        compute_stats=True,
-                        compact_output=True,
-                    )
-                )
-            precip_models_cascade.append(precip_model_cascade)
-        precip_models_cascade = np.array(precip_models_cascade)
+def _compute_cascade_recomposition_nwp(precip_models, recompositor):
+    """If necessary, recompose (NWP) model forecasts and stack cascades."""
+    precip_models_cascade = precip_models
+    precip_models_pm = None
 
-        precip_models, precip_model_cascade = None, None
-    else:
-        precip_models_cascade = precip_models
-        precip_models_pm = None
-
-    # Stack the (NWP) model cascades in separate normalized cascades and return
-    # the means and sigmas.
-    # The normalized model cascade should have the shape:
-    # [n_models, n_timesteps, n_cascade_levels, m, n]
-    precip_models = []
-    mu_models = []
-    sigma_models = []
-    # Stack it per model and combine that
+    # Recompose the (NWP) model cascades to have rainfall fields per
+    # model and time step, which will be used in the probability matching steps.
+    # Recomposed cascade will have shape: [n_models, n_timesteps, m, n]
+    precip_models_pm = []
     for i in range(precip_models_cascade.shape[0]):
-        precip_model, mu_model, sigma_model = blending.utils.stack_cascades(
-            R_d=precip_models_cascade[i, :], donorm=False
-        )
-        precip_models.append(precip_model)
-        mu_models.append(mu_model)
-        sigma_models.append(sigma_model)
-
-    precip_models = np.stack(precip_models)
-    mu_models = np.stack(mu_models)
-    sigma_models = np.stack(sigma_models)
-
-    precip_model, mu_model, sigma_model = None, None, None
-
-    if precip_models_pm is None:
-        # Finally, recompose the (NWP) model cascades to have rainfall fields per
-        # model and time step, which will be used in the probability matching steps.
-        # Recomposed cascade will have shape: [n_models, n_timesteps, m, n]
-        precip_models_pm = []
-        for i in range(precip_models_cascade.shape[0]):
-            precip_model_pm = []
-            for time_step in range(precip_models_cascade.shape[1]):
-                precip_model_pm.append(
-                    recompositor(precip_models_cascade[i, time_step])
-                )
-            precip_models_pm.append(precip_model_pm)
+        precip_model_pm = []
+        for time_step in range(precip_models_cascade.shape[1]):
+            precip_model_pm.append(recompositor(precip_models_cascade[i, time_step]))
+        precip_models_pm.append(precip_model_pm)
 
-        precip_models_pm = np.stack(precip_models_pm)
-        precip_model_pm = None
+    precip_models_pm = np.stack(precip_models_pm)
+    precip_model_pm = None
 
-    return precip_models, mu_models, sigma_models, precip_models_pm
+    return precip_models_pm
 
 
 def _estimate_ar_parameters_radar(
@@ -2476,4 +2457,4 @@ def _determine_max_nr_rainy_cells_nwp(
                 max_rain_pixels_t = t
     precip_noise_input = precip_models_pm[max_rain_pixels_j][max_rain_pixels_t]
 
-    return precip_noise_input
+    return precip_noise_input.astype(np.float64, copy=False)

From 34d90baad8cb014add659e89ba02936fb98a741b Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 11:38:51 +0200
Subject: [PATCH 02/18] apply black

---
 pysteps/blending/steps.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 1a92c27e..4a1c5810 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -51,8 +51,7 @@
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import (binary_dilation, generate_binary_structure,
-                           iterate_structure)
+from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils
@@ -787,7 +786,7 @@ def forecast(
                 starttime = time.time()
 
             if precip_models_cascade is not None:
-                decomp_precip_models = list(precip_models_cascade[:,t])
+                decomp_precip_models = list(precip_models_cascade[:, t])
 
             if precip_models.shape[0] == 1:
                 decomp_precip_models = [

From dabd93ab34a8ca96ebd7af9e0d628278d8e873da Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 12:33:01 +0200
Subject: [PATCH 03/18] use threadpool

---
 pysteps/blending/steps.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 4a1c5810..98d314cb 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -47,11 +47,12 @@
 import time
 from copy import deepcopy
 from functools import partial
-from multiprocessing import Pool
+from multiprocessing.pool import ThreadPool
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
+from scipy.ndimage import (binary_dilation, generate_binary_structure,
+                           iterate_structure)
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils
@@ -801,7 +802,7 @@ def forecast(
                     )
                 ]
             else:
-                with Pool(num_workers) as pool:
+                with ThreadPool(num_workers) as pool:
                     decomp_precip_models = pool.map(
                         partial(
                             decompositor,

From 4e2758a4abd5c861c60b01dbc1c3f863dc2e3df8 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 12:35:16 +0200
Subject: [PATCH 04/18] only decompose if not already done

---
 pysteps/blending/steps.py | 40 +++++++++++++++++++--------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 98d314cb..ca42fad1 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -788,33 +788,33 @@ def forecast(
 
             if precip_models_cascade is not None:
                 decomp_precip_models = list(precip_models_cascade[:, t])
-
-            if precip_models.shape[0] == 1:
-                decomp_precip_models = [
-                    decompositor(
-                        precip_models[0, t, :, :],
-                        bp_filter=bp_filter,
-                        fft_method=fft,
-                        output_domain=domain,
-                        normalize=True,
-                        compute_stats=True,
-                        compact_output=True,
-                    )
-                ]
             else:
-                with ThreadPool(num_workers) as pool:
-                    decomp_precip_models = pool.map(
-                        partial(
-                            decompositor,
+                if precip_models.shape[0] == 1:
+                    decomp_precip_models = [
+                        decompositor(
+                            precip_models[0, t, :, :],
                             bp_filter=bp_filter,
                             fft_method=fft,
                             output_domain=domain,
                             normalize=True,
                             compute_stats=True,
                             compact_output=True,
-                        ),
-                        list(precip_models[:, t, :, :]),
-                    )
+                        )
+                    ]
+                else:
+                    with ThreadPool(num_workers) as pool:
+                        decomp_precip_models = pool.map(
+                            partial(
+                                decompositor,
+                                bp_filter=bp_filter,
+                                fft_method=fft,
+                                output_domain=domain,
+                                normalize=True,
+                                compute_stats=True,
+                                compact_output=True,
+                            ),
+                            list(precip_models[:, t, :, :]),
+                        )
 
             precip_models_cascade_temp = np.array(
                 [decomp["cascade_levels"] for decomp in decomp_precip_models]

From 1a4ab4b18c4eac06630968841add9188a58fda83 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 12:44:12 +0200
Subject: [PATCH 05/18] fix black

---
 pysteps/blending/steps.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index ca42fad1..2f505341 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -51,8 +51,7 @@
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import (binary_dilation, generate_binary_structure,
-                           iterate_structure)
+from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils

From a468411fee839bfa113f3b87fe757dc556abe459 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 17:40:45 +0200
Subject: [PATCH 06/18] update docstring

---
 pysteps/blending/steps.py | 37 +++++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 2f505341..6f8f1717 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -119,24 +119,37 @@ def forecast(
       ordered by timestamp from oldest to newest. The time steps between the
       inputs are assumed to be regular.
     precip_models: array-like
-      Array of shape (n_models,timesteps+1) containing, per timestep (t=0 to
-      lead time here) and per (NWP) model or model ensemble member, a
-      dictionary with a list of cascades obtained by calling a method
-      implemented in :py:mod:`pysteps.cascade.decomposition`. In case of one
-      (deterministic) model as input, add an extra dimension to make sure
-      precip_models is five dimensional prior to calling this function.
-      It is also possible to supply the original (NWP) model forecasts containing
-      rainfall fields as an array of shape (n_models,timestep+1,m,n), which will
-      then be decomposed in this function. Note that for an operational application
-      or for testing with multiple model runs, it is recommended to decompose
-      the model forecasts outside beforehand, as this reduces calculation times.
+      Either raw (NWP) model forecast data or decomposed (NWP) model forecast data.
+      If you supply decomposed data, it needs to be an array of shape
+      (n_models,timesteps+1) containing, per timestep (t=0 to lead time here) and
+      per (NWP) model or model ensemble member, a dictionary with a list of cascades
+      obtained by calling a method implemented in :py:mod:`pysteps.cascade.decomposition`.
+      If you supply the original (NWP) model forecast data, it needs to be an array of shape
+      (n_models,timestep+1,m,n) containing rainfall fields, which will
+      then be decomposed in this function.
+
+      Depending on your use case it can be advantageous to decompose the model
+      forecasts outside beforehand, as this slightly reduces calculation times.
       This is possible with :py:func:`pysteps.blending.utils.decompose_NWP`,
       :py:func:`pysteps.blending.utils.compute_store_nwp_motion`, and
-      :py:func:`pysteps.blending.utils.load_NWP`.
+      :py:func:`pysteps.blending.utils.load_NWP`. However, if you have a lot of (NWP) model
+      members (e.g. 1 model member per nowcast member), this can lead to excessive memory
+      usage.
+
+      To further reduce memory usage, both this array and the ``velocity_models`` array
+      can be given as float32, they will then be converted to float64 before computations
+      to minimize loss in precision
+
+      In case of one (deterministic) model as input, add an extra dimension to make sure
+      precip_models is five dimensional prior to calling this function.
     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. All values are required to be finite.
+
+      To reduce memory usage, both this array and the ``precip_models`` array
+      can be given as float32, they will then be converted to float64 before computations
+      to minimize loss in precision
     velocity_models: array-like
       Array of shape (n_models,timestep,2,m,n) containing the x- and y-components
       of the advection field for the (NWP) model field per forecast lead time.

From a5c9b22af1f39505221bc5e67baf49cb1ed8913c Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 24 Sep 2024 17:44:13 +0200
Subject: [PATCH 07/18] fix one more docstring and rename some variables

---
 pysteps/blending/steps.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 6f8f1717..727b737d 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -595,7 +595,7 @@ def forecast(
         timesteps = nowcast_utils.binned_timesteps(original_timesteps)
         timestep_type = "list"
 
-    # 2.3.1 If precip is below the norain threshold and precip_models_pm is zero,
+    # 2.3.1 If precip is below the norain threshold and precip_models is zero,
     # we consider it as no rain in the domain.
     # The forecast will directly return an array filled with the minimum
     # value present in precip (which equals zero rainfall in the used
@@ -861,7 +861,7 @@ def forecast(
             # the maximum number of (ensemble) members in the input (either the nowcasts or NWP).
             (
                 precip_models_cascade_temp,
-                precip_models_pm_temp,
+                precip_models_temp,
                 velocity_models_temp,
                 mu_models_temp,
                 sigma_models_temp,
@@ -1467,7 +1467,7 @@ def worker(j):
                             R_pm_stacked = np.concatenate(
                                 (
                                     R_pm_ep[None, t_index],
-                                    precip_models_pm_temp,
+                                    precip_models_temp,
                                 ),
                                 axis=0,
                             )
@@ -1475,7 +1475,7 @@ def worker(j):
                             R_pm_stacked = np.concatenate(
                                 (
                                     R_pm_ep[None, t_index],
-                                    precip_models_pm_temp[None, j],
+                                    precip_models_temp[None, j],
                                 ),
                                 axis=0,
                             )
@@ -1492,11 +1492,11 @@ def worker(j):
                                 weights_pm_normalized_mod_only.reshape(
                                     weights_pm_normalized_mod_only.shape[0], 1, 1
                                 )
-                                * precip_models_pm_temp,
+                                * precip_models_temp,
                                 axis=0,
                             )
                         else:
-                            R_pm_blended_mod_only = precip_models_pm_temp[j]
+                            R_pm_blended_mod_only = precip_models_temp[j]
 
                         # The extrapolation components are NaN outside the advected
                         # radar domain. This results in NaN values in the blended
@@ -1587,7 +1587,7 @@ def worker(j):
                         # that for the probability matching.
                         if probmatching_method is not None and resample_distribution:
                             arr1 = R_pm_ep[t_index]
-                            arr2 = precip_models_pm_temp[j]
+                            arr2 = precip_models_temp[j]
                             # resample weights based on cascade level 2.
                             # Areas where one of the fields is nan are not included.
                             R_pm_resampled = probmatching.resample_distributions(
@@ -2090,7 +2090,7 @@ def _compute_cascade_decomposition_radar(
 
 
 def _compute_cascade_recomposition_nwp(precip_models, recompositor):
-    """If necessary, recompose (NWP) model forecasts and stack cascades."""
+    """If necessary, recompose (NWP) model forecasts."""
     precip_models_cascade = precip_models
     precip_models_pm = None
 

From 6dc7efef53c8e1e674c66e176e7460b1543e2155 Mon Sep 17 00:00:00 2001
From: mats-knmi <145579783+mats-knmi@users.noreply.github.com>
Date: Thu, 26 Sep 2024 16:54:02 +0200
Subject: [PATCH 08/18] Apply suggestions from code review

Co-authored-by: Ruben Imhoff <31476760+RubenImhoff@users.noreply.github.com>
---
 pysteps/blending/steps.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 727b737d..9bac47f3 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -125,7 +125,7 @@ def forecast(
       per (NWP) model or model ensemble member, a dictionary with a list of cascades
       obtained by calling a method implemented in :py:mod:`pysteps.cascade.decomposition`.
       If you supply the original (NWP) model forecast data, it needs to be an array of shape
-      (n_models,timestep+1,m,n) containing rainfall fields, which will
+      (n_models,timestep+1,m,n) containing precipitation (or other) fields, which will
       then be decomposed in this function.
 
       Depending on your use case it can be advantageous to decompose the model
@@ -137,19 +137,19 @@ def forecast(
       usage.
 
       To further reduce memory usage, both this array and the ``velocity_models`` array
-      can be given as float32, they will then be converted to float64 before computations
-      to minimize loss in precision
+      can be given as float32. They will then be converted to float64 before computations
+      to minimize loss in precision.
 
       In case of one (deterministic) model as input, add an extra dimension to make sure
-      precip_models is five dimensional prior to calling this function.
+      precip_models is four dimensional prior to calling this function.
     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. All values are required to be finite.
 
-      To reduce memory usage, both this array and the ``precip_models`` array
-      can be given as float32, they will then be converted to float64 before computations
-      to minimize loss in precision
+      To reduce memory usage, this array
+      can be given as float32. They will then be converted to float64 before computations
+      to minimize loss in precision.
     velocity_models: array-like
       Array of shape (n_models,timestep,2,m,n) containing the x- and y-components
       of the advection field for the (NWP) model field per forecast lead time.

From af0cd1571003ae8dfdd2ac78e6044ac8a9b32458 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Thu, 26 Sep 2024 17:40:49 +0200
Subject: [PATCH 09/18] rename precip_models_pm

---
 pysteps/blending/steps.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 9bac47f3..587000ae 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -2089,25 +2089,24 @@ def _compute_cascade_decomposition_radar(
     return R_c, mu_extrapolation, sigma_extrapolation
 
 
-def _compute_cascade_recomposition_nwp(precip_models, recompositor):
+def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor):
     """If necessary, recompose (NWP) model forecasts."""
-    precip_models_cascade = precip_models
-    precip_models_pm = None
+    precip_models = None
 
     # Recompose the (NWP) model cascades to have rainfall fields per
     # model and time step, which will be used in the probability matching steps.
     # Recomposed cascade will have shape: [n_models, n_timesteps, m, n]
-    precip_models_pm = []
+    precip_models = []
     for i in range(precip_models_cascade.shape[0]):
-        precip_model_pm = []
+        precip_model = []
         for time_step in range(precip_models_cascade.shape[1]):
-            precip_model_pm.append(recompositor(precip_models_cascade[i, time_step]))
-        precip_models_pm.append(precip_model_pm)
+            precip_model.append(recompositor(precip_models_cascade[i, time_step]))
+        precip_models.append(precip_model)
 
-    precip_models_pm = np.stack(precip_models_pm)
-    precip_model_pm = None
+    precip_models = np.stack(precip_models)
+    precip_model = None
 
-    return precip_models_pm
+    return precip_models
 
 
 def _estimate_ar_parameters_radar(

From 17c00c73a21c0116470f71449b8510fe98b81682 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Thu, 26 Sep 2024 17:43:03 +0200
Subject: [PATCH 10/18] move comment about memory reduction to velocity_models

---
 pysteps/blending/steps.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 587000ae..3a4965c9 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -146,14 +146,14 @@ def forecast(
       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. All values are required to be finite.
-
-      To reduce memory usage, this array
-      can be given as float32. They will then be converted to float64 before computations
-      to minimize loss in precision.
     velocity_models: array-like
       Array of shape (n_models,timestep,2,m,n) containing the x- and y-components
       of the advection field for the (NWP) model field per forecast lead time.
       All values are required to be finite.
+
+      To reduce memory usage, this array
+      can be given as float32. They will then be converted to float64 before computations
+      to minimize loss in precision.
     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

From 3d1366b4c938d32ca95f2537b446b071620a3fd2 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 1 Oct 2024 08:10:33 +0200
Subject: [PATCH 11/18] fix zerovalue precip_cascade

---
 pysteps/blending/steps.py | 44 +++++++++++++++++++++++++++------------
 1 file changed, 31 insertions(+), 13 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 3a4965c9..07b07a6f 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -644,8 +644,30 @@ def forecast(
     else:
         # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain
         # only nans or infs. If so, fill it with the zero value.
-        if zero_precip_radar:
-            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(precip)
+        if np.any(precip_cascade[np.isfinite(precip_cascade)]):
+            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(precip_cascade)
+        else:
+            done = False
+            for t in timesteps:
+                if done:
+                    break
+                for j in range(precip_models.shape[0]):
+                    if not blending.utils.check_norain(
+                        precip_models, precip_thr, norain_thr
+                    ):
+                        precip_models_cascade_temp = decompositor(
+                            precip_models[j, t, :, :],
+                            bp_filter=bp_filter,
+                            fft_method=fft,
+                            output_domain=domain,
+                            normalize=True,
+                            compute_stats=True,
+                            compact_output=True,
+                        )["cascade_levels"]
+                        precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
+                            precip_models_cascade_temp
+                        )
+                        break
 
         # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP
         # forecast. I.e., velocity (radar) equals velocity_model at the first time
@@ -2426,7 +2448,7 @@ def _init_noise_cascade(
 
 def _fill_nans_infs_nwp_cascade(
     precip_models_cascade,
-    precip_models_pm,
+    precip_models,
     precip_cascade,
     precip,
     mu_models,
@@ -2438,34 +2460,30 @@ def _fill_nans_infs_nwp_cascade(
     min_cascade = np.nanmin(precip_cascade)
     min_precip = np.nanmin(precip)
     precip_models_cascade[~np.isfinite(precip_models_cascade)] = min_cascade
-    precip_models_pm[~np.isfinite(precip_models_pm)] = min_precip
+    precip_models[~np.isfinite(precip_models)] = min_precip
     # Also set any nans or infs in the mean and sigma of the cascade to
     # respectively 0.0 and 1.0
     mu_models[~np.isfinite(mu_models)] = 0.0
     sigma_models[~np.isfinite(sigma_models)] = 0.0
 
-    return precip_models_cascade, precip_models_pm, mu_models, sigma_models
+    return precip_models_cascade, precip_models, mu_models, sigma_models
 
 
-def _determine_max_nr_rainy_cells_nwp(
-    precip_models_pm, precip_thr, n_models, timesteps
-):
+def _determine_max_nr_rainy_cells_nwp(precip_models, precip_thr, n_models, timesteps):
     """Initialize noise based on the NWP field time step where the fraction of rainy cells is highest"""
     if precip_thr is None:
-        precip_thr = np.nanmin(precip_models_pm)
+        precip_thr = np.nanmin(precip_models)
 
     max_rain_pixels = -1
     max_rain_pixels_j = -1
     max_rain_pixels_t = -1
     for j in range(n_models):
         for t in timesteps:
-            rain_pixels = precip_models_pm[j][t][
-                precip_models_pm[j][t] > precip_thr
-            ].size
+            rain_pixels = precip_models[j][t][precip_models[j][t] > precip_thr].size
             if rain_pixels > max_rain_pixels:
                 max_rain_pixels = rain_pixels
                 max_rain_pixels_j = j
                 max_rain_pixels_t = t
-    precip_noise_input = precip_models_pm[max_rain_pixels_j][max_rain_pixels_t]
+    precip_noise_input = precip_models[max_rain_pixels_j][max_rain_pixels_t]
 
     return precip_noise_input.astype(np.float64, copy=False)

From c5718cbf0e6c48070065818ed6e08e1148aadea8 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 1 Oct 2024 08:13:01 +0200
Subject: [PATCH 12/18] use precip_models_cascade if present

---
 pysteps/blending/steps.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 07b07a6f..d00bdbeb 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -646,6 +646,10 @@ def forecast(
         # only nans or infs. If so, fill it with the zero value.
         if np.any(precip_cascade[np.isfinite(precip_cascade)]):
             precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(precip_cascade)
+        elif precip_models_cascade is not None:
+            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
+                precip_models_cascade
+            )
         else:
             done = False
             for t in timesteps:

From f5e5697b8eb5811296c47929efbcb4afac2285d1 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 1 Oct 2024 09:33:51 +0200
Subject: [PATCH 13/18] set done to true

---
 pysteps/blending/steps.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index d00bdbeb..cb3944f7 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -51,7 +51,8 @@
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
+from scipy.ndimage import (binary_dilation, generate_binary_structure,
+                           iterate_structure)
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils
@@ -671,6 +672,7 @@ def forecast(
                         precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
                             precip_models_cascade_temp
                         )
+                        done = True
                         break
 
         # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP

From e478bc0909aeccdb6e8680984b0cc32ecfb64b9f Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Tue, 1 Oct 2024 09:35:40 +0200
Subject: [PATCH 14/18] fix black

---
 pysteps/blending/steps.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index cb3944f7..a5a18b7f 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -51,8 +51,7 @@
 
 import numpy as np
 from scipy.linalg import inv
-from scipy.ndimage import (binary_dilation, generate_binary_structure,
-                           iterate_structure)
+from scipy.ndimage import binary_dilation, generate_binary_structure, iterate_structure
 
 from pysteps import blending, cascade, extrapolation, noise, utils
 from pysteps.nowcasts import utils as nowcast_utils

From 7fd7122da77a47b8870e65bef6dbe8841c647623 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Wed, 2 Oct 2024 17:29:17 +0200
Subject: [PATCH 15/18] revert to old logic, but now with precip_models_cascade
 using new loop

---
 pysteps/blending/steps.py | 55 +++++++++++++++++++--------------------
 1 file changed, 27 insertions(+), 28 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index a5a18b7f..dba88706 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -644,35 +644,34 @@ def forecast(
     else:
         # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain
         # only nans or infs. If so, fill it with the zero value.
-        if np.any(precip_cascade[np.isfinite(precip_cascade)]):
-            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(precip_cascade)
-        elif precip_models_cascade is not None:
-            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
-                precip_models_cascade
-            )
-        else:
-            done = False
-            for t in timesteps:
-                if done:
-                    break
-                for j in range(precip_models.shape[0]):
-                    if not blending.utils.check_norain(
-                        precip_models, precip_thr, norain_thr
-                    ):
-                        precip_models_cascade_temp = decompositor(
-                            precip_models[j, t, :, :],
-                            bp_filter=bp_filter,
-                            fft_method=fft,
-                            output_domain=domain,
-                            normalize=True,
-                            compute_stats=True,
-                            compact_output=True,
-                        )["cascade_levels"]
-                        precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
-                            precip_models_cascade_temp
-                        )
-                        done = True
+        if zero_precip_radar:
+            if precip_models_cascade is not None:
+                precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
+                    precip_models_cascade
+                )
+            else:
+                done = False
+                for t in timesteps:
+                    if done:
                         break
+                    for j in range(precip_models.shape[0]):
+                        if not blending.utils.check_norain(
+                            precip_models, precip_thr, norain_thr
+                        ):
+                            precip_models_cascade_temp = decompositor(
+                                precip_models[j, t, :, :],
+                                bp_filter=bp_filter,
+                                fft_method=fft,
+                                output_domain=domain,
+                                normalize=True,
+                                compute_stats=True,
+                                compact_output=True,
+                            )["cascade_levels"]
+                            precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
+                                precip_models_cascade_temp
+                            )
+                            done = True
+                            break
 
         # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP
         # forecast. I.e., velocity (radar) equals velocity_model at the first time

From 6d51fd4b08a00f402403b83a530bb1a30df6b53e Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Wed, 2 Oct 2024 17:44:15 +0200
Subject: [PATCH 16/18] fix issue with precip_models_cascade

---
 pysteps/blending/steps.py | 50 +++++++++++++++++++--------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index dba88706..71e769ad 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -645,33 +645,33 @@ def forecast(
         # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain
         # only nans or infs. If so, fill it with the zero value.
         if zero_precip_radar:
-            if precip_models_cascade is not None:
-                precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
-                    precip_models_cascade
-                )
-            else:
-                done = False
-                for t in timesteps:
-                    if done:
-                        break
-                    for j in range(precip_models.shape[0]):
-                        if not blending.utils.check_norain(
-                            precip_models, precip_thr, norain_thr
-                        ):
-                            precip_models_cascade_temp = decompositor(
-                                precip_models[j, t, :, :],
-                                bp_filter=bp_filter,
-                                fft_method=fft,
-                                output_domain=domain,
-                                normalize=True,
-                                compute_stats=True,
-                                compact_output=True,
-                            )["cascade_levels"]
+            done = False
+            for t in timesteps:
+                if done:
+                    break
+                for j in range(precip_models.shape[0]):
+                    if not blending.utils.check_norain(
+                        precip_models, precip_thr, norain_thr
+                    ):
+                        if precip_models_cascade is not None:
                             precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
-                                precip_models_cascade_temp
+                                precip_models_cascade[j, t]
                             )
-                            done = True
-                            break
+                            continue
+                        precip_models_cascade_temp = decompositor(
+                            precip_models[j, t, :, :],
+                            bp_filter=bp_filter,
+                            fft_method=fft,
+                            output_domain=domain,
+                            normalize=True,
+                            compute_stats=True,
+                            compact_output=True,
+                        )["cascade_levels"]
+                        precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
+                            precip_models_cascade_temp
+                        )
+                        done = True
+                        break
 
         # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP
         # forecast. I.e., velocity (radar) equals velocity_model at the first time

From 0198a54681b54ec756adecc0888497991bf5349c Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Wed, 2 Oct 2024 19:44:02 +0200
Subject: [PATCH 17/18] fix tests

---
 pysteps/blending/steps.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 71e769ad..00cfcb08 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -650,12 +650,13 @@ def forecast(
                 if done:
                     break
                 for j in range(precip_models.shape[0]):
+                    # Look for a timestep and member with rain so that we have a sensible decomposition
                     if not blending.utils.check_norain(
-                        precip_models, precip_thr, norain_thr
+                        precip_models[j, t], precip_thr, norain_thr
                     ):
                         if precip_models_cascade is not None:
                             precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin(
-                                precip_models_cascade[j, t]
+                                precip_models_cascade[j, t]["cascade_levels"]
                             )
                             continue
                         precip_models_cascade_temp = decompositor(

From 5c14399f67a5b1f7c41b80a3898506400e986c98 Mon Sep 17 00:00:00 2001
From: Mats Veldhuizen <mats.veldhuizen@knmi.nl>
Date: Thu, 3 Oct 2024 17:38:11 +0200
Subject: [PATCH 18/18] Move comment up a little

---
 pysteps/blending/steps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
index 00cfcb08..a6ecc09f 100644
--- a/pysteps/blending/steps.py
+++ b/pysteps/blending/steps.py
@@ -645,12 +645,12 @@ def forecast(
         # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain
         # only nans or infs. If so, fill it with the zero value.
         if zero_precip_radar:
+            # Look for a timestep and member with rain so that we have a sensible decomposition
             done = False
             for t in timesteps:
                 if done:
                     break
                 for j in range(precip_models.shape[0]):
-                    # Look for a timestep and member with rain so that we have a sensible decomposition
                     if not blending.utils.check_norain(
                         precip_models[j, t], precip_thr, norain_thr
                     ):