Nonlinear least squares parameter optimiziation of extracted functions with automatic differentiation

[david-hofmann]

Hi @MilesCranmer ! I've been trying to use automatic differentiation on the data structure provided by Symbolic Regression to optimize parameters of the functions post-search. Haven't been able to sort out all the issues yet thus I would be really grateful for any hints on perhaps simpler ways to achieving this in Julia.

In reference to this discussion #596, here an MWE to show my current attempt:

using SymbolicRegression, MLJBase
using NonlinearSolve

function evaltree_newparams(eqtree, x, params)
    constants = filter(node -> node.degree == 0 && node.constant, eqtree)
    foreach(zip(constants, params)) do (constant, β)
        constant.val = β
    end
    eqtree(x)
end

function lsqfitting(fitfct, (x,y), sigma, β0)
    res(β,(x,y)) = fitfct(x,β) .- y
    prob = NonlinearLeastSquaresProblem(res, β0, (x,y))
    sol = solve(prob, GaussNewton(); maxiters=1000)
    chi2(β) = sum(abs2, res(β,(x,y)) ./ sigma)
    return sol.u, inv(2*ForwardDiff.hessian(chi2,sol.u)), chi2(sol.u)
end

n = 100
nd = 5
X = randn(n, nd);
y = @. cos(X[:, 1] * 2.1 - 0.5) + 0.2 * X[:, 3] * X[:, 2];

model = SRRegressor(binary_operators=[+, -, *, /], unary_operators=[cos], niterations=100);
mach = machine(model, X, y);
fit!(mach, verbosity=0);

equations = report(mach).equations
best = equations[end-1]
eq = get_tree(best)
function_to_fit = (x, β) -> evaltree_newparams(eq, x, β) 

lsqfitting(function_to_fit, (X',y), 1, rand(nd))

It breaks here constant.val = β since The type Float64 exists, but no method is defined for this combination of argument types when trying to construct it.

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{NonlinearFunction{…}, Float64}, Float64, 5})

[MilesCranmer]

See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/ConstantOptimization.jl for how I do constant optimization internally. You could just call that directly

[david-hofmann]

I've been digging into the SR.jl structure more deeply but not been able to make it work just yet. Can you tell me what I am doing wrong here?
I've added the call to optimize_constants() as so:

using SymbolicRegression, MLJBase

n = 100
nd = 5
X = randn(n, nd);
y = @. cos(X[:, 1] * 2.1 - 0.5) + 0.2 * X[:, 3] * X[:, 2];

model = SRRegressor(binary_operators=[+, -, *, /], unary_operators=[cos], niterations=100);
mach = machine(model, X, y);
fit!(mach, verbosity=0);

r = report(mach)
best = r.equations[r.best_idx]

pm = PopMember(best, r.scores[r.best_idx], r.losses[r.best_idx], deterministic=true)
SymbolicRegression.ConstantOptimizationModule.optimize_constants(Dataset(X', y), pm, Options())

Which results in the error:

ERROR: BoundsError: attempt to access Tuple{} at index [0]
Stacktrace:
  [1] getindex(t::Tuple, i::Int64)
    @ Base ./tuple.jl:31
  [2] macro expansion
    @ ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:395 [inlined]
  [3] dispatch_deg1_eval(tree::DynamicExpressions.NodeModule.Node{…}, cX::LinearAlgebra.Adjoint{…}, op_idx::UInt8, operators::OperatorEnum{…}, eval_options::EvalOptions{…})
    @ DynamicExpressions.EvaluateModule ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:375
  [4] _eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::LinearAlgebra.Adjoint{…}, operators::OperatorEnum{…}, eval_options::EvalOptions{…})
    @ DynamicExpressions.EvaluateModule ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:277
  [5] macro expansion
    @ ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:362 [inlined]
  [6] dispatch_deg2_eval(tree::DynamicExpressions.NodeModule.Node{…}, cX::LinearAlgebra.Adjoint{…}, op_idx::UInt8, operators::OperatorEnum{…}, eval_options::EvalOptions{…})
    @ DynamicExpressions.EvaluateModule ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:321
  [7] _eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::LinearAlgebra.Adjoint{…}, operators::OperatorEnum{…}, eval_options::EvalOptions{…})
    @ DynamicExpressions.EvaluateModule ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:282
  [8] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::LinearAlgebra.Adjoint{…}, operators::OperatorEnum{…}; eval_options::EvalOptions{…}, _deprecated_kws::@Kwargs{})
    @ DynamicExpressions.EvaluateModule ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:227
  [9] eval_tree_array
    @ ~/.julia/packages/DynamicExpressions/3sROs/src/Evaluate.jl:200 [inlined]
 [10] eval_tree_array(ex::Expression{…}, cX::LinearAlgebra.Adjoint{…}, operators::OperatorEnum{…}; kws::@Kwargs{…})
    @ DynamicExpressions.ExpressionModule ~/.julia/packages/DynamicExpressions/3sROs/src/Expression.jl:436
 [11] eval_tree_array
    @ ~/.julia/packages/DynamicExpressions/3sROs/src/Expression.jl:429 [inlined]
 [12] #eval_tree_array#1
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/InterfaceDynamicExpressions.jl:68 [inlined]
 [13] eval_tree_array
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/InterfaceDynamicExpressions.jl:56 [inlined]
 [14] eval_tree_dispatch
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/LossFunctions.jl:54 [inlined]
 [15] _eval_loss(tree::Expression{…}, dataset::Dataset{…}, options::Options{…}, regularization::Bool, idx::Nothing)
    @ SymbolicRegression.LossFunctionsModule ~/.julia/packages/SymbolicRegression/xt8Jj/src/LossFunctions.jl:87
 [16] eval_loss(tree::Expression{…}, dataset::Dataset{…}, options::Options{…}; regularization::Bool, idx::Nothing)
    @ SymbolicRegression.LossFunctionsModule ~/.julia/packages/SymbolicRegression/xt8Jj/src/LossFunctions.jl:154
 [17] eval_loss
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/LossFunctions.jl:138 [inlined]
 [18] #_#1
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/ConstantOptimization.jl:115 [inlined]
 [19] Evaluator
    @ ~/.julia/packages/SymbolicRegression/xt8Jj/src/ConstantOptimization.jl:113 [inlined]
 [20] _optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…}, algorithm::Optim.BFGS{…}, optimizer_options::Optim.Options{…}, idx::Nothing)
    @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/SymbolicRegression/xt8Jj/src/ConstantOptimization.jl:75
 [21] dispatch_optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…}, idx::Nothing)
    @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/SymbolicRegression/xt8Jj/src/ConstantOptimization.jl:46
 [22] optimize_constants(dataset::Dataset{…}, member::PopMember{…}, options::Options{…})
    @ SymbolicRegression.ConstantOptimizationModule ~/.julia/packages/SymbolicRegression/xt8Jj/src/ConstantOptimization.jl:27
 [23] top-level scope
    @ ~/Projects/neuroblox/codes/NeurobloxSciML/lib/symbolic_regression/sr_cc.jl:285
Some type information was truncated. Use `show(err)` to see complete types.

Thank you so much!

[MilesCranmer]

Ah it’s because you passed Options() with different operators than the ones you passed to SRRegressor. Need to pass binary_operators and unary_operators. Same exact ones you passed to SRRegressor