Spaces:
Sleeping
Sleeping
File size: 3,376 Bytes
618a3f8 87880d1 618a3f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import subprocess
import tempfile
import textwrap
import unittest
from pathlib import Path
import numpy as np
from .. import PySRRegressor
from .params import (
DEFAULT_NCYCLES,
DEFAULT_NITERATIONS,
DEFAULT_PARAMS,
DEFAULT_POPULATIONS,
)
class TestWarmStart(unittest.TestCase):
def setUp(self):
# Using inspect,
# get default niterations from PySRRegressor, and double them:
self.default_test_kwargs = dict(
progress=False,
model_selection="accuracy",
niterations=DEFAULT_NITERATIONS * 2,
populations=DEFAULT_POPULATIONS * 2,
temp_equation_file=True,
)
self.rstate = np.random.RandomState(0)
self.X = self.rstate.randn(100, 5)
def test_warm_start_from_file(self):
"""Test that we can warm start in another process."""
with tempfile.TemporaryDirectory() as tmpdirname:
model = PySRRegressor(
**self.default_test_kwargs,
unary_operators=["cos"],
)
model.warm_start = True
model.temp_equation_file = False
model.equation_file = Path(tmpdirname) / "equations.csv"
model.deterministic = True
model.multithreading = False
model.random_state = 0
model.procs = 0
model.early_stop_condition = 1e-10
rstate = np.random.RandomState(0)
X = rstate.randn(100, 2)
y = np.cos(X[:, 0]) ** 2
model.fit(X, y)
best_loss = model.equations_.iloc[-1]["loss"]
# Save X and y to a file:
X_file = Path(tmpdirname) / "X.npy"
y_file = Path(tmpdirname) / "y.npy"
np.save(X_file, X)
np.save(y_file, y)
# Now, create a new process and warm start from the file:
result = subprocess.run(
[
"python",
"-c",
textwrap.dedent(
f"""
from pysr import PySRRegressor
import numpy as np
X = np.load("{X_file}")
y = np.load("{y_file}")
print("Loading model from file")
model = PySRRegressor.from_file("{model.equation_file}")
assert model.julia_state_ is not None
model.warm_start = True
model.niterations = 0
model.max_evals = 0
model.ncycles_per_iteration = 0
model.fit(X, y)
best_loss = model.equations_.iloc[-1]["loss"]
assert best_loss <= {best_loss}
"""
),
],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
self.assertEqual(result.returncode, 0)
self.assertIn("Loading model from file", result.stdout.decode())
self.assertIn("Started!", result.stderr.decode())
def runtests():
suite = unittest.TestSuite()
loader = unittest.TestLoader()
suite.addTests(loader.loadTestsFromTestCase(TestWarmStart))
runner = unittest.TextTestRunner()
return runner.run(suite)
|