Symbolic integrals with SymbolicRegression.jl

[MilesCranmer]

Example of solving an integral with SymbolicRegression.jl. Since automatic differentiation is enabled (via Zygote.jl), you can literally just evaluate the forward derivative, and compute the loss against your function. Then the resultant expression == the integral! (Or some approximation of it)

Screen.Recording.2023-08-07.at.15.26.36.mp4

Code:

using MLJ  # for fit/predict
using SymbolicRegression  # for SRRegressor
using Zygote  # For `enable_autodiff=true`
using SymbolicUtils

f(x) = x^2 + 2x + 1 + sin(x)

X = reshape(0.0:0.01:10.0, :, 1)
y = f.(X[:, 1])

function derivative_loss(tree, dataset::Dataset{T,L}, options, idx) where {T,L}
    # Column-major:
    X = idx === nothing ? dataset.X : view(dataset.X, :, idx)
    ∂y = idx === nothing ? dataset.y : view(dataset.y, idx)

    ŷ, ∂ŷ, completed = eval_grad_tree_array(tree, X, options; variable=true)

    !completed && return L(Inf)

    mse = sum(i -> (∂ŷ[i] - ∂y[i])^2, eachindex(∂y)) / length(∂y)
    return mse
end


model = SRRegressor(
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp],
    loss_function=derivative_loss,
    enable_autodiff=true,
    batching=true,
    batch_size=25,
    niterations=100,
)

mach = machine(model, X, y)

fit!(mach)

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

symbolic_eq = node_to_symbolic(eq, model)

[tomaklutfu]

Does y argument to machine accept a named tuple so that we can feed derivative info too? I want to search for both function for y and ensure that the derivative of the searched functions is also reasonable. Or maybe with a vector of a custom element type like Pairs or ForwardDiff's Dual type?

[MilesCranmer]

@tomaklutfu good question. Right now the MLJ interface doesn’t have the option of passing extra data. However, the Dataset type has the extra property which is an arbitrary named tuple, and is set up for allowing you to pass arbitrary extra data needed for a custom loss function. However for this you need to use equation_search directly rather than the MLJ interface which is obviously less convenient...

It shouldn’t be too hard to set up though if you’d be interested in trying? I will eventually get to it but if you make a PR it will happen quicker: I would gladly merge and help with any issues that come up.

[MilesCranmer]

For this to propagate, you will need to add an extra parameter to these functions:

https://github.com/MilesCranmer/SymbolicRegression.jl/blob/9bb4aca909c02f55195d1b33d742f6df398b8157/src/MLJInterface.jl#L126 (plus the two functions before it)

https://github.com/MilesCranmer/SymbolicRegression.jl/blob/9bb4aca909c02f55195d1b33d742f6df398b8157/src/SymbolicRegression.jl#L330

https://github.com/MilesCranmer/SymbolicRegression.jl/blob/9bb4aca909c02f55195d1b33d742f6df398b8157/src/SearchUtils.jl#L307

And also verify that the MLJ interface actually accepts this.

[MilesCranmer]

@tomaklutfu I have added your request to this PR: astroautomata/SymbolicRegression.jl#249. Could you try it out?

For example:

function derivative_loss(tree, dataset::Dataset{T,L}, options, idx) where {T,L}
    # Select from the batch indices, if given
    X = idx === nothing ? dataset.X : view(dataset.X, :, idx)

    # Evaluate both f(x) and f'(x), where f is defined by `tree`
    ŷ, ∂ŷ, completed = eval_grad_tree_array(tree, X, options; variable=true)

    !completed && return L(Inf)

    y = idx === nothing ? dataset.y : view(dataset.y, idx)
    ∂y = idx === nothing ? dataset.extra.∂y : view(dataset.extra.∂y, idx)

    mse_deriv = sum(i -> (∂ŷ[i] - ∂y[i])^2, eachindex(∂y)) / length(∂y)
    mse_value = sum(i -> (ŷ[i] - y[i])^2, eachindex(y)) / length(y)

    return mse_value + mse_deriv
end

where you pass the extra data as a NamedTuple:

    model = SRRegressor(;
        binary_operators=[+, -, *],
        unary_operators=[cos],
        loss_function=derivative_loss,
        enable_autodiff=true,
        batching=true,
        batch_size=25,
        niterations=100,
        early_stop_condition=1e-6,
    )
    mach = machine(model, X, y, (; ∂y=∂y))

[tomaklutfu]

Thanks for quick work. I will try as soon as I can and report the result. I tried making loss function a callable struct and giving derivative data as struct field without fiddling too much but the search is painstakingly slower if used.