This repository contains code for the paper ODEformer: symbolic regression of dynamical systems with transformers.

This package is installable via pip:

pip install odeformer

We include a small notebook that loads a pre-trained model you can play with:

Import the model in a few lines of code:

import odeformer
from odeformer.model import SymbolicTransformerRegressor
dstr = SymbolicTransformerRegressor(from_pretrained=True)
model_args = {'beam_size':50, 'beam_temperature':0.1}
dstr.set_model_args(model_args)

Basic usage:

import numpy as np
from odeformer.metrics import r2_score

times = np.linspace(0, 10, 50)
x = 2.3*np.cos(times+.5)
y = 1.2*np.sin(times+.1)
trajectory = np.stack([x, y], axis=1)

candidates = dstr.fit(times, trajectory)
dstr.print(n_predictions=1)
pred_trajectory = dstr.predict(times, trajectory[0])
print(r2_score(trajectory, pred_trajectory))

To launch a model training with additional arguments (arg1,val1), (arg2,val2): python train.py --arg1 val1 --arg2 val2

All hyper-parameters related to training are specified in parsers.py, and those related to the environment are in envs/environment.py.

To launch evaluation, please use the flag reload_checkpoint to specify in which folder the saved model is located: python evaluate.py --reload_checkpoint XXX

This repository is licensed under MIT licence.