Model Builder refactoring

pymc_experimental/model_builder.py

@@ -70,7 +70,7 @@ def build(self):
         Builds the defined model.
         """
 
-        self.build_model(self.model_config, self.data)
+        self.build_model(self, self.model_config, self.data)
 
     @abstractmethod
     def _data_setter(
@@ -188,8 +188,7 @@ def load(cls, fname):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
-        self = ModelBuilder(idata)
-        self.model = cls(
+        self = cls(
             dict(zip(idata.attrs["model_config_keys"], idata.attrs["model_config_values"])),
             dict(zip(idata.attrs["sample_config_keys"], idata.attrs["sample_config_values"])),
             idata.fit_data.to_dataframe(),

pymc_experimental/tests/test_model_builder.py

@@ -12,10 +12,14 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import hashlib
+import sys
+import tempfile
 
 import numpy as np
 import pandas as pd
 import pymc as pm
+import pytest
 
 from pymc_experimental.model_builder import ModelBuilder
 
@@ -24,12 +28,13 @@ class test_ModelBuilder(ModelBuilder):
     _model_type = "LinearModel"
     version = "0.1"
 
-    def build_model(self, model_config, data=None):
-
+    def build_model(self, model_instance, model_config, data=None):
+        model_instance.model_config = model_config
+        model_instance.data = data
         self.model_config = model_config
         self.data = data
 
-        with pm.Model() as self.model:
+        with pm.Model() as model_instance.model:
             if data is not None:
                 x = pm.MutableData("x", data["input"].values)
                 y_data = pm.MutableData("y_data", data["output"].values)
@@ -83,12 +88,12 @@ def create_sample_input(self):
     @staticmethod
     def initial_build_and_fit(check_idata=True):
         data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
-        model = test_ModelBuilder(model_config, sampler_config, data)
-        model.fit(data=data)
+        model_builder = test_ModelBuilder(model_config, sampler_config, data)
+        model_builder.idata = model_builder.fit(data=data)
         if check_idata:
-            assert model.idata is not None
-            assert "posterior" in model.idata.groups()
-        return model
+            assert model_builder.idata is not None
+            assert "posterior" in model_builder.idata.groups()
+        return model_builder
 
 
 def test_fit():
@@ -101,16 +106,16 @@ def test_fit():
     assert "y_model" in post_pred.keys()
 
 
-"""
 @pytest.mark.skipif(
     sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
 )
 def test_save_load():
-    test_builder = test_ModelBuilder.initial_build_and_fit(False)
+    test_builder = test_ModelBuilder.initial_build_and_fit()
     temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
     test_builder.save(temp.name)
-    test_builder2 = test_ModelBuilder.load(temp.name)
-    assert test_builder.model.idata.groups() == test_builder2.model.idata.groups()
+    test_builder2 = test_ModelBuilder.initial_build_and_fit()
+    test_builder2.model = test_ModelBuilder.load(temp.name)
+    assert test_builder.idata.groups() == test_builder2.idata.groups()
 
     x_pred = np.random.uniform(low=0, high=1, size=100)
     prediction_data = pd.DataFrame({"input": x_pred})
@@ -171,4 +176,3 @@ def test_id():
     ).hexdigest()[:16]
 
     assert model.id == expected_id
-"""