AquaGrid

Pixel-wise AquaCrop on rasters: zarr in, Numba kernels out — CPU or GPU.

Quick start · Config · Architecture · Schema · Testing · Citing · License

Climate as zarr cubes (time, y, x), per-pixel sowing dates (y, x), zarr outputs. Daily-step kernels from AquaCrop-OSPy 3.x are vendored and recompiled with Numba — njit + prange on CPU (one thread per pixel) and numba.cuda on GPU — with bit-exact parity against AquaCrop-OSPy (tests/test_parity.py).

This is not official FAO AquaCrop and not an aquacropos extension. It started in CyMP (Unioeste-LEA).

Current scope: rainfed (no irrigation, no groundwater), one season per pixel. Soil: a single AquaCrop preset or a per-pixel zarr raster (HiHydroSoil hydraulics or sand/silt/clay texture).

Prerequisites

• Python 3.11+
• GPU (optional): NVIDIA GPU with a CUDA driver — Numba talks to the driver; the full CUDA toolkit is not required

Quick start

pip install aquagrid

From Git (no clone, before or besides PyPI):

pip install "git+https://github.com/Paloschi/aquagrid.git"

For development, clone and install editable:

git clone https://github.com/Paloschi/aquagrid.git
cd aquagrid
pip install -e ".[dev]"

# 1. synthetic 10×10 pixels / 540 days
aquagrid synth --out ./data

# 2. config
cp examples/config.example.yaml ./data/config.yaml
# (adjust paths if needed)

# 3. run
aquagrid run --config ./data/config.yaml            # CPU
aquagrid run --config ./data/config.yaml -b gpu     # GPU

Output: output.zarr with final yield/biomass (y, x) and, with save_daily: true, daily series (time, y, x) in group daily. Full schema: docs/zarr-schema.md.

Programmatic

from aquagrid.pipeline import run_grid

run_grid("climate.zarr", "sowing.zarr", "output.zarr",
         crop_name="Maize", soil_name="SandyLoam",
         backend="cpu", save_daily=False)

CLI

YAML config

climate: data/climate.zarr      # cube (time, y, x): tmin, tmax, precip, eto
sowing: data/sowing.zarr        # grid (y, x) int32 YYYYDDD; <=0 = masked
output: data/output.zarr
crop:
  name: Maize                   # any AquaCrop-OSPy crop
  # zarr: data/crop.zarr        # optional (param, y, x) overrides per pixel
soil:
  name: SandyLoam               # AquaCrop preset (xor with zarr below)
  # zarr: data/soil.zarr        # ksat/wcsat/wcpf2/wcpf3 or sand/silt/clay
  # ksat_unit: cm/d
backend: cpu                    # cpu | gpu
options:
  save_daily: false
  tile: 128                     # spatial tile size (pixels)
  max_season_days: 400

Copy from examples/config.example.yaml.

Project structure

aquagrid/
├── docs/                    # zarr I/O schema
├── examples/                # sample YAML config
├── src/aquagrid/
│   ├── cli.py               # Typer: synth / run / bench
│   ├── pipeline.py          # zarr in → tiles → zarr out
│   ├── params.py            # Crop / Soil → kernel arrays
│   ├── soil_grid.py         # PTF + HiHydro layers → compartments
│   ├── bench.py
│   ├── engine/
│   │   ├── cpu.py           # njit + prange
│   │   ├── gpu.py           # numba.cuda
│   │   └── run.py           # shared runner over flat arrays
│   ├── kernels/
│   │   ├── impl.py          # daily step (njit / cuda.jit subset)
│   │   └── loader.py        # compile one source for both backends
│   └── io/                  # schema, validation, synthetic data
└── tests/                   # pytest: parity, grid, soil, io, gpu

Where to look first

Architecture

• kernels/impl.py — AquaCrop daily step ported to a common njit / cuda.jit subset (no allocation in kernels, scalar state per pixel + 1-D compartment views). One source compiled for both backends by kernels/loader.py.
• params.py — flattens aquacrop Crop / Soil into kernel arrays (weather-independent init, including deepening the profile to Zmax + 0.1).
• GDD phenology is computed inside the kernel per pixel (sowing date changes each pixel's thermal accumulation), mirroring aquacrop compute_crop_calendar.

I/O schema

Grids must already be aligned — AquaGrid does not reproject. Details, units, and soil bands: docs/zarr-schema.md.

Benchmark

aquagrid bench --pixels 65536 --days 540 (excluding JIT compile):

One full season (540 days) per pixel. On larger grids the GPU scales better (climate transfer dominates on small grids).

Testing

pytest            # parity vs AquaCrop-OSPy, io, grid driver, gpu (if present)

CI (GitHub Actions) runs on every pull request against main, and again on push to main: Ubuntu and Windows × Python 3.11 / 3.12. Hosted runners are CPU-only; test_gpu.py is skipped.

Contributing

1. Use conventional commits (feat:, fix:, chore:, ci:, docs:).
2. Put tests in the same change as the behaviour they cover.
3. Open a pull request against main — CI must pass before merge.

Citing this project

If you use AquaGrid in research or operational work, please cite this repository:

@software{paloschi_aquagrid,
  author  = {Paloschi, Rennan Andres},
  title   = {AquaGrid: pixel-wise AquaCrop on rasters (Numba CPU/GPU)},
  year    = {2026},
  url     = {https://github.com/Paloschi/aquagrid},
  version = {0.1.0}
}

Also acknowledge AquaCrop-OSPy (and FAO AquaCrop) for the underlying crop water-productivity model that the kernels follow.

License and attribution

MIT. Kernels in src/aquagrid/kernels/impl.py are derived from AquaCrop-OSPy (Apache-2.0).