Speed up the post-spurt 2-pi ambiguity interpolation

conda-env.yml

@@ -6,6 +6,8 @@ dependencies:
   - pip>=21.3  # https://pip.pypa.io/en/stable/reference/build-system/pyproject-toml/#editable-installation
   - git  # for pip install, due to setuptools_scm
   - gdal>=3.3
+  - libgdal-netcdf
+  - libgdal-hdf5
   - h5py>=3.6
   - hdf5!=1.12.2 # https://github.com/SciTools/iris/issues/5187 and https://github.com/pydata/xarray/issues/7549
   - jax>=0.4.19
@@ -16,7 +18,7 @@ dependencies:
   - pyproj>=3.3
   - rasterio>=1.3
   - ruamel.yaml>=0.15
-  - scipy>=1.9 # "scipy 0.16+ is required for linear algebra", numba. 1.9 is the oldest version working with jax=0.4.19
+  - scipy>=1.12 # "scipy 0.16+ is required for linear algebra", numba. 1.9 is the oldest version working with jax=0.4.19
   - threadpoolctl>=3.0
   - tqdm>=4.60
   - typing_extensions>=3.10

src/dolphin/unwrap/_post_process.py

@@ -5,6 +5,7 @@
 import numpy as np
 from numba import njit, stencil
 from numpy.typing import ArrayLike, NDArray
+from scipy import ndimage
 from scipy.interpolate import NearestNDInterpolator
 
 TWOPI = 2 * np.pi
@@ -109,17 +110,34 @@ def interpolate_masked_gaps(
 
     # Calculate ambiguities for valid unwrapped pixels
     valid_pixels = ifg_valid & unw_valid
-    ambiguities = get_2pi_ambiguities(unw[valid_pixels])
+    # Work with a subsampled version for the nearest interpolator to speed it up
+    sub = 4
+    ambiguities = get_2pi_ambiguities(unw[::sub, ::sub][valid_pixels[::sub, ::sub]])
 
     # Get coordinates for valid pixels and pixels to interpolate
-    valid_coords = np.array(np.where(valid_pixels)).T
-    interp_coords = np.array(np.where(interpolate_mask)).T
+    valid_coords = np.array(np.where(valid_pixels[::sub, ::sub])).T
+    interp_coords = np.array(np.where(interpolate_mask[::sub, ::sub])).T
 
     # Create and apply the interpolator
     interpolator = NearestNDInterpolator(valid_coords, ambiguities)
     interpolated_ambiguities = interpolator(interp_coords)
 
     # Apply interpolated ambiguities to the wrapped phase
+    # Fill in the subsampled version, then resize it to full
+    interpolated_sub = np.zeros(ifg[::sub, ::sub].shape, dtype=int)
+    interpolated_sub[interpolate_mask[::sub, ::sub]] = interpolated_ambiguities
+
     unw[interpolate_mask] = np.angle(ifg[interpolate_mask]) + (
-        interpolated_ambiguities * TWOPI
+        TWOPI
+        * np.round(_resize(interpolated_sub, ifg.shape)).astype(int)[interpolate_mask]
     )
+
+
+def _resize(image: ArrayLike, output_shape: tuple[int, int]):
+    """Resize `image` to be the same shape as `output_shape`."""
+    input_shape = image.shape[-2:]
+    factors = np.divide(input_shape, output_shape)
+
+    # Translate modes used by np.pad to those used by scipy.ndimage
+    zoom_factors = [1 / f for f in factors]
+    return ndimage.zoom(image, zoom_factors, grid_mode=True, mode="grid-constant")

src/dolphin/unwrap/_unwrap_3d.py

@@ -155,7 +155,9 @@ def run_with_retry(cmd: list[str], num_retries: int = 3) -> int:
         options.temporal_coherence_threshold,
         unw_filenames=unw_filenames,
     )
-    filled_masked_unw_regions(unw_filenames, ifg_filenames)
+
+    if options.run_ambiguity_interpolation:
+        filled_masked_unw_regions(unw_filenames, ifg_filenames)
     return unw_filenames, conncomp_filenames
 
 

src/dolphin/workflows/config/_unwrap_options.py

@@ -266,6 +266,13 @@ class SpurtOptions(BaseModel, extra="forbid"):
         ge=0.0,
         lt=1.0,
     )
+    run_ambiguity_interpolation: bool = Field(
+        True,
+        description=(
+            "After running spurt, interpolate the values that were masked during"
+            " unwrapping (which are otherwise left as nan)."
+        ),
+    )
     # TODO: do we want to allow a "AND" or "OR" option, so users can pick if they want
     # `good_sim & good_temp_coh`, or `good_sim | good_temp_coh`
     general_settings: SpurtGeneralSettings = Field(default_factory=SpurtGeneralSettings)