pivtools 0.1.3__cp311-cp311-win_amd64.whl

pivtools_cli/piv/piv_backend/base.py

@@ -0,0 +1,825 @@
+from abc import ABC, abstractmethod
+import sys
+from pathlib import Path
+from typing import List
+
+import cv2
+import dask.array as da
+import numpy as np
+import scipy.sparse as sp
+import scipy.sparse.linalg as spla
+from scipy.interpolate import griddata
+from skimage.restoration import inpaint_biharmonic
+
+# Try to import line_profiler for detailed profiling
+try:
+    from line_profiler import profile
+except ImportError:
+    profile = lambda f: f
+
+
+from pivtools_core.config import Config
+
+
+class CrossCorrelator(ABC):
+
+    @abstractmethod
+    def correlate_batch(self, images: np.ndarray, config: Config, vector_masks: List[np.ndarray] = None):
+        pass
+
+    def _window_weight_fun(
+        self,
+        wsize: tuple[int, int],
+        window_type: str,
+        SumWindow: tuple[int, int] = None,
+    ) -> np.ndarray:
+        """
+        Generate a 2D cross-correlation weighting window.
+
+        Parameters
+        ----------
+        wsize : tuple[int, int]
+            Size of the window (rows, cols)
+        window_type : str
+            Type of window: 'blackman', 'square', 'bsingle', 'gaussian', 'gaussianTOP', 'singlepix'
+        SumWindow : tuple[int, int] | None
+            Only used for 'bsingle' and 'singlepix' types to define the outer matrix size.
+
+        Returns
+        -------
+        weight : np.ndarray
+            2D array of shape wsize representing the weighting window
+        """
+        m, n = wsize
+        dist_to_cen_max = np.sqrt(((m + 1) / 2) ** 2 + ((n + 1) / 2) ** 2)
+        weight = np.zeros((m, n), dtype=np.float32)
+
+        if window_type.lower() == "blackman":
+            a0, a1, a2 = 0.42659, 0.49656, 0.076849
+            for ii in range(m):
+                for jj in range(n):
+                    dist_to_cen = np.sqrt(
+                        (ii - (m - 1) / 2) ** 2 + (jj - (n - 1) / 2) ** 2
+                    )
+                    theta = dist_to_cen / dist_to_cen_max * np.pi + np.pi
+                    weight[ii, jj] = a0 - a1 * np.cos(theta) + a2 * np.cos(2 * theta)
+            weight[weight <= 0] = 0.01
+
+        elif window_type.lower() == "square":
+            weight[:] = 1.0
+
+        elif window_type.lower() == "bsingle":
+            if SumWindow is None:
+                raise ValueError("SumWindow must be provided for 'bsingle'")
+            weight = np.ones(SumWindow, dtype=np.float32)
+
+        elif window_type.lower() == "gaussian" or window_type.lower() == "gaussiantop":
+            win_x = np.arange(-m / 2 + 0.5, m / 2, dtype=np.float32)
+            win_y = np.arange(-n / 2 + 0.5, n / 2, dtype=np.float32)
+            xm, ym = np.meshgrid(win_x, win_y, indexing="ij")
+            alpha = 1.0
+            weight = np.exp(
+                -0.5 * ((alpha * xm / (m / 2)) ** 2 + (alpha * ym / (n / 2)) ** 2)
+            )
+
+            if window_type.lower() == "gaussiantop":
+                if m > 6 and n > 6:
+                    # Create checkerboard mask
+                    C_mask = np.tile(
+                        np.array([[1, 0], [0, 0]], dtype=np.float32), (m // 2, n // 2)
+                    )
+                    shift_x, shift_y = int(np.ceil(m / 4)), int(np.ceil(n / 4))
+                    C_mask = np.roll(np.roll(C_mask, shift_x, axis=0), shift_y, axis=1)
+                    weight *= C_mask
+                else:
+                    C_mask = np.zeros((m, n), dtype=np.float32)
+                    r0 = m // 2 - 1
+                    c0 = n // 2 - 1
+                    C_mask[r0 : r0 + 4, c0 : c0 + 4] = 1
+                    weight *= C_mask
+
+        elif window_type.lower() == "singlepix":
+            if SumWindow is None:
+                raise ValueError("SumWindow must be provided for 'singlepix'")
+            weight = np.zeros(SumWindow, dtype=np.float32)
+            start_row = int(np.ceil((SumWindow[0] - m) / 2))
+            end_row = start_row + m
+            start_col = int(np.ceil((SumWindow[1] - n) / 2))
+            end_col = start_col + n
+            weight[start_row:end_row, start_col:end_col] = 1.0
+
+        else:
+            raise ValueError(f"Unrecognized window type '{window_type}'")
+
+        return weight
+    
+    @profile
+    def _inpaint_nans_biharm(self, A):
+        mask = np.isnan(A)
+        A_filled = np.copy(A)
+
+        y, x = np.indices(A.shape)
+        known_points = np.array([y[~mask], x[~mask]]).T
+        known_values = A[~mask]
+        nan_points = np.array([y[mask], x[mask]]).T
+
+        A_filled[mask] = griddata(
+            known_points, known_values, nan_points, method="nearest"
+        )
+
+        return inpaint_biharmonic(A_filled, mask)
+    @profile
+    def _inpaint_nans_griddata(self, A):
+        """
+        Fast NaN inpainting using scipy.interpolate.griddata.
+
+        This uses a two-pass approach:
+        1. 'linear' for smooth, fast interpolation.
+        2. 'nearest' to fill any remaining NaNs (e.g., extrapolation).
+        """
+        mask = np.isnan(A)
+        
+        # If there are no NaNs, return immediately
+        if not mask.any():
+            return A
+
+        y, x = np.indices(A.shape)
+        known_points = np.array([y[~mask], x[~mask]]).T
+        known_values = A[~mask]
+        nan_points = np.array([y[mask], x[mask]]).T
+
+        # If there are no known points, we can't interpolate.
+        # Return the array as-is (or return np.zeros_like(A))
+        if known_points.size == 0:
+            return A 
+
+        # Pass 1: Linear interpolation
+        interp_values = griddata(
+            known_points, known_values, nan_points, method="linear"
+        )
+        
+        # Pass 2: Find where linear failed (still NaN) and fill with nearest
+        nan_mask_pass1 = np.isnan(interp_values)
+        if nan_mask_pass1.any():
+            fallback_values = griddata(
+                known_points, known_values, nan_points[nan_mask_pass1], method="nearest"
+            )
+            interp_values[nan_mask_pass1] = fallback_values
+
+        # Fill the original array
+        A_filled = np.copy(A)
+        A_filled[mask] = interp_values
+        
+        return A_filled
+
+    def _inpaint_nans_matlab(self, A):
+
+        A = np.array(A, dtype=float)
+        n, m = A.shape
+        nm = n * m
+
+        A_flat = A.ravel()
+        nan_mask = np.isnan(A_flat)
+        nan_list = np.where(nan_mask)[0]
+        known_list = np.where(~nan_mask)[0]
+
+        if nan_list.size == 0:
+            return A.copy()
+
+        rows, cols, vals = [], [], []
+
+        def add_entries(center, offsets, weights):
+            for off, w in zip(offsets, weights):
+                r = center
+                c = center + off
+                if 0 <= c < nm:
+                    rows.append(r)
+                    cols.append(c)
+                    vals.append(w)
+
+        for idx in nan_list:
+
+            r = idx % n
+            c = idx // n
+
+            if 2 <= r < n - 2 and 2 <= c < m - 2:
+                offsets = [
+                    -2 * n,
+                    -n - 1,
+                    -n,
+                    -n + 1,
+                    -2,
+                    -1,
+                    0,
+                    1,
+                    2,
+                    n - 1,
+                    n,
+                    n + 1,
+                    2 * n,
+                ]
+                weights = [1, 2, -8, 2, 1, -8, 20, -8, 1, 2, -8, 2, 1]
+                add_entries(idx, offsets, weights)
+
+            else:
+
+                offsets = [-n, 0, n]
+                weights = [1, -2, 1]
+                add_entries(idx, offsets, weights)
+                offsets = [-1, 0, 1]
+                weights = [1, -2, 1]
+                add_entries(idx, offsets, weights)
+
+        fda = sp.csr_matrix((vals, (rows, cols)), shape=(nm, nm))
+
+        rhs = -fda[:, known_list] @ A_flat[known_list]
+
+        k = nan_list
+        fda_sub = fda[k[:, None], k].tocsc()
+        rhs_sub = rhs[k]
+
+        x = spla.spsolve(fda_sub, rhs_sub)
+
+        B_flat = A_flat.copy()
+        B_flat[k] = x
+        return B_flat.reshape((n, m))
+
+    def _inpaint_nans_opencv(self, A):
+        """
+        Inpaint NaNs using OpenCV.
+        """
+        mask = np.isnan(A).astype(np.uint8)
+        A[np.isnan(A)] = 0.0
+
+        A_inpainted = cv2.inpaint(A, mask, inpaintRadius=5, flags=cv2.INPAINT_TELEA)
+
+        return A_inpainted
+
+    def _inpaint_nans(self, A, method=3):
+        """
+        Replicates John D’Errico’s inpaint_nans MATLAB function in Python.
+
+        Parameters
+        ----------
+        A : 2D array_like
+            Input array with NaNs to be filled.
+        method : int, optional
+            Which PDE/finite-difference scheme to use (0 to 5). Default is 0.
+
+        Returns
+        -------
+        B : 2D ndarray
+            Array with NaNs replaced/inpainted.
+        """
+        # Convert input A to a float array in Fortran (column-major) order
+        A = np.array(A, dtype=float, order="F")
+        n, m = A.shape
+        nm = n * m
+
+        # Flatten (column-major) so linear indexing matches MATLAB's
+        A_flat = A.ravel(order="F")
+
+        # Find NaNs
+        isnan = np.isnan(A_flat)
+        nan_list = np.where(isnan)[0]  # these are 0-based linear indices
+        known_list = np.where(~isnan)[0]
+        nan_count = len(nan_list)
+        if nan_count == 0:
+            # No NaNs, nothing to do
+            return A  # already a copy
+
+        # Convert those nan indices to row/col, 1-based analog to MATLAB
+        #   In MATLAB: [nr, nc] = ind2sub([n,m], nan_list)
+        #   We store them as 1-based to match the original boundary logic,
+        #   then keep them in a combined array: [lin_idx, row, col]
+        #   where row,col in [1..n], [1..m].
+        nr = (nan_list % n) + 1
+        nc = (nan_list // n) + 1
+        nan_list_info = np.column_stack([nan_list, nr, nc])
+
+        # --- Helper: build a function to do the solve or least-squares:
+        def sparse_solve(M, rhs):
+            """
+            Solve M x = rhs.
+            Uses direct solve if M is square and has full rank,
+            otherwise use lsqr for least-squares.
+            """
+            # M should be shape (R, C). If R==C, try spsolve:
+            R, C = M.shape
+            if R == C:
+                return spla.spsolve(M.tocsc(), rhs)
+            else:
+                # Use lsqr for least squares
+                sol = spla.lsqr(M, rhs, atol=1e-12, btol=1e-12, iter_lim=10_000)
+                return sol[0]
+
+        # We’ll need a function to eliminate known values from RHS, similarly to MATLAB
+        def eliminate_knowns(fda, known_idx, A_known):
+            """
+            Return the adjusted rhs after applying -fda[:, known_idx] * A_known,
+            i.e. the part that remains for the unknown columns.
+            """
+            rhs = -fda[:, known_idx].dot(A_known)
+            return rhs
+
+        # Switch on 'method' just like the original code
+        if method not in [0, 1, 2, 3, 4, 5]:
+            raise ValueError("method must be one of {0,1,2,3,4,5}.")
+
+        # We will build 'fda' (the finite-difference operator) as a sparse matrix
+        # in each method, then solve or least-squares for the unknowns.
+        B = A_flat.copy()
+
+        # ----- Subfunction from the original code, in Python:
+        def identify_neighbors(n, m, nan_list_3col, talks_to):
+            """
+            Identify neighbors of the NaN pixels, not including the NaNs themselves.
+
+            nan_list_3col: array of shape (N,3): [lin_index, row, col]
+                           with row,col in 1-based coords.
+            talks_to: array of shape (p,2) of row/col offsets
+            """
+            if nan_list_3col.shape[0] == 0:
+                return np.empty((0, 3), dtype=int)
+
+            nan_count_local = nan_list_3col.shape[0]
+            talk_count = talks_to.shape[0]
+
+            # For each offset in talks_to, add to row,col
+            # shape => (nan_count_local * talk_count, 2)
+            repeated = np.repeat(
+                nan_list_3col[:, 1:3], talk_count, axis=0
+            )  # row,col repeated
+            offsets = np.tile(talks_to, (nan_count_local, 1))
+            nn = repeated + offsets  # neighbor row,col (1-based)
+
+            # Filter out-of-bounds neighbors
+            in_bounds = (
+                (nn[:, 0] >= 1) & (nn[:, 0] <= n) & (nn[:, 1] >= 1) & (nn[:, 1] <= m)
+            )
+            nn = nn[in_bounds]
+
+            # Convert (row,col) back to linear index in 0-based
+            #  MATLAB 1-based: lin = row + (col-1)*n
+            #  Python 0-based: lin = (row-1) + (col-1)*n
+            # so if row,col is 1-based, we do:
+            lin_idx = (nn[:, 0] - 1) + (nn[:, 1] - 1) * n
+
+            neighbors_list = np.column_stack([lin_idx, nn])
+
+            # Unique rows and remove those that are in the NaN list
+            neighbors_list = np.unique(neighbors_list, axis=0)
+
+            # Build a set of all nan linear indices for easy filter
+            nan_lin_set = set(nan_list_3col[:, 0].tolist())
+
+            # Keep only those not themselves in the NaN set
+            mask_not_nan = np.array(
+                [(x[0] not in nan_lin_set) for x in neighbors_list], dtype=bool
+            )
+            neighbors_list = neighbors_list[mask_not_nan]
+            return neighbors_list
+
+        # Because the methods differ widely, we implement them in turn:
+        if method == 0 or method == 3:
+            # Methods 0 and 3 are similar to method 1 but build the matrix only
+            # around the nans and their neighbors, then do a del^2 or del^4 solve.
+
+            # If method=0 => del^2, if method=3 => del^4
+            if method == 0:
+                # del^2 using neighbors: up/down/left/right
+                # We only build around the nans + their immediate neighbors
+                #  for 2D or 1D.  If 1D, treat specially.
+
+                if n == 1 or m == 1:
+                    # 1D case
+                    # Identify the "work_list" as nan +/- 1
+                    work_list = np.concatenate([nan_list, nan_list - 1, nan_list + 1])
+                    work_list = work_list[(work_list >= 0) & (work_list < nm)]
+                    work_list = np.unique(work_list)
+
+                    # Build fda
+                    # For each i in work_list, we want i-1, i, i+1 with [1, -2, 1]
+                    # We'll do that in a lil_matrix
+                    fda = sp.lil_matrix((len(work_list), nm), dtype=float)
+                    # Fill row-by-row
+                    for row_idx, i in enumerate(work_list):
+                        # center
+                        fda[row_idx, i] = -2.0
+                        if i - 1 >= 0:
+                            fda[row_idx, i - 1] = 1.0
+                        if i + 1 < nm:
+                            fda[row_idx, i + 1] = 1.0
+
+                    # Eliminate knowns
+                    rhs = eliminate_knowns(fda, known_list, A_flat[known_list])
+
+                    # We only solve for columns in nan_list
+                    unknown_idx = nan_list
+                    # We only keep rows that reference those columns
+                    # i.e. any row with a non-zero in unknown columns
+                    mask_rows = fda[:, unknown_idx].sum(axis=1).A.ravel() != 0
+                    row_sel = np.where(mask_rows)[0]
+
+                    fda_sub = fda[row_sel, :][:, unknown_idx]
+                    rhs_sub = rhs[row_sel]
+
+                    sol = sparse_solve(fda_sub, rhs_sub)
+
+                    # Place solution
+                    B[unknown_idx] = sol
+
+                else:
+                    # 2D case
+                    # Horizontal and vertical neighbors only
+                    talks_to = np.array([[-1, 0], [1, 0], [0, -1], [0, 1]])
+                    neighbors_list = identify_neighbors(n, m, nan_list_info, talks_to)
+                    all_list = np.vstack([nan_list_info, neighbors_list])
+
+                    # Build fda
+                    fda = sp.lil_matrix((nm, nm), dtype=float)
+
+                    # second partials row-wise: (row > 1 & row < n)
+                    L = np.where((all_list[:, 1] > 1) & (all_list[:, 1] < n))[0]
+                    for i in L:
+                        idx = all_list[i, 0]
+                        fda[idx, idx] += -2.0
+                        fda[idx, idx - 1] += 1.0
+                        fda[idx, idx + 1] += 1.0
+
+                    # second partials col-wise: (col > 1 & col < m)
+                    L = np.where((all_list[:, 2] > 1) & (all_list[:, 2] < m))[0]
+                    for i in L:
+                        idx = all_list[i, 0]
+                        fda[idx, idx] += -2.0
+                        fda[idx, idx - n] += 1.0
+                        fda[idx, idx + n] += 1.0
+
+                    # Eliminate knowns
+                    rhs = eliminate_knowns(fda, known_list, A_flat[known_list])
+
+                    # Solve only for relevant rows & columns
+                    unknown_idx = nan_list
+                    mask_rows = fda[:, unknown_idx].sum(axis=1).A.ravel() != 0
+                    row_sel = np.where(mask_rows)[0]
+
+                    fda_sub = fda[row_sel, :][:, unknown_idx]
+                    rhs_sub = rhs[row_sel]
+                    sol = sparse_solve(fda_sub, rhs_sub)
+
+                    B[unknown_idx] = sol
+
+            else:
+                # method == 3 => "better plate" using del^4
+                # We use bigger stencils
+                #  The code is quite extensive, we replicate the logic:
+
+                # neighbors for the center region
+                talks_to = np.array(
+                    [
+                        [-2, 0],
+                        [-1, -1],
+                        [-1, 0],
+                        [-1, 1],
+                        [0, -2],
+                        [0, -1],
+                        [0, 1],
+                        [0, 2],
+                        [1, -1],
+                        [1, 0],
+                        [1, 1],
+                        [2, 0],
+                    ]
+                )
+                neighbors_list = identify_neighbors(n, m, nan_list_info, talks_to)
+                all_list = np.vstack([nan_list_info, neighbors_list])
+
+                fda = sp.lil_matrix((nm, nm), dtype=float)
+
+                # main interior: row>=3 & row<=n-2 & col>=3 & col<=m-2
+                L = np.where(
+                    (all_list[:, 1] >= 3)
+                    & (all_list[:, 1] <= n - 2)
+                    & (all_list[:, 2] >= 3)
+                    & (all_list[:, 2] <= m - 2)
+                )[0]
+                # fill with the big 13-point stencil
+                # Coeffs: [1 2 -8 2 1 -8 20 -8 1 2 -8 2 1],
+                # Offsets in linear indices: [-2n, -(n+1), -n, -(n-1), -2, -1, 0, +1, +2, (n-1), +n, (n+1), +2n]
+                base_offsets = np.array(
+                    [
+                        -2 * n,
+                        -n - 1,
+                        -n,
+                        -n + 1,
+                        -2,
+                        -1,
+                        0,
+                        1,
+                        2,
+                        n - 1,
+                        n,
+                        n + 1,
+                        2 * n,
+                    ]
+                )
+                base_coeffs = np.array(
+                    [1, 2, -8, 2, 1, -8, 20, -8, 1, 2, -8, 2, 1], dtype=float
+                )
+                for i in L:
+                    idx = all_list[i, 0]
+                    for off, coef in zip(base_offsets, base_coeffs):
+                        fda[idx, idx + off] += coef
+
+                # boundaries near row=2 or row=n-1 or col=2 or col=m-1
+                # do a simpler 5-point Laplacian: [1 -4 1], etc.
+                # the original code lumps all boundary expansions.  For brevity, replicate:
+
+                # row=2 or row=n-1 or col=2 or col=m-1
+                L = np.where(
+                    (
+                        ((all_list[:, 1] == 2) | (all_list[:, 1] == n - 1))
+                        & (all_list[:, 2] >= 2)
+                        & (all_list[:, 2] <= m - 1)
+                    )
+                    | (
+                        ((all_list[:, 2] == 2) | (all_list[:, 2] == m - 1))
+                        & (all_list[:, 1] >= 2)
+                        & (all_list[:, 1] <= n - 1)
+                    )
+                )[0]
+                # 5-point: offsets = [-n, -1, 0, +1, +n], coeff = [1,1,-4,1,1]
+                offsets_5 = np.array([-n, -1, 0, 1, n])
+                coeffs_5 = np.array([1, 1, -4, 1, 1], dtype=float)
+                for i in L:
+                    idx = all_list[i, 0]
+                    for off, c in zip(offsets_5, coeffs_5):
+                        fda[idx, idx + off] += c
+
+                # row=1 or row=n, col in [2..m-1]
+                L = np.where(
+                    ((all_list[:, 1] == 1) | (all_list[:, 1] == n))
+                    & (all_list[:, 2] >= 2)
+                    & (all_list[:, 2] <= m - 1)
+                )[0]
+                # 3-point vertical second derivative: offsets = [-n, 0, +n], coeffs = [1, -2, 1]
+                offsets_3v = np.array([-n, 0, n])
+                coeffs_3v = np.array([1, -2, 1], dtype=float)
+                for i in L:
+                    idx = all_list[i, 0]
+                    for off, c in zip(offsets_3v, coeffs_3v):
+                        fda[idx, idx + off] += c
+
+                # col=1 or col=m, row in [2..n-1]
+                L = np.where(
+                    ((all_list[:, 2] == 1) | (all_list[:, 2] == m))
+                    & (all_list[:, 1] >= 2)
+                    & (all_list[:, 1] <= n - 1)
+                )[0]
+                # 3-point horizontal second derivative: offsets = [-1, 0, +1], coeffs = [1, -2, 1]
+                offsets_3h = np.array([-1, 0, 1])
+                coeffs_3h = np.array([1, -2, 1], dtype=float)
+                for i in L:
+                    idx = all_list[i, 0]
+                    for off, c in zip(offsets_3h, coeffs_3h):
+                        fda[idx, idx + off] += c
+
+                # Eliminate knowns
+                rhs = eliminate_knowns(fda, known_list, A_flat[known_list])
+
+                # Solve
+                unknown_idx = nan_list
+                mask_rows = fda[:, unknown_idx].sum(axis=1).A.ravel() != 0
+                row_sel = np.where(mask_rows)[0]
+                fda_sub = fda[row_sel, :][:, unknown_idx]
+                rhs_sub = rhs[row_sel]
+                sol = sparse_solve(fda_sub, rhs_sub)
+
+                B[unknown_idx] = sol
+
+        elif method == 1:
+            # Least squares with del^2 on the entire array
+            # Build the Laplacian operator for all points
+            if n == 1 or m == 1:
+                # 1D
+                # second difference for interior points
+                # row i => i=1..(nm-2) in 0-based => fill [i, i+1, i+2]
+                # but we'll do it more systematically:
+                fda = sp.lil_matrix((nm - 2, nm), dtype=float)
+                for i in range(nm - 2):
+                    fda[i, i] = 1.0
+                    fda[i, i + 1] = -2.0
+                    fda[i, i + 2] = 1.0
+
+            else:
+                # 2D
+                fda = sp.lil_matrix((nm, nm), dtype=float)
+                # Row-second-derivatives for i=2..n-1 => index = i+(j-1)*n
+                # We'll just loop or systematically fill them:
+                for j in range(m):
+                    for i in range(1, n - 1):
+                        idx = i + j * n
+                        # i => row, j => col in 0-based, so fda[idx, idx +/- 1]
+                        fda[idx, idx] += -2.0
+                        fda[idx, idx - 1] += 1.0
+                        fda[idx, idx + 1] += 1.0
+
+                # Column-second-derivatives for j=2..m-1 => index i+(j-1)*n
+                for j in range(1, m - 1):
+                    for i in range(n):
+                        idx = i + j * n
+                        fda[idx, idx] += -2.0
+                        fda[idx, idx - n] += 1.0
+                        fda[idx, idx + n] += 1.0
+
+            # Eliminate knowns
+            rhs = eliminate_knowns(fda, known_list, A_flat[known_list])
+
+            # Solve
+            unknown_idx = nan_list
+            mask_rows = fda[:, unknown_idx].sum(axis=1).A.ravel() != 0
+            row_sel = np.where(mask_rows)[0]
+            fda_sub = fda[row_sel, :][:, unknown_idx]
+            rhs_sub = rhs[row_sel]
+            sol = sparse_solve(fda_sub, rhs_sub)
+
+            B[unknown_idx] = sol
+
+        elif method == 2:
+            # Direct solve for del^2 BVP across holes only
+            if n == 1 or m == 1:
+                raise ValueError(
+                    "Method 2 has problems for 1D input. Use another method."
+                )
+            else:
+                # 2D
+                fda = sp.lil_matrix((nm, nm), dtype=float)
+
+                # second partials on row index
+                L = np.where((nan_list_info[:, 1] > 1) & (nan_list_info[:, 1] < n))[0]
+                for i in L:
+                    idx = nan_list_info[i, 0]
+                    fda[idx, idx] += -2.0
+                    fda[idx, idx - 1] += 1.0
+                    fda[idx, idx + 1] += 1.0
+
+                # second partials on column index
+                L = np.where((nan_list_info[:, 2] > 1) & (nan_list_info[:, 2] < m))[0]
+                for i in L:
+                    idx = nan_list_info[i, 0]
+                    fda[idx, idx] += -2.0
+                    fda[idx, idx - n] += 1.0
+                    fda[idx, idx + n] += 1.0
+
+                # fix boundary corners if they are NaN
+                corners = [0, n - 1, nm - n, nm - 1]  # 0-based corners in Fortran?
+                # Actually, in column-major: top-left = 0, bottom-left = n-1,
+                # top-right = (m-1)*n, bottom-right = nm-1
+                # The original code forces certain patterns if those corners are in nan_list
+                for c in corners:
+                    if c in nan_list:
+                        # replicate the code: fda(c, [c c+...]) = ...
+                        # corner examples from the original
+                        if c == 0:
+                            fda[c, c] = -2.0
+                            fda[c, c + 1] += 1.0
+                            fda[c, c + n] += 1.0
+                        elif c == n - 1:
+                            fda[c, c] = -2.0
+                            fda[c, c - 1] += 1.0
+                            fda[c, c + n] += 1.0
+                        elif c == (m - 1) * n:
+                            fda[c, c] = -2.0
+                            fda[c, c + 1] += 1.0
+                            fda[c, c - n] += 1.0
+                        elif c == nm - 1:
+                            fda[c, c] = -2.0
+                            fda[c, c - 1] += 1.0
+                            fda[c, c - n] += 1.0
+
+                # Eliminate knowns
+                rhs = eliminate_knowns(fda, known_list, A_flat[known_list])
+
+                # Solve directly on the nan_list subset
+                unknown_idx = nan_list
+                fda_sub = fda[unknown_idx, :][:, unknown_idx]
+                rhs_sub = rhs[unknown_idx]
+                sol = sparse_solve(fda_sub, rhs_sub)
+
+                B[unknown_idx] = sol
+
+        elif method == 4:
+            # Spring analogy, only horizontal + vertical neighbors
+            # Diagonals in the original code are not used or are they? Actually,
+            # code has "hv_list=[-1  -1 0; 1 1 0; -n 0 -1; n 0 1]", but that’s
+            # for HV.  (No diagonal in the original code method 4.)
+
+            # We'll build a matrix of "springs"
+            hv_list = np.array([[-1, 0], [1, 0], [0, -1], [0, 1]])  # row/col offsets
+            # But in the original code, it actually used:
+            #   hv_list = [ -1 -1 0
+            #               1  1  0
+            #              -n  0 -1
+            #               n  0  1 ]
+            # That was building pairs of (index, index+...). It's simpler to replicate logic:
+
+            # Let’s explicitly gather pairs of neighbors among the nan_list:
+            springs = []
+            for i in range(nan_count):
+                idx = nan_list_info[i, 0]
+                row_i = nan_list_info[i, 1]
+                col_i = nan_list_info[i, 2]
+                # up/down/left/right neighbors
+                # up => if row_i>1 => idx-n
+                if row_i > 1:
+                    springs.append([idx, idx - 1])  # in col-major, up is -1
+                if row_i < n:
+                    springs.append([idx, idx + 1])  # down is +1
+                if col_i > 1:
+                    springs.append([idx, idx - n])  # left is -n
+                if col_i < m:
+                    springs.append([idx, idx + n])  # right is +n
+
+            # Unique + sort each pair
+            springs = np.array(springs)
+            # Sort rows so that [min,max]
+            springs.sort(axis=1)
+            # Unique
+            springs = np.unique(springs, axis=0)
+
+            # Build the sparse system: each spring => row in the system, [1 -1] for those two columns
+            n_springs = springs.shape[0]
+            S = sp.lil_matrix((n_springs, nm), dtype=float)
+            for i in range(n_springs):
+                i1, i2 = springs[i]
+                S[i, i1] = 1.0
+                S[i, i2] = -1.0
+
+            # Right side after eliminating known
+            rhs = -S[:, known_list].dot(B[known_list])
+
+            # Solve only for nan columns
+            unknown_idx = nan_list
+            S_sub = S[:, unknown_idx]
+            sol = sparse_solve(S_sub, rhs)
+
+            B[unknown_idx] = sol
+
+        elif method == 5:
+            # Average of 8 nearest neighbors
+            # The code builds an operator that enforces B(i) = average(8 neighbors).
+            # That translates to sum_of_neighbors - 8*B(i) = 0, or equivalently
+            # for each i in nan_list, sum_{neighbors} (B(neighbor)) - 8 B(i) = 0.
+
+            fda = sp.lil_matrix((nm, nm), dtype=float)
+
+            def add_avg_equation(center_idx, neighbor_idx):
+                # for eq: B(neighbor_idx) - B(center_idx)
+                fda[center_idx, neighbor_idx] += 1.0
+                fda[center_idx, center_idx] += -1.0
+
+            # We replicate the "if" blocks from MATLAB for the 8 neighbors
+            # We only do it for i in nan_list
+            for i in nan_list:
+                # row,col in 1-based
+                r = (i % n) + 1
+                c = (i // n) + 1
+                # top-left => if r>1,c>1 => i-(n+1)
+                if (r > 1) and (c > 1):
+                    add_avg_equation(i, i - n - 1)
+                # top => if c>1 => i-n
+                if c > 1:
+                    add_avg_equation(i, i - n)
+                # top-right => if r<n,c>1 => i-(n-1)
+                if (r < n) and (c > 1):
+                    add_avg_equation(i, i - n + 1)
+                # left => if r>1 => i-1
+                if r > 1:
+                    add_avg_equation(i, i - 1)
+                # right => if r<n => i+1
+                if r < n:
+                    add_avg_equation(i, i + 1)
+                # bottom-left => if r>1,c<m => i+(n-1)
+                if (r > 1) and (c < m):
+                    add_avg_equation(i, i + n - 1)
+                # bottom => if c<m => i+n
+                if c < m:
+                    add_avg_equation(i, i + n)
+                # bottom-right => if r<n,c<m => i+(n+1)
+                if (r < n) and (c < m):
+                    add_avg_equation(i, i + n + 1)
+
+            # Eliminate known
+            rhs = eliminate_knowns(fda, known_list, B[known_list])
+
+            # Solve for unknown
+            unknown_idx = nan_list
+            fda_sub = fda[unknown_idx, :][:, unknown_idx]
+            rhs_sub = rhs[unknown_idx]
+            sol = sparse_solve(fda_sub, rhs_sub)
+
+            B[unknown_idx] = sol
+
+        # Reshape back to (n,m) in Fortran order
+        B = np.reshape(B, (n, m), order="F")
+        return B