[project]
name = "research_notebook"
version = "0.1.0"
description = "A scientific research project template"
channels = ["conda-forge"]
platforms = ["linux-64", "osx-arm64"]

[dependencies]
python = ">=3.12"
numpy = "*"
scipy = "*"
xarray = "*"
netcdf4 = "*"
matplotlib = "*"
dvc = ">=3"
hydra-core = ">=1.3"

[pypi-dependencies]
research_notebook = { path = ".", editable = true }
hydra-zen = ">=0.13"
omegaconf = ">=2.3"

[feature.dev.dependencies]
pytest = ">=8"
pytest-cov = ">=5"
pre-commit = ">=3.7"
ruff = ">=0.9"
jupytext = ">=1.16"

[feature.dev.pypi-dependencies]
ty = ">=0.0.1a6"

[feature.docs.dependencies]
nodejs = ">=18"

[feature.docs.pypi-dependencies]
mystmd = "*"

[feature.jupyterlab.dependencies]
jupyterlab = ">=4"
ipykernel = ">=6.29"
jupyterlab-git = "*"
jupyterlab-lsp = "*"
jupyter-lsp = "*"
python-lsp-server = "*"
jupyterlab-myst = "*"
ipywidgets = ">=8"
nbdime = "*"

[feature.jupyterlab.pypi-dependencies]
jupyterlab-code-formatter = "*"
jupyterlab-spellchecker = "*"

[feature.marimo.dependencies]
marimo = "*"
duckdb = ">=1"
altair = ">=5"
polars = ">=1"
pyarrow = "*"
anywidget = "*"

[feature.marimo.pypi-dependencies]
vegafusion = {version = ">=2", extras = ["embed"]}
plotly = ">=5"

# ---------------------------------------------------------------------------
# Showcase env: pyrox
# Pulls pyrox from its git main + the [colab] and [optax] extras so every
# notebook under notebooks/pyrox/ and projects/gaussian_processes/notebooks/
# can be re-executed end-to-end. The latter adds cartopy (Iberian maps in
# the Kronecker-GEV notebook); the rest of its deps overlap with pyrox's
# existing transitive closure (matplotlib, scipy via sklearn, optax).
# ---------------------------------------------------------------------------
[feature.pyrox.dependencies]
python = "3.12.*"
ipykernel = ">=6.29"
jupyter = "*"
nbconvert = "*"
# cartopy pulls geos/proj/shapely via conda-forge — much cleaner than PyPI
# wheels for the C-library-heavy geospatial stack. Used by
# projects/gaussian_processes/notebooks/kronecker/01_spain_extremes.ipynb.
cartopy = ">=0.24"

[feature.pyrox.pypi-dependencies]
# Pinned to a specific commit (rather than branch = "main") so the lockfile
# has a stable git ref. Bump this when you want a newer pyrox.
pyrox = { git = "https://github.com/jejjohnson/pyrox", rev = "c122a87544d244192db622731d24b4a648eba853", extras = ["colab", "optax"] }
# numpyro 0.20.1 declares `jax>=0.7` but its ops/provenance.py still imports
# `xla_pmap_p`, which JAX removed at the 0.9 line. Clamp below 0.9 so the
# resolver picks a version whose API matches numpyro's code. Drop this pin
# once numpyro ≥0.21 (with the updated import) lands on PyPI.
jax = ">=0.7,<0.9"
jaxlib = ">=0.7,<0.9"

[feature.pyrox.tasks]
execute-pyrox = "jupyter nbconvert --to notebook --execute --inplace notebooks/pyrox/gp/*.ipynb notebooks/pyrox/masterclass/*.ipynb notebooks/pyrox/ensembles/*.ipynb --ExecutePreprocessor.timeout=600"
execute-gaussian-processes = "jupyter nbconvert --to notebook --execute --inplace projects/gaussian_processes/notebooks/pyroxgp/*.ipynb projects/gaussian_processes/notebooks/kronecker/*.ipynb --ExecutePreprocessor.timeout=1200"

# ---------------------------------------------------------------------------
# Showcase env: coordax
# Coordinate-aware JAX arrays. Notebooks 08-10 use diffrax (ODE/PDE solves)
# and optax + equinox (parameter/state estimation), so those are bundled in
# alongside coordax itself. No git pin needed — coordax is on PyPI.
# ---------------------------------------------------------------------------
[feature.coordax.dependencies]
python = "3.12.*"
ipykernel = ">=6.29"
jupyter = "*"
nbconvert = "*"
jupytext = ">=1.16"

[feature.coordax.pypi-dependencies]
coordax = ">=0.2"
# coordax unconditionally imports chex at package init (coordax/testing.py),
# but does not declare it as a runtime dep. Pin it here until coordax fixes
# the missing declaration or guards the import.
chex = ">=0.1"
jax = ">=0.5"
jaxlib = ">=0.5"
diffrax = ">=0.5"
optax = ">=0.2"
equinox = ">=0.11"

[feature.coordax.tasks]
execute-coordax = "jupyter nbconvert --to notebook --execute --inplace notebooks/coordax/foundations/*.ipynb notebooks/coordax/derivatives/*.ipynb notebooks/coordax/dynamics/*.ipynb --ExecutePreprocessor.timeout=600"

# ---------------------------------------------------------------------------
# Project env: methane-pod
# Self-contained subproject living under projects/methane_pod/. Bundles
# numpyro + jax so MCMC fits and Monte Carlo simulations from the POD +
# TMTPP notebooks can be re-executed end-to-end.
# ---------------------------------------------------------------------------
[feature.methane-pod.dependencies]
python = "3.12.*"
ipykernel = ">=6.29"
jupyter = "*"
nbconvert = "*"
jupytext = ">=1.16"

[feature.methane-pod.pypi-dependencies]
methane_pod = { path = "projects/methane_pod", editable = true }
# Clamp jax below the 0.9 line to match numpyro's xla_pmap_p import (same
# reasoning as the pyrox feature above). Drop this pin when numpyro ≥0.21
# lands on PyPI with the updated import.
jax = ">=0.7,<0.9"
jaxlib = ">=0.7,<0.9"
numpyro = ">=0.16"
equinox = ">=0.11"
corner = ">=2.2"
pandas = ">=2"

[feature.methane-pod.tasks]
execute-methane-pod = "jupyter nbconvert --to notebook --execute --inplace projects/methane_pod/notebooks/*.ipynb --ExecutePreprocessor.timeout=1200"
test-methane-pod = "pytest projects/methane_pod/tests -v -m 'not slow'"
test-methane-pod-all = "pytest projects/methane_pod/tests -v"

# ---------------------------------------------------------------------------
# Project env: plume-simulation
# Self-contained sub-project under projects/plume_simulation/. Bundles
# numpyro + jax for the Bayesian inference layer and xarray for the
# forward-simulation helper.
# ---------------------------------------------------------------------------
[feature.plume-simulation.dependencies]
python = "3.12.*"
ipykernel = ">=6.29"
jupyter = "*"
nbconvert = "*"

[feature.plume-simulation.pypi-dependencies]
plume_simulation = { path = "projects/plume_simulation", editable = true }
# Same numpyro-0.9-line compatibility clamp as the pyrox / methane-pod envs.
jax = ">=0.7,<0.9"
jaxlib = ">=0.7,<0.9"
numpyro = ">=0.16"
xarray = "*"
matplotlib = ">=3.8"
# diffrax is used by the gauss_puff sub-model for time-varying wind
# cumulative-integral solves under the Gaussian puff forward model.
diffrax = ">=0.6"
equinox = ">=0.11"
# gaussx supplies the structured-linear-algebra primitives (LowRankUpdate
# operator + Woodbury-backed solve) used by the radtran matched-filter
# retrieval. With a few-hundred-channel hyperspectral target and a rank-16
# covariance, the operator-based path is ~30× faster than materialising
# Σ⁻¹ via a dense inverse. Pulled by git commit for lockfile reproducibility.
gaussx = { git = "https://github.com/jejjohnson/gaussx", rev = "fd9a2427" }
# finitevolx supplies the Arakawa C-grid FVM operators used by the
# les_fvm Eulerian dispersion sub-model (3D advection-diffusion of a
# passive methane tracer under a prescribed wind field). Pulled from
# git by tag so the lockfile is reproducible — bump when a new tag
# is released.
finitevolx = { git = "https://github.com/jejjohnson/finitevolX", tag = "v0.0.41" }
# hitran-api (HAPI) powers the hapi_lut sub-package: line-by-line Voigt
# absorption cross-section LUTs. Imported lazily inside the generator, so
# the rest of plume_simulation does not require it at import time. The
# hapi_lut notebooks need to download HITRAN line parameters on first run
# (tens of MB per gas, cached under projects/plume_simulation/data/hitran_cache/).
hitran-api = "*"

[feature.plume-simulation.tasks]
execute-plume-simulation = "jupyter nbconvert --to notebook --execute --inplace projects/plume_simulation/notebooks/*.ipynb projects/plume_simulation/notebooks/gauss_puff/*.ipynb projects/plume_simulation/notebooks/les_fvm/*.ipynb --ExecutePreprocessor.timeout=1200"
# The hapi_lut and radtran notebooks are kept out of the CI-targeted task
# above because they depend on the HITRAN line parameters fetched by hapi_lut
# (network) and the uncommitted .nc LUTs under projects/plume_simulation/data/.
# Run locally via the tasks below.
execute-plume-hapi-lut = "jupyter nbconvert --to notebook --execute --inplace projects/plume_simulation/notebooks/hapi_lut/*.ipynb --ExecutePreprocessor.timeout=1800"
execute-plume-radtran = "jupyter nbconvert --to notebook --execute --inplace projects/plume_simulation/notebooks/radtran/*.ipynb --ExecutePreprocessor.timeout=1800"
test-plume-simulation = "pytest projects/plume_simulation/tests -v -m 'not slow'"
test-plume-simulation-all = "pytest projects/plume_simulation/tests -v"

[environments]
default = { features = ["dev"], solve-group = "default" }
docs = { features = ["docs"], solve-group = "default" }
jupyterlab = { features = ["dev", "jupyterlab"], solve-group = "default" }
marimo = { features = ["dev", "marimo"], solve-group = "default" }
pyrox = { features = ["pyrox"], solve-group = "pyrox" }
coordax = { features = ["coordax"], solve-group = "coordax" }
methane-pod = { features = ["methane-pod", "dev"], solve-group = "methane-pod" }
plume-simulation = { features = ["plume-simulation", "dev"], solve-group = "plume-simulation" }

[tasks]
test = "pytest -v"
lint = "ruff check src/"
format = "ruff format src/"
typecheck = "ty check src/"
precommit = "pre-commit run --all-files"
train = "python scripts/train.py"
train-zen = "python scripts/train_zen.py"
evaluate = "python scripts/evaluate.py"
preprocess = "python scripts/preprocess.py"

[feature.docs.tasks]
docs-serve = "myst start"
docs-build = "myst build --html"

[feature.jupyterlab.tasks]
lab = "jupyter lab"

[feature.marimo.tasks]
marimo-edit = "marimo edit"
