How to use template expressions and a custom loss function for coupled ODE system?

[okko-w]

Hi there!

Thanks for creating this nice symbolic regression package! I really enjoy working with it.

Background
I am currently using it for my MSc thesis to test whether symbolic regression can recover ecological ODE models and their bifurcation structure. I’ve had some promising results with a 1D test case, and I’m now extending this to a coupled 2D system. Based on the template expression example from the SymbolicRegression.jl docs https://astroautomata.com/SymbolicRegression.jl/dev/examples/template_expression, I’ve set up a template expression for a 2D case. It should find an expression for each of the states variables V and E (fV, fE) where Ebase is a bifurcation parameter as the third feature and dV, dE are the time derivatives of V and E:

# Dataset is created before

struct Derivs{T}
    dV::T
    dE::T
end

X = select(Dataset, :V, :E, :Ebase)
y = Derivs.(Dataset.dV, Dataset.dE)

template_expr = @template_spec(
    expressions = (fV, fE),
) do V, E, Ebase
    _fV = fV(V, E)
    _fE = fE(V, E, Ebase)

    derivs = [Derivs(_fV.x[i], _fE.x[i]) for i in eachindex(V.x)]
    return ValidVector(derivs, _fV.valid && _fE.valid)
end

Questions

1. Custom loss function

As it seemed to improve performance in the 1D case, I'm trying to add a custom loss function that evaluates whether dV,dE are zero when V, E are zero over an input grid and adds the difference from zero as a weighted penalty to the MAE loss:

function eval_boundary_loss(tree, dataset::Dataset{T,L}, options, idx)::L where {T,L}
    # For batching
    X = idx === nothing ? dataset.X : dataset.X[:, idx]
    y = idx === nothing ? dataset.y : view(dataset.y, idx)

    preds, flag = eval_tree_array(tree, X, options)

    if !flag
          return L(Inf)
    end

    # Compute (MAE) loss
    data_loss = mean(abs(pred.dV - target.dV) + abs(pred.dE - target.dE) for (pred, target) in zip(preds, y))

    # Boundary grid
    E_vals = range(0.0, 5.0, length=20)
    Ebase_vals = range(0.5, 1.5, length=20)

    # Enforce dV = 0 at V = 0
    V_grid = hcat((collect((0.0, E, Eb)) for E in E_vals, Eb in Ebase_vals)...)
    pred_V, flag1 = eval_tree_array(tree, V_grid, options)
    
    if !flag1
          return L(Inf)
    end

    penalty_V = mean(abs.([x.dV for x in pred_V]))

    # Enforce dE = 0 at E = 0
    V_vals = range(0.0, 1.0, length=20)
    E_grid = hcat((collect((V, 0.0, Eb)) for V in V_vals, Eb in Ebase_vals)...)
    pred_E, flag2 = eval_tree_array(tree, E_grid, options)

    if !flag2
          return L(Inf)
    end

    penalty_E = mean(abs.([x.dE for x in pred_E]))

    boundary_loss = penalty_V + penalty_E

    return data_loss + L(0.1) * boundary_loss
end

This code runs, but I’m unsure if I’ve set it up correctly. The template returns a ValidVector, but how is this handled by the loss function sinceeval_tree_array returns a regular array and a flag? Do I unpack them correctly in the data_loss calculation and handle the flags properly? Also, is there a more efficient way to implement this?

2. Nested template expressions

I was wondering whether using a nested template expression makes any sense. So if I know that both fV and fE are likely composed of growth (g) and loss (l) terms, could I define nested templates as:

template_expr = @template_spec(
    expressions = (fV, fE, g1, g2, l1, l2),
) do V, E, Ebase
    _fV = fV(g1(V) - l1(E))
    _fE = fE(g2(V, E)- l2(Ebase))

    derivs = [Derivs(_fV.x[i], _fE.x[i]) for i in eachindex(V.x)]
    return ValidVector(derivs, _fV.valid && _fE.valid)
end

3. Using guesses with templates

I'm also interested in using the new guesses option to further guide the search. How can I include parameter guesses in the template together with the guesses option? From this issue: MilesCranmer/SymbolicRegression.jl#510 of SymbolicRegression.jl, I gather that it is possible to set up a parameterized template and provide initial parameters values for the guesses with a full TemplateExpression instead of the @template_spec macro. I am at a bit of a loss on how to set this up, could you provide some pointers?

Thank you in advance! And apologies if any of this is trivial, I’m still new to Julia and learning the ropes!