10 — Weather & climate on one grid: GOES-19 + CHIRPS¶

# Install AerEO and any required plugins for this notebook (Google Colab)
!pip install -q "aereo[viz]" aereo-search-aws-goes aereo-read-satpy \
    "git+https://github.com/frandorr/aereo-earthlens.git#subdirectory=projects/aereo-earthlens"

# 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
for name in ["job_goes19.yaml", "job_chirps.yaml"]:
    urllib.request.urlretrieve(
        f"{GITHUB_RAW}/examples/config/{name}",
        f"config/{name}",
    )

# AOI files
for name in ["cordoba.geojson"]:
    urllib.request.urlretrieve(
        f"{GITHUB_RAW}/examples/config/aoi/{name}",
        f"config/aoi/{name}",
    )

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

job_goes = ExtractionJob.load_from_config(
    config_dir="config",
    config_name="job_goes19",
    overrides=[
        "name=goes19_cordoba",
        "output_uri=/tmp/aereo_weather/goes19",
        # A stormy afternoon over Córdoba.
        "search.start_datetime=2025-12-25T15:00:00Z",
        "search.end_datetime=2025-12-25T15:10:59Z",
    ],
)

job_chirps = ExtractionJob.load_from_config(
    config_dir="config",
    config_name="job_chirps",
    overrides=["name=chirps_cordoba", "output_uri=/tmp/aereo_weather/chirps"],
)

# GOES reads NetCDF via satpy/HDF5, which is not thread-safe ->
# run with processes (no use_threads), as in the GOES-19 notebook.
local_exec = LocalExecutor(workers=-1, cache=TaskResultCache())

def run(job: ExtractionJob):
    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 artifacts, catalog

goes_artifacts, catalog_goes = run(job_goes)

from datetime import date, timedelta

# One job per day: each day lands as its own rows in the catalog,
# identified by start_time (the Major TOM index way).
chirps_catalogs = []
start = date(2025, 12, 19)
for i in range(7):
    day = start + timedelta(days=i)
    job_day = ExtractionJob.load_from_config(
        config_dir="config",
        config_name="job_chirps",
        overrides=[
            f"name=chirps_{day:%Y%m%d}",
            f"output_uri=/tmp/aereo_weather/chirps/{day:%Y%m%d}",
            f"search.start_datetime={day}T00:00:00Z",
            f"search.end_datetime={day}T23:59:59Z",
        ],
    )
    chirps_artifacts, catalog_day = run(job_day)  # keeps the last (wettest) day
    chirps_catalogs.append(catalog_day)

import geopandas as gpd
import pandas as pd

cat_goes = gpd.read_parquet(catalog_goes).assign(slot="goes19")
cat_chirps = pd.concat(
    [gpd.read_parquet(c) for c in chirps_catalogs], ignore_index=True
).assign(slot="chirps")

merged = (
    pd.concat([cat_goes, cat_chirps], ignore_index=True)
    .sort_values(["grid_cell", "start_time"])
    .reset_index(drop=True)
)

index_path = "/tmp/aereo_weather/majortom_index.parquet"
merged.to_parquet(index_path)

common_cells = sorted(set(cat_goes.grid_cell) & set(cat_chirps.grid_cell))
print(f"{len(merged)} rows -> {index_path}")
print(f"{len(common_cells)} grid cells observed by BOTH sources")
merged[["grid_cell", "slot", "start_time", "uri"]].head(6)

from aereo.viz import plot_artifact_patches

fig, ax = plot_artifact_patches(
    goes_artifacts,
    ds_factor=1,
    stretch="percentile",
    cmap="gray",
    aoi=job_goes.target_aoi,
    aoi_edgecolor="blue",
)

fig, ax = plot_artifact_patches(
    chirps_artifacts,
    ds_factor=1,
    stretch="percentile",
    cmap="Blues",
    aoi=job_chirps.target_aoi,
    aoi_edgecolor="blue",
)

import os

import numpy as np
import rasterio

accum_dir = "/tmp/aereo_weather/chirps_accum"
os.makedirs(accum_dir, exist_ok=True)

accum_rows = []
for cell_id in common_cells:
    rows = cat_chirps[cat_chirps.grid_cell == cell_id].sort_values("start_time")
    arrs, profile = [], None
    for uri in dict.fromkeys(rows.uri):
        with rasterio.open(uri) as src:
            arrs.append(src.read(1).astype("float32"))
            profile = src.profile
    accum = np.nansum(np.stack(arrs), axis=0)  # weekly accumulation, mm
    out_uri = f"{accum_dir}/{cell_id}_week.tif"
    profile.update(count=1, dtype="float32", nodata=None)
    with rasterio.open(out_uri, "w", **profile) as dst:
        dst.write(accum, 1)
    row = rows.iloc[0].copy()
    row["uri"] = out_uri
    accum_rows.append(row)

accum_artifacts = gpd.GeoDataFrame(
    accum_rows,
    geometry=gpd.GeoSeries([r["geometry"] for r in accum_rows], crs=cat_chirps.crs),
)
# plot_artifact_patches needs cell_utm_footprint as a GeoSeries too
accum_artifacts["cell_utm_footprint"] = gpd.GeoSeries(
    accum_artifacts["cell_utm_footprint"].values,
    crs=f"EPSG:{accum_artifacts['cell_utm_crs'].iloc[0]}",
)
print(f"{len(accum_artifacts)} accumulated rasters in {accum_dir}")

fig, ax = plot_artifact_patches(
    accum_artifacts,
    ds_factor=1,
    stretch="percentile",
    cmap="Blues",
    aoi=job_chirps.target_aoi,
    aoi_edgecolor="blue",
)

# Overlay: GOES-19 C07 as base, weekly CHIRPS accumulation on top.
fig, ax = plot_artifact_patches(
    goes_artifacts,
    ds_factor=1,
    stretch="percentile",
    cmap="gray",
    aoi=job_goes.target_aoi,
    aoi_edgecolor="blue",
)
fig, ax = plot_artifact_patches(
    accum_artifacts, ds_factor=1, stretch="percentile", cmap="Blues", ax=ax, alpha=0.4
)