Ascat_Fac.py

import numpy as np
import rasterio as rio
from functools import partial
from statsmodels.nonparametric.smoothers_lowess import lowess


partial_lowess = partial(
    lowess,
    frac=2/3,
    it=3,
    delta=0.0,
    xvals=None,
    is_sorted=False,
    missing='drop',
    return_sorted=True
)

with rio.open("ascat_raw_mask.tif") as ascat_src, rio.open("fac_ais_mask.tif") as fac_src:
    ascat_meta = ascat_src.meta.copy()
    fac_meta = fac_src.meta.copy()
    assert ascat_src.crs == fac_src.crs, "CRS mismatch!"
    ascat_data = ascat_src.read(masked=True)
    fac_data = fac_src.read(masked=True)
    ascat_forward = np.array(ascat_src.transform).reshape((3, 3))
    fac_backward = np.array(~fac_src.transform).reshape((3, 3))

    ascat_mask = np.any(ascat_data.mask, axis=0)
    fac_mask = np.any(fac_data.mask, axis=0)
    fac_mask |= fac_data.mean(axis=0, keepdims=False) > 5

    ascat_row, ascat_col = np.where(~ascat_mask)
    ascat_indexes = np.stack(
        [(ascat_col + 0.5), (ascat_row + 0.5), np.ones_like(ascat_col)],
        axis=-1
    )
    fac_indexes = np.round(
        (
            np.einsum(
                "ij,kj -> ki",
                fac_backward, np.einsum(
                    "ij,kj -> ki",
                    ascat_forward,
                    ascat_indexes
                )
            )[:, :-1] - 0.5
        ),
        decimals=0
    ).astype(int)
    fac_col, fac_row = np.array_split(
        ary=fac_indexes,
        indices_or_sections=fac_indexes.shape[-1],
        axis=-1
    )

    fac_col = np.squeeze(a=fac_col, axis=-1)
    fac_row = np.squeeze(a=fac_row, axis=-1)
    valid = np.logical_and(
        np.logical_and((fac_col >= 0), (fac_col < fac_src.width)),
        np.logical_and((fac_row >= 0), (fac_row < fac_src.height))
    )
    fac_col = fac_col[valid]
    fac_row = fac_row[valid]
    ascat_col = ascat_col[valid]
    ascat_row = ascat_row[valid]

    indices = ~fac_mask[fac_row, fac_col]
    fac_col = fac_col[indices]
    fac_row = fac_row[indices]
    ascat_row = ascat_row[indices]
    ascat_col = ascat_col[indices]

    fac_values = fac_data[:, fac_row, fac_col].filled()
    ascat_values = ascat_data[:, ascat_row, ascat_col].filled()
    smooth = partial_lowess(ascat_values.ravel(), fac_values.ravel())
    smooth = smooth.reshape(
        (ascat_values.shape[0], ascat_values.shape[1], 2)
    )

    fac_pred, ascat_pred = np.array_split(ary=smooth, indices_or_sections=2, axis=-1)
    fac_pimg = np.full(shape=fac_data.shape, fill_value=np.nan, dtype=float)
    ascat_pimg = np.full(shape=ascat_data.shape, fill_value=np.nan, dtype=float)
    fac_pimg[:, fac_row, fac_col] = fac_pred.squeeze(-1)
    ascat_pimg[:, ascat_row, ascat_col] = ascat_pred.squeeze(-1)
    fac_rimg = fac_data.filled() - fac_pimg
    ascat_rimg = ascat_data.filled() - ascat_pimg

    ascat_meta["dtype"] = ascat_rimg.dtype
    ascat_meta["nodata"] = np.nan
    with rio.open("ascat_residual.tif", "w", **ascat_meta) as ascat_dst:
        ascat_dst.write(ascat_rimg)

    fac_meta["dtype"] = fac_rimg.dtype
    fac_meta["nodata"] = np.nan
    with rio.open("fac_residual.tif", "w", **fac_meta) as fac_dst:
        fac_dst.write(fac_rimg)