SurEau solver

sureau_solver


def sureau_solver(
    state:SurEauPlantState, # Plant state. Mutated in place. Read/written: all ψ fields, PLC fields, psi_all_soil.
    fluxes:SurEauPlantFluxes, diag:SurEauPlantDiagnostics, soil:SurEauSoil, params:SurEauVegetationParams,
    dt:float, # Timestep size (seconds).
    opt:SurEauComputationOptions, # Solver options: numerical scheme, scaling flags.
)->SurEauPlantState: # Updated state with potentials at time n+1 and updated PLC.

Time-integrate one small timestep dt [seconds] for the four-compartment plant hydraulic system.

This is the core ODE solver of the SurEau-Ecos model. It solves the four coupled water balance equations (Eqs. 6–9) for ψ at time n+1 in all four plant compartments:

Eq. 6:  C_LApo × dψ_LApo/dt
        + K_SL(ψ_LApo − ψ_SApo)
        + K_LSym(ψ_LApo − ψ_LSym)
        − δ_L × K^cav_L × (ψ^cav_LApo − ψ_LApo) = 0

Eq. 7:  C_SApo × dψ_SApo/dt
        + K_SL(ψ_SApo − ψ_LApo)
        + K_SSym(ψ_SApo − ψ_SSym)
        + Σ K_soil(ψ_SApo − ψ_soil)
        − δ_S × K^cav_S × (ψ^cav_SApo − ψ_SApo) = 0

Eq. 8:  C_LSym × dψ_LSym/dt
        + K_LSym(ψ_LSym − ψ_LApo)
        + E_stom + E_cutiL = 0

Eq. 9:  C_SSym × dψ_SSym/dt
        + K_SSym(ψ_SSym − ψ_SApo)
        + E_cutiS = 0

Three numerical schemes are available (Section 2.5, Appendix C):

Implicit (default, Appendix C2, Eqs. C5–C18):

Eliminates the symplasm unknowns analytically, reducing the 4×4 system 
to a 2×2 in (ψ_LApo, ψ_SApo), then solves in closed form. 
Unconditionally stable. Uses the transpiration linearization (**Eq. 60**) 
to account for stomatal response within the timestep.

Semi-Implicit (Appendix C3, Eqs. C19–C22): Holds neighbor potentials at time n and solves each compartment’s ODE analytically as an exponential relaxation (Eqs. 54–58). Unconditionally stable but decoupled — does not capture simultaneous inter-compartment coupling within one timestep.

Explicit (Section 2.5.1, Eqs. 49–50): Forward Euler. Simple but CFL-limited (Eq. 51): dt must satisfy dt ≤ C/(2 × max(K)). For the apoplasm (tiny C, large K), this forces dt < ~10 ms.

A cavitation self-consistency loop (Block 4–6): cavitation flux depends on ψ^{n+1} (unknown), but ψ^{n+1} depends on cavitation flux. The solver tries multiple (δ_L, δ_S) configurations and checks which is self-consistent with the solution.

Physical analogy: Four interconnected water tanks with different “sponginess” (capacitance C). They’re connected by pipes (conductances K). Water is being sucked out of the leaf tanks (transpiration E). The solver figures out how the water levels (ψ) in all four tanks change over one small time interval.

Parameters:

state: SurEauPlantState object. Mutated in place. The following fields are read and written:

Read:
- psi_LApo, psi_SApo, psi_LSym, psi_SSym: Water potentials at time n (MPa).
- psi_LApo_cav, psi_SApo_cav: Historical minimum potentials (MPa). For cavitation flux (Eqs. 25–27).
- k_LSym, k_SSym, k_SLApo: Plant conductances (mmol m⁻² s⁻¹ MPa⁻¹). From update_kplant.
- C_LSym, C_SSym, C_LApo, C_SApo: Capacitances (mmol m⁻² MPa⁻¹). From update_capacitances.
- PLC_leaf, PLC_stem: Current percent loss of conductivity (%). For computing cavitation conductances (Eq. 26).
- Q_LApo_sat_mmol_per_LA, Q_SApo_sat_mmol_per_LA: Saturated apoplasm water (mmol/m² leaf). For Eqs. 25, 27.
- k_soil_to_stem: Soil-to-stem conductances per layer (mmol m⁻² s⁻¹ MPa⁻¹). From update_kplant(Eq. 20).
Written:
- psi_LApo, psi_SApo, psi_LSym, psi_SSym: Updated potentials at time n+1 (MPa). Clamped to ≤ −0.00001 MPa.
- psi_LApo_cav, psi_SApo_cav: Updated historical minima (only if new record lows are reached).
- PLC_leaf, PLC_stem: Updated PLC (only if new cavitation occurs). Irreversible — only increases.
- psi_all_soil: Conductance-weighted average soil potential (MPa). Diagnostic.
fluxes: SurEauPlantFluxes object. Read only:
- E_lim: Water-limited stomatal transpiration (mmol/m²/s). E^{n+1/2} in Eqs. 8, C17, C21.
- E_prime: dE/dψ (mmol/m²/s/MPa). For Eq. 60 transpiration linearization.
- E_min: Leaf cuticular transpiration (mmol/m²/s). Eqs. 8, C17, C21.
- E_min_S: Stem cuticular transpiration (mmol/m²/s). Eqs. 9, C18, C22.
diag: SurEauPlantDiagnostics object. Written:
- diag_nwhile_cavit: Number of iterations in the cavitation self-consistency loop (1–5). Useful for debugging.
soil: SurEauSoil object. Read only:
- psi_soil: Soil water potentials per layer (MPa). Array.
- k_soil: Not read directly (already folded into state.k_soil_to_stem by update_kplant).
params: SurEauVegetationParams object. Read only:
- slope_VC_leaf, slope_VC_stem: Vulnerability curve slopes (% MPa⁻¹). For compute_PLC_prime (Eq. 26).
- P50_VC_leaf, P50_VC_stem: P50 values (MPa). For compute_PLC (Eq. 15).
dt: Timestep size (seconds). Typically 60–3600 s for implicit/ semi-implicit, < 0.01 s for explicit.
opt: SurEauComputationOptions object. Controls the solver:
- numerical_scheme: "Implicit" (default), "Semi-Implicit", or "Explicit".
- Lsym, Ssym: Scaling flags for symplasm capacitances (default 1.0). Set to 0 to make symplasm infinitely stiff.
- CLapo, CTapo: Scaling flags for apoplasm capacitances.
- Lcav, Scav: Enable/disable cavitation flux (1 or 0).
- Eord: Enable/disable transpiration linearization (Eq. 60). Default 1. Set to 0 for first-order scheme.

Returns:

state: Updated SurEauPlantState with potentials at time n+1, and updated PLC if new cavitation occurred.

Overall solution strategy:

┌─────────────────────────────────────────────────────────┐
│  Step 1: Snapshot state at time n                       │
│  Save ψ^n for all 4 compartments + cavitation history   │
│  Extract conductances K, capacitances C                 │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 2: Read transpiration fluxes                      │  
│  E_lim (Eq. 33), E_prime (Eq. 60), E_min (Eqs. 31–32)   │
│  Frozen for this timestep (computed upstream)           │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 3: Cavitation conductances (Eqs. 25–27)           │
│  K^cav_L = −Q_sat_LApo × PLC'(ψ) / dt    (Eq. 25)       │
│  K^cav_S = −Q_sat_SApo × PLC'(ψ) / dt    (Eq. 27)       │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 4: Build cavitation guess list                    │
│  Both enabled:  (0,0),(1,0),(0,1),(1,1),(0,0)  5 tries  │
│  One enabled:   (0,0),(x,1),(0,0)  or similar  3 tries  │
│  Both disabled: (0,0)                          1 try    │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 5: While loop — try each (δ_L, δ_S) guess         │  
│                                                         │
│  ┌── Implicit ──────────────────────────────────────┐   │
│  │  1. Compute K̃_L, ψ̃_L, Ẽ_L (Eqs. C6–C8)           │   │
│  │  2. Compute K̃_S, ψ̃_S, Ẽ_S (Eqs. C11–C14)         │   │
│  │  3. Solve ψ_LApo^{n+1}       (Eq. C15)           │   │
│  │  4. Back-sub ψ_SApo^{n+1}    (Eq. C16)           │   │
│  └──────────────────────────────────────────────────┘   │
│  ┌── Semi-Implicit ─────────────────────────────────┐   │
│  │  Exponential relaxation: ψ^{n+1} = η·ψ^n+(1−η)·ψ̃ │   │
│  │  ψ_LApo^{n+1} (Eq. C19), ψ_SApo^{n+1} (Eq. C20)  │   │
│  └──────────────────────────────────────────────────┘   │
│  ┌── Explicit ──────────────────────────────────────┐   │
│  │  Forward Euler: ψ^{n+1} = ψ^n + (dt/C)·Σfluxes   │   │
│  │  ψ_LApo^{n+1} (Eq. 50→6), ψ_SApo^{n+1} (Eq. 50→7)│   │
│  └──────────────────────────────────────────────────┘   │
│                                                         │
│  Step 6: Check self-consistency                         │
│  L_ok = (δ_L == (ψ_LApo^{n+1} < ψ^cav_LApo))            │
│  S_ok = (δ_S == (ψ_SApo^{n+1} < ψ^cav_SApo))            │
│  If both ok → exit loop. Else → try next guess.         │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 7: Solve symplasm potentials                      │
│  Implicit:      ψ_LSym^{n+1} (Eq. C17)                  │
│                 ψ_SSym^{n+1} (Eq. C18)                  │
│  Semi-Implicit: ψ_LSym^{n+1} (Eq. C21)                  │
│                 ψ_SSym^{n+1} (Eq. C22)                  │
│  Explicit:      Forward Euler (Eqs. C3–C4)              │
└───────────────────────┬─────────────────────────────────┘
                        │
                        ▼
┌─────────────────────────────────────────────────────────┐
│  Step 8: Update state                                   │
│  Clamp all ψ ≤ −0.00001 MPa (always negative)           │
│  Update PLC if new minimum ψ reached (irreversible)     │
│  Compute weighted soil potential (diagnostic)           │
└─────────────────────────────────────────────────────────┘

Cavitation self-consistency loop detail:

Cavitation flux should only occur when ψ^{n+1} drops below the historical minimum ψ^cav (a new record-low potential). But ψ^{n+1} depends on whether cavitation flux is included. The solver handles this cavitation configuration problem by guessing δ (on/off), solving, and checking:

δ=0, solution ψ ≥ ψ^cav: consistent (assumed off, confirmed off)
δ=0, solution ψ < ψ^cav: inconsistent (assumed off, but should be on)
δ=1, solution ψ < ψ^cav: consistent (assumed on, confirmed on)
δ=1, solution ψ ≥ ψ^cav: inconsistent (assumed on, but the released water “cushioned” the drop — shouldn’t be on)

The search space has 2 × 2 = 4 combinations for (δ_L, δ_S), plus one fallback. When one mechanism is disabled (opt.Lcav=0 or opt.Scav=0), only the other dimension is searched (3 tries). When both are disabled, no search is needed (1 try).

Implicit scheme algebraic structure (Eqs. C5–C18):

The 4×4 system is reduced by eliminating ψ_LSym and ψ_SSym:

Step 1: Substitute Eq. 8 into Eq. 6 → "combined leaf eq." (C5)
        Define effective K̃_L, ψ̃_L, Ẽ_L (Eqs. C6–C8)
        → Simplified: K̃_L(ψ_LApo − ψ̃_L) + K_SL(ψ_LApo − ψ_SApo)
                      + Ẽ_L = 0    (Eq. C9)

Step 2: Substitute Eq. 9 into Eq. 7 → "combined stem eq." (C10)
        Define effective K̃_S, ψ̃_S (Eqs. C11–C12)
        → Simplified: K̃_S(ψ_SApo − ψ̃_S) + K_SL(ψ_SApo − ψ_LApo)
                      + Ẽ_S_raw = 0    (Eq. C13)
        where Ẽ_S_raw = Emin_S / (1 + C_SSym/(K_SSym×dt))

Step 3: Substitute C13 into C9 to get a single equation in ψ_LApo
        → Eq. C15: closed-form ψ_LApo^{n+1}
        The substitution produces a factor K_SL/(K_SL + K̃_S)
        on the stem transpiration, which the paper defines as
        Ẽ_S (Eq. C14) — the reduced form that appears in C15.

Step 4: Back-substitute into C9 for ψ_SApo^{n+1} (Eq. C16)
Step 5: Back-substitute for ψ_LSym^{n+1} (Eq. C17)
Step 6: Back-substitute for ψ_SSym^{n+1} (Eq. C18)

Result: 4 analytical expressions, no matrix inversion needed.

The opt.* scaling flags:

- opt.Lsym  → multiplies C_LSym. Set to 0: leaf symplasm becomes 
infinitely stiff (ψ responds instantly).

- opt.Ssym  → multiplies C_SSym.  Same for stem symplasm.
- opt.CLapo → multiplies C_LApo.  Same for leaf apoplasm.
- opt.CTapo → multiplies C_SApo.  Same for stem apoplasm.
- opt.Lcav  → enables leaf cavitation flux (0 or 1).
- opt.Scav  → enables stem cavitation flux (0 or 1).
- opt.Eord  → enables transpiration linearization (Eq. 60).
            Set to 0: first-order scheme (no dE/dψ correction).

All default to 1 in normal operation.

Equation-to-code mapping:

Code variable          Paper symbol       Equation
─────────────────────────────────────────────────────
K_L_Cav                K^cav_L            Eq. 25
K_S_Cav                K^cav_S            Eq. 27
PLC_prime_L            PLC'_L             Eq. 26
K_L_td                 K̃_L                Eq. C6
Psi_L_td               ψ̃_L                Eq. C7
E_L_td                 Ẽ_L                Eq. C8
K_S_td                 K̃_S                Eq. C11
Psi_S_td               ψ̃_S                Eq. C12
E_S_td                 Ẽ_S (reduced)      Eq. C14
psi_LApo_np1           ψ^{n+1}_LApo       Eq. C15
psi_SApo_np1           ψ^{n+1}_SApo       Eq. C16
psi_LSym_np1           ψ^{n+1}_LSym       Eq. C17
psi_SSym_np1           ψ^{n+1}_SSym       Eq. C18
alpha (semi-imp.)      η                  Eq. 57
Psi_td (semi-imp.)     ψ̃                  Eq. 58

state = SurEauPlantState(
    k_LSym=0.1,
    k_SSym=0.1,
    k_SLApo=0.1,
    C_LSym=0.01,
    C_SSym=0.01,
    C_LApo=0.01,
    C_SApo=0.01,
    psi_LApo=-0.5,
    psi_SApo=-0.3,
    psi_LSym=-0.8,
    psi_SSym=-0.5,
    psi_LApo_cav=-0.5,
    psi_SApo_cav=-0.3,
    k_soil_to_stem=np.array([0.05, 0.05, 0.05]),
    Q_LApo_sat_mmol_per_LA=1.0,
    Q_SApo_sat_mmol_per_LA=1.0,
)

sureau_solver(
    state=state,
    fluxes=SurEauPlantFluxes(),
    diag=SurEauPlantDiagnostics(),
    soil=SurEauSoil(
        psi_soil=np.array([-0.5, -0.5, -0.5]),
        k_soil=np.array([0.1, 0.1, 0.1]),
    ),
    params=SurEauVegetationParams(),
    dt=0.01,
    opt=SurEauComputationOptions(),
)

SurEauPlantState(psi_LApo=np.float64(-0.5104418925604771), psi_SApo=np.float64(-0.35162617503840865), psi_LSym=np.float64(-0.7736765356873161), psi_SSym=np.float64(-0.4865114704580371), psi_LApo_cav=np.float64(-0.5104418925604771), psi_SApo_cav=np.float64(-0.35162617503840865), psi_all_soil=np.float64(-0.5), k_plant=0.0, k_LSym=0.1, k_SSym=0.1, k_RSApo=None, k_SLApo=0.1, k_soil_to_stem=array([0.05, 0.05, 0.05]), C_LSym=0.01, C_SSym=0.01, C_LApo=0.01, C_SApo=0.01, PLC_leaf=np.float64(0.09722358230675984), PLC_stem=np.float64(0.06643099597079465), LAI_pheno=0.0, LAI=0.0, canopy_storage_capacity=0.0, FCC=0.0, LAI_dead=0.0, sum_temperature=0.0, budburst_date=nan, Q_LApo_sat_mmol=0.0, Q_LApo_sat_L=0.0, Q_SApo_sat_mmol=0.0, Q_SApo_sat_L=0.0, Q_LSym_sat_mmol=0.0, Q_LSym_sat_L=0.0, Q_SSym_sat_mmol=0.0, Q_SSym_sat_L=0.0, Q_LApo_sat_mmol_per_LA=1.0, Q_SApo_sat_mmol_per_LA=1.0, Q_LSym_sat_mmol_per_LA=0.0, Q_SSym_sat_mmol_per_LA=0.0, Q_LApo_L=0.0, Q_SApo_L=0.0, Q_LSym_L=0.0, Q_SSym_L=0.0, DM_live_canopy=0.0, DM_dead_canopy=0.0)