reestimate constants of equation

[mikamyrseth]

Hi,
How can I reestimate the constants of an equation?
Lets say I split training data into a training and a validation set to pick the best equation, how would I then reestimate the constant of the specific equation with the whole training set?

[MilesCranmer]

If you get the JAX format with .jax(i), or the PyTorch format with .pytorch(i), for equation i, the constants are tunable parameters, so you can fit them afterwards.

If on the Julia (backend) side, you can call SymbolicRegression.ConstantOptimizationModule.optimize_constants, and pass a different dataset object than the one given for the equation search.

[mikamyrseth]

Great! Is it as simple as doing something like jax.fit() with new data or do I have to do something more? I haven't really used Jax before.

[MilesCranmer]

There are details here: https://astroautomata.com/PySR/options/#exporting-to-numpy-pytorch-and-jax but would need you to write some code with a jax or PyTorch optimization loop (or scipy optimize)

[MilesCranmer]

It might be good if the numpy export had this capability as well (https://github.com/MilesCranmer/PySR/blob/master/pysr/export_numpy.py). Right now it embeds the constants in the equation, but having the constants be tunable would allow for easier integration with scipy optimize.

(By the way, I very much welcome PRs, if you come up with an easier way to set this up!)

[mikamyrseth]

If anyone else is looking for this, a simple Jax implementation with mse could look like this.

import jax.numpy as jnp
import jax.scipy.optimize as jopt

jax_moddel = model.jax(best_equation)
jax_callablle = jax_moddel['callable']
jax_params = jax_moddel['parameters']

def loss(params, x, y):
        return jnp.mean((jax_callablle(x, params) - y) ** 2)

jax_params = jopt.minimize(
    fun=loss,
    x0=jax_params,
    args=(X_new.to_numpy(), Y_new.to_numpy()),
    method="BFGS",
    tol=0.000001
).x

prediction = jax_callablle(X_new.to_numpy(), jax_params)