From extractions to an ML-ready dataset

# Install AerEO and any required plugins for this notebook (Google Colab)
!pip install -q "aereo[viz,pc]"

# Download config files and AOIs from the GitHub repository so this
# notebook can run outside the repo (e.g. Google Colab).
import os
import urllib.request

GITHUB_RAW = "https://raw.githubusercontent.com/frandorr/aereo/main"

os.makedirs("config/aoi", exist_ok=True)

# Config files (job_sentinel2-ndvi.yaml inherits from job_sentinel2.yaml)
for name in ["job_sentinel2.yaml", "job_sentinel2-ndvi.yaml", "job_sentinel1.yaml"]:
    urllib.request.urlretrieve(
        f"{GITHUB_RAW}/examples/config/{name}",
        f"config/{name}",
    )

# AOI files
urllib.request.urlretrieve(
    f"{GITHUB_RAW}/examples/config/aoi/chocon.geojson",
    "config/aoi/chocon.geojson",
)

from aereo.cache import TaskResultCache
from aereo.executors import LocalExecutor
from aereo.pipeline import ExtractionJob

def load_ndvi_job(name: str, start: str, end: str) -> ExtractionJob:
    return ExtractionJob.load_from_config(
        config_dir="config",
        config_name="job_sentinel2-ndvi",
        overrides=[
            f"name={name}",
            f"output_uri=/tmp/aereo_ml/{name}",
            f"search.start_datetime={start}",
            f"search.end_datetime={end}",
        ],
    )

job_t1 = load_ndvi_job("ndvi_t1", "2024-01-01T00:00:00Z", "2024-01-03T23:59:59Z")
job_t2 = load_ndvi_job("ndvi_t2", "2024-01-08T00:00:00Z", "2024-01-20T23:59:59Z")

job_s1 = ExtractionJob.load_from_config(
    config_dir="config",
    config_name="job_sentinel1",
    overrides=["name=s1", "output_uri=/tmp/aereo_ml/s1"],
)

local_exec = LocalExecutor(workers=-1, use_threads=True, cache=TaskResultCache())

def run(job: ExtractionJob) -> str:
    assets = job.search()
    tasks = job.build_tasks(assets)
    artifacts = job.execute(tasks, executor=local_exec)
    catalog = job.write_catalog(artifacts)
    print(f"{job.name}: {len(artifacts)} artifacts -> {catalog}")
    return catalog

catalog_t1 = run(job_t1)

catalog_t2 = run(job_t2)

catalog_s1 = run(job_s1)

import geopandas as gpd
import pandas as pd
import rasterio

from aereo.grid import GridCell

cat_t1 = gpd.read_parquet(catalog_t1).assign(slot="ndvi_t1")
cat_t2 = gpd.read_parquet(catalog_t2).assign(slot="ndvi_t2")
cat_s1 = gpd.read_parquet(catalog_s1).assign(slot="s1")
cat = pd.concat([cat_t1, cat_t2, cat_s1], ignore_index=True)

# Keep only grid cells observed by ALL THREE jobs.
counts = cat.groupby("grid_cell")["slot"].nunique()
common_cells = sorted(counts[counts == 3].index)
samples = cat[cat.grid_cell.isin(common_cells)]

print(f"{len(common_cells)} grid cells observed by all three jobs")
print(f"{len(samples)} catalog rows reference only {samples.uri.nunique()} unique GeoTIFFs\n")
print("Scenes per (grid cell, slot):")
print(samples.groupby(["grid_cell", "slot"]).size().unstack(fill_value=0))

# The extraction window for each cell: GridCell.to_geobox, the same object
# AerEO uses in reproject_mode='grid'. Snapped centre + fixed size
# (grid_dist * (1 + margin)) -> EVERY cell gets the exact same pixel shape.
def geobox_bounds(cell_id: str, resolution: float = 10, margin: float = 10):
    r = samples[samples.grid_cell == cell_id].iloc[0]
    cell = GridCell(id=r.grid_cell, d=r.grid_dist, cell_geometry=r.cell_geometry)
    bb = cell.to_geobox(resolution=resolution, margin=margin).boundingbox
    return (bb.left, bb.bottom, bb.right, bb.top)

cell_bounds = {c: geobox_bounds(c) for c in common_cells}
side = (cell_bounds[common_cells[0]][2] - cell_bounds[common_cells[0]][0]) / 10
print(f"\nEvery cell window: {side:.0f} x {side:.0f} pixels at 10 m (11 km with 10% margin)")

import warnings

import numpy as np
from rasterio.merge import merge

def mosaic(rows: pd.DataFrame, bounds) -> np.ndarray:
    """Mosaic all scenes of one (grid_cell, slot), windowed to the cell's GeoBox."""
    aligned = []
    for uri in dict.fromkeys(rows.uri):  # one shared tif per scene
        with rasterio.open(uri) as src:
            # merge() with explicit bounds reads only the cell window and
            # cuts it to the pixel grid; masked=True maps nodata -> NaN.
            arr, _ = merge([src], bounds=bounds, masked=True)
            aligned.append(np.asarray(arr.filled(np.nan), dtype="float32"))
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)  # all-NaN pixels stay NaN
        return np.nanmean(np.stack(aligned), axis=0)

SLOTS = ["ndvi_t1", "ndvi_t2", "s1"]
CHANNELS = ["ndvi_t1", "ndvi_t2", "vv", "vh"]

def read_cell(cell_id: str) -> np.ndarray:
    """Mosaic all slots of one grid cell -> (channel, y, x), cell-sized."""
    cell = samples[samples.grid_cell == cell_id]
    bounds = cell_bounds[cell_id]
    mosaics = [mosaic(cell[cell.slot == slot], bounds) for slot in SLOTS]
    return np.concatenate(mosaics, axis=0)  # shapes match exactly, thanks to the GeoBox

data = np.stack([read_cell(c) for c in common_cells])  # (sample, channel, y, x)
print(f"tensor: {data.shape}  # (sample, channel, y, x) — every cell exactly the same shape")

import matplotlib.pyplot as plt

def valid_frac(arr: np.ndarray) -> float:
    return float(np.isfinite(arr).mean())

def single_scene(uri: str, bounds) -> np.ndarray:
    """One scene windowed to the cell GeoBox -> (bands, y, x)."""
    with rasterio.open(uri) as src:
        arr, _ = merge([src], bounds=bounds, masked=True)
        return np.asarray(arr.filled(np.nan), dtype="float32")

# Pick the (cell, slot) where mosaicking gained the most coverage.
gains = {}
for cell_id in common_cells:
    cell = samples[samples.grid_cell == cell_id]
    bounds = cell_bounds[cell_id]
    for slot in SLOTS:
        rows = cell[cell.slot == slot]
        uris = list(dict.fromkeys(rows.uri))
        if len(uris) < 2:
            continue
        best_single = max(valid_frac(single_scene(u, bounds)) for u in uris)
        gains[(cell_id, slot)] = valid_frac(mosaic(rows, bounds)) - best_single

cell_id, slot = max(gains, key=gains.get)
bounds = cell_bounds[cell_id]
rows = samples[(samples.grid_cell == cell_id) & (samples.slot == slot)]
uris = list(dict.fromkeys(rows.uri))
best_uri = max(uris, key=lambda u: valid_frac(single_scene(u, bounds)))
before = single_scene(best_uri, bounds)[0]
after = mosaic(rows, bounds)[0]

cmap = "RdYlGn" if slot.startswith("ndvi") else "gray"
if slot == "s1":  # dB for display, as in the final plot
    before = 10 * np.log10(np.clip(before, 1e-6, None))
    after = 10 * np.log10(np.clip(after, 1e-6, None))

fig, axes = plt.subplots(1, 2, figsize=(11, 5), constrained_layout=True)
for ax, arr, title in zip(axes, [before, after], ["best single scene", f"mosaic of {len(uris)} scenes"]):
    im = ax.imshow(arr, cmap=cmap, **({"vmin": -1, "vmax": 1} if cmap == "RdYlGn" else {}))
    ax.set_title(f"{cell_id} — {slot}\n{title} ({valid_frac(arr):.0%} valid)")
    ax.axis("off")
fig.colorbar(im, ax=axes, shrink=0.7)
plt.show()

import xarray as xr

tensor = xr.DataArray(
    data,
    dims=("sample", "channel", "y", "x"),
    coords={"sample": common_cells, "channel": CHANNELS},
)
tensor

# The merged Major TOM index: one row per (grid cell, scene, slot).
# Repeated grid_cell values are intentional — identified by start/end time.
merged = (
    pd.concat([cat_t1, cat_t2, cat_s1], ignore_index=True)
    .sort_values(["grid_cell", "start_time"])
    .reset_index(drop=True)
)

index_path = "/tmp/aereo_ml/majortom_index.parquet"
merged.to_parquet(index_path)
print(f"Major TOM index written to {index_path}")
print(f"{len(merged)} rows, {merged.grid_cell.nunique()} unique cells, "
      f"{merged.uri.nunique()} unique rasters, "
      f"{merged.collection.nunique()} collections: {sorted(merged.collection.unique())}")

# Read it back: the merged, multi-constellation catalog.
gpd.read_parquet(index_path)[["grid_cell", "slot", "collection", "start_time", "uri"]].head(8)

n_show = min(4, len(common_cells))
fig, axes = plt.subplots(n_show, 3, figsize=(11, 3 * n_show), constrained_layout=True)
for i in range(n_show):
    for j, ch in enumerate(["ndvi_t1", "ndvi_t2", "vv"]):
        ax = axes[i, j]
        band = tensor.sel(sample=common_cells[i], channel=ch).values
        if ch == "vv":
            band = 10 * np.log10(np.clip(band, 1e-6, None))  # dB for display
            im = ax.imshow(band, cmap="gray")
        else:
            im = ax.imshow(band, cmap="RdYlGn", vmin=-1, vmax=1)
        ax.set_title(f"{common_cells[i]} — {ch}")
        ax.axis("off")
    fig.colorbar(im, ax=axes[i, :].tolist(), shrink=0.6)
plt.show()

cell_id = "437D_592L"  # pick any cell from the index
rows = merged[merged.grid_cell == cell_id]
bounds = cell_bounds[cell_id]

print(f"{len(rows)} index rows for cell {cell_id}:")
print(rows[["slot", "collection", "start_time"]].to_string(index=False))

def load_cell(rows: pd.DataFrame, bounds) -> dict[str, np.ndarray]:
    """Rebuild one cell's mosaicked channels from its index rows (windowed reads)."""
    out = {}
    for slot in SLOTS:
        aligned = []
        for uri in dict.fromkeys(rows[rows.slot == slot].uri):
            with rasterio.open(uri) as src:
                arr, _ = merge([src], bounds=bounds, masked=True)
                aligned.append(np.asarray(arr.filled(np.nan), dtype="float32"))
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            out[slot] = np.nanmean(np.stack(aligned), axis=0)
    return out

channels = load_cell(rows, bounds)
print("\n", {k: v.shape for k, v in channels.items()})  # cell-sized, not AOI-sized

fig, axes = plt.subplots(1, 3, figsize=(12, 4), constrained_layout=True)
views = [("ndvi_t1", "RdYlGn", (-1, 1)), ("ndvi_t2", "RdYlGn", (-1, 1)), ("s1", "gray", None)]
for ax, (slot, cmap, vminmax) in zip(axes, views):
    arr = channels[slot][0]
    if slot == "s1":
        arr = 10 * np.log10(np.clip(arr, 1e-6, None))  # vv in dB
    kwargs = {"vmin": vminmax[0], "vmax": vminmax[1]} if vminmax else {}
    ax.imshow(arr, cmap=cmap, **kwargs)
    ax.set_title(f"{cell_id} — {slot}")
    ax.axis("off")
plt.show()