Reduce the running-time with custom objective

[unary-code]

Hey,

Here is my custom objective function:

objective = """
import Pkg
Pkg.add("Roots")
Pkg.add("ForwardDiff")
Pkg.add("QuadGK")

using Roots
using ForwardDiff
using QuadGK
function eval_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    MAX_THRESHOLD_ALLOWED = 10000
    

    prediction, flag = eval_tree_array(tree, dataset.X, options)
    !flag && return L(Inf)
    
    # NOTE: reshape([x], 1, 1) should give same dimensions
    # as reshape([x;], 1, 1).
    f(x) = 1 / ((eval_tree_array(tree, reshape([Float32(x)], 1, 1), options))[1][1])
    
    fp(x) = ( -1 / ((eval_tree_array(tree, reshape([Float32(x)], 1, 1), options))[1][1])^2 * (eval_diff_tree_array(tree, reshape([Float32(x)], 1, 1), options, 1))[2][1])
    
    found_a_root = true
    roots = -3
    try
        global roots = find_zero((f, fp), 0.3, Roots.Newton(), maxevals=100)
    catch
        global found_a_root = false
        # This means that "Convergence Failed"
    end
    
    # If the function f(x) represented by the tree "tree"
    # has a value f(x) close to positive infinity or negative infinity
    # when x is in [0, 1], then return L(Inf).
    if found_a_root
        if (roots >= 0 && roots <= 1)
            return L(Inf)
        end
    end
    
    global integral = 1
    err = -3
    try
        my_tuple = quadgk(x -> (eval_tree_array(tree, reshape([Float32(x)], 1, 1), options))[1][1], 0, 1, rtol=1e-8)
        global integral = my_tuple[1]
        global err = my_tuple[2]
    catch
        return L(Inf)
    end
    
    
    
    num_points = length(prediction)
    
    for i in 1:num_points
        cur_expr_val = prediction[i]
        if (cur_expr_val > MAX_THRESHOLD_ALLOWED)
            return L(Inf)
        end
        if (cur_expr_val < 0)
            return L(Inf)
        end
    end
    
    norm_constant = integral
    
    actual_probs = (prediction) / norm_constant
    
    final_loss = exp(-1 * sum(log.(actual_probs)) / (length(prediction)))
    return final_loss
end
"""

Here is how I defined my model:

model = PySRRegressor(
    niterations=2,  # < Increase me for better results
    binary_operators=["+", "*", "-", "/"],
    full_objective=objective,
    enable_autodiff=True
)

First, I verified that this objective function works correctly using the julia helper object to evaluate this objective function on two example trees, given one example dataset X.

jl = pysr.julia_helpers.init_julia(julia_kwargs={"threads": 8, "optimize": 3})

Then, I ran model.fit(X, y):

However, it's been over 25 minutes and the only thing printed out to my terminal is this (which was printed like 1 or 2 minutes after I ran "model.fit(X, y)"):

Compiling Julia backend...

/opt/anaconda3/envs/pysr_env/lib/python3.10/site-packages/pysr/julia_helpers.py:208: UserWarning: Your system's Python library is static (e.g., conda), so precompilation will be turned off. For a dynamic library, try using `pyenv` and installing with `--enable-shared`: https://github.com/pyenv/pyenv/blob/master/plugins/python-build/README.md#building-with---enable-shared.
  warnings.warn(
   Resolving package versions...
  No Changes to `~/.julia/environments/pysr-0.14.3/Project.toml`
  No Changes to `~/.julia/environments/pysr-0.14.3/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/pysr-0.14.3/Project.toml`
  No Changes to `~/.julia/environments/pysr-0.14.3/Manifest.toml`
   Resolving package versions...
  No Changes to `~/.julia/environments/pysr-0.14.3/Project.toml`
  No Changes to `~/.julia/environments/pysr-0.14.3/Manifest.toml`
Started!

On my past runnings of model.fit(X, y), it did finish within 5 minutes or less.
I think the only thing I changed was that I changed "integral = 1" to "global integral = 1", which means that now, I am actually updating the value of "integral" (and "norm_constant") correctly.
Before, the value of "integral" (and "norm_constant") was always 1, because I was not actually updating the value of integral.

So, my main question is:

1. What parameters in the construction of the PySRRegressor model can I change to reduce runtime of model.fit(X, y)?
2. When does model.fit(X, y) know when to stop? Is it when the loss of the guessed function becomes less than a arbitrary small number like less than 0.02 ?
3. Is there some useful metric that I can force model.fit(X, y) to print out while it is running, so that I can know how much more time it is going to take to finish running?

[MilesCranmer]

1. In general it sounds the objective is slowing things down, rather than anything in the hyperparameters. i.e., you should try to speed up the objective, or, if not possible, use more compute. To speed up the objective I recommend reading https://docs.julialang.org/en/v1/manual/performance-tips/.
   • I also note your rtol is super high in quadgk. I would not expect to reach that number if you still want it to be fast at evaluation. I would try to limit the number of steps and return a big loss if it can't reach a moderate rtol.
2. By default this is determined by niterations, which is an outer loop in the evolutionary code. It does not quit, regardless of the loss, until it reaches those number of iterations. For other stopping methods you can see the stopping criteria parameters: https://astroautomata.com/PySR/api/#stopping-criteria
3. It should do this by default – it displays a progress bar when you are running with a time estimate for completion. Not sure why this isn't visible to you – maybe are you running in Jupyter? Unfortunately the progress bar is not yet working in Jupyter, so it instead prints updated expressions every 5 seconds (once the first iteration has passed). Maybe try running in IPython, with progress=true to see the progress bar?

Since you are doing a lot of editing Julia code, it might even make sense to work directly in Julia. At least until you speed things up and can move back to Python.

As it happens I just released an improved API for the backend which makes it much easier to use via Julia sklearn equivalent: MLJ. More info here: https://github.com/MilesCranmer/SymbolicRegression.jl/#mlj-interface

Example:

import SymbolicRegression: SRRegressor
import MLJ: machine, fit!, predict, report

data = (x=randn(100), y=randn(100))
y = @. 2 * cos(data.x * 12) + data.y ^ 2 - 2
# This also works: data = randn(100, 2)

model = SRRegressor(
    niterations=100,
    binary_operators=[+, *, /, -],
    unary_operators=[exp, cos],
    parallelism=:multithreading,
)

mach = machine(model, data, y)

# Train model
fit!(mach)

# View expressions
report(mach)

Just note that the full_objective in PySRRegressor is called loss_function in SRRegressor (sorry; will eventually make this consistent).