Step-by-step: building an AerEO pipeline in Python

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

from aereo.builtins.search import search_stac
from datetime import datetime, timezone
from shapely.geometry import Polygon

from aereo.pipeline import ExtractionJob
from aereo.builtins import read_odc_stac, write_geotiff

# AOI polygon — Chocón reservoir, Argentina (inlined so the notebook has no file dependencies).
aoi_polygon = Polygon(
    [
        (-68.90986824592407, -39.23705421799603),
        (-68.65925870907353, -39.23705421799603),
        (-68.65925870907353, -39.41589522092947),
        (-68.90986824592407, -39.41589522092947),
        (-68.90986824592407, -39.23705421799603),
    ]
)

grid_dist = 10_000

job = ExtractionJob(
    name="sentinel2_sample",
    grid_dist=grid_dist,
    output_uri="/tmp/aereo_extraction",
    read=read_odc_stac,
    write=write_geotiff,
    target_aoi=aoi_polygon,
)

assets = search_stac(
    stac_api_url="https://earth-search.aws.element84.com/v1",
    collections={"sentinel-2-l2a": ["red", "nir"]},
    intersects=aoi_polygon,
    start_datetime=datetime(2024, 1, 1, tzinfo=timezone.utc),
    end_datetime=datetime(2024, 1, 10, tzinfo=timezone.utc),
)

print(f"\u2713 Found {len(assets)} asset rows")
print("Columns:", list(assets.columns))
print("First rows:")
assets[["id", "collection", "channel_id", "crs", "href"]].head()

from aereo.builtins.task_builder import build_grouped_tasks

tasks = build_grouped_tasks(
    search_results=assets,
    job=job,
)
print(f"\u2713 Built {len(tasks)} extraction task(s)")
for i, task in enumerate(tasks):
    print(f"Task {i}: {task}")

# Inspect the first task in detail
task = tasks[0]

print("Grid cells:", task.grid_cells)
print("Task context:", dict(task.task_context))
print(f"Number of asset rows in task: {len(task.assets)}")

# Extraction stages available on the task (delegated from the job)
print("Extraction stages:")
print("  read callable:", callable(task.job.read))
print("  write callable:", callable(task.job.write))

ds_native = job.read(task)

print("Native dataset:")
ds_native
# ds_native["red"].plot()

from pathlib import Path

# Write the native dataset to a temporary path
out_path = Path(job.output_uri) / "demo_native_red.tif"
# crate the parent directory if it doesn't exist
out_path.parent.mkdir(parents=True, exist_ok=True)
written = job.write(ds=ds_native, path=out_path)

print(f"\u2713 Wrote {written}")

from aereo.executors import LocalExecutor

# set job with partial functions for read and write
from functools import partial

# create a write partial function passing rioxarray/to_raster kwargs
write_partial = partial(write_geotiff, compress="LZW", tags={"creator": "aereo"})

job = ExtractionJob(
    name="sentinel2_sample",
    grid_dist=grid_dist,
    output_uri="/tmp/aereo_extraction",
    read=read_odc_stac,
    write=write_partial,
    target_aoi=aoi_polygon,
)

# lets build again the tasks with the new job
tasks = build_grouped_tasks(
    search_results=assets,
    job=job,
)

executor = LocalExecutor(workers=4, use_threads=True)
artifacts = executor(tasks)

print(f"\u2713 Extracted {len(artifacts)} artifact row(s) from {len(tasks)} task(s)")
artifacts[["id", "uri", "grid_cell", "start_time"]].head()

# lets check that chunks are ok
import rioxarray

rioxarray.open_rasterio(artifacts["uri"].iloc[0])

from aereo.viz import plot_artifact_patches

plot_artifact_patches(artifacts, ds_factor=1, cmap="viridis")