Temporal Aggregation for Time-Series Geodata

pip install \
  "xarray>=2023.10" \
  "rioxarray>=0.15" \
  "geopandas>=1.0" \
  "pandas>=2.2" \
  "numpy>=1.26" \
  "dask>=2024.1" \
  "zarr>=2.17" \
  "pydantic>=2.5" \
  "s3fs>=2024.1"

import xarray as xr
import pandas as pd
import numpy as np
from pathlib import Path


def load_and_validate(
    dataset_path: str,
    chunk_dims: dict | None = None,
) -> xr.Dataset:
    """
    Open a multi-temporal Zarr archive, validate the time dimension,
    and normalise timestamps to UTC-naive for consistent resampling.
    """
    if chunk_dims is None:
        chunk_dims = {"time": 1, "y": 512, "x": 512}

    ds = xr.open_zarr(dataset_path, chunks=chunk_dims, consolidated=True)

    if "time" not in ds.dims:
        raise ValueError(
            "Dataset must contain a 'time' coordinate dimension. "
            f"Found dimensions: {list(ds.dims)}"
        )

    # Normalise to UTC-naive so resample() produces deterministic month boundaries.
    if ds.time.dt.tz is not None:
        ds = ds.assign_coords(
            time=ds.time.dt.tz_convert("UTC").dt.tz_localize(None)
        )

    # Detect large temporal gaps that would invalidate monthly composites.
    time_vals = pd.DatetimeIndex(ds.time.values)
    gaps = time_vals[1:] - time_vals[:-1]
    max_gap_days = gaps.max().days
    if max_gap_days > 30:
        import warnings
        warnings.warn(
            f"Largest temporal gap is {max_gap_days} days. "
            "Monthly composites over this gap will have zero valid observations.",
            stacklevel=2,
        )

    return ds


# Sanity check
ds = load_and_validate("s3://my-bucket/sentinel2-daily.zarr")
assert "time" in ds.dims, "Time dimension missing after load"
print(f"Loaded {len(ds.time)} time steps spanning "
      f"{ds.time.values[0]} – {ds.time.values[-1]}")

def apply_cloud_mask(
    ds: xr.Dataset,
    scl_var: str = "SCL",
    valid_classes: tuple[int, ...] = (4, 5, 6),  # vegetation, bare soil, water
) -> xr.Dataset:
    """
    Null pixels whose Scene Classification Layer value is not in valid_classes.
    Returns a Dataset with invalid pixels set to NaN for all non-SCL variables.
    """
    if scl_var not in ds.data_vars:
        raise KeyError(
            f"Expected SCL variable '{scl_var}' not found. "
            f"Available: {list(ds.data_vars)}"
        )

    scl = ds[scl_var]
    # Build boolean valid mask; True where pixel is usable.
    valid_mask = xr.zeros_like(scl, dtype=bool)
    for cls in valid_classes:
        valid_mask = valid_mask | (scl == cls)

    masked_vars = {}
    for var in ds.data_vars:
        if var == scl_var:
            continue
        masked_vars[var] = ds[var].where(valid_mask)

    return ds.assign(masked_vars)


ds_masked = apply_cloud_mask(ds, scl_var="SCL")
# Quick check: masked fraction should be < 70% per tile on average
cloud_fraction = float(ds_masked["ndvi"].isnull().mean())
print(f"Masked pixel fraction (NDVI): {cloud_fraction:.1%}")
assert cloud_fraction < 0.95, "Cloud masking removed >95% of pixels — check SCL classes"

def resample_and_summarise(
    ds: xr.Dataset,
    variables: list[str],
    freq: str = "ME",
    min_obs: int = 3,
) -> xr.Dataset:
    """
    Resample masked variables to `freq` and compute mean, std, min, max,
    and valid-observation count. Months with fewer than min_obs valid
    pixels are set to NaN.

    Parameters
    ----------
    ds : xr.Dataset
        Cloud-masked dataset with NaN for invalid pixels.
    variables : list[str]
        Variable names to aggregate.
    freq : str
        Pandas/xarray frequency alias. Use 'ME' for month-end (xarray >= 2023.10).
    min_obs : int
        Minimum valid observations required; periods below this are nulled.
    """
    missing = [v for v in variables if v not in ds.data_vars]
    if missing:
        raise KeyError(f"Variables not found in dataset: {missing}")

    agg_parts = []
    for var in variables:
        da = ds[var]
        r = da.resample(time=freq)

        obs_count = r.count(dim="time").rename(f"{var}_obs_count")
        da_mean = r.mean(dim="time").rename(f"{var}_mean")
        da_std = r.std(dim="time").rename(f"{var}_std")
        da_min = r.min(dim="time").rename(f"{var}_min")
        da_max = r.max(dim="time").rename(f"{var}_max")

        # Null any period where valid observations fall below threshold.
        sufficient = obs_count >= min_obs
        da_mean = da_mean.where(sufficient)
        da_std = da_std.where(sufficient)

        merged = xr.merge([da_mean, da_std, da_min, da_max, obs_count])
        agg_parts.append(merged)

    aggregated = xr.merge(agg_parts)
    aggregated.attrs.update({
        "temporal_aggregation_freq": freq,
        "variables_aggregated": variables,
        "min_obs_threshold": min_obs,
        "processing_engine": "xarray+dask",
    })
    return aggregated


aggregated = resample_and_summarise(
    ds_masked,
    variables=["ndvi", "lst"],
    freq="ME",
    min_obs=3,
)
print(aggregated)

def fill_temporal_gaps(
    ds: xr.Dataset,
    method: str = "linear",
    max_gap_periods: int = 2,
) -> xr.Dataset:
    """
    Interpolate NaN periods along the time axis.

    Parameters
    ----------
    method : str
        'linear' for continuous variables, 'ffill' for high-autocorrelation vars.
    max_gap_periods : int
        Maximum consecutive NaN periods to fill; longer gaps remain NaN.
    """
    filled = {}
    for var in ds.data_vars:
        if var.endswith("_obs_count"):
            filled[var] = ds[var]  # never interpolate observation counts
            continue
        if method == "linear":
            filled[var] = ds[var].interpolate_na(
                dim="time",
                method="linear",
                max_gap=max_gap_periods,
            )
        elif method == "ffill":
            filled[var] = ds[var].ffill(dim="time", limit=max_gap_periods)
        else:
            raise ValueError(f"Unsupported fill method: '{method}'")

    return ds.assign(filled)


aggregated_filled = fill_temporal_gaps(aggregated, method="linear", max_gap_periods=2)

from pydantic import BaseModel, field_validator
import zarr


class TemporalFeatureSchema(BaseModel):
    feature_name: str
    temporal_resolution: str
    aggregation_type: str
    valid_range: tuple[float, float]
    missingness_threshold: float = 0.15

    @field_validator("valid_range")
    @classmethod
    def check_range_order(cls, v: tuple[float, float]) -> tuple[float, float]:
        if v[0] >= v[1]:
            raise ValueError(
                f"Lower bound {v[0]} must be strictly less than upper bound {v[1]}."
            )
        return v

    @field_validator("missingness_threshold")
    @classmethod
    def check_threshold(cls, v: float) -> float:
        if not 0.0 <= v <= 1.0:
            raise ValueError("missingness_threshold must be in [0, 1].")
        return v


def register_features(
    ds: xr.Dataset,
    output_path: str,
    schemas: list[TemporalFeatureSchema],
    pipeline_version: str = "1.0.0",
) -> None:
    """
    Validate each feature variable against its schema, then write
    the aggregated dataset to a consolidated Zarr store.
    """
    import datetime

    schema_map = {s.feature_name: s for s in schemas}

    for schema in schemas:
        if schema.feature_name not in ds.data_vars:
            raise KeyError(
                f"Feature '{schema.feature_name}' not found in aggregated dataset."
            )
        da = ds[schema.feature_name]
        missingness = float(da.isnull().mean())
        if missingness > schema.missingness_threshold:
            raise ValueError(
                f"Feature '{schema.feature_name}' missingness {missingness:.1%} "
                f"exceeds threshold {schema.missingness_threshold:.1%}."
            )
        lo, hi = schema.valid_range
        out_of_range = float(((da < lo) | (da > hi)).mean())
        if out_of_range > 0.01:
            raise ValueError(
                f"Feature '{schema.feature_name}': {out_of_range:.1%} of values "
                f"are outside valid range [{lo}, {hi}]."
            )

    ds.attrs.update({
        "pipeline_version": pipeline_version,
        "aggregation_timestamp": datetime.datetime.utcnow().isoformat(),
        "source_dataset": "sentinel2-daily.zarr",
    })
    ds.to_zarr(output_path, mode="w", consolidated=True, compute=True)
    print(f"Registered {len(schemas)} feature(s) to {output_path}")


register_features(
    aggregated_filled,
    output_path="s3://my-bucket/features/monthly-composites.zarr",
    schemas=[
        TemporalFeatureSchema(
            feature_name="ndvi_mean",
            temporal_resolution="ME",
            aggregation_type="mean",
            valid_range=(-0.2, 1.0),
            missingness_threshold=0.10,
        ),
        TemporalFeatureSchema(
            feature_name="lst_mean",
            temporal_resolution="ME",
            aggregation_type="mean",
            valid_range=(220.0, 340.0),  # Kelvin range for land surface temperature
            missingness_threshold=0.15,
        ),
    ],
)

def verify_temporal_output(output_path: str, expected_months: int) -> None:
    """
    Load the registered feature store and assert structural and
    statistical correctness.
    """
    result = xr.open_zarr(output_path, consolidated=True)

    # 1. Temporal completeness
    assert len(result.time) == expected_months, (
        f"Expected {expected_months} monthly timesteps, "
        f"got {len(result.time)}"
    )

    # 2. Calendar alignment — all timestamps should be month-end boundaries
    times = pd.DatetimeIndex(result.time.values)
    is_month_end = times == times + pd.offsets.MonthEnd(0)
    assert is_month_end.all(), (
        f"Non-month-end timestamps detected: {times[~is_month_end].tolist()}"
    )

    # 3. NDVI mean plausibility check
    ndvi_mean = float(result["ndvi_mean"].mean())
    assert -0.2 <= ndvi_mean <= 1.0, (
        f"NDVI mean {ndvi_mean:.3f} is outside valid range"
    )

    # 4. Missingness is within acceptable bounds
    missingness = float(result["ndvi_mean"].isnull().mean())
    assert missingness < 0.20, (
        f"NDVI missingness {missingness:.1%} exceeds 20% threshold"
    )

    print(
        f"Verification passed: {expected_months} monthly timesteps, "
        f"NDVI mean = {ndvi_mean:.3f}, missingness = {missingness:.1%}"
    )


verify_temporal_output(
    "s3://my-bucket/features/monthly-composites.zarr",
    expected_months=24,  # 2 years of monthly composites
)