How to create a custom loss function with a Riemann-sum approximation of the integral

[unary-code]

Hi,

I am creating a custom objective function and I want to calculate an approximation of the value of the integral of the guessed function over the domain x=0 to x=1.
Purpose of doing this is so that I can scale my function's values by this amount so that the definite integral of the re-scaled guessed function over the domain x=0 to x=1 becomes 1.

Assume there is only one independent variable, x.

objective = """
function eval_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    # Evaluate numerator:
    prediction, flag = eval_tree_array(tree, dataset.X, options)
    !flag && return L(Inf)
    println("TREE")
    
    num_rectangles = 1000
    
    norm_constant = 0
    prev_expr_val = -1
    for i in 0:num_rectangles:
        cur_float_num = i / (1.0 * num_rectangles)
        abstract_array = [cur_float_num]
        cur_expr_val = eval_tree_array(tree, abstract_array, options)
        
        if i > 0:
            cur_trapezoid_area = (cur_expr_val + prev_expr_val) * / (2.0 * num_rectangles)
            norm_constant += cur_trapezoid_area
        
        prev_expr_val = cur_expr_val
    
    prediction = prediction / norm_constant
    
    print("HELLO WORLD")
    
    return -1 * sum(log.(prediction.*prediction)) / (2*length(prediction))
end
"""

model = PySRRegressor(
    niterations=40,  # < Increase me for better results
    binary_operators=["+", "*", "-", "/"],
    unary_operators=[
        "cos",
        "exp",
        "sin",
        "inv(x) = 1/x",
        # ^ Custom operator (julia syntax)
    ],
    extra_sympy_mappings={"inv": lambda x: 1 / x},
    # ^ Define operator for SymPy as well
    #loss="loss(prediction, target) = (prediction - target)^2",
    full_objective=objective
    # ^ Custom loss function (julia syntax)
)

Is this close to how you would approximate the definite integral of the guessed function (the guessed function is represented by the tree "tree" here) in a custom Julia objective function?
Thank you, very much appreciated.

[MilesCranmer]

This looks okay to me. Some notes:

1. You should probably return early with L(Inf) if any of the cur_expr_val are negative. Or simply treat the square of the expression as the probability to guarantee this.
2. You are mixing Python and Julia syntax. Instead of

 if i > 0:
    cur_trapezoid_area = (cur_expr_val + prev_expr_val) * / (2.0 * num_rectangles)
    norm_constant += cur_trapezoid_area

you should write

 if i > 0
    cur_trapezoid_area = (cur_expr_val + prev_expr_val) / (2.0 * num_rectangles)
    norm_constant += cur_trapezoid_area
end

(also there's an error in the expression)

Similarly for the for loop.

3. Note you can also use external packages in case you don't want to do this from scratch. See https://astroautomata.com/PySR/examples/#7-julia-packages-and-types for an example.
4. It would be faster if you vectorize the calculation along the domain at once. But maybe don't worry about that unless you need to, as it would require a bit of work.

[unary-code]

Hi,

I just ran this objective function:

objective =
function eval_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    # Evaluate numerator:
    prediction, flag = eval_tree_array(tree, dataset.X, options)
    !flag && return L(Inf)
    println("TREE")
    
    num_rectangles = 1000
    
    # You make sure that this object is of type Matrix, not Vector and not Array.
    evenly_spaced_numbers = reshape([LinRange(0, 1, num_rectangles + 1);], num_rectangles + 1, 1)
    println(typeof(evenly_spaced_numbers))
    println(size(evenly_spaced_numbers))
    println(evenly_spaced_numbers[1:5, 1])
    
    prediction_on_evenly_spaced_numbers, flag_on_evenly_spaced_numbers = eval_tree_array(tree, evenly_spaced_numbers, options)
    
    println(length(prediction_on_evenly_spaced_numbers))
    println(typeof(prediction_on_evenly_spaced_numbers))
    println(prediction_on_evenly_spaced_numbers[1])
    
    prediction_on_evenly_spaced_numbers = prediction_on_evenly_spaced_numbers .* prediction_on_evenly_spaced_numbers
    
    println(length(prediction_on_evenly_spaced_numbers))
    println(typeof(prediction_on_evenly_spaced_numbers))
    println(prediction_on_evenly_spaced_numbers[1])
    
    norm_constant = 0
    prev_expr_val = -1
    for i in 0:num_rectangles
        cur_expr_val = prediction_on_evenly_spaced_numbers[i+1]
        
        if (i > 0)
            cur_trapezoid_area = (cur_expr_val + prev_expr_val) / (2.0 * num_rectangles)
            norm_constant += cur_trapezoid_area
        end
        
        prev_expr_val = cur_expr_val
    end
    
    prediction = prediction / norm_constant
    
    println("HELLO WORLD")
    
    return -1 * sum(log.(prediction.*prediction)) / (2*length(prediction))
end

The printed output of these three prints (before prediction_on_evenly_spaced_numbers .* prediction_on_evenly_spaced_numbers)

println(length(prediction_on_evenly_spaced_numbers))
println(typeof(prediction_on_evenly_spaced_numbers))
println(prediction_on_evenly_spaced_numbers[1])

is

1
Vector{Float64}
1000.0

Why is the "prediction_on_evenly_spaced_numbers" output of eval_tree_array(tree, evenly_spaced_numbers, options) only have 1 element?
I would expect it to have 1001 elements, because "evenly_spaced_numbers" has 1001 elements.

Thanks.

[MilesCranmer]

This is a common gotcha. In Julia, similar to Fortran, arrays are stored in column-major order. This means that columns come first, rows come second (i.e., functions operating on matrices usually expect data to come in this format). So for the function

eval_tree_array(tree, X, options)

X has to have columns (features) first, and rows second. So that X[feature, row] is your data.

What you wrote is

evenly_spaced_numbers = reshape([LinRange(0, 1, num_rectangles + 1);], num_rectangles + 1, 1)

which looks to be X[row, feature]. This is row-major order which is used in Python/C++, but incompatible with Julia/Fortran.

(the backend should really try to automatically catch when users do this and throw an error.)

[unary-code]

Oh okay, thanks for this suggestion. I fixed the order so now the code is "reshape([LinRange(0, 1, num_rectangles + 1);], 1, num_rectangles + 1)", but now the issue is that when I call model.fit(X, y), it takes forever (more than 10 minutes).

This was the only change I did, but just for your reference, the objective function is now

objective = """
function eval_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    # Evaluate numerator:
    prediction, flag = eval_tree_array(tree, dataset.X, options)
    !flag && return L(Inf)
    println("TREE")
    
    num_rectangles = 1000
    
    # You make sure that this object is of type Matrix, not Vector and not Array.
    evenly_spaced_numbers = reshape([LinRange(0, 1, num_rectangles + 1);], 1, num_rectangles + 1)
    println(typeof(evenly_spaced_numbers))
    println(size(evenly_spaced_numbers))
    println(evenly_spaced_numbers[1, 1:5])
    
    prediction_on_evenly_spaced_numbers, flag_on_evenly_spaced_numbers = eval_tree_array(tree, evenly_spaced_numbers, options)
    
    println(length(prediction_on_evenly_spaced_numbers))
    println(typeof(prediction_on_evenly_spaced_numbers))
    println(prediction_on_evenly_spaced_numbers[1])
    
    prediction_on_evenly_spaced_numbers = prediction_on_evenly_spaced_numbers .* prediction_on_evenly_spaced_numbers
    
    println(length(prediction_on_evenly_spaced_numbers))
    println(typeof(prediction_on_evenly_spaced_numbers))
    println(prediction_on_evenly_spaced_numbers[1])
    
    norm_constant = 0
    prev_expr_val = -1
    for i in 0:num_rectangles
        cur_expr_val = prediction_on_evenly_spaced_numbers[i+1]
        
        if (i > 0)
            cur_trapezoid_area = (cur_expr_val + prev_expr_val) / (2.0 * num_rectangles)
            norm_constant += cur_trapezoid_area
        end
        
        prev_expr_val = cur_expr_val
    end
    
    prediction = prediction / norm_constant
    
    println("HELLO WORLD")
    
    return -1 * sum(log.(prediction.*prediction)) / (2*length(prediction))
end
"""

What are some possible reasons why model.fit(X, y) takes forever?
Is there a way to get model.fit(X, y) to show you the output of print statements within the custom objective function, before model.fit(X, y) finishes, so that I can understand what is being run?
Is there a way to force model.fit(X, y) to finish, so that you can see the output of your debug print statements (that you put inside the lsos function) so that you can debug?

[MilesCranmer]

Not sure I 100% understand the question but a few things look to be slowing it down:

1. the print statements,

2. because you are doing 1000x more evaluations per expression,

3. perhaps the biggest issue is that there are some type instabilities in your code — eg norm_constant starts off being an integer but then changes to a float. That kind of thing can really slow down Julia code as the compiler can’t do its magic. To fix it you should do norm_constant = T(0) so that it starts off being of type T (which is the same type as the evaluation output. Same issue for prev_expr_val.

[MilesCranmer]

Also the 2.0 should be wrapped with a T as well.