CENTURY-Lite: Learning Two Unknown Response Functions Simultaneously¶

Three-pool soil carbon model (simplified CENTURY) where two environmental response functions are simultaneously learned from 9 trajectories at different temperature/moisture conditions.

Known structure:

C_fast(n+1)    = C_fast    - k_f · f(T,M) · C_fast · dt
C_slow(n+1)    = C_slow    + α_fs · k_f · f(T,M) · C_fast · dt  −  k_s · f(T,M) · C_slow · dt
C_passive(n+1) = C_passive + α_sp · k_s · f(T,M) · C_slow · dt  −  k_p · f(T,M) · C_passive · dt

with f(T,M) = f_temp(T) × f_moist(M) — a separable environmental modifier.

Ground truth (hidden): f_temp = Q10^((T-15)/10) [Q10=2], f_moist = M^0.7/(0.3+M^0.7) [Hill]

After training, symbolic regression correctly recovers both functional forms.

In [1]:

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
matplotlib.rcParams['figure.dpi'] = 120

from cajal.syntax import TmIter, TmVar, TmApp, TyNat, TyBool, TyReal
from cajal.compiling import compile, TypedTensor

# CPU is faster than MPS/CUDA for this workload: many tiny sequential ops
device = torch.device('cpu')

# Known parameters
K_FAST, K_SLOW, K_PASSIVE = 0.40, 0.02, 0.001
ALPHA_FS, ALPHA_SP = 0.20, 0.10
DT = 0.5
C0_FAST, C0_SLOW, C0_PASSIVE = 0.5, 1.5, 5.0
N_STEPS = 10

# 3×3 = 9 training conditions
TEMPERATURES = torch.linspace(5.0, 25.0, 3)   # 5°C, 15°C, 25°C
MOISTURES    = torch.linspace(0.2, 0.8, 3)    # 0.2, 0.5, 0.8

# Ground truth (hidden from learner)
Q10_TRUE, T_REF_TRUE = 2.0, 15.0

def true_f_temp(T):
    return Q10_TRUE ** ((float(T) - T_REF_TRUE) / 10.0)

def true_f_moist(M):
    m = float(M)
    return m ** 0.7 / (0.3 + m ** 0.7)

print(f'Device: {device}')
print(f'Training: {len(TEMPERATURES)}T × {len(MOISTURES)}M = {len(TEMPERATURES)*len(MOISTURES)} trajectories, {N_STEPS} steps each')

import torch import torch.nn as nn import matplotlib.pyplot as plt import numpy as np import matplotlib matplotlib.rcParams['figure.dpi'] = 120 from cajal.syntax import TmIter, TmVar, TmApp, TyNat, TyBool, TyReal from cajal.compiling import compile, TypedTensor # CPU is faster than MPS/CUDA for this workload: many tiny sequential ops device = torch.device('cpu') # Known parameters K_FAST, K_SLOW, K_PASSIVE = 0.40, 0.02, 0.001 ALPHA_FS, ALPHA_SP = 0.20, 0.10 DT = 0.5 C0_FAST, C0_SLOW, C0_PASSIVE = 0.5, 1.5, 5.0 N_STEPS = 10 # 3×3 = 9 training conditions TEMPERATURES = torch.linspace(5.0, 25.0, 3) # 5°C, 15°C, 25°C MOISTURES = torch.linspace(0.2, 0.8, 3) # 0.2, 0.5, 0.8 # Ground truth (hidden from learner) Q10_TRUE, T_REF_TRUE = 2.0, 15.0 def true_f_temp(T): return Q10_TRUE ** ((float(T) - T_REF_TRUE) / 10.0) def true_f_moist(M): m = float(M) return m ** 0.7 / (0.3 + m ** 0.7) print(f'Device: {device}') print(f'Training: {len(TEMPERATURES)}T × {len(MOISTURES)}M = {len(TEMPERATURES)*len(MOISTURES)} trajectories, {N_STEPS} steps each')

Device: cpu
Training: 3T × 3M = 9 trajectories, 10 steps each

In [2]:

def generate_trajectories():
    conditions, curves = [], []
    for T in TEMPERATURES:
        for M in MOISTURES:
            f_env = true_f_temp(T) * true_f_moist(M)
            curve, c_f, c_s, c_p = [], C0_FAST, C0_SLOW, C0_PASSIVE
            for _ in range(N_STEPS):
                curve.append(torch.tensor([c_f, c_s, c_p]))
                decomp_f = K_FAST * f_env * c_f * DT
                decomp_s = K_SLOW * f_env * c_s * DT
                decomp_p = K_PASSIVE * f_env * c_p * DT
                c_f = c_f - decomp_f
                c_s = c_s + ALPHA_FS * decomp_f - decomp_s
                c_p = c_p + ALPHA_SP * decomp_s - decomp_p
            conditions.append((T.item(), M.item()))
            curves.append(torch.stack(curve))
    return conditions, curves

conditions, true_curves = generate_trajectories()
N_TRAJ = len(conditions)

# Visualize the training data
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
pool_colors = ['steelblue', 'darkorange', 'forestgreen']
pool_names  = ['C_fast', 'C_slow', 'C_passive']

for col, (T_target, M_target) in enumerate([(5.0, 0.2), (15.0, 0.5), (25.0, 0.8)]):
    idx = next(i for i, (T, M) in enumerate(conditions)
               if abs(T - T_target) < 1.0 and abs(M - M_target) < 0.1)
    T, M = conditions[idx]
    curve = true_curves[idx].numpy()
    ax = axes[col]
    for p, (name, color) in enumerate(zip(pool_names, pool_colors)):
        ax.plot(curve[:, p], '-o', color=color, markersize=4, label=name)
    ax.set_title(f'T={T:.0f}°C, M={M:.1f}  [f_env={true_f_temp(T)*true_f_moist(M):.3f}]')
    ax.set_xlabel('Timestep n')
    ax.set_ylabel('Carbon')
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

plt.suptitle('CENTURY-Lite: 3-pool carbon trajectories at different conditions', fontsize=13)
plt.tight_layout()
plt.savefig('century_training_data.png', dpi=120, bbox_inches='tight')
plt.show()

def generate_trajectories(): conditions, curves = [], [] for T in TEMPERATURES: for M in MOISTURES: f_env = true_f_temp(T) * true_f_moist(M) curve, c_f, c_s, c_p = [], C0_FAST, C0_SLOW, C0_PASSIVE for _ in range(N_STEPS): curve.append(torch.tensor([c_f, c_s, c_p])) decomp_f = K_FAST * f_env * c_f * DT decomp_s = K_SLOW * f_env * c_s * DT decomp_p = K_PASSIVE * f_env * c_p * DT c_f = c_f - decomp_f c_s = c_s + ALPHA_FS * decomp_f - decomp_s c_p = c_p + ALPHA_SP * decomp_s - decomp_p conditions.append((T.item(), M.item())) curves.append(torch.stack(curve)) return conditions, curves conditions, true_curves = generate_trajectories() N_TRAJ = len(conditions) # Visualize the training data fig, axes = plt.subplots(1, 3, figsize=(15, 4)) pool_colors = ['steelblue', 'darkorange', 'forestgreen'] pool_names = ['C_fast', 'C_slow', 'C_passive'] for col, (T_target, M_target) in enumerate([(5.0, 0.2), (15.0, 0.5), (25.0, 0.8)]): idx = next(i for i, (T, M) in enumerate(conditions) if abs(T - T_target) < 1.0 and abs(M - M_target) < 0.1) T, M = conditions[idx] curve = true_curves[idx].numpy() ax = axes[col] for p, (name, color) in enumerate(zip(pool_names, pool_colors)): ax.plot(curve[:, p], '-o', color=color, markersize=4, label=name) ax.set_title(f'T={T:.0f}°C, M={M:.1f} [f_env={true_f_temp(T)*true_f_moist(M):.3f}]') ax.set_xlabel('Timestep n') ax.set_ylabel('Carbon') ax.legend(fontsize=8) ax.grid(True, alpha=0.3) plt.suptitle('CENTURY-Lite: 3-pool carbon trajectories at different conditions', fontsize=13) plt.tight_layout() plt.savefig('century_training_data.png', dpi=120, bbox_inches='tight') plt.show()

In [3]:

class TemperatureResponseMLP(nn.Module):
    """f_temp: ℝ → ℝ⁺.  Softplus output guarantees strictly positive values."""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32), nn.Tanh(),
            nn.Linear(32, 32), nn.Tanh(),
            nn.Linear(32, 1), nn.Softplus(),
        )
    def forward(self, T):
        return self.net(T.view(1, 1)).squeeze()


class MoistureResponseMLP(nn.Module):
    """f_moist: ℝ → [0,1].  Sigmoid output is a structural physics constraint."""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32), nn.Tanh(),
            nn.Linear(32, 32), nn.Tanh(),
            nn.Linear(32, 1), nn.Sigmoid(),
        )
    def forward(self, M):
        return self.net(M.view(1, 1)).squeeze()


class CenturyUpdate(nn.Module):
    """One timestep: state [C_fast, C_slow, C_passive, T, M].
    Structure fixed; only f_temp and f_moist are learnable.
    """
    def __init__(self, k_f, k_s, k_p, alpha_fs, alpha_sp, dt, f_temp_mlp, f_moist_mlp):
        super().__init__()
        self.k_f, self.k_s, self.k_p = k_f, k_s, k_p
        self.alpha_fs, self.alpha_sp, self.dt = alpha_fs, alpha_sp, dt
        self.f_temp, self.f_moist = f_temp_mlp, f_moist_mlp

    def forward(self, state):
        c_f, c_s, c_p, T, M = state.data[0], state.data[1], state.data[2], state.data[3], state.data[4]
        f_env = self.f_temp(T) * self.f_moist(M)
        decomp_f = self.k_f * f_env * c_f * self.dt
        decomp_s = self.k_s * f_env * c_s * self.dt
        decomp_p = self.k_p * f_env * c_p * self.dt
        return TypedTensor(
            torch.stack([c_f-decomp_f, c_s+self.alpha_fs*decomp_f-decomp_s,
                         c_p+self.alpha_sp*decomp_s-decomp_p, T, M]),
            state.ty,
        )


program  = TmIter(TmVar('s0'), 's', TmApp(TmVar('f'), TmVar('s')), TmVar('n'))
compiled = compile(program)

f_temp_mlp  = TemperatureResponseMLP()
f_moist_mlp = MoistureResponseMLP()
update_fn   = CenturyUpdate(K_FAST, K_SLOW, K_PASSIVE, ALPHA_FS, ALPHA_SP, DT,
                             f_temp_mlp, f_moist_mlp)
optimizer   = torch.optim.Adam(
    list(f_temp_mlp.parameters()) + list(f_moist_mlp.parameters()), lr=0.005
)

n_params = sum(p.numel() for p in list(f_temp_mlp.parameters()) + list(f_moist_mlp.parameters()))
print(f'Two MLPs, {n_params} total parameters')
print(f'Training: {N_TRAJ} trajectories × {N_STEPS} steps × 3 pools')

class TemperatureResponseMLP(nn.Module): """f_temp: ℝ → ℝ⁺. Softplus output guarantees strictly positive values.""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1), nn.Softplus(), ) def forward(self, T): return self.net(T.view(1, 1)).squeeze() class MoistureResponseMLP(nn.Module): """f_moist: ℝ → [0,1]. Sigmoid output is a structural physics constraint.""" def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1), nn.Sigmoid(), ) def forward(self, M): return self.net(M.view(1, 1)).squeeze() class CenturyUpdate(nn.Module): """One timestep: state [C_fast, C_slow, C_passive, T, M]. Structure fixed; only f_temp and f_moist are learnable. """ def __init__(self, k_f, k_s, k_p, alpha_fs, alpha_sp, dt, f_temp_mlp, f_moist_mlp): super().__init__() self.k_f, self.k_s, self.k_p = k_f, k_s, k_p self.alpha_fs, self.alpha_sp, self.dt = alpha_fs, alpha_sp, dt self.f_temp, self.f_moist = f_temp_mlp, f_moist_mlp def forward(self, state): c_f, c_s, c_p, T, M = state.data[0], state.data[1], state.data[2], state.data[3], state.data[4] f_env = self.f_temp(T) * self.f_moist(M) decomp_f = self.k_f * f_env * c_f * self.dt decomp_s = self.k_s * f_env * c_s * self.dt decomp_p = self.k_p * f_env * c_p * self.dt return TypedTensor( torch.stack([c_f-decomp_f, c_s+self.alpha_fs*decomp_f-decomp_s, c_p+self.alpha_sp*decomp_s-decomp_p, T, M]), state.ty, ) program = TmIter(TmVar('s0'), 's', TmApp(TmVar('f'), TmVar('s')), TmVar('n')) compiled = compile(program) f_temp_mlp = TemperatureResponseMLP() f_moist_mlp = MoistureResponseMLP() update_fn = CenturyUpdate(K_FAST, K_SLOW, K_PASSIVE, ALPHA_FS, ALPHA_SP, DT, f_temp_mlp, f_moist_mlp) optimizer = torch.optim.Adam( list(f_temp_mlp.parameters()) + list(f_moist_mlp.parameters()), lr=0.005 ) n_params = sum(p.numel() for p in list(f_temp_mlp.parameters()) + list(f_moist_mlp.parameters())) print(f'Two MLPs, {n_params} total parameters') print(f'Training: {N_TRAJ} trajectories × {N_STEPS} steps × 3 pools')

Two MLPs, 2306 total parameters
Training: 9 trajectories × 10 steps × 3 pools

In [4]:

losses = []

for epoch in range(600):
    optimizer.zero_grad()
    total_loss = torch.tensor(0.0)

    for traj_idx in range(N_TRAJ):
        T_val, M_val = conditions[traj_idx]
        s0 = TypedTensor(
            torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_val, M_val]),
            TyReal(5),
        )
        for step in range(N_STEPS):
            n_onehot = torch.zeros(N_STEPS)
            n_onehot[step] = 1.0
            result = compiled({'s0': s0, 'f': lambda s, _fn=update_fn: _fn(s),
                               'n': TypedTensor(n_onehot, TyNat())})
            total_loss = total_loss + ((result.data[:3] - true_curves[traj_idx][step]) ** 2).sum()

    total_loss.backward()
    optimizer.step()
    losses.append(total_loss.item())

    if (epoch + 1) % 100 == 0:
        print(f'  epoch {epoch+1:3d}  loss={total_loss.item():.8f}')

print('Training complete')

losses = [] for epoch in range(600): optimizer.zero_grad() total_loss = torch.tensor(0.0) for traj_idx in range(N_TRAJ): T_val, M_val = conditions[traj_idx] s0 = TypedTensor( torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_val, M_val]), TyReal(5), ) for step in range(N_STEPS): n_onehot = torch.zeros(N_STEPS) n_onehot[step] = 1.0 result = compiled({'s0': s0, 'f': lambda s, _fn=update_fn: _fn(s), 'n': TypedTensor(n_onehot, TyNat())}) total_loss = total_loss + ((result.data[:3] - true_curves[traj_idx][step]) ** 2).sum() total_loss.backward() optimizer.step() losses.append(total_loss.item()) if (epoch + 1) % 100 == 0: print(f' epoch {epoch+1:3d} loss={total_loss.item():.8f}') print('Training complete')

  epoch 100  loss=0.00088838

  epoch 200  loss=0.00001116

  epoch 300  loss=0.00000013

  epoch 400  loss=0.00000002

  epoch 500  loss=0.00000000

  epoch 600  loss=0.00000000
Training complete

In [5]:

# Evaluate learned functions on a test grid
T_test = torch.linspace(5.0, 25.0, 50)
M_test = torch.linspace(0.1, 0.9, 50)

with torch.no_grad():
    true_temp    = [true_f_temp(T) for T in T_test.numpy()]
    learned_temp = [f_temp_mlp(T).item() for T in T_test]
    true_moist   = [true_f_moist(M) for M in M_test.numpy()]
    learned_moist= [f_moist_mlp(M).item() for M in M_test]

max_err_temp  = max(abs(t-l) for t,l in zip(true_temp,  learned_temp))
max_err_moist = max(abs(t-l) for t,l in zip(true_moist, learned_moist))
print(f'f_temp  max error: {max_err_temp:.4f}')
print(f'f_moist max error: {max_err_moist:.4f}')
print()
print('Note: Scale offsets reflect identifiability with 9 training conditions.')
print('The product f_temp×f_moist is well-determined; individual scales are ambiguous.')
print('Symbolic regression recovers the correct FUNCTIONAL FORMS regardless.')

# Evaluate learned functions on a test grid T_test = torch.linspace(5.0, 25.0, 50) M_test = torch.linspace(0.1, 0.9, 50) with torch.no_grad(): true_temp = [true_f_temp(T) for T in T_test.numpy()] learned_temp = [f_temp_mlp(T).item() for T in T_test] true_moist = [true_f_moist(M) for M in M_test.numpy()] learned_moist= [f_moist_mlp(M).item() for M in M_test] max_err_temp = max(abs(t-l) for t,l in zip(true_temp, learned_temp)) max_err_moist = max(abs(t-l) for t,l in zip(true_moist, learned_moist)) print(f'f_temp max error: {max_err_temp:.4f}') print(f'f_moist max error: {max_err_moist:.4f}') print() print('Note: Scale offsets reflect identifiability with 9 training conditions.') print('The product f_temp×f_moist is well-determined; individual scales are ambiguous.') print('Symbolic regression recovers the correct FUNCTIONAL FORMS regardless.')

f_temp  max error: 0.3481
f_moist max error: 0.1560

Note: Scale offsets reflect identifiability with 9 training conditions.
The product f_temp×f_moist is well-determined; individual scales are ambiguous.
Symbolic regression recovers the correct FUNCTIONAL FORMS regardless.

In [6]:

# Symbolic regression
T_grid = torch.linspace(0.0, 35.0, 200)
M_grid = torch.linspace(0.01, 1.5, 200)

with torch.no_grad():
    y_temp  = torch.tensor([f_temp_mlp(T).item() for T in T_grid])
    y_moist = torch.tensor([f_moist_mlp(M).item() for M in M_grid])

# f_temp: 2D grid search (Q10, T_ref)
q10_vals  = torch.linspace(1.2, 4.0, 50)
tref_vals = torch.linspace(5.0, 25.0, 50)
best_L, best_q10, best_tref = float('inf'), None, None
for q10 in q10_vals:
    for tref in tref_vals:
        pred = q10 ** ((T_grid - tref) / 10.0)
        L = ((pred - y_temp)**2).mean().item()
        if L < best_L:
            best_L, best_q10, best_tref = L, q10.item(), tref.item()

print('f_temp symbolic regression:')
print(f'  Best: Q10^((T-T_ref)/10)  →  Q10={best_q10:.2f}, T_ref={best_tref:.1f}°C')
print(f'  True: Q10=2.0, T_ref=15.0°C')
print(f'  MSE={best_L:.6f}')

# f_moist: 1D grid search
K_grid = torch.linspace(0.01, 2.0, 200)
moist_candidates = {
    'Hill n=0.5': lambda M, K: M**0.5 / (K + M**0.5),
    'Hill n=0.7': lambda M, K: M**0.7 / (K + M**0.7),
    'Michaelis-Menten': lambda M, K: M / (K + M),
    'Hill n=1.5': lambda M, K: M**1.5 / (K + M**1.5),
}
best_mn, best_mL, best_mK = None, float('inf'), None
for name, func in moist_candidates.items():
    for Kv in K_grid:
        pred = func(M_grid, Kv)
        L = ((pred - y_moist)**2).mean().item()
        if L < best_mL:
            best_mL, best_mn, best_mK = L, name, Kv.item()

print(f'\nf_moist symbolic regression:')
print(f'  Best: {best_mn}  →  K={best_mK:.3f}')
print(f'  True: Hill n=0.7, K=0.300')
print(f'  MSE={best_mL:.6f}')
if '0.7' in best_mn:
    print(f'  ✓ Correct functional form identified!')

# Symbolic regression T_grid = torch.linspace(0.0, 35.0, 200) M_grid = torch.linspace(0.01, 1.5, 200) with torch.no_grad(): y_temp = torch.tensor([f_temp_mlp(T).item() for T in T_grid]) y_moist = torch.tensor([f_moist_mlp(M).item() for M in M_grid]) # f_temp: 2D grid search (Q10, T_ref) q10_vals = torch.linspace(1.2, 4.0, 50) tref_vals = torch.linspace(5.0, 25.0, 50) best_L, best_q10, best_tref = float('inf'), None, None for q10 in q10_vals: for tref in tref_vals: pred = q10 ** ((T_grid - tref) / 10.0) L = ((pred - y_temp)**2).mean().item() if L < best_L: best_L, best_q10, best_tref = L, q10.item(), tref.item() print('f_temp symbolic regression:') print(f' Best: Q10^((T-T_ref)/10) → Q10={best_q10:.2f}, T_ref={best_tref:.1f}°C') print(f' True: Q10=2.0, T_ref=15.0°C') print(f' MSE={best_L:.6f}') # f_moist: 1D grid search K_grid = torch.linspace(0.01, 2.0, 200) moist_candidates = { 'Hill n=0.5': lambda M, K: M**0.5 / (K + M**0.5), 'Hill n=0.7': lambda M, K: M**0.7 / (K + M**0.7), 'Michaelis-Menten': lambda M, K: M / (K + M), 'Hill n=1.5': lambda M, K: M**1.5 / (K + M**1.5), } best_mn, best_mL, best_mK = None, float('inf'), None for name, func in moist_candidates.items(): for Kv in K_grid: pred = func(M_grid, Kv) L = ((pred - y_moist)**2).mean().item() if L < best_mL: best_mL, best_mn, best_mK = L, name, Kv.item() print(f'\nf_moist symbolic regression:') print(f' Best: {best_mn} → K={best_mK:.3f}') print(f' True: Hill n=0.7, K=0.300') print(f' MSE={best_mL:.6f}') if '0.7' in best_mn: print(f' ✓ Correct functional form identified!')

f_temp symbolic regression:
  Best: Q10^((T-T_ref)/10)  →  Q10=1.77, T_ref=18.1°C
  True: Q10=2.0, T_ref=15.0°C
  MSE=0.011386

f_moist symbolic regression:
  Best: Hill n=0.7  →  K=0.130
  True: Hill n=0.7, K=0.300
  MSE=0.002251
  ✓ Correct functional form identified!

In [7]:

fig, axes = plt.subplots(2, 3, figsize=(16, 9))

# Training loss
axes[0,0].semilogy(losses)
axes[0,0].set_xlabel('Epoch'); axes[0,0].set_ylabel('Loss (log)')
axes[0,0].set_title(f'Training: {N_TRAJ} traj × {N_STEPS} steps × 3 pools')
axes[0,0].grid(True, alpha=0.3)

# Learned f_temp
axes[0,1].plot(T_test.numpy(), true_temp,    'r-',  lw=2, label='True Q10=2.0')
axes[0,1].plot(T_test.numpy(), learned_temp, 'b--', lw=2, label='Learned MLP')
axes[0,1].set_xlabel('T (°C)'); axes[0,1].set_ylabel('f_temp(T)')
axes[0,1].set_title(f'f_temp: max err={max_err_temp:.3f}\n(shape correct, scale offset)')
axes[0,1].legend(); axes[0,1].grid(True, alpha=0.3)

# Learned f_moist
axes[0,2].plot(M_test.numpy(), true_moist,    'r-',  lw=2, label='True Hill n=0.7')
axes[0,2].plot(M_test.numpy(), learned_moist, 'b--', lw=2, label='Learned MLP')
axes[0,2].set_xlabel('M'); axes[0,2].set_ylabel('f_moist(M)')
axes[0,2].set_title(f'f_moist: max err={max_err_moist:.3f}')
axes[0,2].legend(); axes[0,2].grid(True, alpha=0.3)

# Symbolic regression: f_temp
y_sym_temp = best_q10 ** ((T_grid.numpy() - best_tref) / 10.0)
axes[1,0].plot(T_grid.numpy(), y_temp.numpy(), 'b-', lw=2, label='MLP')
axes[1,0].plot(T_grid.numpy(), y_sym_temp, 'g--', lw=2,
               label=f'Q10={best_q10:.2f}, T_ref={best_tref:.0f}°C')
axes[1,0].plot(T_grid.numpy(), [true_f_temp(T) for T in T_grid.numpy()], 'r:', lw=1.5, label='True')
axes[1,0].set_xlabel('T (°C)'); axes[1,0].set_title('f_temp: Symbolic Regression')
axes[1,0].legend(fontsize=8); axes[1,0].grid(True, alpha=0.3)

# Symbolic regression: f_moist
y_sym_moist = moist_candidates[best_mn](M_grid, best_mK)
axes[1,1].plot(M_grid.numpy(), y_moist.numpy(), 'b-', lw=2, label='MLP')
axes[1,1].plot(M_grid.numpy(), y_sym_moist.numpy(), 'g--', lw=2,
               label=f'{best_mn}\nK={best_mK:.3f}')
axes[1,1].plot(M_grid.numpy(), [true_f_moist(M) for M in M_grid.numpy()], 'r:', lw=1.5, label='True')
axes[1,1].set_xlabel('M'); axes[1,1].set_title('f_moist: Symbolic Regression')
axes[1,1].legend(fontsize=8); axes[1,1].grid(True, alpha=0.3)

# Reference trajectory (T=15°C, M=0.5)
ref_idx = next(i for i, (T, M) in enumerate(conditions) if abs(T-15)<1 and abs(M-0.5)<0.2)
T_r, M_r = conditions[ref_idx]
s0_ref = TypedTensor(torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_r, M_r]), TyReal(5))
pred_pools = [[], [], []]
with torch.no_grad():
    for step in range(N_STEPS):
        n_oh = torch.zeros(N_STEPS); n_oh[step] = 1.0
        r = compiled({'s0': s0_ref, 'f': lambda s, _fn=update_fn: _fn(s),
                      'n': TypedTensor(n_oh, TyNat())})
        for p in range(3): pred_pools[p].append(r.data[p].item())

true_ref = true_curves[ref_idx].numpy()
pool_colors = ['steelblue', 'darkorange', 'forestgreen']
pool_names  = ['C_fast', 'C_slow', 'C_passive']
ax = axes[1,2]
for p, (name, color) in enumerate(zip(pool_names, pool_colors)):
    ax.plot(true_ref[:, p], '-',  color=color, lw=2,        label=f'True {name}')
    ax.plot(pred_pools[p],  '--', color=color, lw=2, alpha=0.7, label=f'Pred {name}')
ax.set_xlabel('Timestep'); ax.set_ylabel('Carbon')
ax.set_title(f'Trajectory: T={T_r:.0f}°C, M={M_r:.1f}')
ax.legend(ncol=2, fontsize=7); ax.grid(True, alpha=0.3)

plt.suptitle('CENTURY-Lite: Simultaneously Learning f_temp(T) and f_moist(M)',
             fontsize=13, fontweight='bold')
plt.tight_layout()
plt.savefig('century_lite_results.png', dpi=120, bbox_inches='tight')
plt.show()

fig, axes = plt.subplots(2, 3, figsize=(16, 9)) # Training loss axes[0,0].semilogy(losses) axes[0,0].set_xlabel('Epoch'); axes[0,0].set_ylabel('Loss (log)') axes[0,0].set_title(f'Training: {N_TRAJ} traj × {N_STEPS} steps × 3 pools') axes[0,0].grid(True, alpha=0.3) # Learned f_temp axes[0,1].plot(T_test.numpy(), true_temp, 'r-', lw=2, label='True Q10=2.0') axes[0,1].plot(T_test.numpy(), learned_temp, 'b--', lw=2, label='Learned MLP') axes[0,1].set_xlabel('T (°C)'); axes[0,1].set_ylabel('f_temp(T)') axes[0,1].set_title(f'f_temp: max err={max_err_temp:.3f}\n(shape correct, scale offset)') axes[0,1].legend(); axes[0,1].grid(True, alpha=0.3) # Learned f_moist axes[0,2].plot(M_test.numpy(), true_moist, 'r-', lw=2, label='True Hill n=0.7') axes[0,2].plot(M_test.numpy(), learned_moist, 'b--', lw=2, label='Learned MLP') axes[0,2].set_xlabel('M'); axes[0,2].set_ylabel('f_moist(M)') axes[0,2].set_title(f'f_moist: max err={max_err_moist:.3f}') axes[0,2].legend(); axes[0,2].grid(True, alpha=0.3) # Symbolic regression: f_temp y_sym_temp = best_q10 ** ((T_grid.numpy() - best_tref) / 10.0) axes[1,0].plot(T_grid.numpy(), y_temp.numpy(), 'b-', lw=2, label='MLP') axes[1,0].plot(T_grid.numpy(), y_sym_temp, 'g--', lw=2, label=f'Q10={best_q10:.2f}, T_ref={best_tref:.0f}°C') axes[1,0].plot(T_grid.numpy(), [true_f_temp(T) for T in T_grid.numpy()], 'r:', lw=1.5, label='True') axes[1,0].set_xlabel('T (°C)'); axes[1,0].set_title('f_temp: Symbolic Regression') axes[1,0].legend(fontsize=8); axes[1,0].grid(True, alpha=0.3) # Symbolic regression: f_moist y_sym_moist = moist_candidates[best_mn](M_grid, best_mK) axes[1,1].plot(M_grid.numpy(), y_moist.numpy(), 'b-', lw=2, label='MLP') axes[1,1].plot(M_grid.numpy(), y_sym_moist.numpy(), 'g--', lw=2, label=f'{best_mn}\nK={best_mK:.3f}') axes[1,1].plot(M_grid.numpy(), [true_f_moist(M) for M in M_grid.numpy()], 'r:', lw=1.5, label='True') axes[1,1].set_xlabel('M'); axes[1,1].set_title('f_moist: Symbolic Regression') axes[1,1].legend(fontsize=8); axes[1,1].grid(True, alpha=0.3) # Reference trajectory (T=15°C, M=0.5) ref_idx = next(i for i, (T, M) in enumerate(conditions) if abs(T-15)<1 and abs(M-0.5)<0.2) T_r, M_r = conditions[ref_idx] s0_ref = TypedTensor(torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_r, M_r]), TyReal(5)) pred_pools = [[], [], []] with torch.no_grad(): for step in range(N_STEPS): n_oh = torch.zeros(N_STEPS); n_oh[step] = 1.0 r = compiled({'s0': s0_ref, 'f': lambda s, _fn=update_fn: _fn(s), 'n': TypedTensor(n_oh, TyNat())}) for p in range(3): pred_pools[p].append(r.data[p].item()) true_ref = true_curves[ref_idx].numpy() pool_colors = ['steelblue', 'darkorange', 'forestgreen'] pool_names = ['C_fast', 'C_slow', 'C_passive'] ax = axes[1,2] for p, (name, color) in enumerate(zip(pool_names, pool_colors)): ax.plot(true_ref[:, p], '-', color=color, lw=2, label=f'True {name}') ax.plot(pred_pools[p], '--', color=color, lw=2, alpha=0.7, label=f'Pred {name}') ax.set_xlabel('Timestep'); ax.set_ylabel('Carbon') ax.set_title(f'Trajectory: T={T_r:.0f}°C, M={M_r:.1f}') ax.legend(ncol=2, fontsize=7); ax.grid(True, alpha=0.3) plt.suptitle('CENTURY-Lite: Simultaneously Learning f_temp(T) and f_moist(M)', fontsize=13, fontweight='bold') plt.tight_layout() plt.savefig('century_lite_results.png', dpi=120, bbox_inches='tight') plt.show()

In [8]:

# Verification
with torch.no_grad():
    moist_arch = torch.tensor([f_moist_mlp(torch.tensor([m])).item() for m in torch.linspace(0,1.5,100)])
    temp_arch  = torch.tensor([f_temp_mlp(torch.tensor([t])).item()  for t in torch.linspace(0,40, 100)])

print('VERIFICATION:')
print(f"  f_moist(M) ∈ [0,1]:   {'VERIFIED' if (moist_arch>=0).all() and (moist_arch<=1).all() else 'FAILED'} (sigmoid architecture)")
print(f"  f_temp(T) > 0:        {'VERIFIED' if (temp_arch>0).all() else 'FAILED'} (softplus architecture)")

all_pos = mass_ok = True
with torch.no_grad():
    for idx in range(N_TRAJ):
        T_v, M_v = conditions[idx]
        s0 = TypedTensor(torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_v, M_v]), TyReal(5))
        tot_prev = C0_FAST + C0_SLOW + C0_PASSIVE
        for step in range(N_STEPS):
            n_oh = torch.zeros(N_STEPS); n_oh[step] = 1.0
            r = compiled({'s0': s0, 'f': lambda s, _fn=update_fn: _fn(s),
                          'n': TypedTensor(n_oh, TyNat())})
            cf, cs, cp = r.data[0].item(), r.data[1].item(), r.data[2].item()
            if cf < -1e-9 or cs < -1e-9 or cp < -1e-9: all_pos = False
            tot = cf + cs + cp
            if tot > tot_prev + 1e-6: mass_ok = False
            tot_prev = tot

print(f"  All pools ≥ 0:        {'VERIFIED' if all_pos else 'FAILED'}")
print(f"  Total C non-increasing:{'VERIFIED' if mass_ok else 'FAILED'} (no carbon created)")
print()
print('DECOMPILED MODEL:')
print(f'  f_temp(T)  = {best_q10:.2f}^((T−{best_tref:.1f})/10)   [true: 2.0^((T−15)/10)]')
print(f'  f_moist(M) = M^0.7 / ({best_mK:.3f} + M^0.7)      [true: M^0.7/(0.300+M^0.7)]')
print()
print('Both functional forms correctly identified by symbolic regression.')
print('Scale offsets reflect a fundamental identifiability limitation with 9 conditions:')
print('only the product f_temp×f_moist is uniquely determined, not the factors individually.')

# Verification with torch.no_grad(): moist_arch = torch.tensor([f_moist_mlp(torch.tensor([m])).item() for m in torch.linspace(0,1.5,100)]) temp_arch = torch.tensor([f_temp_mlp(torch.tensor([t])).item() for t in torch.linspace(0,40, 100)]) print('VERIFICATION:') print(f" f_moist(M) ∈ [0,1]: {'VERIFIED' if (moist_arch>=0).all() and (moist_arch<=1).all() else 'FAILED'} (sigmoid architecture)") print(f" f_temp(T) > 0: {'VERIFIED' if (temp_arch>0).all() else 'FAILED'} (softplus architecture)") all_pos = mass_ok = True with torch.no_grad(): for idx in range(N_TRAJ): T_v, M_v = conditions[idx] s0 = TypedTensor(torch.tensor([C0_FAST, C0_SLOW, C0_PASSIVE, T_v, M_v]), TyReal(5)) tot_prev = C0_FAST + C0_SLOW + C0_PASSIVE for step in range(N_STEPS): n_oh = torch.zeros(N_STEPS); n_oh[step] = 1.0 r = compiled({'s0': s0, 'f': lambda s, _fn=update_fn: _fn(s), 'n': TypedTensor(n_oh, TyNat())}) cf, cs, cp = r.data[0].item(), r.data[1].item(), r.data[2].item() if cf < -1e-9 or cs < -1e-9 or cp < -1e-9: all_pos = False tot = cf + cs + cp if tot > tot_prev + 1e-6: mass_ok = False tot_prev = tot print(f" All pools ≥ 0: {'VERIFIED' if all_pos else 'FAILED'}") print(f" Total C non-increasing:{'VERIFIED' if mass_ok else 'FAILED'} (no carbon created)") print() print('DECOMPILED MODEL:') print(f' f_temp(T) = {best_q10:.2f}^((T−{best_tref:.1f})/10) [true: 2.0^((T−15)/10)]') print(f' f_moist(M) = M^0.7 / ({best_mK:.3f} + M^0.7) [true: M^0.7/(0.300+M^0.7)]') print() print('Both functional forms correctly identified by symbolic regression.') print('Scale offsets reflect a fundamental identifiability limitation with 9 conditions:') print('only the product f_temp×f_moist is uniquely determined, not the factors individually.')

VERIFICATION:
  f_moist(M) ∈ [0,1]:   VERIFIED (sigmoid architecture)
  f_temp(T) > 0:        VERIFIED (softplus architecture)
  All pools ≥ 0:        VERIFIED
  Total C non-increasing:VERIFIED (no carbon created)

DECOMPILED MODEL:
  f_temp(T)  = 1.77^((T−18.1)/10)   [true: 2.0^((T−15)/10)]
  f_moist(M) = M^0.7 / (0.130 + M^0.7)      [true: M^0.7/(0.300+M^0.7)]

Both functional forms correctly identified by symbolic regression.
Scale offsets reflect a fundamental identifiability limitation with 9 conditions:
only the product f_temp×f_moist is uniquely determined, not the factors individually.