PySR / test /test_torch.py
MilesCranmer's picture
Make tests non-random
51a6b05
raw
history blame
1.47 kB
import unittest
import numpy as np
import pandas as pd
from pysr import sympy2torch, get_hof
import torch
import sympy
class TestTorch(unittest.TestCase):
def setUp(self):
np.random.seed(0)
def test_sympy2torch(self):
x, y, z = sympy.symbols('x y z')
cosx = 1.0 * sympy.cos(x) + y
X = torch.tensor(np.random.randn(1000, 3))
true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
torch_module = sympy2torch(cosx, [x, y, z])
self.assertTrue(
np.all(np.isclose(torch_module(X).detach().numpy(), true.detach().numpy()))
)
def test_pipeline(self):
X = np.random.randn(100, 10)
equations = pd.DataFrame({
'Equation': ['1.0', 'cos(x0)', 'square(cos(x0))'],
'MSE': [1.0, 0.1, 1e-5],
'Complexity': [1, 2, 3]
})
equations['Complexity MSE Equation'.split(' ')].to_csv(
'equation_file.csv.bkup', sep='|')
equations = get_hof(
'equation_file.csv', n_features=2, variables_names='x1 x2 x3'.split(' '),
extra_sympy_mappings={}, output_torch_format=True,
multioutput=False, nout=1, selection=[1, 2, 3])
tformat = equations.iloc[-1].torch_format
np.testing.assert_almost_equal(
tformat(torch.tensor(X)).detach().numpy(),
np.square(np.cos(X[:, 1])), #Selection 1st feature
decimal=4
)