Strange behaviour when using custom loss function

[TadeuNP]

Hi,
I was trying to introduce a penalty to solutions that don't satisfy a certain constraint, but for some reason I get very strange results if PySR runs for >= 100 iterations. If running for shorter than that, it seems to output normal results.

The custom loss is included below. The idea is to apply a penalty to the loss if the expression blows up when evaluated near zero.

 objective1 = 
 """ function my_custom_objective(tree, dataset::Dataset{T}, options) where {T<:Real}
    #Applies a penalty if points do not behave as expected.
 
    prediction, flag = eval_tree_array(tree, dataset.X, options)
 
    p1x = zeros(Float32, 1) .+ 0.0005f0
    p1y = zeros(Float32, 1) .+ 0.0005f0
    p1z = zeros(Float32, 1) .+ 0.0005f0
      
    concat = vcat(transpose(p1x[1:1]), transpose(p1y[1:1]),  transpose(p1z[1:1])) 
    p1 = eval_tree_array(tree, concat, options)[1][1] 
 
    
    penalty = 1
 
    if p1 >= 50
        penalty = 1.5 #changing to 1 seems to make it work?
    end 


    prediction_loss = sum((prediction .- dataset.y) .^ 2) / dataset.n * penalty
    return prediction_loss 
end  
"""

The data used should be here: x_data__.txt and y_data__.txt. The target I was aiming for (first column of y_data) is the function $$-0.5X + \frac{1.5}{1+Z^{10}}$$

where $X$ is the first column of x_data and $Z$ is the third column of the same file.

PySR_args = {"model_selection": "best",
        "niterations": 200,  #changing to <= 100 seems to fix it?
        "populations": 25,
        "binary_operators": ["+", "*", "/",], 
        "unary_operators": ["squared(x) = x^2"],
        'nested_constraints':{"squared": {"squared": 4} },
        "weight_optimize": 0.001, 
        "ncyclesperiteration":1000, 
        "warm_start":False, 
        "multithreading":False,  
        "enable_autodiff":False, 
         "extra_sympy_mappings": {"squared": lambda x: x**2 },
        "procs":2,   
        "maxsize":25, 
        "turbo":False} 

I'm also including a Hall of Fame that was seen in such a case: bizarre_hof.csv

I did not open this as a bug report because I'm not really sure if it is my faulty custom loss that is causing the issue or a PySR problem.

Thanks!

[MilesCranmer]

Two quick comments:

1. A simpler way to define the concat array in your code would be:

concat = fill(5f-4, 3, 1)

(which creates an array with shape (3, 1) and value $5 \times 10^{-4}$)

2. The second output of eval_tree_array is important, as it tells you when the evaluation encounters a NaN or not. So I would do:

p, completed = eval_tree_array(tree, concat, options)

penalty = (!completed || p[1] >= 50) ? 1.5 : 1.0

You should also check the first flag and return a large penalty if it is false.

3. These flags should return a constant large value, rather than a multiplicative value against the predictive loss. Because if there is a NaN, your prediction loss might be NaN, which could be bad. So I would do:

if !flag || !completed
    return 1000f0
end

and then just compute the multiplicative penalty from p[1] if it gets past that stage.

[TadeuNP]

I implemented all your suggestions and it seems to be working! Thanks! And thanks for the suggestion about the array!