Implement the full_output option to return forward-backward advection…

… and consistency fields

pysteps/motion/_proesmans.pyx

@@ -29,6 +29,8 @@ def _compute_advection_field(float64 [:, :, :] R, lam, intp num_iter,
     cdef np.ndarray[float64, ndim=4] V_cur = np.zeros((2, 2, m, n))
     cdef np.ndarray[float64, ndim=4] V_next
 
+    cdef np.ndarray[float64, ndim=3] GAMMA = np.empty((2, R.shape[1], R.shape[2]))
+
     for i in range(n_levels-1, -1, -1):
         _proesmans(np.stack([R_p[0][i], R_p[1][i]]), V_cur, num_iter, lam)
 
@@ -41,7 +43,9 @@ def _compute_advection_field(float64 [:, :, :] R, lam, intp num_iter,
             _initialize_next_level(V_cur, V_next)
             V_cur = V_next
 
-    return V_cur[0, :, :, :]
+    _compute_consistency_maps(V_cur, GAMMA)
+
+    return V_cur, GAMMA
 
 @cython.boundscheck(False)
 @cython.wraparound(False)

pysteps/motion/proesmans.py

@@ -19,7 +19,13 @@
 
 @check_input_frames(2, 2)
 def proesmans(
-    input_images, lam=50.0, num_iter=100, num_levels=6, filter_std=0.0, verbose=True,
+    input_images,
+    lam=50.0,
+    num_iter=100,
+    num_levels=6,
+    filter_std=0.0,
+    verbose=True,
+    full_output=False,
 ):
     """Implementation of the anisotropic diffusion method of Proesmans et al.
     (1994).
@@ -40,14 +46,23 @@ def proesmans(
         computing the optical flow.
     verbose : bool, optional
         Verbosity enabled if True (default).
+    full_output : bool, optional
+        If True, the output is a two-element tuple containing the
+        forward-backward advection and consistency fields. The first element
+        is shape (2, 2, m, n), where the index along the first dimension refers
+        to the forward and backward advection fields. The second element is an
+        array of shape (2, m, n), where the index along the first dimension
+        refers to the forward and backward consistency fields.
+        Default : False.
 
     Returns
     -------
     out : ndarray
-        The advection field having shape (2, m, n), where out[0, :, :] contains
-        the x-components of the motion vectors and out[1, :, :] contains the
-        y-components. The velocities are in units of pixels / timestep, where
-        timestep is the time difference between the two input images.
+        If full_output=False, the advection field having shape (2, m, n), where
+        out[0, :, :] contains the x-components of the motion vectors and
+        out[1, :, :] contains the y-components. The velocities are in units of
+        pixels / timestep, where timestep is the time difference between the
+        two input images.
 
     References
     ----------
@@ -69,4 +84,9 @@ def proesmans(
         im[0, :, :] = gaussian_filter(im[0, :, :], filter_std)
         im[1, :, :] = gaussian_filter(im[1, :, :], filter_std)
 
-    return _compute_advection_field(im, lam, num_iter, num_levels)
+    advfield, quality = _compute_advection_field(im, lam, num_iter, num_levels)
+
+    if not full_output:
+        return advfield[0]
+    else:
+        return advfield, quality