Goldstein filtering step for unwrapping workflow

src/dolphin/__init__.py

@@ -5,3 +5,4 @@
 from ._log import *
 from ._show_versions import *
 from ._types import *
+from .goldstein import goldstein

src/dolphin/goldstein.py

@@ -0,0 +1,97 @@
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+def goldstein(phase: ArrayLike, alpha: float, psize: int = 32) -> np.ndarray:
+    """Apply the Goldstein adaptive filter to the given data.
+
+    Parameters
+    ----------
+    phase : np.ndarray
+        2D complex array containing the data to be filtered.
+    alpha : float
+        Filtering parameter for Goldstein algorithm
+        Must be between 0 (no filtering) and 1 (maximum filtering)
+    psize : int, optional
+        edge length of square patch
+        Default = 32
+
+    Returns
+    -------
+        2D numpy array of filtered data.
+
+    """
+
+    def apply_pspec(data: ArrayLike):
+        # NaN is allowed value
+        assert not (alpha < 0), f"Invalid parameter value {alpha} < 0"
+        wgt = np.power(np.abs(data) ** 2, alpha / 2)
+        data = wgt * data
+        return data
+
+    def make_wgt(nxp: int, nyp: int) -> np.ndarray:
+        # Create arrays of horizontal and vertical weights
+        wx = 1.0 - np.abs(np.arange(nxp // 2) - (nxp / 2.0 - 1.0)) / (nxp / 2.0 - 1.0)
+        wy = 1.0 - np.abs(np.arange(nyp // 2) - (nyp / 2.0 - 1.0)) / (nyp / 2.0 - 1.0)
+        # Compute the outer product of wx and wy to create
+        # the top-left quadrant of the weight matrix
+        quadrant = np.outer(wy, wx)
+        # Create a full weight matrix by mirroring the quadrant along both axes
+        wgt = np.block(
+            [
+                [quadrant, np.flip(quadrant, axis=1)],
+                [np.flip(quadrant, axis=0), np.flip(np.flip(quadrant, axis=0), axis=1)],
+            ]
+        )
+        return wgt
+
+    def patch_goldstein_filter(
+        data: ArrayLike, wgt: ArrayLike, psize: int
+    ) -> np.ndarray:
+        """Apply the filter to a single patch of data.
+
+        Parameters
+        ----------
+        data : np.ndarray
+            2D complex array containing the data to be filtered.
+        wgt : np.ndarray
+            weight matrix for summing neighboring data
+        psize : int
+            edge length of square FFT area
+
+        Returns
+        -------
+            2D numpy array of filtered data.
+
+        """
+        # Calculate alpha
+        data = np.fft.fft2(data, s=(psize, psize))
+        data = apply_pspec(data)
+        data = np.fft.ifft2(data, s=(psize, psize))
+        return wgt * data
+
+    def apply_goldstein_filter(data: ArrayLike) -> np.ndarray:
+        # Create an empty array for the output
+        out = np.zeros(data.shape, dtype=np.complex64)
+        # ignore processing for empty chunks
+        if np.all(np.isnan(data)):
+            return data
+        # Create the weight matrix
+        wgt_matrix = make_wgt(psize, psize)
+        # Iterate over windows of the data
+        for i in range(0, data.shape[0] - psize, psize // 2):
+            for j in range(0, data.shape[1] - psize, psize // 2):
+                # Create proocessing windows
+                data_window = data[i : i + psize, j : j + psize]
+                wgt_window = wgt_matrix[: data_window.shape[0], : data_window.shape[1]]
+                # Apply the filter to the window
+                filtered_window = patch_goldstein_filter(data_window, wgt_window, psize)
+                # Add the result to the output array
+                slice_i = slice(i, min(i + psize, out.shape[0]))
+                slice_j = slice(j, min(j + psize, out.shape[1]))
+                out[slice_i, slice_j] += filtered_window[
+                    : slice_i.stop - slice_i.start, : slice_j.stop - slice_j.start
+                ]
+        return out
+
+    return apply_goldstein_filter(phase)

src/dolphin/unwrap/_unwrap.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 from typing import Optional, Sequence, Union
 
+import numpy as np
 from tqdm.auto import tqdm
 
-from dolphin import io
+from dolphin import goldstein, io
 from dolphin._log import get_log, log_runtime
 from dolphin._types import Filename
 from dolphin.utils import DummyProcessPoolExecutor, full_suffix
@@ -18,7 +19,6 @@
     DEFAULT_UNW_NODATA,
     UNW_SUFFIX,
 )
-from ._snaphu_py import unwrap_snaphu_py
 from ._tophu import multiscale_unwrap
 from ._utils import create_combined_mask, set_nodata_values
 
@@ -47,6 +47,8 @@ def run(
     ccl_nodata: int | None = DEFAULT_CCL_NODATA,
     scratchdir: Optional[Filename] = None,
     overwrite: bool = False,
+    run_goldstein: bool = False,
+    alpha: float = 0.5,
 ) -> tuple[list[Path], list[Path]]:
     """Run snaphu on all interferograms in a directory.
 
@@ -97,6 +99,10 @@ def run(
         If None, uses `tophu`'s `/tmp/...` default.
     overwrite : bool, optional, default = False
         Overwrite existing unwrapped files.
+    run_goldstein : bool, optional, default = False
+        Whether to run Goldstein filtering on interferogram
+    alpha : float, optional, default = 0.5
+        Alpha parameter for Goldstein filtering
 
     Returns
     -------
@@ -159,6 +165,8 @@ def run(
                 unw_nodata=unw_nodata,
                 ccl_nodata=ccl_nodata,
                 scratchdir=scratchdir,
+                run_goldstein=run_goldstein,
+                alpha=alpha,
             )
             for ifg_file, out_file, cor_file in zip(in_files, out_files, cor_filenames)
         ]
@@ -190,6 +198,8 @@ def unwrap(
     unw_nodata: float | None = DEFAULT_UNW_NODATA,
     ccl_nodata: int | None = DEFAULT_CCL_NODATA,
     scratchdir: Optional[Filename] = None,
+    run_goldstein: bool = False,
+    alpha: float = 0.5,
 ) -> tuple[Path, Path]:
     """Unwrap a single interferogram using snaphu, isce3, or tophu.
 
@@ -244,6 +254,10 @@ def unwrap(
     scratchdir : Filename, optional
         Path to scratch directory to hold intermediate files.
         If None, uses `tophu`'s `/tmp/...` default.
+    run_goldstein : bool, optional, default = False
+        Whether to run Goldstein filtering on interferogram
+    alpha : float, optional, default = 0.5
+        Alpha parameter for Goldstein filtering
 
     Returns
     -------
@@ -272,12 +286,47 @@ def unwrap(
             output_filename=combined_mask_file,
         )
 
+    if run_goldstein:
+        suf = Path(unw_filename).suffix
+        if suf == ".tif":
+            driver = "GTiff"
+            opts = list(io.DEFAULT_TIFF_OPTIONS)
+        else:
+            driver = "ENVI"
+            opts = list(io.DEFAULT_ENVI_OPTIONS)
+
+        # If we're running Goldstein filtering, the intermediate
+        # filtered/unwrapped rasters are temporary rasters in the scratch dir.
+        filt_ifg_filename = (
+            Path(scratchdir or ".") / Path(ifg_filename).with_suffix(".filt" + suf).name
+        )
+        scratch_unw_filename = Path(unw_filename).with_suffix(".filt.unw" + suf)
+
+        ifg = io.load_gdal(ifg_filename)
+        logger.info(f"Goldstein filtering {ifg_filename} -> {filt_ifg_filename}")
+        filt_ifg = goldstein(ifg, alpha=alpha)
+        logger.info(f"Writing filtered output to {filt_ifg_filename}")
+        io.write_arr(
+            arr=filt_ifg,
+            output_name=filt_ifg_filename,
+            like_filename=ifg_filename,
+            driver=driver,
+            options=opts,
+        )
+        unwrapper_ifg_filename = filt_ifg_filename
+        unwrapper_unw_filename = scratch_unw_filename
+    else:
+        unwrapper_ifg_filename = Path(ifg_filename)
+        unwrapper_unw_filename = Path(unw_filename)
+
     if unwrap_method == UnwrapMethod.SNAPHU:
+        from ._snaphu_py import unwrap_snaphu_py
+
         # Pass everything to snaphu-py
         unw_path, conncomp_path = unwrap_snaphu_py(
-            ifg_filename,
+            unwrapper_ifg_filename,
             corr_filename,
-            unw_filename,
+            unwrapper_unw_filename,
             nlooks,
             ntiles=ntiles,
             tile_overlap=tile_overlap,
@@ -292,9 +341,9 @@ def unwrap(
         )
     else:
         unw_path, conncomp_path = multiscale_unwrap(
-            ifg_filename,
+            unwrapper_ifg_filename,
             corr_filename,
-            unw_filename,
+            unwrapper_unw_filename,
             downsample_factor,
             ntiles=ntiles,
             nlooks=nlooks,
@@ -321,4 +370,22 @@ def unwrap(
         filename=conncomp_path, output_nodata=ccl_nodata, like_filename=ifg_filename
     )
 
+    # Transfer ambiguity numbers from filtered unwrapped interferogram
+    # back to original interferogram
+    if run_goldstein:
+        logger.info(
+            f"Transferring ambiguity numbers from filtered ifg {scratch_unw_filename}"
+        )
+        unw_arr = io.load_gdal(scratch_unw_filename)
+
+        final_arr = np.angle(ifg) + (unw_arr - np.angle(filt_ifg))
+
+        io.write_arr(
+            arr=final_arr,
+            output_name=unw_filename,
+            dtype=np.float32,
+            driver=driver,
+            options=opts,
+        )
+
     return unw_path, conncomp_path

src/dolphin/workflows/config/_common.py

@@ -164,6 +164,12 @@ class UnwrapOptions(BaseModel, extra="forbid"):
             "Whether to run the unwrapping step after wrapped phase estimation."
         ),
     )
+    run_goldstein: bool = Field(
+        False,
+        description=(
+            "Whether to run Goldstein filtering step on wrapped interferogram."
+        ),
+    )
     _directory: Path = PrivateAttr(Path("unwrapped"))
     unwrap_method: UnwrapMethod = UnwrapMethod.SNAPHU
     n_parallel_jobs: int = Field(
@@ -205,6 +211,13 @@ class UnwrapOptions(BaseModel, extra="forbid"):
         description="Statistical cost mode method for SNAPHU.",
     )
 
+    alpha: float = Field(
+        0.5,
+        description=(
+            "(for Goldstein filtering) Power parameter for Goldstein algorithm."
+        ),
+    )
+
     @field_validator("ntiles", "downsample_factor", mode="before")
     @classmethod
     def _to_tuple(cls, v):

src/dolphin/workflows/unwrapping.py

@@ -87,6 +87,8 @@ def run(
         downsample_factor=unwrap_options.downsample_factor,
         unwrap_method=unwrap_options.unwrap_method,
         scratchdir=unwrap_scratchdir,
+        run_goldstein=unwrap_options.run_goldstein,
+        alpha=unwrap_options.alpha,
     )
 
     if add_overviews:

tests/test_unwrap.py

@@ -115,6 +115,21 @@ def test_unwrap_snaphu_nodata(self, tmp_path, list_of_gtiff_ifgs, corr_raster):
         assert np.isnan(io.get_raster_nodata(unw_path))
         assert io.get_raster_nodata(conncomp_path) == 123
 
+    @pytest.mark.parametrize("method", [UnwrapMethod.SNAPHU, UnwrapMethod.PHASS])
+    def test_goldstein(self, tmp_path, list_of_gtiff_ifgs, corr_raster, method):
+        # test other init_method
+        unw_filename = tmp_path / "unwrapped.unw.tif"
+        unw_path, conncomp_path = dolphin.unwrap.unwrap(
+            ifg_filename=list_of_gtiff_ifgs[0],
+            corr_filename=corr_raster,
+            unw_filename=unw_filename,
+            nlooks=1,
+            unwrap_method=method,
+            run_goldstein=True,
+        )
+        assert unw_path.exists()
+        assert conncomp_path.exists()
+
 
 class TestUnwrapRun:
     def test_run_gtiff(self, list_of_gtiff_ifgs, corr_raster):