Use Python Class inside a custom loss function.

[saimani5]

I have a slightly unconventional use case for PySR. Essentially, I want to utilize the genetic algorithm (GA) to optimize complex processes whose input is a 1D trajectory in time (or similar). An example of this could be finding a charging trajectory that charges a battery in the shortest possible time. Here's how I plan to use Symbolic Regression to solve this problem using a custom loss function:

# Define the loss function as a string of Julia code
custom_loss = """

using PythonCall
mymodule = pyimport("BatterySIM")

function custom_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    predictions, completed = eval_tree_array(tree, dataset.X, options)
    if !completed
        return L(Inf)
    end
    
    # Code that should be run in python
    time = mymodule(np.array(predictions, dtype=np.float32)).time()

    loss = pyconvert(Float32, loss_value)
    
    return L(loss)
end
"""

Here, BatterySIM is an example for a class that runs the simulation of a physical system in Python. I have run these simulations across individual CPU cores using embarrassingly parallel computation.

I wrote a custom loss function to test this which calls a python class that evaluates the RMSE between the sampled expression and the ground truth curve store as an attribute inside the class. Every time a python class is called inside the custom loss function, the kernel restarts.

Any help or suggestion on this are highly appreciated. @MilesCranmer, if this peaks your interest, we can discuss this further.

[MilesCranmer]

Sounds cool, are there any issues you’ve run into? Or just starting it?