diff --git a/src/dolphin/phase_link/_core.py b/src/dolphin/phase_link/_core.py index 7ec59e11..0f393985 100644 --- a/src/dolphin/phase_link/_core.py +++ b/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, diff --git a/src/dolphin/phase_link/covariance.py b/src/dolphin/phase_link/covariance.py index b81088d2..d8029cdf 100644 --- a/src/dolphin/phase_link/covariance.py +++ b/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) diff --git a/src/dolphin/shp/__init__.py b/src/dolphin/shp/__init__.py index 11de421f..def1e9da 100644 --- a/src/dolphin/shp/__init__.py +++ b/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) diff --git a/src/dolphin/shp/_gaussian.py b/src/dolphin/shp/_gaussian.py new file mode 100644 index 00000000..015e7e24 --- /dev/null +++ b/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 diff --git a/src/dolphin/workflows/config/_enums.py b/src/dolphin/workflows/config/_enums.py index efb2514d..ac44e431 100644 --- a/src/dolphin/workflows/config/_enums.py +++ b/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" diff --git a/tests/test_phase_link_core.py b/tests/test_phase_link_core.py index df26f74a..89530808 100644 --- a/tests/test_phase_link_core.py +++ b/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, + ) diff --git a/tests/test_phase_link_covariance.py b/tests/test_phase_link_covariance.py index 2765d348..4b98e78c 100644 --- a/tests/test_phase_link_covariance.py +++ b/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) diff --git a/tests/test_shp.py b/tests/test_shp.py index 5f9c97d7..02fa918d 100644 --- a/tests/test_shp.py +++ b/tests/test_shp.py @@ -6,6 +6,7 @@ from dolphin.phase_link import simulate from dolphin.shp import ShpMethod from dolphin.shp._common import remove_unconnected +from dolphin.shp._gaussian import estimate_neighbors as gaussian_estimate_neighbors simulate._seed(1234) @@ -268,3 +269,85 @@ def test_remove_unconnected(): expected[2, 2] = expected[3, 3] = True remove_unconnected(data, inplace=True) assert np.all(data == expected) + + +# Tests for Gaussian multilooking + + +def test_shp_gaussian_smoketest(): + """Basic smoke test for Gaussian multilooking.""" + shape = (5, 50, 50) + slcs = 20 * (np.random.rand(*shape) + 1j * np.random.rand(*shape)) + amp_stack = np.abs(slcs) + mean = np.mean(amp_stack, axis=0) + + # Make sure no errors + weights = shp.estimate_neighbors( + mean=mean, + halfwin_rowcol=(3, 5), + nslc=slcs.shape[0], + alpha=0.05, + method=ShpMethod.GAUSSIAN, + ) + + # Check that weights are returned (not None like RECT) + assert weights is not None + # Check shape: (out_rows, out_cols, window_rows, window_cols) + assert weights.shape == (50, 50, 7, 11) + # Check weights are float, not boolean + assert weights.dtype != np.bool_ + + +def test_shp_gaussian_weights_properties(): + """Test that Gaussian weights have correct properties.""" + halfwin_rowcol = (5, 5) + input_shape = (11, 11) + + weights = gaussian_estimate_neighbors( + halfwin_rowcol=halfwin_rowcol, + input_shape=input_shape, + strides=(1, 1), + ) + + # Check shape + assert weights.shape == (11, 11, 11, 11) + + # All weights should be non-negative + assert np.all(weights >= 0) + + # Center pixel (self) should have weight 0 + window = weights[5, 5] + assert window[5, 5] == 0 + + # Weights should sum to 1 (normalized) + assert np.isclose(np.sum(window), 1.0) + + # Weights should decrease from center + # Check that corners are smaller than near-center pixels + assert window[4, 5] > window[0, 0] + assert window[5, 4] > window[0, 0] + + # Should be symmetric + assert np.isclose(window[4, 5], window[6, 5]) + assert np.isclose(window[5, 4], window[5, 6]) + assert np.isclose(window[0, 0], window[10, 10]) + + +def test_shp_gaussian_with_strides(): + """Test Gaussian weights with strides.""" + halfwin_rowcol = (3, 3) + input_shape = (20, 20) + strides = (2, 2) + + weights = gaussian_estimate_neighbors( + halfwin_rowcol=halfwin_rowcol, + input_shape=input_shape, + strides=strides, + ) + + # Output shape should be reduced by strides + assert weights.shape == (10, 10, 7, 7) + + # All output pixels should have the same weights (Gaussian is position-invariant) + assert np.allclose(weights[0, 0], weights[5, 5]) + assert np.allclose(weights[0, 0], weights[9, 9])