Custom power operator with exponents between 0 and 1 in PySR

[ma-sadeghi]

Hi, thanks for the great package!

We’d like to define a custom power operator in PySR such that the exponent is always constrained to lie between 0 and 1. For example, we want expressions like x^a, where a ∈ [0, 1].

Is there a recommended way to implement such a constraint?

Thanks!

[MilesCranmer]

Maybe something like

function my_loss(ex, dataset::Dataset{T,L}, options) where {T,L}
    INDEX_OF_POW = 5  # Note: Julia indexes at 1, so 5 would be the last index if you have 5 operators
    number_invalids = 0
    # each expression within the template expression
    tree = get_tree(ex)
    for node in tree
        # Each node in this expression tree
        if node.degree == 2 && node.op == INDEX_OF_POW
            right_subtree = node.r
            for inner_node in right_subtree
                # Each node within the right argument of ^
                is_variable = inner_node.degree == 0 && !inner_node.constant
                is_constant = inner_node.degree == 0 && inner_node.constant
                if is_variable
                    # prevent anything other than a constant
                    number_invalids += 1
                elseif is_constant
                    val = inner_node.val
                    bad_constant_value = val > 1 || val < 0
                    if bad_constant_value
                        number_invalids += 1
                    end
                else
                    # expressions also disallowed
                    number_invalids += 1
                end
            end
        end
    end
    if number_invalids > 0
        # Return early, but scale by *number* of invalids (rather than just Inf), so that
        # the genetic algorithm knows the "direction" to evolve the expression
        big_penalty = 10_000
        return L(big_penalty * number_invalids / length(prediction))  # (Divide by length(prediction) so it scales with batch)
    end

    predictions, is_valid = eval_tree_array(ex, dataset.X, options)
    if !is_valid
        return L(Inf)
    end
    squared_error = sum(
        i -> (predictions[i] - dataset.y[i])^2,
        eachindex(predictions)
    )
    mse = squared_error / length(predictions)
    return L(mse)
end

and then pass this to loss_function_expression (as a string in Python, or just my_loss if in Julia)