Metadata-Version: 2.5
Name: agri-jax
Version: 0.0.3
Summary: Differentiable, batch-parallel field-scale crop-soil process models in JAX (CERES-Maize + RZWQM-style soil water and PET).
Project-URL: Homepage, https://github.com/juksentang/agri-jax
Project-URL: Documentation, https://juksentang.github.io/agri-jax
Author-email: Juk-Sen TANG <juksen.tang@mail.mcgill.ca>
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: CERES-Maize,DSSAT,RZWQM,agriculture,crop model,differentiable,jax
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering
Requires-Python: >=3.11
Requires-Dist: equinox>=0.13
Requires-Dist: jax<0.12,>=0.8
Requires-Dist: jaxtyping
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: pyyaml
Requires-Dist: xarray
Provides-Extra: calib
Requires-Dist: arviz; extra == 'calib'
Requires-Dist: blackjax; extra == 'calib'
Requires-Dist: lineax; extra == 'calib'
Requires-Dist: numpyro; extra == 'calib'
Requires-Dist: optax; extra == 'calib'
Requires-Dist: optimistix; extra == 'calib'
Requires-Dist: salib; extra == 'calib'
Provides-Extra: cuda
Requires-Dist: jax[cuda12]; extra == 'cuda'
Provides-Extra: dev
Requires-Dist: beartype; extra == 'dev'
Requires-Dist: chex; extra == 'dev'
Requires-Dist: hypothesis; extra == 'dev'
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: pyright; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-benchmark; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest-regressions; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: towncrier; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-jupyter; extra == 'docs'
Requires-Dist: mkdocs-material; extra == 'docs'
Requires-Dist: mkdocstrings[python]; extra == 'docs'
Provides-Extra: port
Requires-Dist: fparser>=0.2.5; extra == 'port'
Provides-Extra: report
Requires-Dist: jax; extra == 'report'
Requires-Dist: jinja2; extra == 'report'
Requires-Dist: pandas; extra == 'report'
Description-Content-Type: text/markdown

# Agri-JAX

**Differentiable, batch-parallel field-scale crop–soil process models in JAX.**

Agri-JAX ports the CERES-Maize crop module from [DSSAT-CSM](https://github.com/DSSAT/dssat-csm-os) and re-implements RZWQM2-style soil physics (Brooks–Corey Richards equation, Shuttleworth–Wallace potential evapotranspiration, tile drainage) as pure JAX functions, so that a full multi-year field simulation can be run for 10⁵ parameter sets in one `vmap`, differentiated with `jax.grad`, and embedded in gradient-based calibration, Hamiltonian Monte Carlo, ensemble data assimilation, or hybrid process–ML training loops.

The goal is not speed for its own sake. It is to make a *field-scale* model, with real soil physics and management events, as cheap to run in bulk and as differentiable as the regional-scale bucket models that already exist in this space, while staying compatible with the DSSAT and RZWQM parameter files that agronomists have built up over decades.

> **Status: design stage (September 2026).** The package on PyPI is a placeholder that reserves the name and installs an empty module. The design documents below are complete; the proof of concept starts now. Nothing here has been validated against a Fortran model yet. Do not use it for science until the validation report exists.

## Why

| Task | Runs needed | Needs gradients |
|---|---|---|
| Calibration / uncertainty quantification (LHS, MCMC, HMC) | 10⁴–10⁶ | yes (HMC, variational) |
| Regional gridded simulation (pixels × years × scenarios) | 10⁶–10⁸ | no |
| Ensemble data assimilation (EnKF, particle filters) | 10³–10⁴ members in lockstep | no |
| Hybrid process–ML models (model inside a training loop) | one batch per step | yes |

Existing crop models (DSSAT, RZWQM2, APSIM, STICS, WOFOST) are sequential Fortran or C# codes: fast for one field, impossible to batch, and not differentiable. Differentiable modelling has already reshaped hydrology; field-scale crop–soil modelling is still a gap. The differentiable crop models that exist today (diffWOFOST, torchcrop) are regional-scale bucket models. Agri-JAX is field-scale.

## Design in three rules

A process is a pure function of `(state, params, forcing) -> state`. Authors follow three rules and never see `scan`, `vmap`, `jit`, or `lax.cond`:

1. Everything read is in the arguments; everything changed is in the return value.
2. Branch with `jnp.where` / `jnp.select`, never with Python `if` on state.
3. Never write a loop. Time and samples are handled by the runtime.

A model is a state definition plus an ordered list of processes; any process can be swapped. Soil water comes in two interchangeable flavours, `tipping_bucket` (DSSAT) and `richards` (RZWQM), so the same crop module can be compared under both. Crops are a state dimension (`n_crop`), so intercropping is a shape, not a code path.

## Preliminary throughput (skeleton, not the real model)

A computational skeleton with the same per-day shape as the planned model (37-node implicit Richards × 24 sub-steps × 3 Newton iterations, tridiagonal solve, Shuttleworth–Wallace-shaped PET, CERES-shaped crop step, 3287 days) on one NVIDIA H100:

| Precision | 10⁵ nine-year runs | Throughput |
|---|---|---|
| float64 | 243 s | 411 sims/s |
| float32 | 101 s | 993 sims/s |

The Fortran reference (RZWQM2, one CPU core) takes 24 s per run: 667 core-hours and roughly 10 wall-clock hours for the same 10⁵ runs on a cluster. At this batch size the GPU is saturated and time scales with the number of implicit solves per day: halving sub-steps and Newton iterations (12 × 2) gives 82 s for 10⁵ runs and 14 minutes for 10⁶ runs on one H100. How far the scheme can be thinned without losing agreement with the Fortran oracle is the core question of the first paper. See `docs/en/08_throughput_comparison.md`.

## Roadmap

| Phase | Deliverable |
|---|---|
| Proof of concept (4 weeks) | RZWQM water balance + Shuttleworth–Wallace PET + CERES-Maize on one site, validated day by day against the RZWQM2 and DSSAT-CSM binaries; 10⁵-sample `vmap` timing; gradient check and one NUTS run |
| Framework + first models (6 months) | pip package, docs, validation report, differential-testing tools, Sobol and multi-site calibration, intercropping, an economist-facing report API (posterior intervals, marginal effects, elasticities, identifiability diagnostics) with Stata/R front ends |
| Applications | Multi-site joint gradient calibration and parameter identifiability; hybrid process–ML models on held-out flux-tower sites |

## Documentation

English versions live in `docs/en/`, Chinese originals in `docs/zh_cn/`.

| Document | Content |
|---|---|
| [docs/en/02_architecture.md](docs/en/02_architecture.md) | Package layout, State/Params/Forcing pytrees, `@process`, runtime, IO, calibration, report API |
| [docs/en/03_development_plan.md](docs/en/03_development_plan.md) | Week-by-week proof-of-concept plan, later milestones, risks |
| [docs/en/04_porting_and_diff_testing.md](docs/en/04_porting_and_diff_testing.md) | Fortran-to-JAX porting procedure and subroutine-level differential testing |
| [docs/en/05_maintenance_pipeline.md](docs/en/05_maintenance_pipeline.md) | Tooling, lint rules, test tiers, CI, cluster workflow, releases |
| [docs/en/06_open_source_ecosystem.md](docs/en/06_open_source_ecosystem.md) | Survey of related open-source work and chosen dependencies |
| [docs/en/08_throughput_comparison.md](docs/en/08_throughput_comparison.md) | Fortran vs GPU throughput measurements |
| [README.zh.md](README.zh.md) | Original plan (Chinese) |

## Install

```bash
pip install agri-jax        # placeholder 0.0.x: reserves the name, no model code yet
```

## Licensing and provenance

Apache-2.0. CERES-Maize is ported from DSSAT-CSM (BSD-3). Soil water and PET are re-implemented from the published RZWQM2 equations (Ahuja et al., 2000, and related papers); no RZWQM2 source code is included or redistributed. Validation against the RZWQM2 binary is done privately and reported as numbers only; validation against DSSAT-CSM is public and reproducible.

## Citation

See `CITATION.cff`. A design note and the validation report will be posted as preprints.
