How to speed up custom loss function?

[DenisSvirin]

I have custom loss function ( https://gist.github.com/DenisSvirin/64952e83496968cfce9d117328670c04 ).

(edit - pasted below)

include("data/fcc_coords.jl")
include("data/potential.jl")
include("data/Training_r.jl")
include("data/Target_properties.jl")
include("data/Training_E.jl")
using Symbolics

@variables x1 x2 x3

function distance(v1::AbstractVector{T}, v2::AbstractVector{T})::Float64 where
{T<:Real}
    # v1: Cartisean coordinates of atom 1
    # v2: Cartisean coordinates of atom 2
    # -----
    # returns: Euqlidian distance between atom 1 and atom 2

    return sqrt(sum((v1 .- v2).^2))
end

# Define the smoothness penalty using mixed partial derivatives
function smoothness_penalty(x1, x2, x3, tree, options, values)
    # x1, x2, x3: variables of a tree
    # tree: tree of a function
    # options: options
    # values: Dictionary of values of variables
    # -----
    # returns: sum of squares of partial derivatives

    # Transform function from a tree to an experssion
    f_str = string_tree(tree, options)

    # Calculate penalty catching imaginary results, like: (-2) ^ 0.5
    penalty = 0
    try
        f = eval(Meta.parse(f_str))

        # Calculate Hessian of the function
        H = Symbolics.hessian(f, [x1, x2, x3])

        # Calcute final penalty as sum of squares of partial derivatives
        penalty = sum(substitute(H, values) .^ 2).val
    catch
        penalty = 1e6
    end

    return penalty
end

function count_non_zero_values(atom::AbstractVector{T})::Int8 where {T<:Real}
    # atom: Cartisean coordinates of the atom
    # -----
    # returns: Count of non zero coordinates

    return count(!iszero, atom)
end

mutable struct atom_env
    # n: number of atoms with similar coordinates
    # coords: coordinates of an atom [x, y, z]

    n::Float64
    coords::Array{Float64,1}
end

function add_r(d, v::AbstractVector{T}, key::Float64) where {K<:String, V<:Real,
T<:Real}
    # d: dictionary (key: distance to the atom, value: number of
    # atoms with the same distance)
    # key: distance to the atom
    # v: number of atoms with the same distance

    d_keys = keys(d)
    if key in keys(d)
        d[key].n += 2 ^ count_non_zero_values(v)
    else
        d[key] = atom_env(2 ^ count_non_zero_values(v), v)
    end
end

function get_distances(scale_factor::Union{T, Vector{T}}, r_cutoff::T = 7.0) where
{T<:Float64}
    # scale_factor: array of scaling coefficents for each axis ([x, y, z])
    # r_cutoff: Cutoff distance
    # -----
    # returns: dictionary of distances

    distance_dict = Dict()
    atomic_coordinates = relative_coordinates
    if scale_factor isa Number
         scale_factor = [scale_factor, scale_factor, scale_factor]
    end

    for atom_i in atomic_coordinates
        atom = atom_i .* scale_factor
        distance_ij = distance([0., 0., 0.], atom)

        if distance_ij != 0 && distance_ij <= r_cutoff
            add_r(distance_dict, atom, distance_ij)
        end
    end
    return distance_dict
end

# Creating array of distance dictionaries for each r in train_r
dist_arr = []
for i in train_r
    distances_i = get_distances(i, 7.0)
    push!(dist_arr, distances_i)
end

function my_custom_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
    # tree: evaluation tree
    # dataset: training dataset
    # options: options

    r_cutoff = 7.0
    tree.degree != 2 && return L(Inf)

    # Two body part
    E_two_body = tree.l
    E_two_body.degree != 2 && return L(Inf)

    # Three body part
    E_three_body = tree.r

    has_x1_in_E_three_body = any(node -> node.degree==0 && node.constant==false && node.feature==1, E_three_body)
    has_x2_in_E_three_body = any(node -> node.degree==0 && node.constant==false && node.feature==2, E_three_body)
    has_x3_in_E_three_body = any(node -> node.degree==0 && node.constant==false && node.feature==3, E_three_body)

    # Return huge loss if E_three_body doesn't have all three variables
    if !has_x1_in_E_three_body || !has_x2_in_E_three_body || !has_x3_in_E_three_body
        return L(Inf)
    end

    # E repulsion
    tree_rep = E_two_body.l

    # Return huge loss if tree_rep doesn't have at least one variable
    has_variable = any(node -> node.degree==0 && node.constant==false, tree_rep)
    if !has_variable
         return L(Inf)
    end

    # E attraction, power
    tree_att_p = E_two_body.r
    tree_att_p.degree != 2 && return L(Inf)

    # E attraction
    tree_att = tree_att_p.l

    # Return huge loss if tree_att doesn't have at least one variable
    has_variable = any(node -> node.degree==0 && node.constant==false, tree_att)
    if !has_variable
         return L(Inf)
    end

    # power
    power = tree_att_p.r
    !any(node -> node.degree==0 && node.constant==true, power) && return L(Inf)

    total_E_two_body = zeros(length(dist_arr))      # predicted total energy for two body part
    total_E_three_body = zeros(length(dist_arr))    # predicted total energy for three body part
    calculate_two_body_energy = true # flag to calculate two body part only 1 time
    flag = true                      # flag to ensure tree evaluation
    smooth_penalty_E_three_body = 0  # penalty for too high second derivatives

    for (i, dict_i) in enumerate(dist_arr)
        if calculate_two_body_energy
            N = []
            r = [[0, 0, 0]]
            for i in keys(dict_i)
                N = [N; dict_i[i].n]
                r = vcat(r, [[i,i,i]])
            end

            r = T.(hcat(r[2:end, :]...))

            eval_tree_att, finished_tree_att =  eval_tree_array(tree_att, r, options)
            eval_tree_rep, finished_tree_rep =  eval_tree_array(tree_rep, r, options)
            eval_power, finished_power =  eval_tree_array(power, r, options)

            total_E_two_body[i] = sum(N .* eval_tree_rep) - sum(N .* (eval_tree_att.^ 2) .^ eval_power)
            flag *= finished_tree_rep * finished_tree_att * finished_power
            calculate_two_body_energy = false
        end
        triplets = []
        N = []
        for atom_2 in range(start=1, stop=length(dict_i), step=1)
            N_x = []
            distances_x = [[0, 0, 0]]
            for atom_3 in range(start=atom_2, stop=length(dict_i), step=1)
                atom_12 = Float64.(keys(dict_i))[atom_2]
                atom_13 = Float64.(keys(dict_i))[atom_3]
                atom_23 = distance(dict_i[atom_12].coords, dict_i[atom_13].coords)

                if atom_23 != 0  && atom_12 <= r_cutoff && atom_13 <= r_cutoff && atom_23 <= r_cutoff
                    distances_x = vcat(distances_x, [[atom_12, atom_13, atom_23]])
                    N_x = [N_x; dict_i[atom_12].n * dict_i[atom_13].n]
                    vals = Dict(x1 => atom_12, x2 => atom_13, x3 => atom_23)
                    smooth_penalty_E_three_body += dict_i[atom_12].n * dict_i[atom_13].n * smoothness_penalty(x1, x2, x3, E_three_body, options, vals)
                end
            end # end of loop for atom_3

            triplets = vcat(triplets, distances_x[2:end, :])
            N = [N; N_x]

        end # end of loop for atom_2

        r = T.(hcat(triplets...))
        if length(r) != 0
            eval_E_three_body, finished_E_three_body = eval_tree_array(E_three_body, r, options)
            total_E_three_body[i] = sum(N .* eval_E_three_body)
            flag *= finished_E_three_body
        end

        calculate_two_body_energy = true
    end # end of loop for dict_i

    !flag && return L(Inf)

    diffs = total_E_two_body .+ total_E_three_body .- dataset.y

  return sum((diffs .^ 2)) / length(diffs ) + sum((diffs_dimer .^ 2) .* dimer_weights) / length(diffs_dimer) + smooth_penalty_E_three_body
end

I was following tips on your website, but still I had only 4 iterations in 22 hours. Is there way to speed it up?
Here is my model:

model = PySRRegressor(
    model_selection="accuracy",
    niterations=30000,
    maxsize=45,
    populations=32*3,
    ncyclesperiteration=4500,
    binary_operators=["+", "-","*", "/", "^"],
    constraints={'^': (4, 4), '*':(5, 5), '/':(2,2)},
    nested_constraints={"^":{"^":1}},
    full_objective=objective,
    weight_optimize=0.001,
    parsimony=0.0001,
   #procs=32,
   #cluster_manager="slurm",
    turbo = True,
    use_frequency=False,
    warm_start=True)

I have 23 training points and each cycle shouldn't be over 200 iterations.

[MilesCranmer]

I would try to benchmark it with BenchmarkTools.jl (https://github.com/JuliaCI/BenchmarkTools.jl), in pure Julia, with SymbolicRegression.jl. You can do this with, for example:

using SymbolicRegression
using BenchmarkTools

# For generating random trees:
using SymbolicRegression.MutationFunctionsModule: gen_random_tree_fixed_size

#= define objective function =#

# Set up dataset and options
options = Options(binary_operators=[+, *, -, /], unary_operators=[cos, exp])
nrows = 23
nfeatures = 5
T = Float64
dataset = Dataset(
    randn(T, nfeatures, nrows),
    randn(T, nrows)
)

# Benchmark:
@benchmark(
    my_custom_objective(tree, $dataset, $options),
    setup=(
        tree_size = 30;
        tree = gen_random_tree_fixed_size(tree_size, options, nfeatures, T)
    )
)

This will give you a sense of how fast it is. For this dataset size and tree size, the baseline median performance for the default objective is around 1.5 microseconds on my machine:

julia> function default_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
           pred, completed = eval_tree_array(tree, dataset.X, options)
           !completed && return L(Inf)
           return sum(i -> (pred[i] - dataset.y[i])^2, eachindex(pred))
       end
my_custom_objective (generic function with 1 method)

julia> # Benchmark:
       @benchmark(
           default_objective(tree, $dataset, $options),
           setup=(
               tree_size = 30;
               tree = gen_random_tree_fixed_size(tree_size, options, nfeatures, T)
           )
       )
BenchmarkTools.Trial: 10000 samples with 262 evaluations.
 Range (min … max):  267.813 ns … 18.438 μs  ┊ GC (min … max): 0.00% … 90.47%
 Time  (median):       1.507 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.632 μs ±  1.382 μs  ┊ GC (mean ± σ):  7.27% ±  7.91%

                              ▁▂▄▅▅▆█▇▆▇▆▆▄▃▂▁
  ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▂▃▄▅▆▇█████████████████▇▆▅▄▄▃▃▂▂▂▂▂▁▁▁ ▄
  268 ns          Histogram: frequency by time         2.31 μs <

 Memory estimate: 960 bytes, allocs estimate: 4.

So if you are way larger than this (which is what it sounds like) that there is likely some bottleneck in the code. You can continue to use @benchmark to see what changes improve performance.

Then my second comment is you should check out the page https://docs.julialang.org/en/v1/manual/performance-tips/ for common performance "gotchas" in Julia. In particular it seems like some of your code does not declare global variables as const such as @variables x y z, so this is probably one of the things slowing the code down (because Julia can't compile those variables into the function, as they might change in the future). So for example I would replace

@variables x y z

with (you can see what @variables is doing with @macroexpand – it returns [x, y, z] with some definitions for each)

const (x, y, z) = @variables x y z

The Julia discourse is also super helpful for optimizing code snippets: https://discourse.julialang.org/.

Cheers,
Miles

[DenisSvirin]

I'll check it out.
Thank you.