From 8cc3383d4fc3bf9998bdf1aa0484ef8f3e648654 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Tue, 7 Mar 2023 11:16:02 +0100
Subject: [PATCH 01/12] moving build out of constructor

---
 pymc_experimental/model_builder.py            | 6 ++----
 pymc_experimental/tests/test_model_builder.py | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 8c5aa7577..8728820fd 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -63,7 +63,6 @@ def __init__(
         self.sample_config = sampler_config  # parameters for sampling
         self.idata = None  # inference data object
         self.data = data
-        self.build()
 
     def build(self):
         """
@@ -192,6 +191,7 @@ def load(cls, fname):
             idata.fit_data.to_dataframe(),
         )
         self.idata = idata
+        self.build()
         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}'"
@@ -224,10 +224,8 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
 
         if data is not None:
             self.data = data
-            self._data_setter(data)
-
-        if self.basic_RVs == []:
             self.build()
+            self._data_setter(data)
 
         with self:
             self.idata = pm.sample(**self.sample_config)
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 5bdb5d06d..243056c4e 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -84,7 +84,7 @@ def create_sample_input(cls):
     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()
+        model.fit(data=data)
         if check_idata:
             assert model.idata is not None
             assert "posterior" in model.idata.groups()

From 2e4507d09f1b8d8be82761b51bf5eb0e3f42ef4c Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Fri, 24 Mar 2023 15:26:08 +0100
Subject: [PATCH 02/12] removing inheritance from pm.Model, adding instance of
 a model as a class property

---
 pymc_experimental/model_builder.py            | 22 ++++---
 pymc_experimental/tests/test_model_builder.py | 64 ++++++++++---------
 2 files changed, 46 insertions(+), 40 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 8728820fd..85c9cdafe 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -14,6 +14,7 @@
 
 
 import hashlib
+from abc import abstractmethod
 from pathlib import Path
 from typing import Dict, Union
 
@@ -23,7 +24,7 @@
 import pymc as pm
 
 
-class ModelBuilder(pm.Model):
+class ModelBuilder:
     """
     ModelBuilder can be used to provide an easy-to-use API (similar to scikit-learn) for models
     and help with deployment.
@@ -69,9 +70,9 @@ def build(self):
         Builds the defined model.
         """
 
-        with self:
-            self.build_model(self.model_config, self.data)
+        self.build_model(self.model_config, self.data)
 
+    @abstractmethod
     def _data_setter(
         self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]], x_only: bool = True
     ):
@@ -98,8 +99,10 @@ def _data_setter(
 
         raise NotImplementedError
 
-    @classmethod
-    def create_sample_input(cls):
+    # need a discussion if it's really needed.
+    @staticmethod
+    @abstractmethod
+    def create_sample_input():
         """
         Needs to be implemented by the user in the inherited class.
         Returns examples for data, model_config, sampler_config.
@@ -168,7 +171,7 @@ def load(cls, fname):
 
         Returns
         -------
-        Returns the inference data that is loaded from local system.
+        Returns an instance of pm.Model, that is loaded from local data.
 
         Raises
         ------
@@ -185,7 +188,8 @@ def load(cls, fname):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
-        self = cls(
+        self = ModelBuilder(idata)
+        self.model = 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(),
@@ -197,7 +201,7 @@ def load(cls, fname):
                 f"The file '{fname}' does not contain an inference data of the same model or configuration as '{self._model_type}'"
             )
 
-        return self
+        return self.model
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
         """
@@ -227,7 +231,7 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
             self.build()
             self._data_setter(data)
 
-        with self:
+        with self.model:
             self.idata = pm.sample(**self.sample_config)
             self.idata.extend(pm.sample_prior_predictive())
             self.idata.extend(pm.sample_posterior_predictive(self.idata))
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 243056c4e..0f1b4b58e 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -13,14 +13,9 @@
 #   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
 
@@ -30,25 +25,30 @@ class test_ModelBuilder(ModelBuilder):
     version = "0.1"
 
     def build_model(self, model_config, data=None):
-        if data is not None:
-            x = pm.MutableData("x", data["input"].values)
-            y_data = pm.MutableData("y_data", data["output"].values)
-
-        # prior parameters
-        a_loc = model_config["a_loc"]
-        a_scale = model_config["a_scale"]
-        b_loc = model_config["b_loc"]
-        b_scale = model_config["b_scale"]
-        obs_error = model_config["obs_error"]
-
-        # priors
-        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)
-
-        # observed data
-        if data is not None:
-            y_model = pm.Normal("y_model", a + b * x, obs_error, shape=x.shape, observed=y_data)
+
+        self.model_config = model_config
+        self.data = data
+
+        with pm.Model() as self.model:
+            if data is not None:
+                x = pm.MutableData("x", data["input"].values)
+                y_data = pm.MutableData("y_data", data["output"].values)
+
+            # prior parameters
+            a_loc = model_config["a_loc"]
+            a_scale = model_config["a_scale"]
+            b_loc = model_config["b_loc"]
+            b_scale = model_config["b_scale"]
+            obs_error = model_config["obs_error"]
+
+            # priors
+            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)
+
+            # observed data
+            if data is not None:
+                y_model = pm.Normal("y_model", a + b * x, obs_error, shape=x.shape, observed=y_data)
 
     def _data_setter(self, data: pd.DataFrame):
         with self.model:
@@ -57,7 +57,7 @@ def _data_setter(self, data: pd.DataFrame):
                 pm.set_data({"y_data": data["output"].values})
 
     @classmethod
-    def create_sample_input(cls):
+    def create_sample_input(self):
         x = np.linspace(start=0, stop=1, num=100)
         y = 5 * x + 3
         y = y + np.random.normal(0, 1, len(x))
@@ -101,20 +101,21 @@ 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():
-    model = test_ModelBuilder.initial_build_and_fit(False)
+    test_builder = test_ModelBuilder.initial_build_and_fit(False)
     temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
-    model.save(temp.name)
-    model2 = test_ModelBuilder.load(temp.name)
-    assert model.idata.groups() == model2.idata.groups()
+    test_builder.save(temp.name)
+    test_builder2 = test_ModelBuilder.load(temp.name)
+    assert test_builder.model.idata.groups() == test_builder2.model.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)
+    pred1 = test_builder.predict(prediction_data)
+    pred2 = test_builder2.predict(prediction_data)
     assert pred1["y_model"].shape == pred2["y_model"].shape
     temp.close()
 
@@ -170,3 +171,4 @@ def test_id():
     ).hexdigest()[:16]
 
     assert model.id == expected_id
+"""

From 064e3aeb9ecb829192cbb0a47cb846bb72b6abcb Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Sun, 26 Mar 2023 20:23:48 +0200
Subject: [PATCH 03/12] ModelBuilder no longer extends pm.Model, pm.Model now
 is an attribute

---
 pymc_experimental/model_builder.py            |  5 ++--
 pymc_experimental/tests/test_model_builder.py | 30 +++++++++++--------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 85c9cdafe..1b6e5cd85 100644
--- a/pymc_experimental/model_builder.py
+++ b/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(),
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 0f1b4b58e..0337ce668 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/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
-"""

From d3a74ec8e40ce695b361712e039506e115d7f65c Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Mon, 27 Mar 2023 18:40:08 +0200
Subject: [PATCH 04/12] typehinting and refactoring load function

---
 pymc_experimental/model_builder.py            | 29 ++++++++++---------
 pymc_experimental/tests/test_model_builder.py |  7 ++---
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 1b6e5cd85..1242ce820 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -28,8 +28,6 @@ class ModelBuilder:
     """
     ModelBuilder can be used to provide an easy-to-use API (similar to scikit-learn) for models
     and help with deployment.
-
-    Extends the pymc.Model class.
     """
 
     _model_type = "BaseClass"
@@ -65,7 +63,7 @@ def __init__(
         self.idata = None  # inference data object
         self.data = data
 
-    def build(self):
+    def build(self) -> None:
         """
         Builds the defined model.
         """
@@ -136,7 +134,7 @@ def create_sample_input():
 
         raise NotImplementedError
 
-    def save(self, fname):
+    def save(self, fname: str) -> None:
         """
         Saves inference data of the model.
 
@@ -160,8 +158,9 @@ def save(self, fname):
         self.idata.to_netcdf(file)
 
     @classmethod
-    def load(cls, fname):
+    def load(cls, fname: str):
         """
+        Creates a ModelBuilder instance from a file,
         Loads inference data for the model.
 
         Parameters
@@ -171,7 +170,7 @@ def load(cls, fname):
 
         Returns
         -------
-        Returns an instance of pm.Model, that is loaded from local data.
+        Returns an instance of ModelBuilder.
 
         Raises
         ------
@@ -188,21 +187,23 @@ def load(cls, fname):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
-        self = cls(
+        model_builder = 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(),
         )
-        self.idata = idata
-        self.build()
-        if self.id != idata.attrs["id"]:
+        model_builder.idata = idata
+        model_builder.build()
+        if model_builder.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.model
+        return model_builder
 
-    def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
+    def fit(
+        self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
+    ) -> az.InferenceData:
         """
         As the name suggests fit can be used to fit a model using the data that is passed as a parameter.
         Sets attrs to inference data of the model.
@@ -248,7 +249,7 @@ 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,
-    ):
+    ) -> dict:
         """
         Uses model to predict on unseen data and return point prediction of all the samples
 
@@ -288,7 +289,7 @@ def predict(
     def predict_posterior(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
-    ):
+    ) -> Dict[str, np.array]:
         """
         Uses model to predict samples on unseen data.
 
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 0337ce668..f26e91569 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -28,7 +28,7 @@ class test_ModelBuilder(ModelBuilder):
     _model_type = "LinearModel"
     version = "0.1"
 
-    def build_model(self, model_instance, model_config, data=None):
+    def build_model(self, model_instance: ModelBuilder, model_config: dict, data: dict = None):
         model_instance.model_config = model_config
         model_instance.data = data
         self.model_config = model_config
@@ -86,7 +86,7 @@ def create_sample_input(self):
         return data, model_config, sampler_config
 
     @staticmethod
-    def initial_build_and_fit(check_idata=True):
+    def initial_build_and_fit(check_idata=True) -> ModelBuilder:
         data, model_config, sampler_config = test_ModelBuilder.create_sample_input()
         model_builder = test_ModelBuilder(model_config, sampler_config, data)
         model_builder.idata = model_builder.fit(data=data)
@@ -113,8 +113,7 @@ def test_save_load():
     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.initial_build_and_fit()
-    test_builder2.model = test_ModelBuilder.load(temp.name)
+    test_builder2 = 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)

From 9c481aa6a17efaf3a7b314aa3bcd3af88bc7db79 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Thu, 30 Mar 2023 13:04:19 +0100
Subject: [PATCH 05/12] data changed to required, changed sampler_config to
 optional. Added extend_idata to prediction functions, adapted test file

---
 pymc_experimental/model_builder.py            | 50 ++++++++++++-------
 pymc_experimental/tests/test_model_builder.py |  9 +++-
 2 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index 9f861cc7f..a779e4756 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -37,8 +37,8 @@ class ModelBuilder:
     def __init__(
         self,
         model_config: Dict,
-        sampler_config: Dict,
-        data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
+        data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]],
+        sampler_config: Dict = None,
     ):
         """
         Initializes model configuration and sampler configuration for the model
@@ -47,10 +47,10 @@ def __init__(
         ----------
         model_config : Dictionary
             dictionary of parameters that initialise model configuration. Generated by the user defined create_sample_input method.
-        sampler_config : Dictionary
-            dictionary of parameters that initialise sampler configuration. Generated by the user defined create_sample_input method.
         data : Dictionary
             It is the data we need to train the model on.
+        sampler_config : Dictionary
+            dictionary of parameters that initialise sampler configuration. Generated by the user defined create_sample_input method.
         Examples
         --------
         >>> class LinearModel(ModelBuilder):
@@ -60,9 +60,11 @@ def __init__(
 
         super().__init__()
         self.model_config = model_config  # parameters for priors etc.
-        self.sample_config = sampler_config  # parameters for sampling
-        self.idata = None  # inference data object
+        self.sampler_config = sampler_config  # parameters for sampling
         self.data = data
+        self.idata = (
+            None  # inference data object placeholder, idata is generated during build execution
+        )
 
     def build(self) -> None:
         """
@@ -188,16 +190,20 @@ def load(cls, fname: str):
 
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
+        if "sampler_config" in idata.attrs:
+            sampler_config = json.loads(idata.attrs["sampler_config"])
+        else:
+            sampler_config = None
         model_builder = cls(
-            json.loads(idata.attrs["model_config"]),
-            json.loads(idata.attrs["sampler_config"]),
-            idata.fit_data.to_dataframe(),
+            model_config=json.loads(idata.attrs["model_config"]),
+            sampler_config=sampler_config,
+            data=idata.fit_data.to_dataframe(),
         )
         model_builder.idata = idata
         model_builder.build()
         if model_builder.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}'"
+                f"The file '{fname}' does not contain an inference data of the same model or configuration as '{cls._model_type}'"
             )
 
         return model_builder
@@ -206,7 +212,7 @@ def fit(
         self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
     ) -> az.InferenceData:
         """
-        As the name suggests fit can be used to fit a model using the data that is passed as a parameter.
+        Fit a model using the data passed as a parameter.
         Sets attrs to inference data of the model.
 
         Parameter
@@ -233,14 +239,18 @@ def fit(
             self._data_setter(data)
 
         with self.model:
-            self.idata = pm.sample(**self.sample_config)
+            if self.sampler_config is not None:
+                self.idata = pm.sample(**self.sampler_config)
+            else:
+                self.idata = pm.sample()
             self.idata.extend(pm.sample_prior_predictive())
             self.idata.extend(pm.sample_posterior_predictive(self.idata))
 
         self.idata.attrs["id"] = self.id
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sampler_config"] = json.dumps(self.sample_config)
+        if self.sampler_config is not None:
+            self.idata.attrs["sampler_config"] = json.dumps(self.sampler_config)
         self.idata.attrs["model_config"] = json.dumps(self.model_config)
         self.idata.add_groups(fit_data=self.data.to_xarray())
         return self.idata
@@ -248,6 +258,7 @@ def fit(
     def predict(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
+        extend_idata: bool = True,
     ) -> dict:
         """
         Uses model to predict on unseen data and return point prediction of all the samples
@@ -256,6 +267,8 @@ def predict(
         ---------
         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.
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to True.
 
         Returns
         -------
@@ -277,7 +290,8 @@ def predict(
 
         with self.model:  # sample with new input data
             post_pred = pm.sample_posterior_predictive(self.idata)
-
+            if extend_idata:
+                self.idata.extend(post_pred)
         # reshape output
         post_pred = self._extract_samples(post_pred)
         for key in post_pred:
@@ -288,6 +302,7 @@ def predict(
     def predict_posterior(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
+        extend_idata: bool = True,
     ) -> Dict[str, np.array]:
         """
         Uses model to predict samples on unseen data.
@@ -296,8 +311,8 @@ def predict_posterior(
         ---------
         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.
-        point_estimate : bool
-            Adds point like estimate used as mean passed as
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to True.
 
         Returns
         -------
@@ -319,6 +334,8 @@ def predict_posterior(
 
         with self.model:  # sample with new input data
             post_pred = pm.sample_posterior_predictive(self.idata)
+            if extend_idata:
+                self.idata.extend(post_pred)
 
         # reshape output
         post_pred = self._extract_samples(post_pred)
@@ -359,5 +376,4 @@ def id(self) -> str:
         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())
         return hasher.hexdigest()[:16]
diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 1e89cd6e3..21c6890d2 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -28,10 +28,17 @@ class test_ModelBuilder(ModelBuilder):
     _model_type = "LinearModel"
     version = "0.1"
 
-    def build_model(self, model_instance: ModelBuilder, model_config: dict, data: dict = None):
+    def build_model(
+        self,
+        model_instance: ModelBuilder,
+        model_config: dict,
+        data: dict = None,
+        sampler_config: dict = None,
+    ):
         model_instance.model_config = model_config
         model_instance.data = data
         self.model_config = model_config
+        self.sampler_config = sampler_config
         self.data = data
 
         with pm.Model() as model_instance.model:

From b4e634944964a670cb74112e3a57257d1f73aaea Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Thu, 30 Mar 2023 14:18:57 +0100
Subject: [PATCH 06/12] changed default sampler_config from None to {}

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

diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py
index a779e4756..191b1b1e7 100644
--- a/pymc_experimental/model_builder.py
+++ b/pymc_experimental/model_builder.py
@@ -59,6 +59,8 @@ def __init__(
         """
 
         super().__init__()
+        if sampler_config is None:
+            sampler_config = {}
         self.model_config = model_config  # parameters for priors etc.
         self.sampler_config = sampler_config  # parameters for sampling
         self.data = data
@@ -193,7 +195,7 @@ def load(cls, fname: str):
         if "sampler_config" in idata.attrs:
             sampler_config = json.loads(idata.attrs["sampler_config"])
         else:
-            sampler_config = None
+            sampler_config = {}
         model_builder = cls(
             model_config=json.loads(idata.attrs["model_config"]),
             sampler_config=sampler_config,
@@ -239,7 +241,7 @@ def fit(
             self._data_setter(data)
 
         with self.model:
-            if self.sampler_config is not None:
+            if self.sampler_config:
                 self.idata = pm.sample(**self.sampler_config)
             else:
                 self.idata = pm.sample()
@@ -249,7 +251,7 @@ def fit(
         self.idata.attrs["id"] = self.id
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        if self.sampler_config is not None:
+        if self.sampler_config:
             self.idata.attrs["sampler_config"] = json.dumps(self.sampler_config)
         self.idata.attrs["model_config"] = json.dumps(self.model_config)
         self.idata.add_groups(fit_data=self.data.to_xarray())

From 68bc5185a7c9dbf9e63776cef47d729d4bf8af55 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Sun, 26 Mar 2023 20:23:48 +0200
Subject: [PATCH 07/12] ModelBuilder no longer extends pm.Model, pm.Model now
 is an attribute

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

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 21c6890d2..d35db29de 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -120,7 +120,8 @@ def test_save_load():
     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)
+    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)

From ac4ec1a171d34e9b9b283f5c4184284bfa383686 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Mon, 27 Mar 2023 18:40:08 +0200
Subject: [PATCH 08/12] typehinting and refactoring load function

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

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index d35db29de..21c6890d2 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -120,8 +120,7 @@ def test_save_load():
     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.initial_build_and_fit()
-    test_builder2.model = test_ModelBuilder.load(temp.name)
+    test_builder2 = 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)

From 7b8c2533561f4c5e25d8821198e0e995d4351d61 Mon Sep 17 00:00:00 2001
From: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>
Date: Thu, 30 Mar 2023 12:18:34 +0200
Subject: [PATCH 09/12] Update API reference listings (#105)

* Update API reference listings
* improve api page+some fixes
* update to pymc theme
* add pymc theme to requirements
* fix indent in docstring
---------

Co-authored-by: Oriol (ZBook) <oriol.abril.pla@gmail.com>
---
 docs/_templates/autosummary/base.rst          |  5 ++
 docs/_templates/autosummary/class.rst         | 29 +++++++++++
 docs/api_reference.rst                        | 51 ++++++++++++++-----
 docs/conf.py                                  | 25 +++++----
 docs/index.rst                                | 11 +---
 pymc_experimental/distributions/continuous.py | 10 ++--
 pymc_experimental/marginal_model.py           |  1 +
 requirements-docs.txt                         |  1 +
 8 files changed, 96 insertions(+), 37 deletions(-)
 create mode 100644 docs/_templates/autosummary/base.rst
 create mode 100644 docs/_templates/autosummary/class.rst

diff --git a/docs/_templates/autosummary/base.rst b/docs/_templates/autosummary/base.rst
new file mode 100644
index 000000000..3fe9858d7
--- /dev/null
+++ b/docs/_templates/autosummary/base.rst
@@ -0,0 +1,5 @@
+{{ objname | escape | underline }}
+
+.. currentmodule:: {{ module }}
+
+.. auto{{ objtype }}:: {{ objname }}
diff --git a/docs/_templates/autosummary/class.rst b/docs/_templates/autosummary/class.rst
new file mode 100644
index 000000000..574b853a5
--- /dev/null
+++ b/docs/_templates/autosummary/class.rst
@@ -0,0 +1,29 @@
+{{ objname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+
+   {% block methods %}
+   .. automethod:: __init__
+
+   {% if methods %}
+   .. rubric:: {{ _('Methods') }}
+
+   .. autosummary::
+   {% for item in methods %}
+      ~{{ name }}.{{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
+
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: {{ _('Attributes') }}
+
+   .. autosummary::
+   {% for item in attributes %}
+      ~{{ name }}.{{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
diff --git a/docs/api_reference.rst b/docs/api_reference.rst
index b47e912fa..2d85c3ab4 100644
--- a/docs/api_reference.rst
+++ b/docs/api_reference.rst
@@ -4,23 +4,50 @@ API Reference
 This reference provides detailed documentation for all modules, classes, and
 methods in the current release of PyMC experimental.
 
+.. currentmodule:: pymc_experimental
+.. autosummary::
+   :toctree: generated/
 
-.. toctree::
-  :maxdepth: 2
+   marginal_model.MarginalModel
+   model_builder.ModelBuilder
 
+Inference
+=========
 
-:mod:`pymc_experimental.distributions`
-=============================
+.. currentmodule:: pymc_experimental.inference
+.. autosummary::
+   :toctree: generated/
 
-.. automodule:: pymc_experimental.distributions.histogram_utils
-   :members: histogram_approximation
+   fit
 
 
-:mod:`pymc_experimental.utils`
-=============================
+Distributions
+=============
 
-.. automodule:: pymc_experimental.utils.spline
-   :members: bspline_interpolation
+.. currentmodule:: pymc_experimental.distributions
+.. autosummary::
+   :toctree: generated/
 
-.. automodule:: pymc_experimental.utils.prior
-   :members: prior_from_idata
+   GenExtreme
+   histogram_utils.histogram_approximation
+
+
+Gaussian Processess
+===================
+
+.. currentmodule:: pymc_experimental.gp
+.. autosummary::
+   :toctree: generated/
+
+   latent_approx.HSGP
+
+
+Utils
+=====
+
+.. currentmodule:: pymc_experimental.utils
+.. autosummary::
+   :toctree: generated/
+
+   spline.bspline_interpolation
+   prior.prior_from_idata
diff --git a/docs/conf.py b/docs/conf.py
index f45b3407c..e831458f2 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -56,16 +56,18 @@
 # ones.
 extensions = [
     "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
     "sphinx.ext.viewcode",
     "sphinx.ext.napoleon",
     "sphinx.ext.mathjax",
     "nbsphinx",
+    "matplotlib.sphinxext.plot_directive",
 ]
 
 nbsphinx_execute = "never"
 
 # Add any paths that contain templates here, relative to this directory.
-# templates_path = ["_templates"]
+templates_path = ["_templates"]
 
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
@@ -81,7 +83,7 @@
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+# language = None
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
@@ -97,7 +99,7 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = "pydata_sphinx_theme"
+html_theme = "pymc_sphinx_theme"
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
@@ -110,14 +112,17 @@
     "show_toc_level": 2,
     "navigation_depth": 4,
     "search_bar_text": "Search the docs...",
-    "icon_links": [
-        {
-            "name": "GitHub",
-            "url": "https://github.com/pymc-devs/pymc-experimental",
-            "icon": "fab fa-github-square",
-        },
-    ],
+    "use_search_override": False,
+    "logo": {"text": project},
 }
+html_context = {
+    "github_user": "pymc-devs",
+    "github_repo": "pymc-experimental",
+    "github_version": "main",
+    "doc_path": "docs",
+    "default_mode": "light",
+}
+
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
diff --git a/docs/index.rst b/docs/index.rst
index b5dfef6fd..28cf4c9a3 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -38,16 +38,7 @@ Contributors
 ============
 See the `GitHub contributor page <https://github.com/pymc-devs/pymc-experimental/graphs/contributors>`_.
 
-Contents
-========
-
 .. toctree::
-   :maxdepth: 4
+   :hidden:
 
    api_reference
-
-Indices
-=======
-
-* :ref:`genindex`
-* :ref:`modindex`
diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py
index e214f0081..2de72c769 100644
--- a/pymc_experimental/distributions/continuous.py
+++ b/pymc_experimental/distributions/continuous.py
@@ -60,7 +60,7 @@ def rng_fn(
 
 class GenExtreme(Continuous):
     r"""
-        Univariate Generalized Extreme Value log-likelihood
+    Univariate Generalized Extreme Value log-likelihood
 
     The cdf of this distribution is
 
@@ -119,13 +119,13 @@ class GenExtreme(Continuous):
 
     Parameters
     ----------
-    mu: float
+    mu : float
         Location parameter.
-    sigma: float
+    sigma : float
         Scale parameter (sigma > 0).
-    xi: float
+    xi : float
         Shape parameter
-    scipy: bool
+    scipy : bool
         Whether or not to use the Scipy interpretation of the shape parameter
         (defaults to `False`).
 
diff --git a/pymc_experimental/marginal_model.py b/pymc_experimental/marginal_model.py
index aed93adb5..599509bfb 100644
--- a/pymc_experimental/marginal_model.py
+++ b/pymc_experimental/marginal_model.py
@@ -54,6 +54,7 @@ class MarginalModel(Model):
     Marginalize over a single variable
 
     .. code-block:: python
+
         import pymc as pm
         from pymc_experimental import MarginalModel
 
diff --git a/requirements-docs.txt b/requirements-docs.txt
index 0b024fea9..185288a8b 100644
--- a/requirements-docs.txt
+++ b/requirements-docs.txt
@@ -1,3 +1,4 @@
 nbsphinx>=0.4.2
 pydata-sphinx-theme>=0.6.3
 sphinx>=4
+git+https://github.com/pymc-devs/pymc-sphinx-theme

From c023abe51b1ffc6ab6be0c3a54a5632b10b8d0e1 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Fri, 24 Mar 2023 15:26:08 +0100
Subject: [PATCH 10/12] removing inheritance from pm.Model, adding instance of
 a model as a class property

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

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 21c6890d2..1f710dcf8 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -19,7 +19,6 @@
 import numpy as np
 import pandas as pd
 import pymc as pm
-import pytest
 
 from pymc_experimental.model_builder import ModelBuilder
 
@@ -113,6 +112,7 @@ 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."
 )
@@ -182,3 +182,4 @@ def test_id():
     ).hexdigest()[:16]
 
     assert model.id == expected_id
+"""

From c6c5cb593cb5caa0d100dce96d217657a9310f98 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Sun, 26 Mar 2023 20:23:48 +0200
Subject: [PATCH 11/12] ModelBuilder no longer extends pm.Model, pm.Model now
 is an attribute

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

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index 1f710dcf8..d35db29de 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -19,6 +19,7 @@
 import numpy as np
 import pandas as pd
 import pymc as pm
+import pytest
 
 from pymc_experimental.model_builder import ModelBuilder
 
@@ -112,7 +113,6 @@ 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."
 )
@@ -120,7 +120,8 @@ def test_save_load():
     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)
+    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)
@@ -182,4 +183,3 @@ def test_id():
     ).hexdigest()[:16]
 
     assert model.id == expected_id
-"""

From b6cbda2ef40183440f1c064b0c150e49f5ea46b3 Mon Sep 17 00:00:00 2001
From: Michal Raczycki <michalr265@gmail.com>
Date: Mon, 27 Mar 2023 18:40:08 +0200
Subject: [PATCH 12/12] typehinting and refactoring load function

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

diff --git a/pymc_experimental/tests/test_model_builder.py b/pymc_experimental/tests/test_model_builder.py
index d35db29de..21c6890d2 100644
--- a/pymc_experimental/tests/test_model_builder.py
+++ b/pymc_experimental/tests/test_model_builder.py
@@ -120,8 +120,7 @@ def test_save_load():
     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.initial_build_and_fit()
-    test_builder2.model = test_ModelBuilder.load(temp.name)
+    test_builder2 = 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)