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 or 3.12
• 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 and status (y, x) and, with save_daily: true, daily series (time, y, x) in group daily (NaN outside that pixel's season). Full schema, including status codes: 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       # (time, y, x): tmin, tmax, precip, eto
sowing: data/sowing.zarr         # (y, x) int32 YYYYDDD; <=0 = not simulated
sowing_var: sowing               # variable name in the sowing store
initial_water_content: FC       # FC | WP | SAT, either soil mode
output: data/output.zarr

crop:
  name: Maize                    # any AquaCrop-OSPy crop
  zarr: null                     # optional (param, y, x) cube; null skips it
soil:
  name: SandyLoam                # preset; remove when zarr is a path
  zarr: null                     # hydraulic or texture raster; null uses name
  ksat_unit: cm/d                # float hydraulic Ksat: cm/d or mm/d
  scale_factors: {}              # hydraulic only, e.g. {ksat: 10}; else ignored

backend: cpu                     # cpu | gpu
options:
  save_daily: false
  tile: 128                      # pixels on each side of a tile
  parallel: true                 # CPU thread per pixel; GPU ignores this
  evap_time_steps: 20            # soil-evaporation substeps per day
  max_season_days: 400           # cap; reaching it records status 0

That block is the full set of keys run_from_config reads. Copy it from examples/config.example.yaml.

initial_water_content accepts FC, WP, or SAT for a preset and for a soil raster. Provide exactly one of soil.name and soil.zarr. soil.scale_factors and soil.ksat_unit apply only when soil.zarr is a hydraulic raster (ksat, wcsat, wcpf2, wcpf3); a preset or a texture raster ignores them. options.parallel is the CPU thread-per-pixel switch; the GPU backend ignores it. Status codes: docs/zarr-schema.md.

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 set NUMBA_ENABLE_CUDASIM=1, so test_gpu.py compiles the CUDA kernel in the simulator instead of being skipped. A runner with a GPU should leave that variable unset and run the same test on the device. Coverage of pipeline, io, params, soil_grid and engine fails the job below 90% (kernels/impl.py stays outside that percentage; parity covers it).

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.3.1}
}

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).