Building complex TemplateExpression structures

[MilesCranmer]

Copying @gm89uk from #756 for visibility in case others are wanting to learn more about the internals of TemplateExpressions –

Sorry @MilesCranmer a bit stuck, and with it being a new feature there isn't much yet in the forums to go off.

I'm working in Julia rather than PySR to permit global variables to be used in a custom loss function otherwise Essentially I have a categorial variable that I know with domain knowledge, has a linear effect on y with just one parameter.

I am interested in finding the expression that finds y, but without the noise of the parameter from the categorial variable. Although parametric expressions seems ideal, the actual problem is more complicated and template expressions is an otherwise better fit (if I can get it to work). I input the y and category variables as part of X so that I could access them in structure. The categorial variable is basically a vector of Ints, one for each group.

#=
structure = TemplateStructure{(:f, :c)}(
    ((; f, c), (x1, x2, x3, x4, x5, x6, y, category)) -> begin
        transformed_x1 = c(x1, x3, x4)
        o = f(transformed_x1, x2, x3, x4, x5, x6)
        residuals = y .- o
        unique_category = unique(category)
        mean_residuals_by_category = Dict(
            group => mean(residuals[category .== group]) for group in unique_category
        )
        o .-= [mean_residuals_by_category[category[i]] for i in eachindex(o)]
        return o
    end
)
=#


model = SRRegressor(
    niterations=1000000,
    binary_operators=[+,-,*,/],
    maxsize=40,
    bumper=true,
    turbo=true,
    populations=18,
    expression_type = TemplateExpression,
    expression_options = (; structure),
    population_size=30,
    parsimony = 0.01,
    batching=true,
)
mach = machine(model, X, y)
fit!(mach)

Although I think my approach is completely wrong, as structure is dealing with valid_vectors, I am really at a loss of how to get something like this to work.

[MilesCranmer]

I think you are almost there! The next thing would be to unpack the ValidVector with .x - which gives you the raw data.

Maybe this would work?

structure = TemplateStructure{(:f, :c)}(
    ((; f, c), (x1, x2, x3, x4, x5, x6, y, category)) -> begin
        transformed_x1 = c(x1, x3, x4)
        o = f(transformed_x1, x2, x3, x4, x5, x6)
        residuals = y - o  # Changed
        unique_category = unique(category.x)  # Changed
        _mean(_x) = sum(_x) / length(_x)  # Just in case it doesn't exist
        mean_residuals_by_category = Dict(
            group => _mean(residuals.x[category.x .== group])  # Changed
            for group in unique_category
        )
        offsets = [mean_residuals_by_category[category[i]] for i in eachindex(o.x)]  # Changed
        return ValidVector(o.x .- offsets, o.valid && residuals.valid)
    end
)

Basically each x1, x2, x3, ... is a ValidVector. This has a .x field which stores the raw data vector, and a .valid which stores whether the evaluation encountered any NaNs/Infs. This is because invalid values are very common in evolution, and it's faster to pass them through as bools rather than throw errors and catch them.

The ValidVector doubles a convenience wrapper around Vector to automatically perform vectorized operations. So you can just write - between two ValidVector objects, like x1 - x2 and it will do x1.x .- x2.x and x1.valid && x2.valid.

When you evaluate an expression like c(x1, x3, x4), the result is another ValidVector.

Changes:

• y .- o into y - o.
• category into category.x
• residuals into residuals.x within the Dict
• Added the _mean function (mean doesn't exist in Julia due to the ambiguity between sample/population means)
• Changed the o .-= offsets into a o.x .- offsets
• I also did an explicit ValidVector as a return value just as a demonstration, and so you could include the .valid part of the residual ValidVector (though I think here it's probably not necessary, but just a good habit!)

[gm89uk]

I think I have it working now:
In your code, I renamed DynamicAutodiff to DynamicDiff.

So for example code to check monotonicity of x1 to x5 to y:

using DynamicDiff: DynamicDiff as DA
using SymbolicRegression.TemplateExpressionModule: ArgumentRecorder
DA.D(f::ArgumentRecorder, _::Integer) = f
structure = TemplateStructure{(:f,)}(
  ((; f), (x1, x2, x3, x4, x5, x6, y, category)) -> begin
    o = f(x1, x2, x3, x4, x5, x6)
    if !o.valid
        return ValidVector(fill(1e9, length(o.x)), false)
    end
    grad_matrix = hcat(
        DA.D(f, 1)(x1, x2, x3, x4, x5, x6).x, 
        DA.D(f, 2)(x1, x2, x3, x4, x5, x6).x,  
        DA.D(f, 3)(x1, x2, x3, x4, x5, x6).x,
        DA.D(f, 4)(x1, x2, x3, x4, x5, x6).x,  
        DA.D(f, 5)(x1, x2, x3, x4, x5, x6).x, 
    )
   #e.g. x6 does not need monotonicity check
    if any(row -> !(all(row .>= 0) || all(row .<= 0)), eachrow(grad_matrix))
        return ValidVector(fill(1e9, length(o.x)), false)
    end
    residuals = o - y
    unique_groups = unique(Int.(category.x))
    _mean(_x) = sum(_x) / length(_x)
    mean_residuals_by_group = Dict(group => _mean(residuals.x[category.x .== group]) for group in unique_groups)
    offsets = [mean_residuals_by_group[category.x[i]] for i in eachindex(o.x)]
    return ValidVector(o.x .- offsets, o.valid && residuals.valid)
  end
)

This works nicely, although it is a bit slower than using a custom loss function and using eval_grad_tree_array(tree, X, options; variable=true), it is much faster than looping around with eval_diff_tree_array(tree, X, options,varNum).

Hopefully I have the right idea there, unless there is a quicker way, and thanks again for working on this so fast.

EDIT: early exit if o is invalid to avoid unnecessary processing. I'm wondering if operators are limited to +,-,*,/, monotonicity can probably be determined symbolically!

[gm89uk]

On another note relevant to complex TemplateExpressions, users may standardise their variables on input, but domain knowledge means that those functions are unlikely to work with standardised variables. The workaround for me has been to feed in the variables I need as non-standardised variables:

((; f), (x1, x2, x3, x4, x5, x6, y, category,x1_NS,x2_NS,x3_NS)) -> begin #NonStandardised

then using the _NS variables within the evaluation. This works, but I suppose is not memory efficient. I was wondering if you had any suggestions regarding that aspect?

[MilesCranmer]

I think you could potentially use eval_grad_tree_array within the TemplateStructure though I haven't tried. You can extract the tree with DE.get_tree(f) (where DE is DynamicExpressions)¹ and then DE.get_operators(f) to get the operators². Then the call is:

tree = DE.get_tree(f)
operators = DE.get_operators(f)
X = stack((x1.x, x2.x, #= ... =#), dims=1)
result, grad, complete = eval_grad_tree_array(tree, X, operators)

Footnotes

1. All it's doing is calling f.tree. The reason to prefer the get_tree is that it's more general in case there is some other type of expression used here in the future. ↩

2. All it is doing is calling f.metadata.operators. ↩

[MilesCranmer]

The main reason the D operator was introduced was to do higher order differentiation efficiently. But if you just want a regular gradient then maybe eval_grad_tree_array is faster.

[MilesCranmer]

Also for questions about the efficiency stuff, maybe create a new discussion thread? I'm realising that this reply thread is long enough it might not be easily searchable for other users