Calculate solid earth tide correction for disp-s1

solid_earth_tides.py

import numpy as np
import pysolid
import rasterio
from dolphin import Filename, io
from opera_utils import get_zero_doppler_time
from rasterio.warp import transform_bounds
from scipy.interpolate import RegularGridInterpolator

from disp_s1._reference import ReferencePoint


def calculate_solid_earth_tides_correction(
    unw_filename: Filename,
    reference_cslc_file: Filename,
    secondary_cslc_file: Filename,
    los_east_file: Filename,
    los_north_file: Filename,
    reference_point: ReferencePoint | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the relative displacement correction for solid earth tides.

    This function calculates the solid earth tides correction using pysolid,
    based on the geographical extent of the unwrapped interferogram and
    the timing information from the reference/secondary CSLC files.

    The output is in meters, relative in time (reference to secondary) and
    space (if the `reference_point` is provided)

    Parameters
    ----------
    unw_filename : str or Path
        Path to sample unwrapped interferogram file.
        This is only used to extract the geographic bounds/CRS for the area of interest.
    reference_cslc_file : str or Path
        Path to the reference CSLC file.
    secondary_cslc_file : str or Path
        Path to the secondary CSLC file.
    los_east_file : str or Path
        Path to the LOS East component raster.
    los_north_file : str or Path
        Path to the LOS North component raster.
    ref_point : ReferencePoint
        (Row, column) index of the spatial reference point to use
        If not provided, output is unreferenced.

    Returns
    -------
    solid_earth_tide : np.ndarray
        2D numpy array containing the correction (in meters) for solid earth tides:
        1. The correction relative to the reference time.
        2. The correction with the reference point subtracted.

    """
    # Get lat/lon boundaries from the unwrapped interferogram
    with rasterio.open(unw_filename) as src:
        bounds = src.bounds
        crs = src.crs
        transform = src.transform
        height, width = src.shape

        # If the CRS is not in lat/lon, transform bounds to EPSG:4326
        if not crs.is_geographic:
            bounds = transform_bounds(crs, "EPSG:4326", *bounds)

    # Extract timing information from the CSLC files
    reference_start_time = get_zero_doppler_time(reference_cslc_file, type_="start")
    secondary_start_time = get_zero_doppler_time(secondary_cslc_file, type_="start")

    # Create the atr object for pysolid
    # We'll use a coarse grid of approximately 2.5 km x 2.5 km
    margin_degrees = 0.1
    height_small = max(25, height // 100)
    width_small = max(25, width // 100)
    atr = {
        # Ensure at least 25 points
        "LENGTH": height_small,
        "WIDTH": width_small,
        "X_FIRST": bounds.left - margin_degrees,
        "Y_FIRST": bounds.top + margin_degrees,  # y decreases as we go south
        "X_STEP": (bounds.right - bounds.left + 2 * margin_degrees) / width_small,
        "Y_STEP": -(bounds.top - bounds.bottom + 2 * margin_degrees) / height_small,
    }

    # Run pySolid to get SET in ENU coordinate system for both times
    set_east, set_north, set_up = pysolid.calc_solid_earth_tides_grid(
        secondary_start_time, atr, display=False, verbose=True
    )
    set_east_ref, set_north_ref, set_up_ref = pysolid.calc_solid_earth_tides_grid(
        reference_start_time, atr, display=False, verbose=True
    )

    # Generate the lat/lon arrays for the SET geogrid
    lat_geo_array = np.linspace(
        atr["Y_FIRST"],
        atr["Y_FIRST"] + atr["Y_STEP"] * atr["LENGTH"],
        num=atr["LENGTH"],
    )
    lon_geo_array = np.linspace(
        atr["X_FIRST"], atr["X_FIRST"] + atr["X_STEP"] * atr["WIDTH"], num=atr["WIDTH"]
    )

    # Create a grid of coordinates for the original unwrapped interferogram
    y, x = np.mgrid[0:height, 0:width]
    if crs.is_geographic:
        lon, lat = rasterio.transform.xy(transform, y, x)
    else:
        # If the original CRS is not geographic, we need to transform the coordinates
        lon, lat = rasterio.warp.transform(crs, "EPSG:4326", x.flatten(), y.flatten())
        lon = lon.reshape(x.shape)
        lat = lat.reshape(y.shape)

    # Interpolate SET components to the unwrapped interferogram grid
    interp_east = RegularGridInterpolator(
        (lat_geo_array, lon_geo_array), set_east - set_east_ref
    )
    interp_north = RegularGridInterpolator(
        (lat_geo_array, lon_geo_array), set_north - set_north_ref
    )
    interp_up = RegularGridInterpolator(
        (lat_geo_array, lon_geo_array), set_up - set_up_ref
    )

    set_east_interp = interp_east((lat, lon))
    set_north_interp = interp_north((lat, lon))
    set_up_interp = interp_up((lat, lon))

    # Load LOS components
    los_east: np.ma.MaskedArray = io.load_gdal(los_east_file, masked=True)
    los_north: np.ma.MaskedArray = io.load_gdal(los_north_file, masked=True)
    los_up: np.ma.MaskedArray = np.sqrt(1 - los_east**2 - los_north**2)

    # Calculate the LOS projection
    set_los = (
        los_east * set_east_interp
        + los_north * set_north_interp
        + los_up * set_up_interp
    )

    if reference_point is None:
        return set_los

    # Subtract the reference point
    ref_row, ref_col = reference_point
    return set_los - set_los[ref_row, ref_col]