Spaces:
Running
Running
MilesCranmer
commited on
Commit
•
47dbec6
1
Parent(s):
b4cb407
Add feature selection for PyTorch
Browse files- test/test_torch.py +18 -0
test/test_torch.py
CHANGED
@@ -4,6 +4,7 @@ import pandas as pd
|
|
4 |
from pysr import sympy2torch, PySRRegressor
|
5 |
import torch
|
6 |
import sympy
|
|
|
7 |
|
8 |
|
9 |
class TestTorch(unittest.TestCase):
|
@@ -137,3 +138,20 @@ class TestTorch(unittest.TestCase):
|
|
137 |
np.sin(X[:, 1]),
|
138 |
decimal=4,
|
139 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
4 |
from pysr import sympy2torch, PySRRegressor
|
5 |
import torch
|
6 |
import sympy
|
7 |
+
from functools import partial
|
8 |
|
9 |
|
10 |
class TestTorch(unittest.TestCase):
|
|
|
138 |
np.sin(X[:, 1]),
|
139 |
decimal=4,
|
140 |
)
|
141 |
+
|
142 |
+
def test_feature_selection(self):
|
143 |
+
X = pd.DataFrame({f"k{i}": np.random.randn(1000) for i in range(10, 21)})
|
144 |
+
y = X["k15"] ** 2 + np.cos(X["k20"])
|
145 |
+
|
146 |
+
model = PySRRegressor(
|
147 |
+
unary_operators=["cos"],
|
148 |
+
select_k_features=3,
|
149 |
+
early_stop_condition=1e-5,
|
150 |
+
)
|
151 |
+
model.fit(X.values, y.values)
|
152 |
+
torch_module = model.pytorch()
|
153 |
+
|
154 |
+
np_output = model.predict(X.values)
|
155 |
+
torch_output = torch_module(torch.tensor(X.values)).detach().numpy()
|
156 |
+
|
157 |
+
np.testing.assert_almost_equal(np_output, torch_output, decimal=4)
|