Water- and N-limited production (winter wheat)¶

This example runs torchcrop's Lintul5 model in its water- and nitrogen-limited mode over the 18-location Brandenburg winter-wheat dataset, then plots the full state and diagnostic trajectories of a few locations against development stage (DVS).

Sitting between the water-limited and the full water + NPK-limited runs, this mode lets two stresses act on growth — soil water (TRANRF) and the nitrogen nutrition index (NNI). The phosphorus and potassium indices (PNI, KNI) are held at 1, so P and K can never limit growth here. The two management levers are supplied externally, day by day, as explicit drivers passed to model(...):

Driver	Shape	Unit	Overrides
irrigation	[B, T]	mm d⁻¹	the soil_params.irri mode (0/1/2)
fertilizer	[B, T, 3]	g m⁻² d⁻¹ (N, P, K)	the ferntab/ferptab/ferktab look-ups

Each tensor is zero on most days and carries a dose only on the scheduled application days. The external irrigation short-circuits the internal IRRI logic, and the external fertilizer replaces the table-driven applications — the scale_factor_fer* scale factors and the recovery fractions nrf/prf/krf still apply to the supplied amounts. Because only nitrogen is limiting in this mode, the management plan focuses on N supply.

It is configured with one switch:

Switch	Value	Effect
crop_params.iopt	3	Run mode = water + N limited

The dataset is autumn-sown (IDPL = 270) while the weather series start on 1 January, so the model is driven with start_doy=1 and relies on its internal sowing latch for the pre-sowing spin-up.

Import libraries¶

In [1]:

import os
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from torch.utils.data import Dataset

from torchcrop import (
    CropParameters,
    SoilParameters,
    SiteParameters,
    Lintul5Model,
    WeatherDriver,
)

plt.rcParams["font.family"] = "DeJavu Serif"
plt.rcParams["font.serif"] = "Times New Roman"

DATA_DIR = Path("../data", "brandenburg", "torchcrop")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

import os from pathlib import Path import numpy as np import pandas as pd import torch import matplotlib.pyplot as plt from torch.utils.data import Dataset from torchcrop import ( CropParameters, SoilParameters, SiteParameters, Lintul5Model, WeatherDriver, ) plt.rcParams["font.family"] = "DeJavu Serif" plt.rcParams["font.serif"] = "Times New Roman" DATA_DIR = Path("../data", "brandenburg", "torchcrop") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}")

Using device: cpu

Prepare dataset and dataloader¶

In [2]:

class TorchCropDataset(Dataset):
    """Reads per-location weather, soil, and site data for torchcrop.

    Directory layout (under ``DATA_DIR``)::

        weather/<location>.csv   daily forcing, one file per location
        soil/soil.csv            one row per location, indexed by ``location``
        site/site.csv            one row per location, indexed by ``location``

    ``__getitem__`` returns a single sample as a dict with:

        * ``weather`` : ``[T, 8]`` float tensor in torchcrop channel order
          (DOY, davtmp, tmin, tmax, irrad, rain, vp, wind)
        * ``soil``    : dict of scalar floats keyed by ``SoilParameters`` field
        * ``site``    : dict of scalar floats keyed by ``SiteParameters`` field

    Use `collate_torchcrop` as the DataLoader ``collate_fn`` to assemble a
    batch into ready-to-run ``WeatherDriver`` / ``SoilParameters`` /
    ``SiteParameters`` objects.
    """

    # CSV weather columns, in the exact order torchcrop expects them.
    WEATHER_COLS = [
        "Date",          # -> doy
        "TempMean",      # -> davtmp (recomputed as (TempMin+TempMax)/2 below)
        "TempMin",       # -> tmin
        "TempMax",       # -> tmax
        "Radiation",     # -> irrad
        "Precipitation", # -> rain
        "VapPressure",   # -> vp
        "Windspeed",     # -> wind
    ]

    # soil.csv column -> SoilParameters field name
    SOIL_MAP = {
        "SMDRY": "wcad",       # air-dry volumetric water content [m³ m⁻³]
        "SMW": "wcwp",         # water content at wilting point [m³ m⁻³]
        "SMFC": "wcfc",        # water content at field capacity [m³ m⁻³]
        "SMO": "wcst",         # water content at saturation [m³ m⁻³]
        "CRAIRC": "crairc",    # critical air content for aeration [m³ m⁻³]
        "SMI": "wci",          # initial root-zone moisture content [m³ m⁻³]
        "SMLOWI": "wci_lower", # initial lower-zone (sub-soil) moisture [m³ m⁻³]
        "RDMSO": "rdmso",      # max rooting depth allowed by the soil [m]
        "RUNFR": "runfr",      # fraction of rainfall lost to surface runoff [0–1]
        "CFEV": "cfev",        # soil-evaporation correction factor [-] (1–4)
        "KSUB": "ksub",        # max percolation rate to deeper layers [mm d⁻¹]
        "NMINS": "nmini",      # initial mineralisable soil organic N [g N m⁻²]
        "PMINS": "pmini",      # initial mineralisable soil organic P [g P m⁻²]
        "KMINS": "kmini",      # initial mineralisable soil organic K [g K m⁻²]
        "RTNMINS": "rtnmins",  # daily N mineralisation fraction [d⁻¹]
        "RTPMINS": "rtpmins",  # daily P mineralisation fraction [d⁻¹]
        "RTKMINS": "rtkmins",  # daily K mineralisation fraction [d⁻¹]
    }

    # site.csv column -> SiteParameters field name
    SITE_MAP = {
        "LATITUDE": "latitude",
        "ALTITUDE": "altitude",
        "IDPL": "idpl",
        "CO2": "co2",
    }

    def __init__(self, weather_dir, soil_dir, site_dir, dtype=torch.float32):
        super().__init__()
        self.weather_dir = Path(weather_dir)
        self.dtype = dtype
        self.soil_data = pd.read_csv(
            os.path.join(soil_dir, "soil.csv")
        ).set_index("location")
        self.site_data = pd.read_csv(
            os.path.join(site_dir, "site.csv")
        ).set_index("location")
        # Locations present in both soil and site tables.
        self.locations = self.soil_data.index.intersection(self.site_data.index)

    def __len__(self):
        return len(self.locations)

    def _load_weather(self, location):
        df = pd.read_csv(
            self.weather_dir / f"{location}.csv", parse_dates=["Date"]
        )
        df = df[self.WEATHER_COLS].copy()
        # SIMPLACE drives phenology with TMPA = (TMIN+TMAX)/2 (Phenology.java:513),
        # not the measured TempMean; match it to reproduce the reference.
        df["TempMean"] = (df["TempMin"] + df["TempMax"]) / 2.0
        df["Radiation"] = df["Radiation"] / 1000.0   # kJ -> MJ m-2 d-1
        df["Date"] = df["Date"].dt.dayofyear          # date -> day-of-year
        return torch.as_tensor(df.values, dtype=self.dtype)  # [T, 8]

    def __getitem__(self, idx):
        location = self.locations[idx]

        weather = self._load_weather(location)

        srow = self.soil_data.loc[location]
        soil = {field: float(srow[col]) for col, field in self.SOIL_MAP.items()}

        trow = self.site_data.loc[location]
        site = {field: float(trow[col]) for col, field in self.SITE_MAP.items()}

        return {"weather": weather, "soil": soil, "site": site}


def collate_torchcrop(batch, dtype=torch.float32):
    """Collate samples into batched torchcrop driver/parameter objects.

    Args:
        batch: list of samples produced by `TorchCropDataset.__getitem__`.
        dtype: target dtype for the assembled tensors.

    Returns:
        Tuple ``(weather, soil_params, site_params)`` where ``weather`` is a
        `WeatherDriver` of shape ``[B, T, 8]`` and ``soil_params`` /
        ``site_params`` are batched dataclasses with ``[B]`` scalar fields.
    """
    # Weather: [B, T, 8]
    weather = torch.stack([s["weather"] for s in batch], dim=0).to(dtype)
    weather = WeatherDriver(weather)

    # Soil / site: stack each field across the batch into a [B] tensor.
    soil_fields = batch[0]["soil"].keys()
    soil_kwargs = {
        f: torch.tensor([s["soil"][f] for s in batch], dtype=dtype)
        for f in soil_fields
    }
    soil_params = SoilParameters(**soil_kwargs)

    site_fields = batch[0]["site"].keys()
    site_kwargs = {
        f: torch.tensor([s["site"][f] for s in batch], dtype=dtype)
        for f in site_fields
    }
    site_params = SiteParameters(**site_kwargs)

    return weather, soil_params, site_params

class TorchCropDataset(Dataset): """Reads per-location weather, soil, and site data for torchcrop. Directory layout (under ``DATA_DIR``):: weather/<location>.csv daily forcing, one file per location soil/soil.csv one row per location, indexed by ``location`` site/site.csv one row per location, indexed by ``location`` ``__getitem__`` returns a single sample as a dict with: * ``weather`` : ``[T, 8]`` float tensor in torchcrop channel order (DOY, davtmp, tmin, tmax, irrad, rain, vp, wind) * ``soil`` : dict of scalar floats keyed by ``SoilParameters`` field * ``site`` : dict of scalar floats keyed by ``SiteParameters`` field Use `collate_torchcrop` as the DataLoader ``collate_fn`` to assemble a batch into ready-to-run ``WeatherDriver`` / ``SoilParameters`` / ``SiteParameters`` objects. """ # CSV weather columns, in the exact order torchcrop expects them. WEATHER_COLS = [ "Date", # -> doy "TempMean", # -> davtmp (recomputed as (TempMin+TempMax)/2 below) "TempMin", # -> tmin "TempMax", # -> tmax "Radiation", # -> irrad "Precipitation", # -> rain "VapPressure", # -> vp "Windspeed", # -> wind ] # soil.csv column -> SoilParameters field name SOIL_MAP = { "SMDRY": "wcad", # air-dry volumetric water content [m³ m⁻³] "SMW": "wcwp", # water content at wilting point [m³ m⁻³] "SMFC": "wcfc", # water content at field capacity [m³ m⁻³] "SMO": "wcst", # water content at saturation [m³ m⁻³] "CRAIRC": "crairc", # critical air content for aeration [m³ m⁻³] "SMI": "wci", # initial root-zone moisture content [m³ m⁻³] "SMLOWI": "wci_lower", # initial lower-zone (sub-soil) moisture [m³ m⁻³] "RDMSO": "rdmso", # max rooting depth allowed by the soil [m] "RUNFR": "runfr", # fraction of rainfall lost to surface runoff [0–1] "CFEV": "cfev", # soil-evaporation correction factor [-] (1–4) "KSUB": "ksub", # max percolation rate to deeper layers [mm d⁻¹] "NMINS": "nmini", # initial mineralisable soil organic N [g N m⁻²] "PMINS": "pmini", # initial mineralisable soil organic P [g P m⁻²] "KMINS": "kmini", # initial mineralisable soil organic K [g K m⁻²] "RTNMINS": "rtnmins", # daily N mineralisation fraction [d⁻¹] "RTPMINS": "rtpmins", # daily P mineralisation fraction [d⁻¹] "RTKMINS": "rtkmins", # daily K mineralisation fraction [d⁻¹] } # site.csv column -> SiteParameters field name SITE_MAP = { "LATITUDE": "latitude", "ALTITUDE": "altitude", "IDPL": "idpl", "CO2": "co2", } def __init__(self, weather_dir, soil_dir, site_dir, dtype=torch.float32): super().__init__() self.weather_dir = Path(weather_dir) self.dtype = dtype self.soil_data = pd.read_csv( os.path.join(soil_dir, "soil.csv") ).set_index("location") self.site_data = pd.read_csv( os.path.join(site_dir, "site.csv") ).set_index("location") # Locations present in both soil and site tables. self.locations = self.soil_data.index.intersection(self.site_data.index) def __len__(self): return len(self.locations) def _load_weather(self, location): df = pd.read_csv( self.weather_dir / f"{location}.csv", parse_dates=["Date"] ) df = df[self.WEATHER_COLS].copy() # SIMPLACE drives phenology with TMPA = (TMIN+TMAX)/2 (Phenology.java:513), # not the measured TempMean; match it to reproduce the reference. df["TempMean"] = (df["TempMin"] + df["TempMax"]) / 2.0 df["Radiation"] = df["Radiation"] / 1000.0 # kJ -> MJ m-2 d-1 df["Date"] = df["Date"].dt.dayofyear # date -> day-of-year return torch.as_tensor(df.values, dtype=self.dtype) # [T, 8] def __getitem__(self, idx): location = self.locations[idx] weather = self._load_weather(location) srow = self.soil_data.loc[location] soil = {field: float(srow[col]) for col, field in self.SOIL_MAP.items()} trow = self.site_data.loc[location] site = {field: float(trow[col]) for col, field in self.SITE_MAP.items()} return {"weather": weather, "soil": soil, "site": site} def collate_torchcrop(batch, dtype=torch.float32): """Collate samples into batched torchcrop driver/parameter objects. Args: batch: list of samples produced by `TorchCropDataset.__getitem__`. dtype: target dtype for the assembled tensors. Returns: Tuple ``(weather, soil_params, site_params)`` where ``weather`` is a `WeatherDriver` of shape ``[B, T, 8]`` and ``soil_params`` / ``site_params`` are batched dataclasses with ``[B]`` scalar fields. """ # Weather: [B, T, 8] weather = torch.stack([s["weather"] for s in batch], dim=0).to(dtype) weather = WeatherDriver(weather) # Soil / site: stack each field across the batch into a [B] tensor. soil_fields = batch[0]["soil"].keys() soil_kwargs = { f: torch.tensor([s["soil"][f] for s in batch], dtype=dtype) for f in soil_fields } soil_params = SoilParameters(**soil_kwargs) site_fields = batch[0]["site"].keys() site_kwargs = { f: torch.tensor([s["site"][f] for s in batch], dtype=dtype) for f in site_fields } site_params = SiteParameters(**site_kwargs) return weather, soil_params, site_params

In [3]:

dataset = TorchCropDataset(
    weather_dir=DATA_DIR / "weather",
    soil_dir=DATA_DIR / "soil",
    site_dir=DATA_DIR / "site",
)

print(f"Dataset: {len(dataset)} locations")

dataset = TorchCropDataset( weather_dir=DATA_DIR / "weather", soil_dir=DATA_DIR / "soil", site_dir=DATA_DIR / "site", ) print(f"Dataset: {len(dataset)} locations")

Dataset: 18 locations

Run the model — water + N limited, with external irrigation & fertilizer¶

The production level is controlled by crop_params.iopt:

iopt	Mode
1	Potential (no stress)
2	Water-limited
3	Water + N limited
4	Water + N + P + K limited

For the water- and N-limited run we set iopt = 3, so the soil-water factor TRANRF and the nitrogen nutrition index NNI can reduce growth, while the phosphorus and potassium indices (PNI, KNI) are forced to 1. Rather than letting the model irrigate and fertilise itself, we hand it two explicit management schedules:

• irrigation [B, T] — 20 mm every 15 days from 1 Apr to 30 Jul 2021, covering the spring–summer water demand. Supplying this tensor overrides soil_params.irri for every day, so the crop receives only rainfall plus these scheduled doses of irrigation.
• fertilizer [B, T, 3] — a small basal N dressing at sowing plus three split N top-dressings through spring (the last axis is ordered N, P, K; the P and K columns stay zero because neither limits growth in this mode).

Both schedules are broadcast identically across all 18 locations here, but each row of the batch could carry its own plan.

In [4]:

# Winter-wheat crop parameters in the water + N-limited mode (IOPT=3)
crop_params = CropParameters(crop_name="wheat")
crop_params.iopt = torch.tensor(3.0)

# Assemble all 18 locations into one batch.
samples = [dataset[i] for i in range(len(dataset))]
weather, soil_params, site_params = collate_torchcrop(samples)
B, T = weather.batch_size, weather.n_days

# Winter-wheat crop parameters in the water + N-limited mode (IOPT=3) crop_params = CropParameters(crop_name="wheat") crop_params.iopt = torch.tensor(3.0) # Assemble all 18 locations into one batch. samples = [dataset[i] for i in range(len(dataset))] weather, soil_params, site_params = collate_torchcrop(samples) B, T = weather.batch_size, weather.n_days

In [5]:

# Build the external management schedules
SIM_START = pd.Timestamp("2020-01-01")

def day_index(date):
    """Day index (0-based) of a calendar date within the weather series."""
    return (pd.Timestamp(date) - SIM_START).days

# Irrigation [B, T] in mm d-1: 20 mm every 15 days across the 2021 season.
irrigation = torch.zeros(B, T)
for date in pd.date_range("2021-04-01", "2021-07-30", freq="15D"):
    irrigation[:, day_index(date)] = 20.0

# Fertiliser [B, T, 3] in g m-2 d-1, last axis ordered (N, P, K).
# Only nitrogen is limiting in IOPT=3, so the P and K columns stay zero.
fertilizer = torch.zeros(B, T, 3)
fert_plan = [
    # (date,        N,   P,   K)  -- g m-2 applied on that day
    ("2020-09-26", 2.0, 0.0, 0.0),  # at sowing
    ("2021-03-01", 8.0, 0.0, 0.0),  # early spring
    ("2021-03-31", 6.0, 0.0, 0.0),  # stem elongation
    ("2021-04-30", 5.0, 0.0, 0.0),  # flag leaf
]
for date, n, p, k in fert_plan:
    fertilizer[:, day_index(date), :] = torch.tensor([n, p, k])

# Build the external management schedules SIM_START = pd.Timestamp("2020-01-01") def day_index(date): """Day index (0-based) of a calendar date within the weather series.""" return (pd.Timestamp(date) - SIM_START).days # Irrigation [B, T] in mm d-1: 20 mm every 15 days across the 2021 season. irrigation = torch.zeros(B, T) for date in pd.date_range("2021-04-01", "2021-07-30", freq="15D"): irrigation[:, day_index(date)] = 20.0 # Fertiliser [B, T, 3] in g m-2 d-1, last axis ordered (N, P, K). # Only nitrogen is limiting in IOPT=3, so the P and K columns stay zero. fertilizer = torch.zeros(B, T, 3) fert_plan = [ # (date, N, P, K) -- g m-2 applied on that day ("2020-09-26", 2.0, 0.0, 0.0), # at sowing ("2021-03-01", 8.0, 0.0, 0.0), # early spring ("2021-03-31", 6.0, 0.0, 0.0), # stem elongation ("2021-04-30", 5.0, 0.0, 0.0), # flag leaf ] for date, n, p, k in fert_plan: fertilizer[:, day_index(date), :] = torch.tensor([n, p, k])

In [6]:

# Move every tensor input onto the device.
crop_params = crop_params.to(device=device)
soil_params = soil_params.to(device=device)
site_params = site_params.to(device=device)
weather = weather.to(device=device)
irrigation = irrigation.to(device=device)
fertilizer = fertilizer.to(device=device)

model = Lintul5Model(crop_params, soil_params, site_params).eval().to(device)
with torch.no_grad():
    output = model(
        weather, start_doy=1, irrigation=irrigation, fertilizer=fertilizer
    )

# Per-location summary at maturity.
results = pd.DataFrame(
    {
        "location": list(dataset.locations),
        "yield_g_m2": np.round(output.yield_.tolist(), 2),
        "adjusted_yield_g_m2": np.round(output.adjusted_yield.tolist(), 2),
        "max_lai": np.round(output.lai.max(dim=1).values.tolist(), 2),
        "final_dvs": output.dvs[:, -1].tolist(),
    }
)
results

# Move every tensor input onto the device. crop_params = crop_params.to(device=device) soil_params = soil_params.to(device=device) site_params = site_params.to(device=device) weather = weather.to(device=device) irrigation = irrigation.to(device=device) fertilizer = fertilizer.to(device=device) model = Lintul5Model(crop_params, soil_params, site_params).eval().to(device) with torch.no_grad(): output = model( weather, start_doy=1, irrigation=irrigation, fertilizer=fertilizer ) # Per-location summary at maturity. results = pd.DataFrame( { "location": list(dataset.locations), "yield_g_m2": np.round(output.yield_.tolist(), 2), "adjusted_yield_g_m2": np.round(output.adjusted_yield.tolist(), 2), "max_lai": np.round(output.lai.max(dim=1).values.tolist(), 2), "final_dvs": output.dvs[:, -1].tolist(), } ) results

Out[6]:

	location	yield_g_m2	adjusted_yield_g_m2	max_lai	final_dvs
0	0	370.91	370.91	1.80	2.0
1	1	472.66	472.66	1.93	2.0
2	2	479.52	479.52	1.84	2.0
3	3	447.08	447.08	1.86	2.0
4	4	418.43	418.43	1.79	2.0
5	5	419.37	419.37	1.86	2.0
6	6	441.44	441.44	1.99	2.0
7	7	432.89	432.89	2.00	2.0
8	8	429.86	429.86	2.02	2.0
9	9	402.99	402.99	1.81	2.0
10	10	421.89	421.89	1.88	2.0
11	11	454.77	454.77	1.84	2.0
12	12	377.91	377.91	1.86	2.0
13	13	374.68	374.68	1.55	2.0
14	14	403.86	403.86	1.76	2.0
15	15	440.84	440.84	1.94	2.0
16	16	393.57	393.57	1.83	2.0
17	17	345.69	345.69	1.42	2.0

Visualizing the results¶

The model returns two per-day trajectories that we visualise the same way:

• State variables — output.states, a list of ModelState holding the model's stored pools (biomass, LAI, soil water, nutrient pools, …).
• Diagnostic variables — output.diagnostics, a list of DiagnosticState holding the intermediate fluxes and stress factors computed each day (TRANRF, RUE, partitioning fractions, transpiration, …).

Each weather series starts on 1 January, but the crop is autumn-sown (IDPL = 270); masking to the sowing latch makes every curve start at the day of sowing. The short near-vertical segment at DVS 0 is the sowing→emergence interval; DVS then climbs 0 → 1 (anthesis) → 2 (maturity).

In [7]:

from matplotlib.lines import Line2D

# Pick five random (reproducible) locations
rng = np.random.default_rng(42)
selected_idx = rng.choice(len(dataset.locations), size=5, replace=False)
locations = [dataset.locations[i] for i in selected_idx]
n_loc = len(locations)
colors = [
    "#1f77b4",  # blue
    "#d62728",  # red
    "#2ca02c",  # green
    "#ff7f0e",  # orange
    "#9467bd",  # purple
]

def mask_until_dvs2(sown_mask, dvs_vals):
    """Boolean mask selecting sown days up to and including maturity (DVS >= 2).

    Args:
        sown_mask: ``[L]`` bool array, ``True`` once the crop is sown.
        dvs_vals:  ``[L]`` development-stage values.

    Returns:
        ``[L]`` bool mask: sown days with index <= the first day DVS reaches 2.
    """
    mature = np.where(dvs_vals >= 2.0)[0]
    cutoff = mature[0] if len(mature) else len(dvs_vals)
    return sown_mask & (np.arange(len(dvs_vals)) <= cutoff)


def plot_trajectories(traj, dvs, sown, field_names, meta, *, title, ncols=4):
    """Plot every channel of ``traj`` against DVS, one line per location.

    A single helper used for both state and diagnostic figures so they share
    the same style, grid layout, labelling and legend.

    Args:
        traj:        ``[B, L, C]`` stacked trajectory array.
        dvs:         ``[B, L]`` development stage (x-axis).
        sown:        ``[B, L]`` sowing latch (> 0 once sown).
        field_names: length-``C`` list naming the channels of ``traj``.
        meta:        ``{field: (label, unit)}`` used for the panel titles.
        title:       figure super-title.
        ncols:       number of subplot columns.
    """
    nrows = int(np.ceil(len(field_names) / ncols))
    fig, axes = plt.subplots(
        nrows, ncols, figsize=(4.2 * ncols, 2.4 * nrows), sharex=True
    )
    axes = axes.ravel()

    for ax, fld in zip(axes, field_names):
        j = field_names.index(fld)
        for c, b in enumerate(selected_idx):
            mask = mask_until_dvs2(sown[b] > 0, dvs[b])
            ax.plot(dvs[b][mask], traj[b, mask, j], color=colors[c], lw=1.2, alpha=0.9)
        label, unit = meta.get(fld, (fld, ""))
        ax.set_title(f"{label} [{unit}]" if unit else label, fontsize=11)
        ax.grid(alpha=0.3)
        ax.margins(x=0.02)

    # Hide unused panels; label the x-axis on the bottom row of plotted panels.
    for ax in axes[len(field_names):]:
        ax.axis("off")
    for ax in axes[max(0, len(field_names) - ncols):len(field_names)]:
        ax.set_xlabel("DVS")

    handles = [Line2D([0], [0], color=colors[i], lw=2) for i in range(n_loc)]
    fig.legend(
        handles,
        [str(loc) for loc in locations],
        title="Location",
        loc="lower center",
        ncol=n_loc,
        fontsize=12,
        bbox_to_anchor=(0.5, 0.01),
        frameon=False,
    )
    fig.suptitle(title, fontsize=13)
    fig.tight_layout(rect=[0, 0.04, 1, 0.99])
    plt.show()

from matplotlib.lines import Line2D # Pick five random (reproducible) locations rng = np.random.default_rng(42) selected_idx = rng.choice(len(dataset.locations), size=5, replace=False) locations = [dataset.locations[i] for i in selected_idx] n_loc = len(locations) colors = [ "#1f77b4", # blue "#d62728", # red "#2ca02c", # green "#ff7f0e", # orange "#9467bd", # purple ] def mask_until_dvs2(sown_mask, dvs_vals): """Boolean mask selecting sown days up to and including maturity (DVS >= 2). Args: sown_mask: ``[L]`` bool array, ``True`` once the crop is sown. dvs_vals: ``[L]`` development-stage values. Returns: ``[L]`` bool mask: sown days with index <= the first day DVS reaches 2. """ mature = np.where(dvs_vals >= 2.0)[0] cutoff = mature[0] if len(mature) else len(dvs_vals) return sown_mask & (np.arange(len(dvs_vals)) <= cutoff) def plot_trajectories(traj, dvs, sown, field_names, meta, *, title, ncols=4): """Plot every channel of ``traj`` against DVS, one line per location. A single helper used for both state and diagnostic figures so they share the same style, grid layout, labelling and legend. Args: traj: ``[B, L, C]`` stacked trajectory array. dvs: ``[B, L]`` development stage (x-axis). sown: ``[B, L]`` sowing latch (> 0 once sown). field_names: length-``C`` list naming the channels of ``traj``. meta: ``{field: (label, unit)}`` used for the panel titles. title: figure super-title. ncols: number of subplot columns. """ nrows = int(np.ceil(len(field_names) / ncols)) fig, axes = plt.subplots( nrows, ncols, figsize=(4.2 * ncols, 2.4 * nrows), sharex=True ) axes = axes.ravel() for ax, fld in zip(axes, field_names): j = field_names.index(fld) for c, b in enumerate(selected_idx): mask = mask_until_dvs2(sown[b] > 0, dvs[b]) ax.plot(dvs[b][mask], traj[b, mask, j], color=colors[c], lw=1.2, alpha=0.9) label, unit = meta.get(fld, (fld, "")) ax.set_title(f"{label} [{unit}]" if unit else label, fontsize=11) ax.grid(alpha=0.3) ax.margins(x=0.02) # Hide unused panels; label the x-axis on the bottom row of plotted panels. for ax in axes[len(field_names):]: ax.axis("off") for ax in axes[max(0, len(field_names) - ncols):len(field_names)]: ax.set_xlabel("DVS") handles = [Line2D([0], [0], color=colors[i], lw=2) for i in range(n_loc)] fig.legend( handles, [str(loc) for loc in locations], title="Location", loc="lower center", ncol=n_loc, fontsize=12, bbox_to_anchor=(0.5, 0.01), frameon=False, ) fig.suptitle(title, fontsize=13) fig.tight_layout(rect=[0, 0.04, 1, 0.99]) plt.show()

State variables vs DVS¶

Stack the per-day ModelState snapshots into a [B, T+1, C] array and plot every stored state variable against DVS with the shared helper.

In [8]:

# Label + unit for each ModelState field.
STATE_META = {
    "tsum": ("Thermal time (post-emergence)", "°C d"),
    "tsump": ("Thermal time (post-sowing)", "°C d"),
    "vern": ("Vernalisation days", "d"),
    "wlv": ("Green leaf DW", "g m⁻²"),
    "wlvd": ("Dead leaf DW", "g m⁻²"),
    "wst": ("Stem DW", "g m⁻²"),
    "wstd": ("Dead stem DW", "g m⁻²"),
    "wrt": ("Root DW", "g m⁻²"),
    "wrtd": ("Dead root DW", "g m⁻²"),
    "wso": ("Storage organ DW", "g m⁻²"),
    "lai": ("Leaf area index", "m² m⁻²"),
    "rootd": ("Rooting depth", "m"),
    "wa": ("Root-zone water", "mm"),
    "wa_lower": ("Lower-zone water", "mm"),
    "dslr": ("Days since last rain", "d"),
    "dsos": ("Days of oxygen shortage", "d"),
    "anlv": ("N in leaves", "g N m⁻²"),
    "anst": ("N in stems", "g N m⁻²"),
    "anrt": ("N in roots", "g N m⁻²"),
    "anso": ("N in storage organs", "g N m⁻²"),
    "aplv": ("P in leaves", "g P m⁻²"),
    "apst": ("P in stems", "g P m⁻²"),
    "aprt": ("P in roots", "g P m⁻²"),
    "apso": ("P in storage organs", "g P m⁻²"),
    "aklv": ("K in leaves", "g K m⁻²"),
    "akst": ("K in stems", "g K m⁻²"),
    "akrt": ("K in roots", "g K m⁻²"),
    "akso": ("K in storage organs", "g K m⁻²"),
    "nmin": ("Soil organic N pool", "g N m⁻²"),
    "pmin": ("Soil organic P pool", "g P m⁻²"),
    "kmin": ("Soil organic K pool", "g K m⁻²"),
    "nmint": ("Soil inorganic N pool", "g N m⁻²"),
    "pmint": ("Soil inorganic P pool", "g P m⁻²"),
    "kmint": ("Soil inorganic K pool", "g K m⁻²"),
    "nlossl": ("N lost via dead leaves", "g N m⁻²"),
    "nlossr": ("N lost via dead roots", "g N m⁻²"),
    "nlosss": ("N lost via dead stems", "g N m⁻²"),
    "plossl": ("P lost via dead leaves", "g P m⁻²"),
    "plossr": ("P lost via dead roots", "g P m⁻²"),
    "plosss": ("P lost via dead stems", "g P m⁻²"),
    "klossl": ("K lost via dead leaves", "g K m⁻²"),
    "klossr": ("K lost via dead roots", "g K m⁻²"),
    "klosss": ("K lost via dead stems", "g K m⁻²"),
    "tran_cum": ("Cumulative transpiration", "mm"),
    "evap_cum": ("Cumulative soil evaporation", "mm"),
    "rain_cum": ("Cumulative rainfall", "mm"),
    "irrig_cum": ("Cumulative irrigation", "mm"),
    "runoff_cum": ("Cumulative runoff", "mm"),
    "drain_cum": ("Cumulative drainage", "mm"),
    "nuptr_cum": ("Cumulative N uptake", "g N m⁻²"),
    "puptr_cum": ("Cumulative P uptake", "g P m⁻²"),
    "kuptr_cum": ("Cumulative K uptake", "g K m⁻²"),
    "nfixtr_cum": ("Cumulative N fixation", "g N m⁻²"),
    "parint_cum": ("Cumulative intercepted PAR", "MJ m⁻²"),
    "gtotal_cum": ("Cumulative gross assimilate", "g DM m⁻²"),
}

# Stack the per-day ModelState snapshots into a [B, T+1, C] array.
field_names = output.states[0].field_names
state_traj = torch.stack([s.stack() for s in output.states], dim=1).cpu().numpy()  # [B, T+1, C]
dvs_traj = output.dvs.cpu().numpy()  # [B, T+1]
sown_traj = torch.stack([s.sown for s in output.states], dim=1).cpu().numpy()  # [B, T+1]

# DVS itself and the binary sowing latch are not informative state curves.
plot_fields = [f for f in field_names if f not in ("dvs", "sown")]
plot_idx = [field_names.index(f) for f in plot_fields]

plot_trajectories(
    state_traj[:, :, plot_idx],
    dvs_traj,
    sown_traj,
    plot_fields,
    STATE_META,
    title="Water- and N-limited winter wheat — state variables vs DVS (5 locations, DVS 0→2)\n",
)

# Label + unit for each ModelState field. STATE_META = { "tsum": ("Thermal time (post-emergence)", "°C d"), "tsump": ("Thermal time (post-sowing)", "°C d"), "vern": ("Vernalisation days", "d"), "wlv": ("Green leaf DW", "g m⁻²"), "wlvd": ("Dead leaf DW", "g m⁻²"), "wst": ("Stem DW", "g m⁻²"), "wstd": ("Dead stem DW", "g m⁻²"), "wrt": ("Root DW", "g m⁻²"), "wrtd": ("Dead root DW", "g m⁻²"), "wso": ("Storage organ DW", "g m⁻²"), "lai": ("Leaf area index", "m² m⁻²"), "rootd": ("Rooting depth", "m"), "wa": ("Root-zone water", "mm"), "wa_lower": ("Lower-zone water", "mm"), "dslr": ("Days since last rain", "d"), "dsos": ("Days of oxygen shortage", "d"), "anlv": ("N in leaves", "g N m⁻²"), "anst": ("N in stems", "g N m⁻²"), "anrt": ("N in roots", "g N m⁻²"), "anso": ("N in storage organs", "g N m⁻²"), "aplv": ("P in leaves", "g P m⁻²"), "apst": ("P in stems", "g P m⁻²"), "aprt": ("P in roots", "g P m⁻²"), "apso": ("P in storage organs", "g P m⁻²"), "aklv": ("K in leaves", "g K m⁻²"), "akst": ("K in stems", "g K m⁻²"), "akrt": ("K in roots", "g K m⁻²"), "akso": ("K in storage organs", "g K m⁻²"), "nmin": ("Soil organic N pool", "g N m⁻²"), "pmin": ("Soil organic P pool", "g P m⁻²"), "kmin": ("Soil organic K pool", "g K m⁻²"), "nmint": ("Soil inorganic N pool", "g N m⁻²"), "pmint": ("Soil inorganic P pool", "g P m⁻²"), "kmint": ("Soil inorganic K pool", "g K m⁻²"), "nlossl": ("N lost via dead leaves", "g N m⁻²"), "nlossr": ("N lost via dead roots", "g N m⁻²"), "nlosss": ("N lost via dead stems", "g N m⁻²"), "plossl": ("P lost via dead leaves", "g P m⁻²"), "plossr": ("P lost via dead roots", "g P m⁻²"), "plosss": ("P lost via dead stems", "g P m⁻²"), "klossl": ("K lost via dead leaves", "g K m⁻²"), "klossr": ("K lost via dead roots", "g K m⁻²"), "klosss": ("K lost via dead stems", "g K m⁻²"), "tran_cum": ("Cumulative transpiration", "mm"), "evap_cum": ("Cumulative soil evaporation", "mm"), "rain_cum": ("Cumulative rainfall", "mm"), "irrig_cum": ("Cumulative irrigation", "mm"), "runoff_cum": ("Cumulative runoff", "mm"), "drain_cum": ("Cumulative drainage", "mm"), "nuptr_cum": ("Cumulative N uptake", "g N m⁻²"), "puptr_cum": ("Cumulative P uptake", "g P m⁻²"), "kuptr_cum": ("Cumulative K uptake", "g K m⁻²"), "nfixtr_cum": ("Cumulative N fixation", "g N m⁻²"), "parint_cum": ("Cumulative intercepted PAR", "MJ m⁻²"), "gtotal_cum": ("Cumulative gross assimilate", "g DM m⁻²"), } # Stack the per-day ModelState snapshots into a [B, T+1, C] array. field_names = output.states[0].field_names state_traj = torch.stack([s.stack() for s in output.states], dim=1).cpu().numpy() # [B, T+1, C] dvs_traj = output.dvs.cpu().numpy() # [B, T+1] sown_traj = torch.stack([s.sown for s in output.states], dim=1).cpu().numpy() # [B, T+1] # DVS itself and the binary sowing latch are not informative state curves. plot_fields = [f for f in field_names if f not in ("dvs", "sown")] plot_idx = [field_names.index(f) for f in plot_fields] plot_trajectories( state_traj[:, :, plot_idx], dvs_traj, sown_traj, plot_fields, STATE_META, title="Water- and N-limited winter wheat — state variables vs DVS (5 locations, DVS 0→2)\n", )

Diagnostic variables vs DVS¶

The diagnostics are the intermediate fluxes and stress factors computed each day — stress factors (TRANRF, NNI, …), photosynthesis drivers (RUE, GTOTAL), light interception, phenology drivers, water/nutrient fluxes and partitioning fractions.

In [9]:

# Label + unit for each DiagnosticState field.
DIAG_META = {
    # Stress factors
    "tranrf": ("Water-stress factor (TRANRF)", "[-]"),
    "rdry": ("Drought reduction (RDRY)", "[-]"),
    "rwet": ("Waterlogging reduction (RWET)", "[-]"),
    "nstress": ("NPK nutrition index (min)", "[-]"),
    "nni": ("Nitrogen nutrition index (NNI)", "[-]"),
    "pni": ("Phosphorus nutrition index (PNI)", "[-]"),
    "kni": ("Potassium nutrition index (KNI)", "[-]"),
    "leaf_heat_factor": ("Leaf heat stress factor", "[-]"),
    "combined_stress": ("Combined stress factor", "[-]"),
    "co2_factor": ("CO₂ transpiration factor", "[-]"),
    # Photosynthesis / growth drivers
    "gtotal": ("Gross assimilate (GTOTAL)", "g DM m⁻² d⁻¹"),
    "rue": ("Radiation use efficiency (RUE)", "g MJ⁻¹"),
    "rtmco": ("Temperature × CO₂ correction", "[-]"),
    # Light / canopy
    "frac_intercepted": ("PAR interception fraction", "[-]"),
    "parint": ("PAR intercepted", "MJ m⁻² d⁻¹"),
    # Phenology drivers
    "dtsu": ("Effective thermal time (DTSU)", "°C d d⁻¹"),
    "photofac": ("Photoperiod factor", "[-]"),
    "vernfac": ("Vernalisation factor", "[-]"),
    # Water fluxes
    "tran": ("Actual transpiration (TRAN)", "mm d⁻¹"),
    "evap": ("Soil evaporation (EVAP)", "mm d⁻¹"),
    "runoff": ("Surface runoff", "mm d⁻¹"),
    "drain": ("Deep drainage", "mm d⁻¹"),
    "rirr": ("Effective irrigation", "mm d⁻¹"),
    "smact": ("Root-zone moisture content", "m³ m⁻³"),
    "smactl": ("Lower-zone moisture content", "m³ m⁻³"),
    # Nutrient fluxes
    "nuptr": ("Daily N uptake", "g N m⁻² d⁻¹"),
    "puptr": ("Daily P uptake", "g P m⁻² d⁻¹"),
    "kuptr": ("Daily K uptake", "g K m⁻² d⁻¹"),
    "nfixtr": ("Daily N fixation", "g N m⁻² d⁻¹"),
    "n_demand": ("N demand", "g N m⁻² d⁻¹"),
    "p_demand": ("P demand", "g P m⁻² d⁻¹"),
    "k_demand": ("K demand", "g K m⁻² d⁻¹"),
    # Partitioning fractions
    "fr": ("Root allocation fraction", "[-]"),
    "fl": ("Leaf allocation fraction", "[-]"),
    "fs": ("Stem allocation fraction", "[-]"),
    "fo": ("Storage organ allocation fraction", "[-]"),
}

# Stack the per-day DiagnosticState snapshots into a [B, T, C] array.
diag_field_names = output.diagnostics[0].field_names
diag_traj = torch.stack(
    [s.stack() for s in output.diagnostics], dim=1
).cpu().numpy()  # [B, T, C]

plot_trajectories(
    diag_traj,
    dvs_traj[:, :-1],
    sown_traj[:, :-1],
    diag_field_names,
    DIAG_META,
    title="Water- and N-limited winter wheat — diagnostic variables vs DVS (5 locations, DVS 0→2)\n",
)

# Label + unit for each DiagnosticState field. DIAG_META = { # Stress factors "tranrf": ("Water-stress factor (TRANRF)", "[-]"), "rdry": ("Drought reduction (RDRY)", "[-]"), "rwet": ("Waterlogging reduction (RWET)", "[-]"), "nstress": ("NPK nutrition index (min)", "[-]"), "nni": ("Nitrogen nutrition index (NNI)", "[-]"), "pni": ("Phosphorus nutrition index (PNI)", "[-]"), "kni": ("Potassium nutrition index (KNI)", "[-]"), "leaf_heat_factor": ("Leaf heat stress factor", "[-]"), "combined_stress": ("Combined stress factor", "[-]"), "co2_factor": ("CO₂ transpiration factor", "[-]"), # Photosynthesis / growth drivers "gtotal": ("Gross assimilate (GTOTAL)", "g DM m⁻² d⁻¹"), "rue": ("Radiation use efficiency (RUE)", "g MJ⁻¹"), "rtmco": ("Temperature × CO₂ correction", "[-]"), # Light / canopy "frac_intercepted": ("PAR interception fraction", "[-]"), "parint": ("PAR intercepted", "MJ m⁻² d⁻¹"), # Phenology drivers "dtsu": ("Effective thermal time (DTSU)", "°C d d⁻¹"), "photofac": ("Photoperiod factor", "[-]"), "vernfac": ("Vernalisation factor", "[-]"), # Water fluxes "tran": ("Actual transpiration (TRAN)", "mm d⁻¹"), "evap": ("Soil evaporation (EVAP)", "mm d⁻¹"), "runoff": ("Surface runoff", "mm d⁻¹"), "drain": ("Deep drainage", "mm d⁻¹"), "rirr": ("Effective irrigation", "mm d⁻¹"), "smact": ("Root-zone moisture content", "m³ m⁻³"), "smactl": ("Lower-zone moisture content", "m³ m⁻³"), # Nutrient fluxes "nuptr": ("Daily N uptake", "g N m⁻² d⁻¹"), "puptr": ("Daily P uptake", "g P m⁻² d⁻¹"), "kuptr": ("Daily K uptake", "g K m⁻² d⁻¹"), "nfixtr": ("Daily N fixation", "g N m⁻² d⁻¹"), "n_demand": ("N demand", "g N m⁻² d⁻¹"), "p_demand": ("P demand", "g P m⁻² d⁻¹"), "k_demand": ("K demand", "g K m⁻² d⁻¹"), # Partitioning fractions "fr": ("Root allocation fraction", "[-]"), "fl": ("Leaf allocation fraction", "[-]"), "fs": ("Stem allocation fraction", "[-]"), "fo": ("Storage organ allocation fraction", "[-]"), } # Stack the per-day DiagnosticState snapshots into a [B, T, C] array. diag_field_names = output.diagnostics[0].field_names diag_traj = torch.stack( [s.stack() for s in output.diagnostics], dim=1 ).cpu().numpy() # [B, T, C] plot_trajectories( diag_traj, dvs_traj[:, :-1], sown_traj[:, :-1], diag_field_names, DIAG_META, title="Water- and N-limited winter wheat — diagnostic variables vs DVS (5 locations, DVS 0→2)\n", )