diff --git a/.gitignore b/.gitignore index baf33ed16..0e3d3464f 100644 --- a/.gitignore +++ b/.gitignore @@ -70,4 +70,4 @@ dmypy.json site/ # Jupyter Notebook load data -.ipynb_checkpoints \ No newline at end of file +.ipynb_checkpoints diff --git a/autora/skl/bms.py b/autora/skl/bms.py index 9e24be187..547e2361e 100644 --- a/autora/skl/bms.py +++ b/autora/skl/bms.py @@ -70,6 +70,7 @@ def __init__( self.X_: Optional[np.ndarray] = None self.y_: Optional[np.ndarray] = None self.model_: Tree = Tree() + self.models_: List[Tree] = [Tree()] self.loss_: float = np.inf self.cache_: List = [] self.variables: List = [] @@ -110,6 +111,7 @@ def fit(self, X: np.ndarray, y: np.ndarray, num_param: int = 1) -> BMSRegressor: prior_par=self.prior_par, ) self.model_, self.loss_, self.cache_ = utils.run(self.pms, self.epochs) + self.models_ = list(self.pms.trees.values()) _logger.info("BMS fitting finished") self.X_, self.y_ = X, y diff --git a/autora/theorist/bms/parallel.py b/autora/theorist/bms/parallel.py index e9ac6cdfc..182c943ef 100644 --- a/autora/theorist/bms/parallel.py +++ b/autora/theorist/bms/parallel.py @@ -49,7 +49,7 @@ def __init__( Ts.sort() self.Ts = [str(T) for T in Ts] self.trees = { - "1": Tree( + "1.0": Tree( ops=ops, variables=deepcopy(variables), parameters=deepcopy(parameters), @@ -60,7 +60,7 @@ def __init__( BT=1, ) } - self.t1 = self.trees["1"] + self.t1 = self.trees["1.0"] for BT in [T for T in self.Ts if T != 1]: treetmp = Tree( ops=ops, @@ -87,7 +87,7 @@ def mcmc_step(self, verbose=False, p_rr=0.05, p_long=0.45) -> None: for T, tree in list(self.trees.items()): # MCMC step tree.mcmc_step(verbose=verbose, p_rr=p_rr, p_long=p_long) - self.t1 = self.trees["1"] + self.t1 = self.trees["1.0"] # ------------------------------------------------------------------------- def tree_swap(self) -> Tuple[Optional[str], Optional[str]]: @@ -119,7 +119,7 @@ def tree_swap(self) -> Tuple[Optional[str], Optional[str]]: self.trees[self.Ts[nT2]] = t1 t1.BT = BT2 t2.BT = BT1 - self.t1 = self.trees["1"] + self.t1 = self.trees["1.0"] return self.Ts[nT1], self.Ts[nT2] else: return None, None @@ -140,7 +140,7 @@ def anneal(self, n=1000, factor=5) -> None: t.BT *= factor for kk in range(n): print( - "# Annealing heating at %g: %d / %d" % (self.trees["1"].BT, kk, n), + "# Annealing heating at %g: %d / %d" % (self.trees["1.0"].BT, kk, n), file=sys.stderr, ) self.mcmc_step() @@ -150,7 +150,8 @@ def anneal(self, n=1000, factor=5) -> None: t.BT = float(BT) for kk in range(2 * n): print( - "# Annealing cooling at %g: %d / %d" % (self.trees["1"].BT, kk, 2 * n), + "# Annealing cooling at %g: %d / %d" + % (self.trees["1.0"].BT, kk, 2 * n), file=sys.stderr, ) self.mcmc_step() diff --git a/tests/test_bms_mcmc.py b/tests/test_bms_mcmc.py index c1aeaaff7..e28846068 100644 --- a/tests/test_bms_mcmc.py +++ b/tests/test_bms_mcmc.py @@ -45,7 +45,7 @@ def test_tree_mcmc_stepping( ) # MCMC - t.mcmc(burnin=200, thin=10, samples=samples, verbose=True) + t.mcmc(burnin=200, thin=10, samples=samples, verbose=False) # Predict print(t.predict(x)) diff --git a/tests/test_bms_multi_model_output.py b/tests/test_bms_multi_model_output.py new file mode 100644 index 000000000..3d3962800 --- /dev/null +++ b/tests/test_bms_multi_model_output.py @@ -0,0 +1,25 @@ +import numpy as np +import pytest + +from autora.skl.bms import BMSRegressor +from autora.theorist.bms import Tree + + +@pytest.fixture +def curve_to_fit(): + x = np.linspace(-10, 10, 100).reshape(-1, 1) + y = (x**3.0) + (2.0 * x**2.0) + (17.0 * x) - 1 + return x, y + + +def test_bms_models(curve_to_fit): + x, y = curve_to_fit + regressor = BMSRegressor(epochs=100) + + regressor.fit(x, y) + + print(regressor.models_) + + assert len(regressor.models_) == len(regressor.ts) # Currently hardcoded + for model in regressor.models_: + assert isinstance(model, Tree)