Extract a function with parameters after fitting (Julia)

[avik-pal]

I am trying to reproduce https://docs.kidger.site/diffrax/examples/symbolic_regression/ on the Julia end. Specifically this part

symbolic_fn = Stack([sympy2jax.SymbolicModule(expr) for expr in best_expressions])
symbolic_model = eqx.tree_at(lambda m: m.func.mlp, model, symbolic_fn)  # noqa: F821
optim = optax.adam(fine_tuning_lr)
opt_state = optim.init(eqx.filter(symbolic_model, eqx.is_inexact_array))

After fitting, is there an API which I can use to extract a function and parameters that can be fine-tuned?

[MilesCranmer]

AbstractExpressionNode{T} implements Optim.optimize if this is what you are looking for?

https://github.com/SymbolicML/DynamicExpressions.jl/blob/4bd83e97d23ebb3fbde57f720b0265c5ecc33881/ext/DynamicExpressionsOptimExt.jl#L84-L113

So if you have a function that takes an expression tree like Node{T}, you can directly use Optim.optimize on it. I'm not sure this answers your question though, could you be more specific?

[avik-pal]

I think I can work through the optim version to do this.

Effectively, I am trying to see if I can take a fitted symbolic expression, extract f and p, and then use it inside a lux layer.

my_ude = Chain(
	LuxSRWrapper(fitted_tree),  # I need to have a parameter vector for this to work
	Dense(),
)

[MilesCranmer]

Sounds super cool!!

Here's a bit more info in the form of an example:

julia> using SymbolicRegression, MLJBase

julia> X = randn(100, 5);

julia> y = @. cos(X[:, 1] * 2.1 - 0.5) + 0.2 * X[:, 3] * X[:, 2];

julia> model = SRRegressor(binary_operators=[+, -, *, /], unary_operators=[cos], niterations=100);

julia> mach = machine(model, X, y);

julia> fit!(mach, verbosity=0);

julia> equations = report(mach).equations
8-element Vector{DynamicExpressions.EquationModule.Node{Float64}}:
 0.05336346662643239
 x₁ * 0.21680232405649386
 cos(x₁ * 2.0617285536753673)
 cos((x₁ * 2.035173184625832) + -0.5033370794240668)
 cos((x₁ * -2.0106626799718876) - -0.4949017834844345) - 0.04205315632795699
 cos((0.009450215065161867 / x₁) - (x₁ + (x₁ - 0.5057275123455831)))
 cos(-0.499999999947676 - (x₁ * -2.09999999993789)) + (x₃ * (x₂ * 0.19999999996546464))
 ((((x₃ - cos(0.20000000003991894)) + cos(0.20000000003991894)) * 0.20000000003991894) * x₂) + cos(-0.4999999999840714 - (x₁ * -2.100000000063916))

julia> best = equations[end-1]
cos(-0.499999999947676 - (x₁ * -2.09999999993789)) + (x₃ * (x₂ * 0.19999999996546464))

julia> constants = filter(node -> node.degree == 0 && node.constant, best)
3-element Vector{DynamicExpressions.EquationModule.Node{Float64}}:
 -0.499999999947676
 -2.09999999993789
 0.19999999996546464

julia> foreach(constants) do constant
           constant.val *= 2
       end

julia> best
cos(-0.999999999895352 - (x₁ * -4.19999999987578)) + (x₃ * (x₂ * 0.3999999999309293))

So you can see we doubled all the values.

We can then do

julia> best(X', mach.fitresult.options)

to evaluate it. Note the X', because MLJ assumes [row, feature] while SymbolicRegression uses [feature, row] internally.

A lot of the Base operations on collections are implemented for Node types so it's easy to walk through the expression trees: https://symbolicml.org/DynamicExpressions.jl/dev/utils/

[avik-pal]

Awesome, thanks for this. This should be good to get me started.

[MilesCranmer]

Awesome! Happy to help get this going, it sounds super cool. If you get a work-in-progress PR feel free to tag me in it and I can help add suggestions.

[MilesCranmer]

Also as a heads-up, Zygote doesn't work out-of-the-box on the evaluation method. So instead there is a hand-rolled forward mode differentiation: https://symbolicml.org/DynamicExpressions.jl/dev/eval/#Derivatives which use Zygote on the individual operators. (Note that I don't have second-order differentiation implemented though)

(I need to get around to rewriting this in ChainRulesCore.jl but haven't yet.)

[leelew]

Is there any method to count constant using Python API?