Python variables in julia custom loss function

[AndrewHutani]

Hi,

I am trying to create a custom loss function based on lower and upper bounds of the dataset. I have these bounds in a python variable, but I am not sure how to use this variable in the julia loss function or whether it would be even possible. My current implementation (with some omissions) is as follows:

upper_bounds = np.max(y,axis=0)
lower_bounds = np.min(y, axis=0)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.10)

model = PySRRegressor(
    niterations=100,  # < Increase me for better results
    binary_operators=["+", "*", '-', '/'],
    unary_operators=[
        "cos",
        "sin",
    ],
    procs=5,
    populations=30,
    model_selection="best",
    maxsize = 15,
    ncyclesperiteration = 10000,
    nested_constraints={"sin": {"sin": 0, "cos": 0}, "cos": {"sin": 0, "cos": 0}},
    weight_optimize = 0.001,
    turbo = True,
    progress = True,
    loss="""
        function bounded_loss(x, y, lower_bounds, upper_bounds)
        loss = abs(x - y)
        for i in 1:{num_features}
            if x[i] < lower_bounds[i]
                loss += 1 * (lower_bounds[i] - x[i])^2
            elseif x[i] > upper_bounds[i]
                loss += 1 * (x[i] - upper_bounds[i])^2
            end
        end
        return loss
    end
    """

The parameters I want to use in the custom loss function are lower_bounds and upper_bounds. Any help would be greatly appreciated!

[MilesCranmer]

Hi @AndrewHutani,

This is definitely possible. However loss expects elementwise loss only (e.g., loss="my_loss(prediction, target) = (prediction - target)^2" for a single scalar prediction and single scalar target). Instead you should use full_objective. The example of this is here: https://astroautomata.com/PySR/api/#the-objective. Do you want to try with that and see if you can get it to work?

To get your Python lower and upper bounds into the loss, I would do something with string interpolation and just put them into the string like:

full_objective=f"""
    function eval_loss(tree, dataset, options)
        lower_bounds = [{", ".join(map(str, np.min(y, axis=0)))}]
        upper_bounds = [{", ".join(map(str, np.max(y, axis=0)))}]
        prediction, flag = eval_tree_array(tree, dataset.X, options)
        if !flag
            return eltype(dataset.X)(Inf)
        end
        return sum((prediction .- dataset.y) .^ 2) / dataset.n
    end
"""

(But modifying the final line with the sum to use your custom loss)

Cheers!
Miles

[MilesCranmer]

No, eval_tree_array returns a 1D vector, not a 2D matrix. The objective function is called independently for each column of y.

One way you could indicate which column of y is being calculated is by passing through information in the first row of X and y. i.e., y[0, :] == [1, 2, 3, 4] where you have 4 columns in that example. Then eval_loss could simply see what y[0] is (in Julia this would be y[1]) and use that to find the corresponding upper/lower limit? (Ensuring that you run convert(Int, y[1]) to get an integral result.)

[AndrewHutani]

I'm a bit confused, because when I just used prediction .- dataset.y I got an error saying something along the lines of "could not broadcast 3839 to 6" (3839 is the size of the total dataset and 6 is the number of outputs). I thought this meant that eval_tree_array returns a (3839, 6) array.

Assuming dataset.X is a the complete training set of (3839, 6) (It is complete coincidence that there are the same amount of input features as output features), what size would the prediction be? This could maybe help me understand what exactly prediction is in this case.

[MilesCranmer]

Assuming dataset.X is a the complete training set of (3839, 6) (It is complete coincidence that there are the same amount of input features as output features), what size would prediction be

it would be size (6,). In Julia, the features come first and rows come second.

[MilesCranmer]

(If you pass in (3800, 6) data to PySR, it would automatically be transposed to (6, 3800) in Julia)

[AndrewHutani]

If prediction would be size (6,), I am even more confused since using

...
        if sum(prediction .- lower_bounds) < 0
            loss += 100
        end
        return loss
    end
    """,

results in the following error (only lower bounds are checked to isolate the error):

RuntimeError: <PyCall.jlwrap (in a Julia function called from Python)
JULIA: DimensionMismatch: arrays could not be broadcast to a common size; got a dimension with lengths 3839 and 6

This lead me to believe that prediction is not of size 6, but I am a bit lost on what exactly is happening here.