Dimensional constraints in custom 'loss_function_expression'

[ibengtsson]

Hi!

I know that dimensional constraints aren't working for template expressions (and I understand why it isn't trivial or even makes sense to make it available for the general case), but would it be possible to make use of it in a custom loss function?

In my case for instance, I look for equations of the form:

y = f(...) * exp(g(...) / h(...)),

so I know that f should have the same dimension as y, whilst g and h should have the same dimension dividing to 1.

Any advice on where to alter the backend to make use of this information, if possible, would be greatly appreciated!

[MilesCranmer]

The way it works internally is in LossFunctions.jl and DimensionalAnalysis.jl. There is a part in LossFunctions that looks like this:

    if regularization
        loss_val += dimensional_regularization(tree, dataset, options)
    end

which calls out to

function dimensional_regularization(
    tree::Union{AbstractExpression{T},AbstractExpressionNode{T}},
    dataset::Dataset{T,L},
    options::AbstractOptions,
) where {T<:DATA_TYPE,L<:LOSS_TYPE}
    if !violates_dimensional_constraints(tree, dataset, options)
        return zero(L)
    end
    return convert(L, something(options.dimensional_constraint_penalty, 1000))
end

which then is unpacked using these:

function violates_dimensional_constraints(
    tree::AbstractExpression, dataset::Dataset, options::AbstractOptions
)
    return violates_dimensional_constraints(get_tree(tree), dataset, options)
end
function violates_dimensional_constraints(
    tree::AbstractExpressionNode, dataset::Dataset, options::AbstractOptions
)
    X = dataset.X
    return violates_dimensional_constraints(
        tree, dataset.X_units, dataset.y_units, (@view X[:, 1]), options
    )
end

So I think you could just pass each f, g, and h to the function violates_dimensional_constraints and specify the units you are requesting for each of them. The (@view X[:, 1]) is just used as an example input by the way.

This function sits in SymbolicRegression.DimensionalAnalysisModule.

So it might look like:

function my_loss(ex::TemplateExpression, dataset::Dataset{T,L}, options) where {T,L}
    prediction, complete = eval_tree_array(ex, dataset, options)
    if !complete
        return L(Inf)
    end

    f = ex.trees.f
    g = ex.trees.g
    h = ex.trees.h
    
    checker(args...) = SymbolicRegression.DimensionalAnalysisModule.violates_dimensional_constraints(args...)

    f_violates = checker(
        f,
        [u"km/s", u"m/s", u"kg"],  #= X_units; the input dimensions to f =#
        [u"1"], #= y_units; The expected output dimensions =#
        ones(T, 3),  #= Example input with 3 features =#
        options
    )
    g_violates = checker(
        g,
        [u"K"],  #= Different input units! =#
        nothing,  #= No requirement for output units =#
        ones(T, 1),  #= Example with 1 input feature =#
        options
    )

Note that if f takes 3 arguments in the template expression, you would pass 3 units to the violates_dimensional_constraints function. And so on.

Note that the u"km/s" is DynamicQuantities syntax. So you would need to load that package.

These could be different between f, g, and h. You can set y_units=nothing if you don't care about the output units.

You would then pass this to loss_function_expression.

If you want to deal with the units of f flowing into the units of g, you will need to see how the unit propagation works here: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/c3a974d8dfebf3ead2413fc670d9b86ef4a1c414/src/DimensionalAnalysis.jl#L200-L221. A "wildcard" basically just means that there is flexibility about units. e.g., the expression [const] * x1 has a wildcard, because that constant could take any form. But the expression [const] * x1 + x2 does not have a wildcard - the units are fixed by x2.

Hope this helps!

[ibengtsson]

Thank you so much for the detailed reply, I'll give it a go!