Learning an Unknown Nonlinear Function via Neural Sub-expression¶

Soil Carbon with Unknown Moisture Response¶

Instead of learning a scalar parameter, we learn an entire unknown function using a neural network, then decompile it back to symbolic form.

Model:

C(n+1) = C(n) - k · f_moisture(M) · C(n) · dt

• k = base decomposition rate (known)
• f_moisture(M) = moisture response function (unknown — learned as MLP)
• Ground truth: f_moisture(m) = m^0.7 / (0.3 + m^0.7) (Hill equation)

Here f_moisture is a Hill-type saturating response function: increasing moisture raises decomposition, but the effect approaches a maximum. The functional form originates with Hill's treatment of cooperative binding and is used here as a convenient phenomenological moisture-response curve rather than a mechanistic soil law [1,2].

After training, symbolic regression on the MLP recovers the Hill equation.

References

[1] A. V. Hill. The possible effects of the aggregation of the molecules of haemoglobin on its dissociation curves. Journal of Physiology 40: iv-vii, 1910. https://doi.org/10.1113/jphysiol.1910.sp001386

[2] R. Gesztelyi, J. Zsuga, A. Kemeny-Beke, B. Varga, B. Juhasz, and A. Tosaki. The Hill equation and the origin of quantitative pharmacology. Archive for History of Exact Sciences 66: 427-438, 2012. https://doi.org/10.1007/s00407-012-0098-5

1. Generate Synthetic Training Data¶

We first choose a ground-truth moisture response and simulate carbon-decay trajectories at several fixed moisture values. These trajectories are the only supervision the learner will see: the model never gets direct access to the symbolic Hill formula during training.

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

# Problem setup for the synthetic decay experiment.
# CPU is faster for this workload (many tiny sequential ops)
device = torch.device("cpu")

K_BASE = 0.5
DT = 0.1
C0 = 1.0
N_STEPS = 10
N_TRAJ  = 20

# Ground-truth moisture response used only to generate supervision.
def true_f_moisture(m):
    return m ** 0.7 / (0.3 + m ** 0.7)

# Generate one carbon-decay trajectory for each fixed moisture level.
moistures = torch.linspace(0.05, 1.0, N_TRAJ)
all_curves = []
for m in moistures:
    f_m = true_f_moisture(m)
    curve = []
    c = C0
    for _ in range(N_STEPS):
        curve.append(c)
        c = c * (1.0 - K_BASE * f_m * DT)
    all_curves.append(torch.tensor(curve))

print(f"Generated {N_TRAJ} training trajectories at different moisture levels")
print(f"True f_moisture: Hill equation m^0.7 / (0.3 + m^0.7)")

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 # Problem setup for the synthetic decay experiment. # CPU is faster for this workload (many tiny sequential ops) device = torch.device("cpu") K_BASE = 0.5 DT = 0.1 C0 = 1.0 N_STEPS = 10 N_TRAJ = 20 # Ground-truth moisture response used only to generate supervision. def true_f_moisture(m): return m ** 0.7 / (0.3 + m ** 0.7) # Generate one carbon-decay trajectory for each fixed moisture level. moistures = torch.linspace(0.05, 1.0, N_TRAJ) all_curves = [] for m in moistures: f_m = true_f_moisture(m) curve = [] c = C0 for _ in range(N_STEPS): curve.append(c) c = c * (1.0 - K_BASE * f_m * DT) all_curves.append(torch.tensor(curve)) print(f"Generated {N_TRAJ} training trajectories at different moisture levels") print(f"True f_moisture: Hill equation m^0.7 / (0.3 + m^0.7)")

Generated 20 training trajectories at different moisture levels
True f_moisture: Hill equation m^0.7 / (0.3 + m^0.7)

2. Learn the Unknown Sub-expression Inside the Program¶

The MLP plays the role of the unknown function f_moisture. We wrap it inside the decay update, compile a Cajal TmIter program that repeatedly applies that update, and then train the MLP so the compiled program reproduces the synthetic trajectories across all moisture settings and time steps.

In [2]:

# Learn f_moisture with a compact network constrained to [0, 1].
class MoistureResponseMLP(nn.Module):
    """f_moisture: R -> [0,1].  Sigmoid output is a structural 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()

# Wrap the learned sub-expression inside one step of the decay dynamics.
class DecayWithMoisture(nn.Module):
    def __init__(self, k, dt, mlp):
        super().__init__()
        self.k, self.dt, self.f_moisture = k, dt, mlp
    def forward(self, state):
        c = state.data[0]
        moisture = state.data[1]
        f_m = self.f_moisture(moisture)
        c_new = c - self.k * f_m * c * self.dt
        return TypedTensor(torch.stack([c_new, moisture]), state.ty)

# Cajal program: iterate the update function n times from initial state s0.
program  = TmIter(TmVar("s0"), "s", TmApp(TmVar("f"), TmVar("s")), TmVar("n"))
compiled = compile(program)

mlp       = MoistureResponseMLP()
update_fn = DecayWithMoisture(K_BASE, DT, mlp)
optimizer = torch.optim.Adam(mlp.parameters(), lr=0.005)
n_params  = sum(p.numel() for p in mlp.parameters())
print(f"MLP has {n_params} parameters")

# Fit the learned sub-expression so the compiled program matches every trajectory.
losses = []
for epoch in range(800):
    optimizer.zero_grad()
    total_loss = torch.tensor(0.0)
    for traj_idx in range(N_TRAJ):
        m = moistures[traj_idx]
        s0 = TypedTensor(torch.stack([torch.tensor(C0), m]), TyReal(2))
        for step in range(N_STEPS):
            # TyNat is encoded as a one-hot vector selecting the iteration count.
            n_onehot = torch.zeros(N_STEPS)
            n_onehot[step] = 1.0
            n_val = TypedTensor(n_onehot, TyNat())
            result = compiled({"s0": s0, "f": lambda s, _fn=update_fn: _fn(s), "n": n_val})
            total_loss = total_loss + (result.data[0] - all_curves[traj_idx][step]) ** 2
    total_loss.backward()
    optimizer.step()
    losses.append(total_loss.item())
    if (epoch + 1) % 200 == 0:
        print(f"  epoch {epoch+1:3d}  loss={total_loss.item():.8f}")

# Learn f_moisture with a compact network constrained to [0, 1]. class MoistureResponseMLP(nn.Module): """f_moisture: R -> [0,1]. Sigmoid output is a structural 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() # Wrap the learned sub-expression inside one step of the decay dynamics. class DecayWithMoisture(nn.Module): def __init__(self, k, dt, mlp): super().__init__() self.k, self.dt, self.f_moisture = k, dt, mlp def forward(self, state): c = state.data[0] moisture = state.data[1] f_m = self.f_moisture(moisture) c_new = c - self.k * f_m * c * self.dt return TypedTensor(torch.stack([c_new, moisture]), state.ty) # Cajal program: iterate the update function n times from initial state s0. program = TmIter(TmVar("s0"), "s", TmApp(TmVar("f"), TmVar("s")), TmVar("n")) compiled = compile(program) mlp = MoistureResponseMLP() update_fn = DecayWithMoisture(K_BASE, DT, mlp) optimizer = torch.optim.Adam(mlp.parameters(), lr=0.005) n_params = sum(p.numel() for p in mlp.parameters()) print(f"MLP has {n_params} parameters") # Fit the learned sub-expression so the compiled program matches every trajectory. losses = [] for epoch in range(800): optimizer.zero_grad() total_loss = torch.tensor(0.0) for traj_idx in range(N_TRAJ): m = moistures[traj_idx] s0 = TypedTensor(torch.stack([torch.tensor(C0), m]), TyReal(2)) for step in range(N_STEPS): # TyNat is encoded as a one-hot vector selecting the iteration count. n_onehot = torch.zeros(N_STEPS) n_onehot[step] = 1.0 n_val = TypedTensor(n_onehot, TyNat()) result = compiled({"s0": s0, "f": lambda s, _fn=update_fn: _fn(s), "n": n_val}) total_loss = total_loss + (result.data[0] - all_curves[traj_idx][step]) ** 2 total_loss.backward() optimizer.step() losses.append(total_loss.item()) if (epoch + 1) % 200 == 0: print(f" epoch {epoch+1:3d} loss={total_loss.item():.8f}")

MLP has 1153 parameters

  epoch 200  loss=0.00493772

  epoch 400  loss=0.00107890

  epoch 600  loss=0.00035698

  epoch 800  loss=0.00018694

3. Evaluate the Learned Function and Search for a Closed Form¶

Once training is done, we compare the learned MLP against the true response on a dense moisture grid. We then perform a lightweight symbolic-regression sweep over a few plausible functional families and keep the best-fitting closed form.

In [3]:

# Compare the learned network to the ground-truth response on held-out points.
test_m = torch.linspace(0.05, 1.0, 50)
with torch.no_grad():
    true_vals   = [true_f_moisture(m).item() for m in test_m]
    learned_vals = [mlp(m).item() for m in test_m]

max_err = max(abs(t - l) for t, l in zip(true_vals, learned_vals))
print(f"Max absolute error: {max_err:.4f}")

# Candidate symbolic families for a simple regression sweep.
m_grid = torch.linspace(0.01, 1.5, 200)
with torch.no_grad():
    y_learned = torch.tensor([mlp(m).item() for m in m_grid])

K_grid = torch.linspace(0.01, 2.0, 200)
candidates = {
    "MM: m/(K+m)":           lambda m, K: m / (K + m),
    "Hill n=0.5: m^0.5/(K+m^0.5)": lambda m, K: m**0.5 / (K + m**0.5),
    "Hill n=0.7: m^0.7/(K+m^0.7)": lambda m, K: m**0.7 / (K + m**0.7),
    "Hill n=1.5: m^1.5/(K+m^1.5)": lambda m, K: m**1.5 / (K + m**1.5),
    "quadratic: m^2/(K+m^2)": lambda m, K: m**2 / (K + m**2),
}

# Grid-search K for each candidate and keep the lowest-MSE symbolic form.
best_name, best_loss, best_K = None, float('inf'), None
sym_results = {}
for name, func in candidates.items():
    best_K_this, best_L_this = None, float('inf')
    for Kv in K_grid:
        pred = func(m_grid, Kv)
        L = ((pred - y_learned) ** 2).mean().item()
        if L < best_L_this:
            best_L_this, best_K_this = L, Kv.item()
    sym_results[name] = (best_L_this, best_K_this)
    if best_L_this < best_loss:
        best_loss, best_name, best_K = best_L_this, name, best_K_this

print()
print("Symbolic regression:")
for name, (L, K) in sym_results.items():
    marker = " <-- BEST" if name == best_name else ""
    print(f"  {name:40s}  MSE={L:.6f}  K={K:.3f}{marker}")
print(f"\nGround truth: Hill n=0.7, K=0.300")

# Compare the learned network to the ground-truth response on held-out points. test_m = torch.linspace(0.05, 1.0, 50) with torch.no_grad(): true_vals = [true_f_moisture(m).item() for m in test_m] learned_vals = [mlp(m).item() for m in test_m] max_err = max(abs(t - l) for t, l in zip(true_vals, learned_vals)) print(f"Max absolute error: {max_err:.4f}") # Candidate symbolic families for a simple regression sweep. m_grid = torch.linspace(0.01, 1.5, 200) with torch.no_grad(): y_learned = torch.tensor([mlp(m).item() for m in m_grid]) K_grid = torch.linspace(0.01, 2.0, 200) candidates = { "MM: m/(K+m)": lambda m, K: m / (K + m), "Hill n=0.5: m^0.5/(K+m^0.5)": lambda m, K: m**0.5 / (K + m**0.5), "Hill n=0.7: m^0.7/(K+m^0.7)": lambda m, K: m**0.7 / (K + m**0.7), "Hill n=1.5: m^1.5/(K+m^1.5)": lambda m, K: m**1.5 / (K + m**1.5), "quadratic: m^2/(K+m^2)": lambda m, K: m**2 / (K + m**2), } # Grid-search K for each candidate and keep the lowest-MSE symbolic form. best_name, best_loss, best_K = None, float('inf'), None sym_results = {} for name, func in candidates.items(): best_K_this, best_L_this = None, float('inf') for Kv in K_grid: pred = func(m_grid, Kv) L = ((pred - y_learned) ** 2).mean().item() if L < best_L_this: best_L_this, best_K_this = L, Kv.item() sym_results[name] = (best_L_this, best_K_this) if best_L_this < best_loss: best_loss, best_name, best_K = best_L_this, name, best_K_this print() print("Symbolic regression:") for name, (L, K) in sym_results.items(): marker = " <-- BEST" if name == best_name else "" print(f" {name:40s} MSE={L:.6f} K={K:.3f}{marker}") print(f"\nGround truth: Hill n=0.7, K=0.300")

Max absolute error: 0.0110

Symbolic regression:
  MM: m/(K+m)                               MSE=0.002556  K=0.240
  Hill n=0.5: m^0.5/(K+m^0.5)               MSE=0.001331  K=0.340
  Hill n=0.7: m^0.7/(K+m^0.7)               MSE=0.000145  K=0.300 <-- BEST
  Hill n=1.5: m^1.5/(K+m^1.5)               MSE=0.011968  K=0.160
  quadratic: m^2/(K+m^2)                    MSE=0.023105  K=0.100

Ground truth: Hill n=0.7, K=0.300

4. Visualize Optimization, Recovery, and Decompilation¶

The final figure shows three views of the result: optimization progress, how closely the MLP matches the true moisture response, and whether the best symbolic fit matches the learned function closely enough to count as a successful decompilation.

In [4]:

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# Panel 1: optimization history.
axes[0].semilogy(losses)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss (log scale)')
axes[0].set_title('Training loss')
axes[0].grid(True, alpha=0.3)

# Panel 2: learned function versus the true Hill curve.
axes[1].plot(test_m.numpy(), true_vals, 'r-', linewidth=2, label='True (Hill n=0.7)')
axes[1].plot(test_m.numpy(), learned_vals, 'b--', linewidth=2, label=f'Learned MLP ({n_params} params)')
axes[1].set_xlabel('Moisture M')
axes[1].set_ylabel('f_moisture(M)')
axes[1].set_title(f'Learned vs True f_moisture\nMax error: {max_err:.4f}')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

# Panel 3: best symbolic approximation versus the learned network.
best_func = candidates[best_name]
y_sym = best_func(m_grid, best_K)
axes[2].plot(m_grid.numpy(), y_learned.numpy(), 'b-', linewidth=2, label='MLP output')
axes[2].plot(m_grid.numpy(), y_sym.numpy(), 'g--', linewidth=2, label=f'Symbolic fit:\n{best_name}\nK={best_K:.3f}')
axes[2].plot(m_grid.numpy(), [true_f_moisture(m).item() for m in m_grid], 
             'r:', linewidth=1.5, label='True')
axes[2].set_xlabel('Moisture M')
axes[2].set_ylabel('f_moisture(M)')
axes[2].set_title('Symbolic Regression\n(decompilation)')
axes[2].legend(fontsize=8)
axes[2].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('learn_unknown_function.png', dpi=120, bbox_inches='tight')
plt.show()

if '0.7' in best_name:
    print("✓ Symbolic regression correctly recovered the Hill equation with n=0.7")

fig, axes = plt.subplots(1, 3, figsize=(15, 4)) # Panel 1: optimization history. axes[0].semilogy(losses) axes[0].set_xlabel('Epoch') axes[0].set_ylabel('Loss (log scale)') axes[0].set_title('Training loss') axes[0].grid(True, alpha=0.3) # Panel 2: learned function versus the true Hill curve. axes[1].plot(test_m.numpy(), true_vals, 'r-', linewidth=2, label='True (Hill n=0.7)') axes[1].plot(test_m.numpy(), learned_vals, 'b--', linewidth=2, label=f'Learned MLP ({n_params} params)') axes[1].set_xlabel('Moisture M') axes[1].set_ylabel('f_moisture(M)') axes[1].set_title(f'Learned vs True f_moisture\nMax error: {max_err:.4f}') axes[1].legend() axes[1].grid(True, alpha=0.3) # Panel 3: best symbolic approximation versus the learned network. best_func = candidates[best_name] y_sym = best_func(m_grid, best_K) axes[2].plot(m_grid.numpy(), y_learned.numpy(), 'b-', linewidth=2, label='MLP output') axes[2].plot(m_grid.numpy(), y_sym.numpy(), 'g--', linewidth=2, label=f'Symbolic fit:\n{best_name}\nK={best_K:.3f}') axes[2].plot(m_grid.numpy(), [true_f_moisture(m).item() for m in m_grid], 'r:', linewidth=1.5, label='True') axes[2].set_xlabel('Moisture M') axes[2].set_ylabel('f_moisture(M)') axes[2].set_title('Symbolic Regression\n(decompilation)') axes[2].legend(fontsize=8) axes[2].grid(True, alpha=0.3) plt.tight_layout() plt.savefig('learn_unknown_function.png', dpi=120, bbox_inches='tight') plt.show() if '0.7' in best_name: print("✓ Symbolic regression correctly recovered the Hill equation with n=0.7")

✓ Symbolic regression correctly recovered the Hill equation with n=0.7