From f6bfb6eab180f55b762c3320e62d2778c4a48aef Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:07:30 +0000 Subject: [PATCH 1/8] Add correlation in the interface --- pysteps/motion/interface.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pysteps/motion/interface.py b/pysteps/motion/interface.py index eac81aa46..1668f0129 100644 --- a/pysteps/motion/interface.py +++ b/pysteps/motion/interface.py @@ -30,6 +30,7 @@ from pysteps.motion.lucaskanade import dense_lucaskanade from pysteps.motion.proesmans import proesmans from pysteps.motion.vet import vet +from pysteps.motion.correlation import correlation _methods = dict() _methods["constant"] = constant @@ -38,6 +39,7 @@ _methods["darts"] = DARTS _methods["proesmans"] = proesmans _methods["vet"] = vet +_methods["correlation"] = correlation _methods[None] = lambda precip, *args, **kw: np.zeros( (2, precip.shape[1], precip.shape[2]) ) @@ -72,6 +74,9 @@ def get_method(name): | | Laroche and Zawadzki (1995) and | | | Germann and Zawadzki (2002) | +-------------------+------------------------------------------------------+ + | correlation | implementation of the classical correlation based | + | | method used in Haiden et al (2011) | + +-------------------+------------------------------------------------------+ +--------------------------------------------------------------------------+ | Methods implemented in C (these require separate compilation and linkage)| From db61111fd5a70b81acbbdbfd5442246f0836ba12 Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:08:07 +0000 Subject: [PATCH 2/8] Introduce Numba correlation-like motion vector computation --- pysteps/motion/correlation.py | 313 ++++++++++++++++++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 pysteps/motion/correlation.py diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py new file mode 100644 index 000000000..3afb44614 --- /dev/null +++ b/pysteps/motion/correlation.py @@ -0,0 +1,313 @@ +# -*- coding: utf-8 -*- +""" +pysteps.motion.correlation +======================== + +Implementation of the classical correlation based motion vector. + +.. autosummary:: + :toctree: ../generated/ + + correlation +""" + +import numpy as np +import math + +from pysteps.decorators import check_input_frames +from pysteps.utils.cleansing import detect_outliers +from pysteps import utils + +#To delete +import pdb + +#Check if numba can be imported +try: + from numba import jit + + NUMBA_IMPORTED = True +except ImportError: + NUMBA_IMPORTED = False + + +@check_input_frames(2) +def correlation( + input_images, + settings={}, + interp_method="idwinterp2d", + interp_kwargs=None, + dense=True, + nr_std_outlier=3, + k_outlier=30, + verbose=False, +): + + """ + Run the correlation based optical flow method and interpolate the motion + vectors. This is using NUMBA to speed up the correlation computation. + + The sparse motion vectors are finally interpolated to return the whole + motion field. + + Parameters + ---------- + input_images: ndarray_ + Array of shape (T, m, n) containing a sequence of *T* two-dimensional + input images of shape (m, n). The indexing order in **input_images** is + assumed to be (time, latitude, longitude). + + *T* = 2 is the minimum required number of images. + With *T* > 2, all the resulting obtained vectors are pooled together for + the final interpolation on a regular grid. + + settings: dict, optional + Optional dictionary containing keyword arguments for the correlation + based algorithm. + + interp_method: {"idwinterp2d", "rbfinterp2d"}, optional + Name of the interpolation method to use. See interpolation methods in + :py:mod:`pysteps.utils.interpolate`. + + interp_kwargs: dict, optional + Optional dictionary containing keyword arguments for the interpolation + algorithm. See the documentation of :py:mod:`pysteps.utils.interpolate`. + + dense: bool, optional + If True, return the three-dimensional array (2, m, n) containing + the dense x- and y-components of the motion field. + + If False, return the sparse motion vectors as 2-D **xy** and **uv** + arrays, where **xy** defines the vector positions, **uv** defines the + x and y direction components of the vectors. + + nr_std_outlier: int, optional + Maximum acceptable deviation from the mean in terms of number of + standard deviations. Any sparse vector with a deviation larger than + this threshold is flagged as outlier and excluded from the + interpolation. + See the documentation of + :py:func:`pysteps.utils.cleansing.detect_outliers`. + + k_outlier: int or None, optional + The number of nearest neighbors used to localize the outlier detection. + If set to None, it employs all the data points (global detection). + See the documentation of + :py:func:`pysteps.utils.cleansing.detect_outliers`. + + verbose: bool, optional + If set to True, print some information about the program. + + Returns + ------- + out: ndarray_ or tuple + If **dense=True** (the default), return 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. + Return a zero motion field of shape (2, m, n) when no motion is + detected. + + If **dense=False**, it returns a tuple containing the 2-dimensional + arrays **xy** and **uv**, where x, y define the vector locations, + u, v define the x and y direction components of the vectors. + Return two empty arrays when no motion is detected. + + References + ---------- + Haiden, T., A. Kann, C. Wittmann, G. Pistotnik, B. Bica, and C. Gruber, 2011: The + Integrated Nowcasting through Comprehensive Analysis (INCA) System and Its + Validation over the Eastern Alpine Region. Wea. Forecasting, 26, 166–183. + """ + + input_images = input_images.copy() + + if interp_kwargs is None: + interp_kwargs = dict() + + if verbose: + print("Computing the motion field with the classical correlation-based method.") + t0 = time.time() + + nr_fields = input_images.shape[0] + domain_size = (input_images.shape[1], input_images.shape[2]) + + PMIN = settings.get('PMIN', 0.05) + nthin = settings.get('nthin', 15) + + NI = domain_size[0] + NJ = domain_size[1] + + NX = int(np.ceil(NI / nthin)) + NY = int(np.ceil(NJ / nthin)) + + xgrid = np.arange(domain_size[1]) + ygrid = np.arange(domain_size[0]) + X, Y = xgrid[::nthin], ygrid[::nthin] + XX, YY =np.meshgrid(X,Y, indexing='ij') + + U_t = [] + V_t = [] + + for n in range(nr_fields - 1): + # extract consecutive images + prvs_img = input_images[n, :, :].copy() + next_img = input_images[n + 1, :, :].copy() + + # Removing values below the PMIN threshold + prvs_img[prvs_img < PMIN] = 0. + next_img[next_img < PMIN] = 0. + + Uana = np.full((NJ, NI),np.nan)[::nthin, ::nthin] + Vana = np.full((NJ, NI),np.nan)[::nthin, ::nthin] + + Uana, Vana = compute_motion_numba( + next_img, prvs_img, + NI, NJ, nthin, Uana, Vana) + + U_t.append(Uana) + V_t.append(Vana) + + Uana = np.nanmean(np.array(U_t),axis=0) + Vana = np.nanmean(np.array(V_t),axis=0) + + ## Choose computed boxes + boole_ = (np.isfinite(Uana) & np.isfinite(Vana)) + + if np.sum(boole_) == 0: + return np.zeros((2, domain_size[0], domain_size[1])) + + uv = np.vstack([Uana[boole_],Vana[boole_]]).T + xy = np.vstack([XX[boole_],YY[boole_]]).T + + # detect and remove outliers + outliers = detect_outliers(uv, nr_std_outlier, xy, k_outlier, verbose) + xy = xy[~outliers, :] + uv = uv[~outliers, :] + + if verbose: + print("--- LK found %i sparse vectors ---" % xy.shape[0]) + + # return sparse vectors if required + if not dense: + return xy, uv + + # # interpolation + interpolation_method = utils.get_method(interp_method) + uvgrid = interpolation_method(xy, uv, xgrid, ygrid, **interp_kwargs) + + if verbose: + print("--- total time: %.2f seconds ---" % (time.time() - t0)) + + return uvgrid + + +@jit(nopython=True) +def compute_motion_numba(image1, image2, NI, NJ, nthin, + IS_tmp, JS_tmp): + + """ + Compute the motion between two images using a correlation-based method optimized with Numba. + + Parameters: + ----------- + image1 : 2D numpy array + The first input image for motion computation. + image2 : 2D numpy array + The second input image for motion computation. + NI : int + The number of rows in the input images. + NJ : int + The number of columns in the input images. + nthin : int + The thinning factor for the grid used in motion computation. + IS_tmp : 2D numpy array + The output array to store the computed motion in the x-direction. + JS_tmp : 2D numpy array + The output array to store the computed motion in the y-direction. + + Returns: + -------- + IS_tmp : 2D numpy array + The updated motion array in the x-direction after computation. + JS_tmp : 2D numpy array + The updated motion array in the y-direction after computation. + + Notes: + ------ + - This function computes the motion by correlating patches of the first image (`image1`) + with patches of the second image (`image2`). + - The computation is performed on a grid defined by the `nthin` parameter to reduce + computational complexity. + - The function uses a fixed-size search window (`nsh`) and a quantization parameter (`nqu`) + to define the neighborhood for correlation computation. + - The correlation is calculated within a defined window and the maximum correlation value is + used to determine the motion vector. + - The function utilizes several optimization techniques to ensure efficient computation + and is compiled with Numba for further speed-up. + + """ + + ## List of parameters from the original paper (TODO: Optimization as input values in func.) + nsh = 20 + nqu = 45 + nsq = nsh + nqu + di = 1 + PAREAMIN = 1.0 + rr_dpa_min = 0.03 + nn = (2 * nqu / di + 1) ** 2. + + for i in range(0, NI, nthin): + if (i >= nsq) and (i < NI - nsq): + for j in range(0, NJ, nthin): + if (j >= nsq) and (j < NJ - nsq): + ii1 = max(i - nqu, 0) + ii2 = min(i + nqu, NI - 1) + jj1 = max(j - nqu, 0) + jj2 = min(j + nqu, NJ - 1) + + sy = 0. + sy2 = 0. + + for ii in range(ii1, ii2+1, di): + for jj in range(jj1, jj2+1, di): + sy = sy + image1[jj, ii] + sy2 = sy2 + image1[jj, ii] ** 2. + + sigy = sy2 - sy ** 2. / nn + isho = -99 + jsho = -99 + + if (sigy > 0.) and (sy > PAREAMIN): + corqx = 0.1 + for ish in range(-nsh, nsh + 1): + for jsh in range(-nsh, nsh + 1): + if ( math.sqrt((ish)**2. + (jsh)**2.) > nsh ): continue + sx = 0. + sx2 = 0. + sxy = 0. + for ii in range(ii1, ii2+1, di): + for jj in range(jj1, jj2+1, di): + ind_x = min(max(0, jj + jsh), NJ - 1) + ind_y = min(max(0, ii + ish), NI - 1) + sx += image2[ind_x, ind_y] + sx2 += image2[ind_x, ind_y] ** 2. + sxy += image2[ind_x, ind_y] * image1[jj, ii] + + sigx = sx2 - sx ** 2. / nn + if sigx > 0.: + corq = (sxy -sx * sy / nn) ** 2. / (sigx * sigy) + if corq > corqx: + corqx = corq + isho = ish + jsho = jsh + + if (isho != -99) and (corqx > 0.3) and (corqx * sy > 1.) and (sigy / sy >= 0.08): + if (abs(isho) < nsh) and (abs(jsho) < nsh): + IS_tmp[int(j/nthin),int(i/nthin)] = -isho + JS_tmp[int(j/nthin),int(i/nthin)] = -jsho + else: + IS_tmp[int(j/nthin),int(i/nthin)] = 0 + JS_tmp[int(j/nthin),int(i/nthin)] = 0 + + return IS_tmp, JS_tmp \ No newline at end of file From 07ddbf1fc9ac7bd48f98c0ac9d82b7df935aecfd Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:08:37 +0000 Subject: [PATCH 3/8] Include correlation tests --- pysteps/tests/test_motion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pysteps/tests/test_motion.py b/pysteps/tests/test_motion.py index 8d198960a..eb188fa57 100644 --- a/pysteps/tests/test_motion.py +++ b/pysteps/tests/test_motion.py @@ -252,6 +252,7 @@ def test_optflow_method_convergence( no_precip_args_values = [ ("lk", 2), ("lk", 3), + ("correlation", 2), ("vet", 2), ("vet", 3), ("darts", 9), @@ -293,6 +294,7 @@ def test_no_precipitation(optflow_method_name, num_times): ) input_tests_args_values = [ ("lk", 2, np.inf), + ("correlation", 2, np.inf), ("vet", 2, 3), ("darts", 9, 9), ("proesmans", 2, 2), From 0b75db80fab71c9083a6f7baed53962d28d6b376 Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:19:52 +0000 Subject: [PATCH 4/8] black modification of formatting --- pysteps/motion/correlation.py | 113 +++++++++++++++++----------------- 1 file changed, 58 insertions(+), 55 deletions(-) diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py index 3afb44614..ae088e1e3 100644 --- a/pysteps/motion/correlation.py +++ b/pysteps/motion/correlation.py @@ -18,10 +18,10 @@ from pysteps.utils.cleansing import detect_outliers from pysteps import utils -#To delete +# To delete import pdb -#Check if numba can be imported +# Check if numba can be imported try: from numba import jit @@ -115,8 +115,8 @@ def correlation( References ---------- - Haiden, T., A. Kann, C. Wittmann, G. Pistotnik, B. Bica, and C. Gruber, 2011: The - Integrated Nowcasting through Comprehensive Analysis (INCA) System and Its + Haiden, T., A. Kann, C. Wittmann, G. Pistotnik, B. Bica, and C. Gruber, 2011: The + Integrated Nowcasting through Comprehensive Analysis (INCA) System and Its Validation over the Eastern Alpine Region. Wea. Forecasting, 26, 166–183. """ @@ -132,8 +132,8 @@ def correlation( nr_fields = input_images.shape[0] domain_size = (input_images.shape[1], input_images.shape[2]) - PMIN = settings.get('PMIN', 0.05) - nthin = settings.get('nthin', 15) + PMIN = settings.get("PMIN", 0.05) + nthin = settings.get("nthin", 15) NI = domain_size[0] NJ = domain_size[1] @@ -144,7 +144,7 @@ def correlation( xgrid = np.arange(domain_size[1]) ygrid = np.arange(domain_size[0]) X, Y = xgrid[::nthin], ygrid[::nthin] - XX, YY =np.meshgrid(X,Y, indexing='ij') + XX, YY = np.meshgrid(X, Y, indexing="ij") U_t = [] V_t = [] @@ -155,30 +155,28 @@ def correlation( next_img = input_images[n + 1, :, :].copy() # Removing values below the PMIN threshold - prvs_img[prvs_img < PMIN] = 0. - next_img[next_img < PMIN] = 0. + prvs_img[prvs_img < PMIN] = 0.0 + next_img[next_img < PMIN] = 0.0 - Uana = np.full((NJ, NI),np.nan)[::nthin, ::nthin] - Vana = np.full((NJ, NI),np.nan)[::nthin, ::nthin] + Uana = np.full((NJ, NI), np.nan)[::nthin, ::nthin] + Vana = np.full((NJ, NI), np.nan)[::nthin, ::nthin] - Uana, Vana = compute_motion_numba( - next_img, prvs_img, - NI, NJ, nthin, Uana, Vana) + Uana, Vana = compute_motion_numba(next_img, prvs_img, NI, NJ, nthin, Uana, Vana) U_t.append(Uana) V_t.append(Vana) - Uana = np.nanmean(np.array(U_t),axis=0) - Vana = np.nanmean(np.array(V_t),axis=0) + Uana = np.nanmean(np.array(U_t), axis=0) + Vana = np.nanmean(np.array(V_t), axis=0) ## Choose computed boxes - boole_ = (np.isfinite(Uana) & np.isfinite(Vana)) + boole_ = np.isfinite(Uana) & np.isfinite(Vana) - if np.sum(boole_) == 0: + if np.sum(boole_) == 0: return np.zeros((2, domain_size[0], domain_size[1])) - uv = np.vstack([Uana[boole_],Vana[boole_]]).T - xy = np.vstack([XX[boole_],YY[boole_]]).T + uv = np.vstack([Uana[boole_], Vana[boole_]]).T + xy = np.vstack([XX[boole_], YY[boole_]]).T # detect and remove outliers outliers = detect_outliers(uv, nr_std_outlier, xy, k_outlier, verbose) @@ -203,12 +201,11 @@ def correlation( @jit(nopython=True) -def compute_motion_numba(image1, image2, NI, NJ, nthin, - IS_tmp, JS_tmp): +def compute_motion_numba(image1, image2, NI, NJ, nthin, IS_tmp, JS_tmp): """ Compute the motion between two images using a correlation-based method optimized with Numba. - + Parameters: ----------- image1 : 2D numpy array @@ -225,27 +222,27 @@ def compute_motion_numba(image1, image2, NI, NJ, nthin, The output array to store the computed motion in the x-direction. JS_tmp : 2D numpy array The output array to store the computed motion in the y-direction. - + Returns: -------- IS_tmp : 2D numpy array The updated motion array in the x-direction after computation. JS_tmp : 2D numpy array The updated motion array in the y-direction after computation. - + Notes: ------ - - This function computes the motion by correlating patches of the first image (`image1`) + - This function computes the motion by correlating patches of the first image (`image1`) with patches of the second image (`image2`). - - The computation is performed on a grid defined by the `nthin` parameter to reduce + - The computation is performed on a grid defined by the `nthin` parameter to reduce computational complexity. - - The function uses a fixed-size search window (`nsh`) and a quantization parameter (`nqu`) + - The function uses a fixed-size search window (`nsh`) and a quantization parameter (`nqu`) to define the neighborhood for correlation computation. - - The correlation is calculated within a defined window and the maximum correlation value is + - The correlation is calculated within a defined window and the maximum correlation value is used to determine the motion vector. - - The function utilizes several optimization techniques to ensure efficient computation + - The function utilizes several optimization techniques to ensure efficient computation and is compiled with Numba for further speed-up. - + """ ## List of parameters from the original paper (TODO: Optimization as input values in func.) @@ -255,7 +252,7 @@ def compute_motion_numba(image1, image2, NI, NJ, nthin, di = 1 PAREAMIN = 1.0 rr_dpa_min = 0.03 - nn = (2 * nqu / di + 1) ** 2. + nn = (2 * nqu / di + 1) ** 2.0 for i in range(0, NI, nthin): if (i >= nsq) and (i < NI - nsq): @@ -266,48 +263,54 @@ def compute_motion_numba(image1, image2, NI, NJ, nthin, jj1 = max(j - nqu, 0) jj2 = min(j + nqu, NJ - 1) - sy = 0. - sy2 = 0. + sy = 0.0 + sy2 = 0.0 - for ii in range(ii1, ii2+1, di): - for jj in range(jj1, jj2+1, di): + for ii in range(ii1, ii2 + 1, di): + for jj in range(jj1, jj2 + 1, di): sy = sy + image1[jj, ii] - sy2 = sy2 + image1[jj, ii] ** 2. + sy2 = sy2 + image1[jj, ii] ** 2.0 - sigy = sy2 - sy ** 2. / nn + sigy = sy2 - sy**2.0 / nn isho = -99 jsho = -99 - if (sigy > 0.) and (sy > PAREAMIN): + if (sigy > 0.0) and (sy > PAREAMIN): corqx = 0.1 for ish in range(-nsh, nsh + 1): for jsh in range(-nsh, nsh + 1): - if ( math.sqrt((ish)**2. + (jsh)**2.) > nsh ): continue - sx = 0. - sx2 = 0. - sxy = 0. - for ii in range(ii1, ii2+1, di): - for jj in range(jj1, jj2+1, di): + if math.sqrt((ish) ** 2.0 + (jsh) ** 2.0) > nsh: + continue + sx = 0.0 + sx2 = 0.0 + sxy = 0.0 + for ii in range(ii1, ii2 + 1, di): + for jj in range(jj1, jj2 + 1, di): ind_x = min(max(0, jj + jsh), NJ - 1) ind_y = min(max(0, ii + ish), NI - 1) sx += image2[ind_x, ind_y] - sx2 += image2[ind_x, ind_y] ** 2. + sx2 += image2[ind_x, ind_y] ** 2.0 sxy += image2[ind_x, ind_y] * image1[jj, ii] - sigx = sx2 - sx ** 2. / nn - if sigx > 0.: - corq = (sxy -sx * sy / nn) ** 2. / (sigx * sigy) + sigx = sx2 - sx**2.0 / nn + if sigx > 0.0: + corq = (sxy - sx * sy / nn) ** 2.0 / (sigx * sigy) if corq > corqx: corqx = corq isho = ish jsho = jsh - if (isho != -99) and (corqx > 0.3) and (corqx * sy > 1.) and (sigy / sy >= 0.08): + if ( + (isho != -99) + and (corqx > 0.3) + and (corqx * sy > 1.0) + and (sigy / sy >= 0.08) + ): if (abs(isho) < nsh) and (abs(jsho) < nsh): - IS_tmp[int(j/nthin),int(i/nthin)] = -isho - JS_tmp[int(j/nthin),int(i/nthin)] = -jsho + IS_tmp[int(j / nthin), int(i / nthin)] = -isho + JS_tmp[int(j / nthin), int(i / nthin)] = -jsho else: - IS_tmp[int(j/nthin),int(i/nthin)] = 0 - JS_tmp[int(j/nthin),int(i/nthin)] = 0 + IS_tmp[int(j / nthin), int(i / nthin)] = 0 + JS_tmp[int(j / nthin), int(i / nthin)] = 0 - return IS_tmp, JS_tmp \ No newline at end of file + return IS_tmp, JS_tmp From ad9b10b6bb5351df605ce0696dfefde25cd4a9d0 Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:36:13 +0000 Subject: [PATCH 5/8] Some changes to avoid the error by lacking NUMBA library --- pysteps/motion/correlation.py | 244 ++++++++++++++++++---------------- 1 file changed, 130 insertions(+), 114 deletions(-) diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py index ae088e1e3..c0d3c066d 100644 --- a/pysteps/motion/correlation.py +++ b/pysteps/motion/correlation.py @@ -120,6 +120,11 @@ def correlation( Validation over the Eastern Alpine Region. Wea. Forecasting, 26, 166–183. """ + if not NUMBA_IMPORTED: + raise ImportError( + "Numba is not installed. The correlation function cannot be loaded." + ) + input_images = input_images.copy() if interp_kwargs is None: @@ -200,117 +205,128 @@ def correlation( return uvgrid -@jit(nopython=True) -def compute_motion_numba(image1, image2, NI, NJ, nthin, IS_tmp, JS_tmp): - - """ - Compute the motion between two images using a correlation-based method optimized with Numba. - - Parameters: - ----------- - image1 : 2D numpy array - The first input image for motion computation. - image2 : 2D numpy array - The second input image for motion computation. - NI : int - The number of rows in the input images. - NJ : int - The number of columns in the input images. - nthin : int - The thinning factor for the grid used in motion computation. - IS_tmp : 2D numpy array - The output array to store the computed motion in the x-direction. - JS_tmp : 2D numpy array - The output array to store the computed motion in the y-direction. - - Returns: - -------- - IS_tmp : 2D numpy array - The updated motion array in the x-direction after computation. - JS_tmp : 2D numpy array - The updated motion array in the y-direction after computation. - - Notes: - ------ - - This function computes the motion by correlating patches of the first image (`image1`) - with patches of the second image (`image2`). - - The computation is performed on a grid defined by the `nthin` parameter to reduce - computational complexity. - - The function uses a fixed-size search window (`nsh`) and a quantization parameter (`nqu`) - to define the neighborhood for correlation computation. - - The correlation is calculated within a defined window and the maximum correlation value is - used to determine the motion vector. - - The function utilizes several optimization techniques to ensure efficient computation - and is compiled with Numba for further speed-up. - - """ - - ## List of parameters from the original paper (TODO: Optimization as input values in func.) - nsh = 20 - nqu = 45 - nsq = nsh + nqu - di = 1 - PAREAMIN = 1.0 - rr_dpa_min = 0.03 - nn = (2 * nqu / di + 1) ** 2.0 - - for i in range(0, NI, nthin): - if (i >= nsq) and (i < NI - nsq): - for j in range(0, NJ, nthin): - if (j >= nsq) and (j < NJ - nsq): - ii1 = max(i - nqu, 0) - ii2 = min(i + nqu, NI - 1) - jj1 = max(j - nqu, 0) - jj2 = min(j + nqu, NJ - 1) - - sy = 0.0 - sy2 = 0.0 - - for ii in range(ii1, ii2 + 1, di): - for jj in range(jj1, jj2 + 1, di): - sy = sy + image1[jj, ii] - sy2 = sy2 + image1[jj, ii] ** 2.0 - - sigy = sy2 - sy**2.0 / nn - isho = -99 - jsho = -99 - - if (sigy > 0.0) and (sy > PAREAMIN): - corqx = 0.1 - for ish in range(-nsh, nsh + 1): - for jsh in range(-nsh, nsh + 1): - if math.sqrt((ish) ** 2.0 + (jsh) ** 2.0) > nsh: - continue - sx = 0.0 - sx2 = 0.0 - sxy = 0.0 - for ii in range(ii1, ii2 + 1, di): - for jj in range(jj1, jj2 + 1, di): - ind_x = min(max(0, jj + jsh), NJ - 1) - ind_y = min(max(0, ii + ish), NI - 1) - sx += image2[ind_x, ind_y] - sx2 += image2[ind_x, ind_y] ** 2.0 - sxy += image2[ind_x, ind_y] * image1[jj, ii] - - sigx = sx2 - sx**2.0 / nn - if sigx > 0.0: - corq = (sxy - sx * sy / nn) ** 2.0 / (sigx * sigy) - if corq > corqx: - corqx = corq - isho = ish - jsho = jsh - - if ( - (isho != -99) - and (corqx > 0.3) - and (corqx * sy > 1.0) - and (sigy / sy >= 0.08) - ): - if (abs(isho) < nsh) and (abs(jsho) < nsh): - IS_tmp[int(j / nthin), int(i / nthin)] = -isho - JS_tmp[int(j / nthin), int(i / nthin)] = -jsho - else: - IS_tmp[int(j / nthin), int(i / nthin)] = 0 - JS_tmp[int(j / nthin), int(i / nthin)] = 0 - - return IS_tmp, JS_tmp +if NUMBA_IMPORTED: + + @jit(nopython=True) + def compute_motion_numba(image1, image2, NI, NJ, nthin, IS_tmp, JS_tmp): + + """ + Compute the motion between two images using a correlation-based method optimized with Numba. + + Parameters: + ----------- + image1 : 2D numpy array + The first input image for motion computation. + image2 : 2D numpy array + The second input image for motion computation. + NI : int + The number of rows in the input images. + NJ : int + The number of columns in the input images. + nthin : int + The thinning factor for the grid used in motion computation. + IS_tmp : 2D numpy array + The output array to store the computed motion in the x-direction. + JS_tmp : 2D numpy array + The output array to store the computed motion in the y-direction. + + Returns: + -------- + IS_tmp : 2D numpy array + The updated motion array in the x-direction after computation. + JS_tmp : 2D numpy array + The updated motion array in the y-direction after computation. + + Notes: + ------ + - This function computes the motion by correlating patches of the first image (`image1`) + with patches of the second image (`image2`). + - The computation is performed on a grid defined by the `nthin` parameter to reduce + computational complexity. + - The function uses a fixed-size search window (`nsh`) and a quantization parameter (`nqu`) + to define the neighborhood for correlation computation. + - The correlation is calculated within a defined window and the maximum correlation value is + used to determine the motion vector. + - The function utilizes several optimization techniques to ensure efficient computation + and is compiled with Numba for further speed-up. + + """ + + ## List of parameters from the original paper (TODO: Optimization as input values in func.) + nsh = 20 + nqu = 45 + nsq = nsh + nqu + di = 1 + PAREAMIN = 1.0 + rr_dpa_min = 0.03 + nn = (2 * nqu / di + 1) ** 2.0 + + for i in range(0, NI, nthin): + if (i >= nsq) and (i < NI - nsq): + for j in range(0, NJ, nthin): + if (j >= nsq) and (j < NJ - nsq): + ii1 = max(i - nqu, 0) + ii2 = min(i + nqu, NI - 1) + jj1 = max(j - nqu, 0) + jj2 = min(j + nqu, NJ - 1) + + sy = 0.0 + sy2 = 0.0 + + for ii in range(ii1, ii2 + 1, di): + for jj in range(jj1, jj2 + 1, di): + sy = sy + image1[jj, ii] + sy2 = sy2 + image1[jj, ii] ** 2.0 + + sigy = sy2 - sy**2.0 / nn + isho = -99 + jsho = -99 + + if (sigy > 0.0) and (sy > PAREAMIN): + corqx = 0.1 + for ish in range(-nsh, nsh + 1): + for jsh in range(-nsh, nsh + 1): + if math.sqrt((ish) ** 2.0 + (jsh) ** 2.0) > nsh: + continue + sx = 0.0 + sx2 = 0.0 + sxy = 0.0 + for ii in range(ii1, ii2 + 1, di): + for jj in range(jj1, jj2 + 1, di): + ind_x = min(max(0, jj + jsh), NJ - 1) + ind_y = min(max(0, ii + ish), NI - 1) + sx += image2[ind_x, ind_y] + sx2 += image2[ind_x, ind_y] ** 2.0 + sxy += image2[ind_x, ind_y] * image1[jj, ii] + + sigx = sx2 - sx**2.0 / nn + if sigx > 0.0: + corq = (sxy - sx * sy / nn) ** 2.0 / ( + sigx * sigy + ) + if corq > corqx: + corqx = corq + isho = ish + jsho = jsh + + if ( + (isho != -99) + and (corqx > 0.3) + and (corqx * sy > 1.0) + and (sigy / sy >= 0.08) + ): + if (abs(isho) < nsh) and (abs(jsho) < nsh): + IS_tmp[int(j / nthin), int(i / nthin)] = -isho + JS_tmp[int(j / nthin), int(i / nthin)] = -jsho + else: + IS_tmp[int(j / nthin), int(i / nthin)] = 0 + JS_tmp[int(j / nthin), int(i / nthin)] = 0 + + return IS_tmp, JS_tmp + +else: + + def compute_motion_numba(image1, image2, NI, NJ, nthin, IS_tmp, JS_tmp): + raise RuntimeError( + "Numba is not available, so the compute_motion_numba function cannot be used." + ) From 482f6b6f335f72707c3a6c5c350212d13f6a8f61 Mon Sep 17 00:00:00 2001 From: aitaten Date: Tue, 23 Jul 2024 15:40:50 +0000 Subject: [PATCH 6/8] Black formatting with same version as in the original repository --- pysteps/motion/correlation.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py index c0d3c066d..a859bf920 100644 --- a/pysteps/motion/correlation.py +++ b/pysteps/motion/correlation.py @@ -41,7 +41,6 @@ def correlation( k_outlier=30, verbose=False, ): - """ Run the correlation based optical flow method and interpolate the motion vectors. This is using NUMBA to speed up the correlation computation. @@ -209,7 +208,6 @@ def correlation( @jit(nopython=True) def compute_motion_numba(image1, image2, NI, NJ, nthin, IS_tmp, JS_tmp): - """ Compute the motion between two images using a correlation-based method optimized with Numba. From 457289cfe3fb5fdc88b09350c672f5df04d58db3 Mon Sep 17 00:00:00 2001 From: aitaten <88198920+aitaten@users.noreply.github.com> Date: Fri, 2 Aug 2024 20:05:40 +0200 Subject: [PATCH 7/8] Correct the LK wording in this message Co-authored-by: Daniele Nerini --- pysteps/motion/correlation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py index a859bf920..048920d4e 100644 --- a/pysteps/motion/correlation.py +++ b/pysteps/motion/correlation.py @@ -188,7 +188,8 @@ def correlation( uv = uv[~outliers, :] if verbose: - print("--- LK found %i sparse vectors ---" % xy.shape[0]) + print("--- found %i sparse vectors ---" % xy.shape[0]) + # return sparse vectors if required if not dense: From eefd3c799a2438b5bb9ba225c17181543a2775b9 Mon Sep 17 00:00:00 2001 From: aitaten <88198920+aitaten@users.noreply.github.com> Date: Fri, 2 Aug 2024 20:34:34 +0200 Subject: [PATCH 8/8] Change ImportError for correct exception Co-authored-by: Daniele Nerini --- pysteps/motion/correlation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pysteps/motion/correlation.py b/pysteps/motion/correlation.py index 048920d4e..008875fbf 100644 --- a/pysteps/motion/correlation.py +++ b/pysteps/motion/correlation.py @@ -120,7 +120,7 @@ def correlation( """ if not NUMBA_IMPORTED: - raise ImportError( + raise MissingOptionalDependency( "Numba is not installed. The correlation function cannot be loaded." )