From 667277ee1bd6bdf65dc59f80f59ea53b6a7231d2 Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Fri, 14 Oct 2022 15:07:25 +0530
Subject: [PATCH 01/19] added predict_posterior

---
 pymc_experimental/model_builder.py | 50 +++++++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 8 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index c0583c19..06332da5 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -245,10 +245,49 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
     def predict(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
-        point_estimate: bool = True,
     ):
         """
-        Uses model to predict on unseen data.
+        Uses model to predict on unseen data and return point prediction of all the samples
+
+        Parameters
+        ---------
+        data_prediction : Dictionary of string and either of numpy array, pandas dataframe or pandas Series
+            It is the data we need to make prediction on using the model.
+
+        Returns
+        -------
+        returns dictionary of sample's mean of posterior predict.
+
+        Examples
+        --------
+        >>> data, model_config, sampler_config = LinearModel.create_sample_input()
+        >>> model = LinearModel(model_config, sampler_config)
+        >>> idata = model.fit(data)
+        >>> x_pred = []
+        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        # point predict
+        >>> pred_mean = model.predict(prediction_data)
+        """
+
+        if data_prediction is not None:  # set new input data
+            self._data_setter(data_prediction)
+
+        with self.model:  # sample with new input data
+            post_pred = pm.sample_posterior_predictive(self.idata)
+
+        # reshape output
+        post_pred = self._extract_samples(post_pred)
+        for key in post_pred:
+            post_pred[key] = post_pred[key].mean(axis=0)
+
+        return post_pred
+
+    def predict_posterior(
+        self,
+        data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
+    ):
+        """
+        Uses model to predict samples on unseen data.
 
         Parameters
         ---------
@@ -268,10 +307,8 @@ def predict(
         >>> idata = model.fit(data)
         >>> x_pred = []
         >>> prediction_data = pd.DataFrame({'input':x_pred})
-        # only point estimate
-        >>> pred_mean = model.predict(prediction_data)
         # samples
-        >>> pred_samples = model.predict(prediction_data, point_estimate=False)
+        >>> pred_mean = model.predict_posterior(prediction_data)
         """
 
         if data_prediction is not None:  # set new input data
@@ -282,9 +319,6 @@ def predict(
 
         # reshape output
         post_pred = self._extract_samples(post_pred)
-        if point_estimate:  # average, if point-like estimate desired
-            for key in post_pred:
-                post_pred[key] = post_pred[key].mean(axis=0)
 
         return post_pred
 

From d539e5c3000fe836a810789ee54c0f5e89cd973e Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Sat, 22 Oct 2022 05:15:25 +0530
Subject: [PATCH 02/19] changed tests

---
 pymc_experimental/model_builder.py            |  2 +-
 pymc_experimental/tests/test_model_builder.py | 32 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 06332da5..b309cff9 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -287,7 +287,7 @@ def predict_posterior(
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
     ):
         """
-        Uses model to predict samples on unseen data.
+        Uses model to predict samples on unseen data.`
 
         Parameters
         ---------
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index d1df5535..c011e44d 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -135,3 +135,35 @@ def test_predict():
         y_test = pm.sample_posterior_predictive(idata)
 
         assert str(model_2.idata.groups) == str(idata.groups)
+
+
+def test_predict_posterior():
+    x_pred = np.random.uniform(low=0, high=1, size=100)
+    prediction_data = pd.DataFrame({"input": x_pred})
+    data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
+    model_2 = test_ModelBuilder(model_config, sampler_config, data)
+    model_2.idata = model_2.fit()
+    model_2.predict_posterior(prediction_data)
+    with pm.Model() as model:
+        x = np.linspace(start=0, stop=1, num=100)
+        y = 5 * x + 3
+        x = pm.MutableData("x", x)
+        y_data = pm.MutableData("y_data", y)
+        a_loc = 7
+        a_scale = 3
+        b_loc = 5
+        b_scale = 3
+        obs_error = 2
+
+        a = pm.Normal("a", a_loc, sigma=a_scale)
+        b = pm.Normal("b", b_loc, sigma=b_scale)
+        obs_error = pm.HalfNormal("σ_model_fmc", obs_error)
+
+        y_model = pm.Normal("y_model", a + b * x, obs_error, observed=y_data)
+
+        idata = pm.sample(tune=10, draws=20, chains=3, cores=1)
+        idata.extend(pm.sample_prior_predictive())
+        idata.extend(pm.sample_posterior_predictive(idata))
+        y_test = pm.sample_posterior_predictive(idata)
+
+        assert str(model_2.idata.groups) == str(idata.groups)

From b3695da921ec471a2b8a4d88560f945b29d1c1d3 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Sat, 22 Oct 2022 18:23:47 +0200
Subject: [PATCH 03/19] Update pymc_experimental/model_builder.py

---
 pymc_experimental/model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index b309cff9..06332da5 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -287,7 +287,7 @@ def predict_posterior(
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
     ):
         """
-        Uses model to predict samples on unseen data.`
+        Uses model to predict samples on unseen data.
 
         Parameters
         ---------

From b1a68443629f51028b03b27c972ecd9ec1a14f1b Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Sun, 15 Jan 2023 22:10:11 +0530
Subject: [PATCH 04/19] updated load method

---
 pymc_experimental/model_builder.py | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 06332da5..6ae50e99 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -157,8 +157,7 @@ def save(self, fname):
         file = Path(str(fname))
         self.idata.to_netcdf(file)
 
-    @classmethod
-    def load(cls, fname):
+    def load(self, fname):
         """
         Loads inference data for the model.
 
@@ -189,16 +188,14 @@ def load(cls, fname):
         idata = data
         # Since there is an issue with attrs getting saved in netcdf format which will be fixed in future the following part of code is commented
         # Link of issue -> https://github.com/arviz-devs/arviz/issues/2109
-        # if model.idata.attrs is not None:
-        #     if model.idata.attrs['id'] == self.idata.attrs['id']:
-        #         self = model
-        #         self.idata = data
-        #         return self
-        #     else:
-        #         raise ValueError(
-        #             f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
-        #         )
-        return idata
+        if idata.attrs is not None:
+            if self.id() == idata.attrs["id"]:
+                self.idata = idata
+            else:
+                raise ValueError(
+                    f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
+                )
+        # return idata
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
         """
@@ -238,8 +235,8 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["id"] = self.id()
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sample_conifg"] = self.sample_config
-        self.idata.attrs["model_config"] = self.model_config
+        self.idata.attrs["sample_conifg"] = tuple(self.sample_config)
+        self.idata.attrs["model_config"] = tuple(self.model_config)
         return self.idata
 
     def predict(

From eb0469ddd252507552cbb306f9d018146aa3d3e8 Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Sun, 15 Jan 2023 22:12:18 +0530
Subject: [PATCH 05/19] updated load method

---
 pymc_experimental/model_builder.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 6ae50e99..c82d8930 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -186,8 +186,6 @@ def load(self, fname):
         filepath = Path(str(fname))
         data = az.from_netcdf(filepath)
         idata = data
-        # Since there is an issue with attrs getting saved in netcdf format which will be fixed in future the following part of code is commented
-        # Link of issue -> https://github.com/arviz-devs/arviz/issues/2109
         if idata.attrs is not None:
             if self.id() == idata.attrs["id"]:
                 self.idata = idata

From b576b38f6ec19cae2bd08dc09667fd8d3b1dd39a Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Sun, 15 Jan 2023 22:45:30 +0530
Subject: [PATCH 06/19] added required changes

---
 pymc_experimental/model_builder.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index c82d8930..d9e4a33d 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -157,7 +157,8 @@ def save(self, fname):
         file = Path(str(fname))
         self.idata.to_netcdf(file)
 
-    def load(self, fname):
+    @classmethod
+    def load(cls, self, fname):
         """
         Loads inference data for the model.
 
@@ -193,7 +194,9 @@ def load(self, fname):
                 raise ValueError(
                     f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
                 )
-        # return idata
+        self = cls(idata.attrs["sample_config"], idata.attrs["model_config"])
+        self.idata = idata
+        return self
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
         """
@@ -233,7 +236,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["id"] = self.id()
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sample_conifg"] = tuple(self.sample_config)
+        self.idata.attrs["sample_config"] = tuple(self.sample_config)
         self.idata.attrs["model_config"] = tuple(self.model_config)
         return self.idata
 

From 6ebce2f5abcb3c04c5669f391509dce2514972ac Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Sun, 15 Jan 2023 22:46:24 +0530
Subject: [PATCH 07/19] added required changes

---
 pymc_experimental/model_builder.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index d9e4a33d..21bb27f6 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -189,13 +189,12 @@ def load(cls, self, fname):
         idata = data
         if idata.attrs is not None:
             if self.id() == idata.attrs["id"]:
+                self = cls(idata.attrs["sample_config"], idata.attrs["model_config"])
                 self.idata = idata
             else:
                 raise ValueError(
                     f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
                 )
-        self = cls(idata.attrs["sample_config"], idata.attrs["model_config"])
-        self.idata = idata
         return self
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):

From 51d1ff034e3c9223b71ba34f36329d00bd3f2c16 Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Thu, 19 Jan 2023 21:58:54 +0530
Subject: [PATCH 08/19] fixed laod function

---
 pymc_experimental/model_builder.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 21bb27f6..b2da0484 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -158,7 +158,7 @@ def save(self, fname):
         self.idata.to_netcdf(file)
 
     @classmethod
-    def load(cls, self, fname):
+    def load(cls, fname):
         """
         Loads inference data for the model.
 
@@ -185,16 +185,13 @@ def load(cls, self, fname):
         """
 
         filepath = Path(str(fname))
-        data = az.from_netcdf(filepath)
-        idata = data
-        if idata.attrs is not None:
-            if self.id() == idata.attrs["id"]:
-                self = cls(idata.attrs["sample_config"], idata.attrs["model_config"])
-                self.idata = idata
-            else:
-                raise ValueError(
-                    f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
-                )
+        idata = az.from_netcdf(filepath)
+        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.data)
+        self.idata=idata
+        if self.id() != idata.attrs["id"]:
+            raise ValueError(
+                f"The route '{fname}' does not contain an inference data of the same model '{self._model_type}'"
+            )
         return self
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
@@ -235,8 +232,11 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["id"] = self.id()
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sample_config"] = tuple(self.sample_config)
-        self.idata.attrs["model_config"] = tuple(self.model_config)
+        self.idata.attrs["sample_config_keys"] = tuple(self.sample_config.keys())
+        self.idata.attrs["sample_config_values"] = tuple(self.sample_config.values())
+        self.idata.attrs["model_config_keys"] = tuple(self.model_config.keys())
+        self.idata.attrs["model_config_values"] = tuple(self.model_config.values())
+        self.idata.add_groups(data = self.data.to_xarray())
         return self.idata
 
     def predict(
@@ -352,5 +352,5 @@ def id(self):
         hasher.update(str(self.model_config.values()).encode())
         hasher.update(self.version.encode())
         hasher.update(self._model_type.encode())
-        hasher.update(str(self.sample_config.values()).encode())
+        # hasher.update(str(self.sample_config.values()).encode())
         return hasher.hexdigest()[:16]

From d6183b0ac8a10789711eeb95ef53cd398823c202 Mon Sep 17 00:00:00 2001
From: S <Shashankkirtania@gmail.com>
Date: Thu, 19 Jan 2023 21:59:35 +0530
Subject: [PATCH 09/19] fixed laod functio

---
 pymc_experimental/model_builder.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index b2da0484..e0fb5af2 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -186,8 +186,12 @@ def load(cls, fname):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
-        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.data)
-        self.idata=idata
+        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.data,
+        )
+        self.idata = idata
         if self.id() != idata.attrs["id"]:
             raise ValueError(
                 f"The route '{fname}' does not contain an inference data of the same model '{self._model_type}'"
@@ -236,7 +240,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["sample_config_values"] = tuple(self.sample_config.values())
         self.idata.attrs["model_config_keys"] = tuple(self.model_config.keys())
         self.idata.attrs["model_config_values"] = tuple(self.model_config.values())
-        self.idata.add_groups(data = self.data.to_xarray())
+        self.idata.add_groups(data=self.data.to_xarray())
         return self.idata
 
     def predict(

From 93bee6a0181900fdf0a4a33003d04294dbe1e6aa Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Fri, 20 Jan 2023 11:47:38 +0100
Subject: [PATCH 10/19] Update pymc_experimental/model_builder.py

---
 pymc_experimental/model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index e0fb5af2..bcb8791e 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -194,7 +194,7 @@ def load(cls, fname):
         self.idata = idata
         if self.id() != idata.attrs["id"]:
             raise ValueError(
-                f"The route '{fname}' does not contain an inference data of the same model '{self._model_type}'"
+                f"The file '{fname}' does not contain an inference data of the same model or configuration as '{self._model_type}'"
             )
         return self
 

From 35f661d41b1613bf760d92e5351e74399ff18354 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Fri, 20 Jan 2023 11:48:14 +0100
Subject: [PATCH 11/19] Update pymc_experimental/model_builder.py

---
 pymc_experimental/model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index bcb8791e..7edce989 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -240,7 +240,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["sample_config_values"] = tuple(self.sample_config.values())
         self.idata.attrs["model_config_keys"] = tuple(self.model_config.keys())
         self.idata.attrs["model_config_values"] = tuple(self.model_config.values())
-        self.idata.add_groups(data=self.data.to_xarray())
+        self.idata.add_groups(constant_data=self.data.to_xarray())
         return self.idata
 
     def predict(

From 218761b603c9a3141463454a6c4b80f7aa3fd9cd Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Mon, 23 Jan 2023 12:17:37 +0100
Subject: [PATCH 12/19] Apply suggestions from code review

---
 pymc_experimental/model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 7edce989..be3730ba 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -240,7 +240,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["sample_config_values"] = tuple(self.sample_config.values())
         self.idata.attrs["model_config_keys"] = tuple(self.model_config.keys())
         self.idata.attrs["model_config_values"] = tuple(self.model_config.values())
-        self.idata.add_groups(constant_data=self.data.to_xarray())
+        self.idata.add_groups(fit_data=self.data.to_xarray())
         return self.idata
 
     def predict(

From 38b8883abe358b3057afaa51e10c82808d474e52 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 19:43:42 +0100
Subject: [PATCH 13/19] Restructure tests. Fix load().

---
 pymc_experimental/model_builder.py            |   7 +-
 pymc_experimental/tests/test_model_builder.py | 105 ++++--------------
 2 files changed, 27 insertions(+), 85 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index be3730ba..ac5cfb73 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -1,4 +1,4 @@
-#   Copyright 2022 The PyMC Developers
+#   Copyright 2023 The PyMC Developers
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
 #   you may not use this file except in compliance with the License.
@@ -189,13 +189,14 @@ def load(cls, fname):
         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.data,
+            idata.fit_data.to_dataframe(),
         )
         self.idata = idata
         if self.id() != idata.attrs["id"]:
             raise ValueError(
                 f"The file '{fname}' does not contain an inference data of the same model or configuration as '{self._model_type}'"
             )
+
         return self
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
@@ -307,7 +308,7 @@ def predict_posterior(
         >>> model = LinearModel(model_config, sampler_config)
         >>> idata = model.fit(data)
         >>> x_pred = []
-        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        >>> prediction_data = pd.DataFrame({'input': x_pred})
         # samples
         >>> pred_mean = model.predict_posterior(prediction_data)
         """
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index c011e44d..3ee03d55 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -1,4 +1,4 @@
-#   Copyright 2022 The PyMC Developers
+#   Copyright 2023 The PyMC Developers
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
 #   you may not use this file except in compliance with the License.
@@ -13,6 +13,8 @@
 #   limitations under the License.
 
 
+import tempfile
+
 import numpy as np
 import pandas as pd
 import pymc as pm
@@ -77,93 +79,32 @@ def create_sample_input(cls):
 
 
 def test_fit():
-    with pm.Model() as model:
-        x = np.linspace(start=0, stop=1, num=100)
-        y = 5 * x + 3
-        x = pm.MutableData("x", x)
-        y_data = pm.MutableData("y_data", y)
-
-        a_loc = 7
-        a_scale = 3
-        b_loc = 5
-        b_scale = 3
-        obs_error = 2
-
-        a = pm.Normal("a", a_loc, sigma=a_scale)
-        b = pm.Normal("b", b_loc, sigma=b_scale)
-        obs_error = pm.HalfNormal("σ_model_fmc", obs_error)
-
-        y_model = pm.Normal("y_model", a + b * x, obs_error, observed=y_data)
-
-        idata = pm.sample(tune=100, draws=200, chains=1, cores=1, target_accept=0.5)
-        idata.extend(pm.sample_prior_predictive())
-        idata.extend(pm.sample_posterior_predictive(idata))
-
     data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
-    model_2 = test_ModelBuilder(model_config, sampler_config, data)
-    model_2.idata = model_2.fit()
-    assert str(model_2.idata.groups) == str(idata.groups)
-
+    model = test_ModelBuilder(model_config, sampler_config, data)
+    model.fit()
+    assert model.idata is not None
+    assert "posterior" in model.idata.groups()
 
-def test_predict():
     x_pred = np.random.uniform(low=0, high=1, size=100)
     prediction_data = pd.DataFrame({"input": x_pred})
-    data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
-    model_2 = test_ModelBuilder(model_config, sampler_config, data)
-    model_2.idata = model_2.fit()
-    model_2.predict(prediction_data)
-    with pm.Model() as model:
-        x = np.linspace(start=0, stop=1, num=100)
-        y = 5 * x + 3
-        x = pm.MutableData("x", x)
-        y_data = pm.MutableData("y_data", y)
-        a_loc = 7
-        a_scale = 3
-        b_loc = 5
-        b_scale = 3
-        obs_error = 2
-
-        a = pm.Normal("a", a_loc, sigma=a_scale)
-        b = pm.Normal("b", b_loc, sigma=b_scale)
-        obs_error = pm.HalfNormal("σ_model_fmc", obs_error)
-
-        y_model = pm.Normal("y_model", a + b * x, obs_error, observed=y_data)
+    pred = model.predict(prediction_data)
+    assert "y_model" in pred.keys()
+    post_pred = model.predict_posterior(prediction_data)
+    assert "y_model" in post_pred.keys()
 
-        idata = pm.sample(tune=10, draws=20, chains=3, cores=1)
-        idata.extend(pm.sample_prior_predictive())
-        idata.extend(pm.sample_posterior_predictive(idata))
-        y_test = pm.sample_posterior_predictive(idata)
-
-        assert str(model_2.idata.groups) == str(idata.groups)
 
+def test_save_load():
+    data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
+    model = test_ModelBuilder(model_config, sampler_config, data)
+    temp = tempfile.TemporaryFile()
+    model.fit()
+    model.save(temp.name)
+    model2 = test_ModelBuilder.load(temp.name)
+    assert model.idata.groups() == model2.idata.groups()
 
-def test_predict_posterior():
     x_pred = np.random.uniform(low=0, high=1, size=100)
     prediction_data = pd.DataFrame({"input": x_pred})
-    data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
-    model_2 = test_ModelBuilder(model_config, sampler_config, data)
-    model_2.idata = model_2.fit()
-    model_2.predict_posterior(prediction_data)
-    with pm.Model() as model:
-        x = np.linspace(start=0, stop=1, num=100)
-        y = 5 * x + 3
-        x = pm.MutableData("x", x)
-        y_data = pm.MutableData("y_data", y)
-        a_loc = 7
-        a_scale = 3
-        b_loc = 5
-        b_scale = 3
-        obs_error = 2
-
-        a = pm.Normal("a", a_loc, sigma=a_scale)
-        b = pm.Normal("b", b_loc, sigma=b_scale)
-        obs_error = pm.HalfNormal("σ_model_fmc", obs_error)
-
-        y_model = pm.Normal("y_model", a + b * x, obs_error, observed=y_data)
-
-        idata = pm.sample(tune=10, draws=20, chains=3, cores=1)
-        idata.extend(pm.sample_prior_predictive())
-        idata.extend(pm.sample_posterior_predictive(idata))
-        y_test = pm.sample_posterior_predictive(idata)
-
-        assert str(model_2.idata.groups) == str(idata.groups)
+    pred1 = model.predict(prediction_data)
+    pred2 = model2.predict(prediction_data)
+    assert pred1["y_model"].shape == pred2["y_model"].shape
+    temp.close()

From 80fae7b21bd78d318eaa7964cc914cf06b553cea Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 19:55:55 +0100
Subject: [PATCH 14/19] Bump isort.

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1904521d..848ac5af 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,7 +11,7 @@ repos:
         args: [--branch, main]
     -   id: trailing-whitespace
 - repo: https://github.com/PyCQA/isort
-  rev: 5.10.1
+  rev: 5.12.0
   hooks:
     - id: isort
       name: isort

From 50d76f27207b9c5bdc58f7faa55ee538228e7d8f Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 20:29:02 +0100
Subject: [PATCH 15/19] Skip saving on windows.

---
 pymc_experimental/tests/test_model_builder.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 3ee03d55..56dadda6 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -13,11 +13,13 @@
 #   limitations under the License.
 
 
+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
 
@@ -93,6 +95,9 @@ def test_fit():
     assert "y_model" in post_pred.keys()
 
 
+@pytest.mark.xfail(
+    sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
+)
 def test_save_load():
     data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
     model = test_ModelBuilder(model_config, sampler_config, data)

From 85a1f375dcb7a26a8bde5cbcb4902db43fac6b53 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 20:34:11 +0100
Subject: [PATCH 16/19] Different approach.

---
 pymc_experimental/tests/test_model_builder.py | 32 +++++++++----------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 56dadda6..e712c38f 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -12,14 +12,12 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-
-import sys
+import os
 import tempfile
 
 import numpy as np
 import pandas as pd
 import pymc as pm
-import pytest
 
 from pymc_experimental.model_builder import ModelBuilder
 
@@ -95,21 +93,21 @@ def test_fit():
     assert "y_model" in post_pred.keys()
 
 
-@pytest.mark.xfail(
-    sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
-)
 def test_save_load():
     data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
     model = test_ModelBuilder(model_config, sampler_config, data)
-    temp = tempfile.TemporaryFile()
     model.fit()
-    model.save(temp.name)
-    model2 = test_ModelBuilder.load(temp.name)
-    assert model.idata.groups() == model2.idata.groups()
-
-    x_pred = np.random.uniform(low=0, high=1, size=100)
-    prediction_data = pd.DataFrame({"input": x_pred})
-    pred1 = model.predict(prediction_data)
-    pred2 = model2.predict(prediction_data)
-    assert pred1["y_model"].shape == pred2["y_model"].shape
-    temp.close()
+    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
+    try:
+        model.save(temp.name)
+        model2 = test_ModelBuilder.load(temp.name)
+        assert model.idata.groups() == model2.idata.groups()
+
+        x_pred = np.random.uniform(low=0, high=1, size=100)
+        prediction_data = pd.DataFrame({"input": x_pred})
+        pred1 = model.predict(prediction_data)
+        pred2 = model2.predict(prediction_data)
+        assert pred1["y_model"].shape == pred2["y_model"].shape
+    finally:
+        temp.close()
+        os.unlink(file.name)

From eef7d44f0146d3b96c4de3905b54c1c4a8e66446 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 20:51:40 +0100
Subject: [PATCH 17/19] Different approach.

---
 pymc_experimental/tests/test_model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index e712c38f..c8196af8 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -110,4 +110,4 @@ def test_save_load():
         assert pred1["y_model"].shape == pred2["y_model"].shape
     finally:
         temp.close()
-        os.unlink(file.name)
+        os.unlink(temp.name)

From 7a5705c0a82670e1c937b486db9d8fb5f8d022ce Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 21:15:36 +0100
Subject: [PATCH 18/19] Different approach.

---
 pymc_experimental/tests/test_model_builder.py | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index c8196af8..f627f404 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -12,12 +12,14 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-import os
+
+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
 
@@ -93,21 +95,21 @@ def test_fit():
     assert "y_model" in post_pred.keys()
 
 
+@pytest.mark.xfail(
+    sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
+)
 def test_save_load():
     data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
     model = test_ModelBuilder(model_config, sampler_config, data)
     model.fit()
     temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
-    try:
-        model.save(temp.name)
-        model2 = test_ModelBuilder.load(temp.name)
-        assert model.idata.groups() == model2.idata.groups()
-
-        x_pred = np.random.uniform(low=0, high=1, size=100)
-        prediction_data = pd.DataFrame({"input": x_pred})
-        pred1 = model.predict(prediction_data)
-        pred2 = model2.predict(prediction_data)
-        assert pred1["y_model"].shape == pred2["y_model"].shape
-    finally:
-        temp.close()
-        os.unlink(temp.name)
+    model.save(temp.name)
+    model2 = test_ModelBuilder.load(temp.name)
+    assert model.idata.groups() == model2.idata.groups()
+
+    x_pred = np.random.uniform(low=0, high=1, size=100)
+    prediction_data = pd.DataFrame({"input": x_pred})
+    pred1 = model.predict(prediction_data)
+    pred2 = model2.predict(prediction_data)
+    assert pred1["y_model"].shape == pred2["y_model"].shape
+    temp.close()

From f4d872a63aae6e41028927a9acb5e39667b449de Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Wed, 15 Feb 2023 21:44:57 +0100
Subject: [PATCH 19/19] Different approach.

---
 pymc_experimental/tests/test_model_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index f627f404..be29ad9d 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -95,7 +95,7 @@ def test_fit():
     assert "y_model" in post_pred.keys()
 
 
-@pytest.mark.xfail(
+@pytest.mark.skipif(
     sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
 )
 def test_save_load():