Add Gaussian multilooking option for phase linking

src/dolphin/phase_link/_core.py

@@ -49,7 +49,12 @@ class PhaseLinkOutput(NamedTuple):
     """
 
     shp_counts: np.ndarray
-    """Number of neighbor pixels used in adaptive multilooking."""
+    """Number of effective looks used in multilooking.
+
+    For boolean SHP methods (GLRT, KS), this is the count of neighbor pixels.
+    For float weights (Gaussian), this is the effective number of looks (ENL)
+    via Kish's formula: ENL = (sum(w))^2 / sum(w^2).
+    """
 
     eigenvalues: np.ndarray
     """The smallest (largest) eigenvalue resulting from EMI (EVD)."""
@@ -362,8 +367,14 @@ def run_cpl(
     # Get the SHP counts for each pixel (if not using Rect window)
     if neighbor_arrays is None:
         shp_counts = jnp.zeros(temp_coh.shape, dtype=np.int16)
-    else:
+    elif neighbor_arrays.dtype == np.bool_:
         shp_counts = jnp.sum(neighbor_arrays, axis=(-2, -1))
+    else:
+        # For float weights (e.g. Gaussian), compute effective number of looks
+        # ENL = (sum(w))^2 / sum(w^2)  (Kish, 1965, Survey Sampling)
+        w_sum = jnp.sum(neighbor_arrays, axis=(-2, -1))
+        w_sq_sum = jnp.sum(neighbor_arrays**2, axis=(-2, -1))
+        shp_counts = jnp.round(w_sum**2 / w_sq_sum).astype(jnp.int16)
 
     return PhaseLinkOutput(
         cpx_phase=cpx_phase_reshaped,

src/dolphin/phase_link/covariance.py

@@ -113,26 +113,55 @@ def coh_mat_single(
 ) -> Array:
     """Given (n_slc, n_samps) SLC samples, get the (nslc, nslc) coherence matrix.
 
-    Note this requires `slc_samples` to be transposed from `coh_mat_single`.
+    Supports both boolean masks (for GLRT/KS adaptive methods) and float weights
+    (for Gaussian multilooking).
+
+    Parameters
+    ----------
+    slc_samples : ArrayLike
+        SLC samples with shape (n_slc, n_samples).
+    neighbor_mask : ArrayLike, optional
+        Either a boolean mask or float weights with shape (n_samples,).
+        For boolean: True means include the sample.
+        For float: values are weights (should be non-negative).
+
+    Returns
+    -------
+    Array
+        Coherence matrix with shape (n_slc, n_slc).
+
     """
     _nslc, nsamps = slc_samples.shape
 
     if neighbor_mask is None:
         neighbor_mask = jnp.ones(nsamps, dtype=jnp.bool_)
+
     valid_samples_mask = ~jnp.isnan(slc_samples)
-    combined_mask = valid_samples_mask & neighbor_mask[None, :]
 
-    # Mask the slc samples
+    # Handle boolean mask (GLRT/KS) vs float weights (Gaussian)
+    # For boolean: use masking (zero out non-neighbors)
+    # For float: use weighted samples with sqrt(weights) for proper weighted coherence
     # note that it's not possible to change the size based on the mask
     # https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#dynamic-shapes
-    masked_slc = jnp.where(combined_mask, slc_samples, 0)
+
+    # Convert neighbor_mask to float weights for unified handling
+    # Boolean True -> 1.0, False -> 0.0; float stays as-is
+    weights = jnp.asarray(neighbor_mask, dtype=jnp.float32)
+
+    # Zero out invalid samples and apply sqrt(weights) for proper weighted coherence
+    # C_ij = sum(w * z_i * conj(z_j)) / sqrt(sum(w * |z_i|^2) * sum(w * |z_j|^2))
+    # Using sqrt(w) on samples gives the correct weighted formula
+    sample_weights = jnp.where(valid_samples_mask, weights[None, :], 0.0)
+    weighted_slc = slc_samples * jnp.sqrt(sample_weights)
+    # Zero out invalid samples
+    weighted_slc = jnp.where(valid_samples_mask, weighted_slc, 0.0)
 
     # Compute cross-correlation
-    numer = jnp.dot(masked_slc, jnp.conj(masked_slc.T))
+    numer = jnp.dot(weighted_slc, jnp.conj(weighted_slc.T))
 
     # Compute amplitudes so we normalize the covariance to a coherence matrix
     # a1 is shape (nslc,)
-    amp_vec = jnp.sum(jnp.abs(masked_slc) ** 2, axis=1)
+    amp_vec = jnp.sum(jnp.abs(weighted_slc) ** 2, axis=1)
     # Form outer product of amplitudes for each slc
     power_mat = amp_vec[:, None] * amp_vec[None, :]
     amp_mat = jnp.sqrt(power_mat)

src/dolphin/shp/__init__.py

@@ -9,7 +9,7 @@
 from dolphin import Strides
 from dolphin.workflows import ShpMethod
 
-from . import _glrt, _ks
+from . import _gaussian, _glrt, _ks
 
 logger = logging.getLogger("dolphin")
 
@@ -28,6 +28,7 @@ def estimate_neighbors(
     is_sorted: bool = False,
     method: ShpMethod = ShpMethod.GLRT,
     prune_disconnected: bool = False,
+    input_shape: Optional[tuple[int, int]] = None,
 ) -> np.ndarray:
     """Estimate the statistically similar neighbors of each pixel.
 
@@ -58,17 +59,23 @@ def estimate_neighbors(
         If True, keeps only SHPs that are 8-connected to the current pixel.
         Otherwise, any pixel within the window may be considered an SHP, even
         if it is not directly connected.
+    input_shape : tuple[int, int], optional
+        Shape of the input image (rows, cols). Required for GAUSSIAN method
+        if mean/var/amp_stack are not provided.
 
     Returns
     -------
     Optional[np.ndarray]
         Array of estimated statistically similar neighbors.
+        For GLRT/KS methods, this is a boolean array.
+        For GAUSSIAN method, this is a float array of weights.
 
     Raises
     ------
     ValueError
         - nslc is not provided for GLRT method
         - amp_stack is not provided for the KS method.
+        - input_shape cannot be inferred for GAUSSIAN method.
         - `method` not a valid `ShpMethod`
 
     """
@@ -109,6 +116,24 @@ def estimate_neighbors(
             alpha=alpha,
             is_sorted=is_sorted,
         )
+    elif method.lower() == ShpMethod.GAUSSIAN:
+        logger.debug("Using Gaussian weighting for multilooking")
+        # Infer input_shape if not provided
+        if input_shape is None:
+            if mean is not None:
+                input_shape = np.asarray(mean).shape
+            elif var is not None:
+                input_shape = np.asarray(var).shape
+            elif amp_stack is not None:
+                input_shape = np.asarray(amp_stack).shape[1:]
+            else:
+                msg = "input_shape must be provided for GAUSSIAN method"
+                raise ValueError(msg)
+        neighbor_arrays = _gaussian.estimate_neighbors(
+            halfwin_rowcol=halfwin_rowcol,
+            input_shape=input_shape,
+            strides=tuple(strides),
+        )
     else:
         msg = f"SHP method {method} is not implemented"
         raise ValueError(msg)

src/dolphin/shp/_gaussian.py

@@ -0,0 +1,80 @@
+"""Gaussian weighting for multilooking."""
+
+from __future__ import annotations
+
+from functools import partial
+
+import jax.numpy as jnp
+from jax import Array, jit
+
+from dolphin.utils import compute_out_shape
+
+
+@partial(
+    jit,
+    static_argnames=["halfwin_rowcol", "strides", "input_shape"],
+)
+def estimate_neighbors(
+    halfwin_rowcol: tuple[int, int],
+    input_shape: tuple[int, int],
+    strides: tuple[int, int] = (1, 1),
+) -> Array:
+    """Generate Gaussian weights for multilooking.
+
+    Unlike GLRT or KS methods, Gaussian multilooking uses a fixed weight
+    pattern (same for all pixels) based on a 2D Gaussian window.
+
+    Parameters
+    ----------
+    halfwin_rowcol : tuple[int, int]
+        Half the size of the window in (row, col) dimensions.
+    input_shape : tuple[int, int]
+        Shape of the input image (rows, cols).
+    strides : tuple[int, int]
+        The (row, col) strides to use for the sliding window.
+        By default (1, 1), meaning output size equals input size.
+
+    Returns
+    -------
+    weights : Array, 4D
+        Float array of Gaussian weights for each pixel in the window.
+        Shape is (out_rows, out_cols, window_rows, window_cols).
+        Weights are normalized so they sum to 1 (excluding the center pixel).
+
+    """
+    rows, cols = input_shape
+    half_row, half_col = halfwin_rowcol
+
+    out_rows, out_cols = compute_out_shape((rows, cols), strides)
+
+    window_rsize = 2 * half_row + 1
+    window_csize = 2 * half_col + 1
+
+    # Create a 2D Gaussian window
+    # Use sigma = half_window / 2 so that the window covers ~2 sigma
+    sigma_row = half_row / 2.0 if half_row > 0 else 0.5
+    sigma_col = half_col / 2.0 if half_col > 0 else 0.5
+
+    # Create coordinate grids centered at 0
+    row_coords = jnp.arange(window_rsize) - half_row
+    col_coords = jnp.arange(window_csize) - half_col
+
+    # Compute 2D Gaussian: exp(-(r^2/(2*sr^2) + c^2/(2*sc^2)))
+    row_gauss = jnp.exp(-0.5 * (row_coords / sigma_row) ** 2)
+    col_gauss = jnp.exp(-0.5 * (col_coords / sigma_col) ** 2)
+    gaussian_window = jnp.outer(row_gauss, col_gauss)
+
+    # Set center pixel to 0 (don't include self in weighting, matching GLRT behavior)
+    gaussian_window = gaussian_window.at[half_row, half_col].set(0.0)
+
+    # Normalize so weights sum to 1
+    gaussian_window = gaussian_window / jnp.sum(gaussian_window)
+
+    # Broadcast to all output pixels (same weights for each pixel)
+    # Shape: (out_rows, out_cols, window_rows, window_cols)
+    weights = jnp.broadcast_to(
+        gaussian_window[None, None, :, :],
+        (out_rows, out_cols, window_rsize, window_csize),
+    )
+
+    return weights

src/dolphin/workflows/config/_enums.py

@@ -12,6 +12,7 @@ class ShpMethod(str, Enum):
     GLRT = "glrt"
     KS = "ks"
     RECT = "rect"
+    GAUSSIAN = "gaussian"
     # Alias for no SHP search
     NONE = "rect"
 

tests/test_phase_link_core.py

@@ -6,6 +6,7 @@
 
 from dolphin._types import HalfWindow, Strides
 from dolphin.phase_link import _core, covariance, simulate
+from dolphin.shp._gaussian import estimate_neighbors as gaussian_estimate_neighbors
 from dolphin.utils import gpu_is_available
 
 GPU_AVAILABLE = gpu_is_available() and os.environ.get("NUMBA_DISABLE_JIT") != "1"
@@ -181,3 +182,62 @@ def test_run_phase_linking_ps_fill(slc_samples, use_max_ps, strides):
     )
 
     assert pl_out.temp_coh[out_idx, out_idx] == 1
+
+
+def test_run_phase_linking_gaussian_vs_rect(C_truth):
+    """Test that Gaussian multilooking gives similar results to RECT for synthetic data.
+
+    For homogeneous synthetic data (all from same distribution), Gaussian and RECT
+    should give similar phase estimates since both methods use the full window
+    (just with different weighting).
+    """
+    C, truth = C_truth
+    ns = 11 * 11
+    slc_samples = simulate.simulate_neighborhood_stack(C, ns)
+    slc_stack = slc_samples.reshape(NUM_ACQ, 11, 11)
+
+    half_window = HalfWindow(x=5, y=5)
+    strides = Strides(x=1, y=1)
+
+    # Run phase linking with RECT (no neighbor_arrays = uniform weighting)
+    pl_out_rect = _core.run_phase_linking(
+        slc_stack,
+        half_window=half_window,
+        strides=strides,
+        neighbor_arrays=None,  # RECT method
+    )
+
+    # Run phase linking with Gaussian weights
+    gaussian_weights = gaussian_estimate_neighbors(
+        halfwin_rowcol=(half_window.y, half_window.x),
+        input_shape=(11, 11),
+        strides=(strides.y, strides.x),
+    )
+    pl_out_gaussian = _core.run_phase_linking(
+        slc_stack,
+        half_window=half_window,
+        strides=strides,
+        neighbor_arrays=np.array(gaussian_weights),
+    )
+
+    # Both methods should recover the truth reasonably well
+    # Compare at the center pixel (full window available)
+    rect_phase = np.angle(pl_out_rect.cpx_phase[:, 5, 5])
+    gaussian_phase = np.angle(pl_out_gaussian.cpx_phase[:, 5, 5])
+
+    # Check both methods are close to truth
+    err_deg_threshold = 15  # degrees
+    assert np.degrees(simulate.rmse(truth, rect_phase)) < err_deg_threshold
+    assert np.degrees(simulate.rmse(truth, gaussian_phase)) < err_deg_threshold
+
+    # Gaussian and RECT should give similar results for homogeneous data
+    # Allow some tolerance since weighting differs
+    phase_diff = np.angle(np.exp(1j * (rect_phase - gaussian_phase)))
+    assert np.degrees(np.std(phase_diff)) < err_deg_threshold
+
+    # Temporal coherence should be similar
+    npt.assert_allclose(
+        pl_out_rect.temp_coh[5, 5],
+        pl_out_gaussian.temp_coh[5, 5],
+        atol=0.1,
+    )

tests/test_phase_link_covariance.py

@@ -152,3 +152,81 @@ def test_estimate_stack_covariance_neighbors_masked(slcs):
 
     C_neighbors = covariance.coh_mat_single(slc_samples, neighbor_mask=neighbor_mask)
     npt.assert_allclose(C_nan, C_neighbors)
+
+
+def test_coh_mat_single_float_weights(slcs):
+    """Test that float weights work correctly for Gaussian multilooking."""
+    num_slc, _rows, cols = slcs.shape
+    slc_samples = slcs.reshape(num_slc, -1)
+    nsamps = slc_samples.shape[1]
+
+    # Test 1: uniform weights should give same result as no weights
+    uniform_weights = np.ones(nsamps, dtype=np.float32) / nsamps
+    C_uniform = covariance.coh_mat_single(slc_samples, neighbor_mask=uniform_weights)
+
+    # With all True boolean mask
+    bool_mask = np.ones(nsamps, dtype=np.bool_)
+    C_bool = covariance.coh_mat_single(slc_samples, neighbor_mask=bool_mask)
+
+    # The results should be similar (coherence is normalized)
+    npt.assert_allclose(C_uniform, C_bool, rtol=1e-5)
+
+    # Test 2: zero weight on some pixels should reduce their influence
+    # Create weights that zero out the first row
+    weights_partial = np.ones(nsamps, dtype=np.float32)
+    weights_partial[:cols] = 0  # Zero out first row
+    weights_partial = weights_partial / weights_partial.sum()
+
+    C_partial = covariance.coh_mat_single(slc_samples, neighbor_mask=weights_partial)
+
+    # Create equivalent boolean mask
+    bool_partial = np.ones(nsamps, dtype=np.bool_)
+    bool_partial[:cols] = False
+
+    C_bool_partial = covariance.coh_mat_single(slc_samples, neighbor_mask=bool_partial)
+
+    # Results should be similar when weights are 0 or 1
+    npt.assert_allclose(C_partial, C_bool_partial, rtol=1e-5)
+
+
+def test_estimate_stack_covariance_gaussian_weights():
+    """Test covariance estimation with Gaussian weights."""
+    from dolphin.shp._gaussian import estimate_neighbors as gaussian_estimate_neighbors
+
+    # Create a simple test stack
+    nslc, rows, cols = 5, 20, 20
+    slc_stack = np.random.rand(nslc, rows, cols) + 1j * np.random.rand(nslc, rows, cols)
+    slc_stack = slc_stack.astype(np.complex64)
+
+    half_window = HalfWindow(x=3, y=3)
+    strides = Strides(x=1, y=1)
+
+    # Get Gaussian weights
+    gaussian_weights = gaussian_estimate_neighbors(
+        halfwin_rowcol=(half_window.y, half_window.x),
+        input_shape=(rows, cols),
+        strides=(strides.y, strides.x),
+    )
+
+    # Compute covariance with Gaussian weights
+    C_gaussian = covariance.estimate_stack_covariance(
+        slc_stack,
+        half_window=half_window,
+        strides=strides,
+        neighbor_arrays=gaussian_weights,
+    )
+
+    # Check output shape
+    assert C_gaussian.shape == (rows, cols, nslc, nslc)
+
+    # Check that diagonal is 1 (coherence with self)
+    for r in range(rows):
+        for c in range(cols):
+            diag = np.diag(C_gaussian[r, c])
+            npt.assert_allclose(np.abs(diag), 1.0, rtol=1e-5)
+
+    # Check Hermitian symmetry
+    for r in range(rows):
+        for c in range(cols):
+            C = C_gaussian[r, c]
+            npt.assert_allclose(C, C.conj().T, rtol=1e-5)