Fix vet padding bug #163 (PR #164)

This PR fixes issue #163.

- Fix the issue with the pysteps.motion.vet function's padding argument when MaskedArrays are used as input.
- Add a short test to check that pysteps.motion.vet correctly pad Masked and Non-Masked input images.

Author: aperezhortal

pysteps/motion/vet.py

@@ -513,20 +513,25 @@ def debug_print(*args, **kwargs):
             + "Supported values: {0}".format(str(valid_indexing))
         )
 
-    # Get mask
-    if isinstance(input_images, MaskedArray):
-        mask = numpy.ma.getmaskarray(input_images)
-    else:
-        # Mask invalid data
-        if padding > 0:
-            padding_tuple = ((0, 0), (padding, padding), (padding, padding))
+    # Convert input_images to a MaskedArray if it is a regular ndarray
+    if not isinstance(input_images, MaskedArray):
+        input_images = numpy.ma.masked_invalid(input_images)
 
-            input_images = numpy.pad(
-                input_images, padding_tuple, "constant", constant_values=numpy.nan
-            )
+    mask = numpy.ma.getmaskarray(input_images)
 
-        input_images = numpy.ma.masked_invalid(input_images)
-        mask = numpy.ma.getmaskarray(input_images)
+    if padding > 0:
+        padding_tuple = ((0, 0), (padding, padding), (padding, padding))
+
+        input_images_data = numpy.pad(
+            numpy.ma.getdata(input_images),
+            padding_tuple,
+            "constant",
+            constant_values=numpy.nan,
+        )
+
+        mask = numpy.pad(mask, padding_tuple, "constant", constant_values=True)
+
+        input_images = numpy.ma.MaskedArray(data=input_images_data, mask=mask)
 
     input_images.data[mask] = 0  # Remove any Nan from the raw data
 

pysteps/tests/test_motion.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import pytest
+from functools import partial
 from scipy.ndimage import uniform_filter
 
 import pysteps as stp
@@ -324,6 +325,40 @@ def test_input_shape_checks(
             motion_method(np.zeros((frames, image_size, image_size)), verbose=False)
 
 
+def test_vet_padding():
+    """
+    Test that the padding functionality in vet works correctly with ndarrays and
+    masked arrays.
+    """
+
+    _, precip_obs = _create_observations(
+        reference_field.copy(), "linear_y", num_times=2
+    )
+
+    # Use a small region to speed up the test
+    precip_obs = precip_obs[:, 200:427, 250:456]
+    # precip_obs.shape == (227 , 206)
+    # 227 is a prime number ; 206 = 2*103
+    # Using this shape will force vet to internally pad the input array for the sector's
+    # blocks to divide exactly the input shape.
+    # NOTE: This "internal padding" is different from the padding keyword being test next.
+
+    for padding in [0, 3, 10]:
+        # No padding
+        vet_method = partial(
+            motion.get_method("vet"),
+            verbose=False,
+            sectors=((16, 4, 2), (16, 4, 2)),
+            options=dict(maxiter=5),
+            padding=padding
+            # We use only a few iterations since
+            # we don't care about convergence in this test
+        )
+
+        assert precip_obs.shape == vet_method(precip_obs).shape
+        assert precip_obs.shape == vet_method(np.ma.masked_invalid(precip_obs)).shape
+
+
 def test_vet_cost_function():
     """
     Test that the vet cost_function computation gives always the same result