{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {},
   "source": [
    "# Composable Callables & Full‑Stack `GrowthModel` — Practical Guide\n",
    "\n",
    "This notebook shows how to:\n",
    "\n",
    "1. **Wrap scientific models as `Callable`s** (e.g., Elfving 2010 DBH increment, Söderberg 1986 heights,\n",
    "   Fridman–Ståhl 2006 mortality, and Edgren–Nylinder 1949 taper + a `PriceList`).\n",
    "2. Build a **full‑stack `GrowthModel`** that runs all components inside `grow()`.\n",
    "3. Use the **factory** to adapt Angle‑Count stands (Bitterlich) to safe working inventories.\n",
    "4. Run the **DSL** controller (triggers/schedules).\n",
    "5. Broadcast aggregate runs with the **ContextEnsemble** (Python/Numba/JAX engines).\n",
    "\n",
    "> All types and utilities come from `pyforestry.base.simulation` & `pyforestry.base.helpers`.\n",
    "Everything lives in `base/` as per the project conventions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acae54e37e7d407bbb7b55eff062a284",
   "metadata": {},
   "source": [
    "## Notebook Objectives\n",
    "- Explain composable callable patterns and how they connect to `GrowthModel` runtime orchestration.\n",
    "- Provide runnable, copy-safe snippets that work in the docs build environment.\n",
    "\n",
    "## Prerequisites\n",
    "- Python environment with `pyforestry` installed from this repository.\n",
    "- Execute cells in order; random components should use fixed seeds where shown.\n",
    "\n",
    "## Sources\n",
    "- Simulation architecture in `pyforestry.base.simulation` and `pyforestry.simulation`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
   "metadata": {},
   "source": [
    "## 1) Imports & quick recap\n",
    "\n",
    "- `GrowthModel` is the **factory + behavior**.\n",
    "- `SimulationContext` is the **sandbox** (all mutation + history live here).\n",
    "- `run_pipeline` is the **controller**: an ordered list of steps over the whole stand, one period at a time.\n",
    "- A **policy** (`Callable[[ctx], Sequence[Action]]`) is all management is -- triggers, schedules and rulesets are policies with an `if`.\n",
    "- `AdapterRegistry` provides **Angle‑Count → pseudo tree‑list/spatial/diameter‑class** adapters.\n",
    "- `ContextEnsemble` broadcasts aggregate steps through a **batch engine** (Python/Numba/JAX)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8dd0d8092fe74a7c96281538738b07e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from typing import Any, Callable, Dict, Optional\n",
    "\n",
    "from pyforestry.base.helpers import (\n",
    "    PICEA_ABIES,\n",
    "    PINUS_SYLVESTRIS,\n",
    "    AngleCount,\n",
    "    CircularPlot,\n",
    "    Stand,\n",
    "    Tree,\n",
    ")\n",
    "from pyforestry.base.simulation import (\n",
    "    ContextEnsemble,\n",
    "    GrowthModel,\n",
    "    Requirements,\n",
    "    Action,\n",
    "    GrowthStep,\n",
    "    ManagementStep,\n",
    "    SimulationContext,\n",
    "    run_pipeline,\n",
    "    when,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72eea5119410473aa328ad9291626812",
   "metadata": {},
   "source": [
    "## 2) Define the scientific components as `Callable`s\n",
    "\n",
    "We keep each study/model as a function. In production, you’ll call your real implementations here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8edb47106e1a46a883d545849b8ab81b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Signatures (type hints are for clarity; not strictly required)\n",
    "ElfvingDbhIncrementFn = Callable[[Any, float, Optional[float], Any, float, Dict[str, Any]], float]\n",
    "SoderbergHeightFn     = Callable[[Any, float, Any, Optional[float]], float]\n",
    "FridmanStahlSurvivalFn= Callable[[Any, float, Optional[float], Any, float, Dict[str, Any]], float]\n",
    "EdgrenNylinderVolFn   = Callable[[Any, float, float, Any], float]\n",
    "PriceFn               = Callable[[Any, float, Any], float]\n",
    "\n",
    "# --- Replace these with your calibrated implementations ---\n",
    "def elfving_dbh_increment(sp, dbh_cm, h_m, site, dt, state):\n",
    "    # Δdbh in cm over dt years (placeholder)\n",
    "    return 0.25 * dt\n",
    "\n",
    "def soderberg_height(sp, dbh_cm, site, age):\n",
    "    # height in m from DBH (placeholder)\n",
    "    return max(1.3, 1.3 + 0.6 * (dbh_cm ** 0.5))\n",
    "\n",
    "def fridman_stahl_survival(sp, dbh_cm, h_m, site, dt, state):\n",
    "    # survival fraction in [0,1] for the step (placeholder)\n",
    "    return max(0.0, min(1.0, 1.0 - 0.006 * dt))\n",
    "\n",
    "def edgren_nylinder_volume(sp, dbh_cm, h_m, site):\n",
    "    # taper-based whole-stem volume in m3 (placeholder)\n",
    "    return 0.00007854 * (dbh_cm ** 2) * h_m\n",
    "\n",
    "def price_list(sp, vol_m3, site):\n",
    "    # SEK per m3 (placeholder)\n",
    "    return 500.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10185d26023b46108eb7d9f57d49d2b3",
   "metadata": {},
   "source": [
    "## 3) A full‑stack `GrowthModel` that chains the callables inside `grow()`\n",
    "\n",
    "This model prefers `tree_list`/`spatial` inventories. If the input Stand uses **Angle‑Count**, we’ll ask the factory to adapt to a pseudo tree‑list so per-tree operations work safely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8763a12b2bbd4a93a75aff182afb95dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class FullStackCallableModel(GrowthModel):\n",
    "    dbh_increment_fn: ElfvingDbhIncrementFn\n",
    "    height_fn:       SoderbergHeightFn\n",
    "    survival_fn:     FridmanStahlSurvivalFn\n",
    "    volume_fn:       EdgrenNylinderVolFn\n",
    "    price_fn:        PriceFn\n",
    "    remove_zero_weight: bool = True\n",
    "\n",
    "    def requirements(self) -> Requirements:\n",
    "        # We need per-tree inventories to run taper/price properly.\n",
    "        return Requirements(inventory=\"tree_list\")\n",
    "\n",
    "    def update_step(self, ctx: SimulationContext, dt: float) -> None:\n",
    "        self.grow(ctx, dt)\n",
    "\n",
    "    def grow(self, ctx: SimulationContext, dt: float) -> None:\n",
    "        if ctx.mode not in (\"tree_list\", \"spatial\"):\n",
    "            raise RuntimeError(f\"{self.__class__.__name__} requires per-tree mode; got {ctx.mode}.\")\n",
    "        ctx.state[\"years_since_thin\"] = ctx.state.get(\"years_since_thin\", 0.0) + dt\n",
    "        step_value = 0.0\n",
    "        for p in ctx.plots:\n",
    "            for t in p.trees:\n",
    "                sp = getattr(t, \"species\", None)\n",
    "                if sp is None:\n",
    "                    continue\n",
    "                dbh = float(getattr(t, \"diameter_cm\", 0.0) or 0.0)\n",
    "                h   = float(getattr(t, \"height_m\",   0.0) or 0.0)\n",
    "                age = getattr(t, \"age\", None)\n",
    "\n",
    "                # 1) DBH increment (Elfving 2010)\n",
    "                ddbh = float(self.dbh_increment_fn(sp, dbh, (h if h > 0 else None), ctx.site, dt, ctx.state))\n",
    "                dbh = max(0.0, dbh + ddbh)\n",
    "                t.diameter_cm = dbh\n",
    "\n",
    "                # 2) Height update (Söderberg 1986)\n",
    "                h = float(self.height_fn(sp, dbh, ctx.site, age))\n",
    "                t.height_m = h\n",
    "\n",
    "                # 3) Mortality (Fridman–Ståhl 2006)\n",
    "                surv = float(self.survival_fn(sp, dbh, h, ctx.site, dt, ctx.state))\n",
    "                surv = 0.0 if surv < 0.0 else (1.0 if surv > 1.0 else surv)\n",
    "                t.weight_n = float(getattr(t, \"weight_n\", 1.0)) * surv\n",
    "\n",
    "                # 4) Value (Edgren–Nylinder 1949 taper + PriceList)\n",
    "                if dbh > 0.0 and h > 0.0 and t.weight_n > 0.0:\n",
    "                    vol_m3 = float(self.volume_fn(sp, dbh, h, ctx.site))\n",
    "                    price  = float(self.price_fn(sp, vol_m3, ctx.site))\n",
    "                    step_value += vol_m3 * price * float(t.weight_n)\n",
    "\n",
    "        if self.remove_zero_weight:\n",
    "            for p in ctx.plots:\n",
    "                p.trees = [t for t in p.trees if float(getattr(t, \"weight_n\", 0.0) or 0.0) > 1e-9]\n",
    "\n",
    "        ctx.state[\"last_step_value_SEK_per_ha\"] = step_value\n",
    "        ctx.state[\"cum_value_SEK_per_ha\"] = ctx.state.get(\"cum_value_SEK_per_ha\", 0.0) + step_value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7623eae2785240b9bd12b16a66d81610",
   "metadata": {},
   "source": [
    "## 4) Build a Stand and run the model through a pipeline\n",
    "\n",
    "We’ll use a small synthetic tree‑list Stand. The `GrowthModel` factory builds a sandboxed `SimulationContext` for us and `run_pipeline` steps it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7cdc8c89c7104fffa095e18ddfef8986",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cumulative value (SEK/ha): 7058.081815293103\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>t</th>\n",
       "      <th>op</th>\n",
       "      <th>details</th>\n",
       "      <th>ba_total</th>\n",
       "      <th>n_total</th>\n",
       "      <th>qmd_total_cm</th>\n",
       "      <th>model_state</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>6.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>17.706657</td>\n",
       "      <td>530.494635</td>\n",
       "      <td>20.614977</td>\n",
       "      <td>{'t': 6.0, 'years_since_thin': 6.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>7.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>18.029392</td>\n",
       "      <td>527.311667</td>\n",
       "      <td>20.864689</td>\n",
       "      <td>{'t': 7.0, 'years_since_thin': 7.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>8.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>18.352762</td>\n",
       "      <td>524.147797</td>\n",
       "      <td>21.114407</td>\n",
       "      <td>{'t': 8.0, 'years_since_thin': 8.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>9.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>18.676717</td>\n",
       "      <td>521.002910</td>\n",
       "      <td>21.364132</td>\n",
       "      <td>{'t': 9.0, 'years_since_thin': 9.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>10.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>19.001208</td>\n",
       "      <td>517.876893</td>\n",
       "      <td>21.613863</td>\n",
       "      <td>{'t': 10.0, 'years_since_thin': 10.0, 'last_dt...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      t           op                        details   ba_total     n_total  \\\n",
       "5   6.0  update_step  {'dt': 1.0, 'management': {}}  17.706657  530.494635   \n",
       "6   7.0  update_step  {'dt': 1.0, 'management': {}}  18.029392  527.311667   \n",
       "7   8.0  update_step  {'dt': 1.0, 'management': {}}  18.352762  524.147797   \n",
       "8   9.0  update_step  {'dt': 1.0, 'management': {}}  18.676717  521.002910   \n",
       "9  10.0  update_step  {'dt': 1.0, 'management': {}}  19.001208  517.876893   \n",
       "\n",
       "   qmd_total_cm                                        model_state  \n",
       "5     20.614977  {'t': 6.0, 'years_since_thin': 6.0, 'last_dt':...  \n",
       "6     20.864689  {'t': 7.0, 'years_since_thin': 7.0, 'last_dt':...  \n",
       "7     21.114407  {'t': 8.0, 'years_since_thin': 8.0, 'last_dt':...  \n",
       "8     21.364132  {'t': 9.0, 'years_since_thin': 9.0, 'last_dt':...  \n",
       "9     21.613863  {'t': 10.0, 'years_since_thin': 10.0, 'last_dt...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# --- A tiny tree-list stand ---\n",
    "stand = Stand(\n",
    "    area_ha=1.0,\n",
    "    plots=[\n",
    "        CircularPlot(id=1, area_m2=200.0, trees=[\n",
    "            Tree(species=PICEA_ABIES, diameter_cm=20.0, height_m=15.0, weight_n=6),\n",
    "            Tree(species=PICEA_ABIES, diameter_cm=18.0, height_m=13.0, weight_n=5),\n",
    "        ])\n",
    "    ],\n",
    ")\n",
    "\n",
    "model = FullStackCallableModel(\n",
    "    dbh_increment_fn=elfving_dbh_increment,\n",
    "    height_fn=soderberg_height,\n",
    "    survival_fn=fridman_stahl_survival,\n",
    "    volume_fn=edgren_nylinder_volume,\n",
    "    price_fn=price_list,\n",
    ")\n",
    "\n",
    "ok, missing = model.can_build(stand, allow_adapters=True, mode_hint=\"tree_list\")\n",
    "assert ok, f\"missing: {missing}\"\n",
    "ctx = model.build_context(stand, mode_hint=\"tree_list\")\n",
    "\n",
    "run_pipeline(ctx, (GrowthStep(),), years=10.0, step=1.0)\n",
    "print(\"Cumulative value (SEK/ha):\", ctx.state.get(\"cum_value_SEK_per_ha\", 0.0))\n",
    "ctx.to_pandas().tail()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b118ea5561624da68c537baed56e602f",
   "metadata": {},
   "source": [
    "## 5) Angle‑Count stands → pseudo tree‑lists (safe adapter)\n",
    "\n",
    "If the Stand carries Bitterlich tallies, you can still run per‑tree logic by **opting in** to a pseudo tree‑list at build time (your original Stand stays immutable)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "938c804e27f84196a10c8828c723f798",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inventory origin: angle_count_pseudo_tree_list\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>t</th>\n",
       "      <th>op</th>\n",
       "      <th>details</th>\n",
       "      <th>ba_total</th>\n",
       "      <th>n_total</th>\n",
       "      <th>qmd_total_cm</th>\n",
       "      <th>model_state</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>49.769666</td>\n",
       "      <td>4.970000</td>\n",
       "      <td>357.074823</td>\n",
       "      <td>{'t': 1.0, 'years_since_thin': 1.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>49.540345</td>\n",
       "      <td>4.940180</td>\n",
       "      <td>357.324823</td>\n",
       "      <td>{'t': 2.0, 'years_since_thin': 2.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0, 'management': {}}</td>\n",
       "      <td>49.312033</td>\n",
       "      <td>4.910539</td>\n",
       "      <td>357.574823</td>\n",
       "      <td>{'t': 3.0, 'years_since_thin': 3.0, 'last_dt':...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     t           op                        details   ba_total   n_total  \\\n",
       "0  1.0  update_step  {'dt': 1.0, 'management': {}}  49.769666  4.970000   \n",
       "1  2.0  update_step  {'dt': 1.0, 'management': {}}  49.540345  4.940180   \n",
       "2  3.0  update_step  {'dt': 1.0, 'management': {}}  49.312033  4.910539   \n",
       "\n",
       "   qmd_total_cm                                        model_state  \n",
       "0    357.074823  {'t': 1.0, 'years_since_thin': 1.0, 'last_dt':...  \n",
       "1    357.324823  {'t': 2.0, 'years_since_thin': 2.0, 'last_dt':...  \n",
       "2    357.574823  {'t': 3.0, 'years_since_thin': 3.0, 'last_dt':...  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# --- Bitterlich tally stand ---\n",
    "ac_plot = CircularPlot(id=\"ac1\", area_m2=10000.0, AngleCount=[\n",
    "    AngleCount(\n",
    "        ba_factor=10.0,\n",
    "        species=[PICEA_ABIES, PINUS_SYLVESTRIS],\n",
    "        value=[3.0, 2.0],\n",
    "        # Tallies expand into a pseudo tree list only when the tallied trees'\n",
    "        # diameters were recorded: stems/ha is derived from them.\n",
    "        diameters_cm=[[24.0, 21.0, 19.0], [27.0, 23.0]],\n",
    "    )\n",
    "])\n",
    "ac_stand = Stand(area_ha=1.0, plots=[ac_plot])\n",
    "\n",
    "ok, missing = model.can_build(ac_stand, allow_adapters=True, mode_hint=\"tree_list\")\n",
    "assert ok, missing\n",
    "ctx_ac = model.build_context(ac_stand, mode_hint=\"tree_list\")  # auto AC→pseudo tree-list\n",
    "\n",
    "run_pipeline(ctx_ac, (GrowthStep(),), years=5.0, step=1.0)\n",
    "print(\"Inventory origin:\", ctx_ac.attrs.get(\"inventory_origin\"))\n",
    "ctx_ac.to_pandas().head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504fb2a444614c0babb325280ed9130a",
   "metadata": {},
   "source": [
    "## 6) Using triggers/schedules (DSL)\n",
    "\n",
    "You can still add management actions with the DSL. Here we add a simple **post** trigger that prints a message whenever QMD exceeds a threshold (placeholder for a thinning action)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "59bbdb311c014d738909a11f9e486628",
   "metadata": {},
   "outputs": [],
   "source": [
    "def qmd_exceeds(ctx: SimulationContext, threshold_cm: float = 22.0) -> bool:\n",
    "    return float(ctx.metrics[\"QMD\"][\"TOTAL\"]) > threshold_cm\n",
    "\n",
    "def announce(ctx: SimulationContext):\n",
    "    print(f\"t={ctx.state['t']:.1f} → QMD now {float(ctx.metrics['QMD']['TOTAL']):.2f} cm\")\n",
    "\n",
    "# A trigger is a policy with an `if`: `when(predicate, action)`.\n",
    "watch = when(\n",
    "    lambda c: qmd_exceeds(c, 22.0),\n",
    "    Action(name=\"qmd_watch\", apply=announce),\n",
    ")\n",
    "\n",
    "ctx2 = model.build_context(stand, mode_hint=\"tree_list\")\n",
    "# Placing the ManagementStep after the GrowthStep is what `check_phase=\"post\"`\n",
    "# used to mean -- the ordering is now the pipeline, so you can read it.\n",
    "run_pipeline(\n",
    "    ctx2,\n",
    "    (GrowthStep(), ManagementStep(watch)),\n",
    "    years=5.0,\n",
    "    step=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b43b363d81ae4b689946ece5c682cd59",
   "metadata": {},
   "source": [
    "## 7) Vectorized aggregate runs with `ContextEnsemble`\n",
    "\n",
    "If you also keep an **aggregate** approximation of your model, you can opt in to the batch engine by\n",
    "implementing `has_batch_engine()` and `batch_grow_step(ba, n, dt, fert_mask)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8a65eabff63a45729fe45fb5ade58bdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>t</th>\n",
       "      <th>op</th>\n",
       "      <th>details</th>\n",
       "      <th>ba_total</th>\n",
       "      <th>n_total</th>\n",
       "      <th>qmd_total_cm</th>\n",
       "      <th>model_state</th>\n",
       "      <th>context_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0}</td>\n",
       "      <td>16.260098</td>\n",
       "      <td>546.700000</td>\n",
       "      <td>19.459967</td>\n",
       "      <td>{'t': 1.0, 'years_since_thin': 0.0, 'last_dt':...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0}</td>\n",
       "      <td>16.747901</td>\n",
       "      <td>543.419800</td>\n",
       "      <td>19.809227</td>\n",
       "      <td>{'t': 2.0, 'years_since_thin': 0.0, 'last_dt':...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>3.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0}</td>\n",
       "      <td>17.250338</td>\n",
       "      <td>540.159281</td>\n",
       "      <td>20.164755</td>\n",
       "      <td>{'t': 3.0, 'years_since_thin': 0.0, 'last_dt':...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>4.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0}</td>\n",
       "      <td>17.767848</td>\n",
       "      <td>536.918326</td>\n",
       "      <td>20.526664</td>\n",
       "      <td>{'t': 4.0, 'years_since_thin': 0.0, 'last_dt':...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>5.0</td>\n",
       "      <td>update_step</td>\n",
       "      <td>{'dt': 1.0}</td>\n",
       "      <td>18.300884</td>\n",
       "      <td>533.696816</td>\n",
       "      <td>20.895068</td>\n",
       "      <td>{'t': 5.0, 'years_since_thin': 0.0, 'last_dt':...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     t           op      details   ba_total     n_total  qmd_total_cm  \\\n",
       "5  1.0  update_step  {'dt': 1.0}  16.260098  546.700000     19.459967   \n",
       "6  2.0  update_step  {'dt': 1.0}  16.747901  543.419800     19.809227   \n",
       "7  3.0  update_step  {'dt': 1.0}  17.250338  540.159281     20.164755   \n",
       "8  4.0  update_step  {'dt': 1.0}  17.767848  536.918326     20.526664   \n",
       "9  5.0  update_step  {'dt': 1.0}  18.300884  533.696816     20.895068   \n",
       "\n",
       "                                         model_state  context_id  \n",
       "5  {'t': 1.0, 'years_since_thin': 0.0, 'last_dt':...           1  \n",
       "6  {'t': 2.0, 'years_since_thin': 0.0, 'last_dt':...           1  \n",
       "7  {'t': 3.0, 'years_since_thin': 0.0, 'last_dt':...           1  \n",
       "8  {'t': 4.0, 'years_since_thin': 0.0, 'last_dt':...           1  \n",
       "9  {'t': 5.0, 'years_since_thin': 0.0, 'last_dt':...           1  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class AggregateBANModel(GrowthModel):\n",
    "    def __init__(self, ba_rel_per_year=0.03, mort_per_year=0.006):\n",
    "        self.ba_rel = ba_rel_per_year\n",
    "        self.mort = mort_per_year\n",
    "    def requirements(self) -> Requirements:\n",
    "        return Requirements(inventory=\"aggregate\")\n",
    "    def has_batch_engine(self) -> bool:\n",
    "        return True\n",
    "    def batch_grow_step(self, ba: np.ndarray, n: np.ndarray, dt: float, fert_mask: np.ndarray):\n",
    "        # Simple BA relative growth and mortality on N (vectorized)\n",
    "        ba2 = ba * (1.0 + self.ba_rel * dt)\n",
    "        n2  = n  * (1.0 - self.mort * dt)\n",
    "        return ba2, n2\n",
    "    def update_step(self, ctx: SimulationContext, dt: float) -> None:\n",
    "        # Fallback scalar (rarely used when batched)\n",
    "        ba = float(ctx.metrics[\"BasalArea\"][\"TOTAL\"]) * (1.0 + self.ba_rel * dt)\n",
    "        n  = float(ctx.metrics[\"Stems\"][\"TOTAL\"]) * (1.0 - self.mort * dt)\n",
    "        ctx.set_aggregate_metrics(ba_total=ba, stems_total=n)\n",
    "\n",
    "agg_model = AggregateBANModel()\n",
    "\n",
    "# Build a few contexts in aggregate mode (force with mode_hint)\n",
    "stands = [stand, stand]  # reuse same stand for brevity\n",
    "agg_contexts = []\n",
    "for s in stands:\n",
    "    ok, _ = agg_model.can_build(s, mode_hint=\"aggregate\")\n",
    "    ctxa = agg_model.build_context(s, mode_hint=\"aggregate\")\n",
    "    agg_contexts.append(ctxa)\n",
    "\n",
    "ens = ContextEnsemble(contexts=agg_contexts, model=agg_model)\n",
    "for _ in range(5):\n",
    "    ens.update_step(1.0)\n",
    "\n",
    "ens.to_pandas().tail()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3933fab20d04ec698c2621248eb3be0",
   "metadata": {},
   "source": [
    "## 8) Wrap‑up\n",
    "\n",
    "- Keep **each scientific component** testable as a `Callable`.\n",
    "- The **full‑stack model** chains them in one place (`grow`).\n",
    "- The **factory** adapts inventories (Angle‑Count → pseudo tree‑list) when you ask for per‑tree modes.\n",
    "- The **DSL** orchestrates **when** things happen; history is logged on every step.\n",
    "- For many aggregate scenarios, opt‑in to the **batch engine** by exposing `batch_grow_step`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyforestry",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}