Extrapolation fixes (#290)

* Fix incorrect replacement of nan values with zeros when interp_order>1

* Remove unnecessary encoding line

* Set allow_nonfinite_values automatically based on the input data

Author: pulkkins

pysteps/extrapolation/semilagrangian.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 """
 pysteps.extrapolation.semilagrangian
 ====================================
@@ -152,6 +151,8 @@ def extrapolate(
             precip = precip.copy()
             precip[~mask_finite] = 0.0
             mask_finite = mask_finite.astype(float)
+        else:
+            mask_finite = np.ones(precip.shape)
 
     prefilter = True if interp_order > 1 else False
 
@@ -241,16 +242,15 @@ def interpolate_motion(displacement, velocity_inc, td):
                 )
                 precip_warped[mask_warped < 0.5] = minval
 
-                if allow_nonfinite_values:
-                    mask_warped = ip.map_coordinates(
-                        mask_finite,
-                        coords_warped,
-                        mode=map_coordinates_mode,
-                        cval=0,
-                        order=1,
-                        prefilter=False,
-                    )
-                    precip_warped[mask_warped < 0.5] = np.nan
+                mask_warped = ip.map_coordinates(
+                    mask_finite,
+                    coords_warped,
+                    mode=map_coordinates_mode,
+                    cval=0,
+                    order=1,
+                    prefilter=False,
+                )
+                precip_warped[mask_warped < 0.5] = np.nan
 
             precip_extrap.append(np.reshape(precip_warped, precip.shape))
 

pysteps/nowcasts/anvil.py

@@ -196,6 +196,10 @@ def forecast(
 
     # transform the input fields to Lagrangian coordinates by extrapolation
     extrapolator = extrapolation.get_method(extrap_method)
+    extrap_kwargs["allow_nonfinite_values"] = (
+        True if np.any(~np.isfinite(vil)) else False
+    )
+
     res = list()
 
     def worker(vil, i):
@@ -205,7 +209,6 @@ def worker(vil, i):
                 vil[i, :],
                 velocity,
                 vil.shape[0] - 1 - i,
-                allow_nonfinite_values=True,
                 **extrap_kwargs,
             )[-1],
         )
@@ -364,12 +367,16 @@ def worker(vil, i):
                     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],
-                    allow_nonfinite_values=True,
                     **extrap_kwargs,
                 )
                 r_f.append(r_f_ep[0])
@@ -384,7 +391,6 @@ def worker(vil, i):
                 None,
                 velocity,
                 [t_diff_prev],
-                allow_nonfinite_values=True,
                 **extrap_kwargs,
             )
             t_prev = t + 1

pysteps/nowcasts/extrapolation.py

@@ -12,6 +12,7 @@
 """
 
 import time
+import numpy as np
 
 from pysteps import extrapolation
 
@@ -73,6 +74,10 @@ def forecast(
     else:
         extrap_kwargs = extrap_kwargs.copy()
 
+    extrap_kwargs["allow_nonfinite_values"] = (
+        True if np.any(~np.isfinite(precip)) else False
+    )
+
     if measure_time:
         print(
             "Computing extrapolation nowcast from a "

pysteps/nowcasts/linda.py

@@ -207,6 +207,8 @@ def forecast(
         feature_kwargs = dict()
     if extrap_kwargs is None:
         extrap_kwargs = dict()
+    else:
+        extrap_kwargs = extrap_kwargs.copy()
 
     if localization_window_radius is None:
         localization_window_radius = 0.2 * np.min(precip_fields.shape[1:])
@@ -276,6 +278,10 @@ def forecast(
             vel_pert_kwargs["vp_par"] = vp_par
             vel_pert_kwargs["vp_perp"] = vp_perp
 
+    extrap_kwargs["allow_nonfinite_values"] = (
+        True if np.any(~np.isfinite(precip_fields)) else False
+    )
+
     fct_gen = _linda_deterministic_init(
         precip_fields,
         advection_field,

pysteps/nowcasts/sprog.py

@@ -214,7 +214,7 @@ def forecast(
 
     extrap_kwargs = extrap_kwargs.copy()
     extrap_kwargs["xy_coords"] = xy_coords
-    extrap_kwargs["allow_nonfinite_values"] = True
+    extrap_kwargs["allow_nonfinite_values"] = True if np.any(~np.isfinite(R)) else False
 
     # advect the previous precipitation fields to the same position with the
     # most recent one (i.e. transform them into the Lagrangian coordinates)

pysteps/nowcasts/sseps.py

@@ -316,6 +316,8 @@ def forecast(
     R = R[-(ar_order + 1) :, :, :].copy()
     extrap_kwargs = extrap_kwargs.copy()
     extrap_kwargs["xy_coords"] = xy_coords
+    extrap_kwargs["allow_nonfinite_values"] = True if np.any(~np.isfinite(R)) else False
+
     res = []
     f = lambda R, i: extrapolator_method(
         R[i, :, :], V, ar_order - i, "min", **extrap_kwargs

pysteps/nowcasts/steps.py

@@ -402,7 +402,8 @@ def forecast(
     # most recent one (i.e. transform them into the Lagrangian coordinates)
     extrap_kwargs = extrap_kwargs.copy()
     extrap_kwargs["xy_coords"] = xy_coords
-    extrap_kwargs["allow_nonfinite_values"] = True
+    extrap_kwargs["allow_nonfinite_values"] = True if np.any(~np.isfinite(R)) else False
+
     res = list()
 
     def f(R, i):