pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/model_builder.py

@@ -0,0 +1,759 @@
+#   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.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+
+import hashlib
+import json
+import warnings
+
+from abc import abstractmethod
+from pathlib import Path
+from typing import Any
+
+import arviz as az
+import numpy as np
+import pandas as pd
+import pymc as pm
+import xarray as xr
+
+from pymc.util import RandomState
+
+# If scikit-learn is available, use its data validator
+try:
+    from sklearn.utils.validation import check_array, check_X_y
+# If scikit-learn is not available, return the data unchanged
+except ImportError:
+
+    def check_X_y(X, y, **kwargs):
+        return X, y
+
+    def check_array(X, **kwargs):
+        return X
+
+
+class ModelBuilder:
+    """
+    ModelBuilder can be used to provide an easy-to-use API (similar to scikit-learn) for models
+    and help with deployment.
+    """
+
+    _model_type = "BaseClass"
+    version = "None"
+
+    def __init__(
+        self,
+        model_config: dict | None = None,
+        sampler_config: dict | None = None,
+    ):
+        """
+        Initializes model configuration and sampler configuration for the model
+
+        Parameters
+        ----------
+        data : Dictionary, optional
+            It is the data we need to train the model on.
+        model_config : Dictionary, optional
+            dictionary of parameters that initialise model configuration. Class-default defined by the user default_model_config method.
+        sampler_config : Dictionary, optional
+            dictionary of parameters that initialise sampler configuration. Class-default defined by the user default_sampler_config method.
+
+        Examples
+        --------
+        >>> class MyModel(ModelBuilder):
+        >>>     ...
+        >>> model = MyModel(model_config, sampler_config)
+        """
+        sampler_config = (
+            self.get_default_sampler_config() if sampler_config is None else sampler_config
+        )
+        self.sampler_config = sampler_config
+        model_config = self.get_default_model_config() if model_config is None else model_config
+
+        self.model_config = model_config  # parameters for priors etc.
+        self.model = None  # Set by build_model
+        self.idata: az.InferenceData | None = None  # idata is generated during fitting
+        self.is_fitted_ = False
+
+    def _validate_data(self, X, y=None):
+        if y is not None:
+            return check_X_y(X, y, accept_sparse=False, y_numeric=True, multi_output=False)
+        else:
+            return check_array(X, accept_sparse=False)
+
+    @abstractmethod
+    def _data_setter(
+        self,
+        X: np.ndarray | pd.DataFrame,
+        y: np.ndarray | pd.DataFrame | list = None,
+    ) -> None:
+        """
+        Sets new data in the model.
+
+        Parameters
+        ----------
+        X : array, shape (n_obs, n_features)
+            The training input samples.
+        y : array, shape (n_obs,)
+            The target values (real numbers).
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> def _data_setter(self, data : pd.DataFrame):
+        >>>     with self.model:
+        >>>         pm.set_data({'x': X['x'].values})
+        >>>         try: # if y values in new data
+        >>>             pm.set_data({'y_data': y.values})
+        >>>         except: # dummies otherwise
+        >>>             pm.set_data({'y_data': np.zeros(len(data))})
+
+        """
+
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def output_var(self):
+        """
+        Returns the name of the output variable of the model.
+
+        Returns
+        -------
+        output_var : str
+            Name of the output variable of the model.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_default_model_config() -> dict:
+        """
+        Returns a class default config dict for model builder if no model_config is provided on class initialization
+        Useful for understanding structure of required model_config to allow its customization by users
+
+        Examples
+        --------
+        >>>     @staticmethod
+        >>>     def default_model_config():
+        >>>         return {
+        >>>             'a' : {
+        >>>                 'loc': 7,
+        >>>                 'scale' : 3
+        >>>             },
+        >>>             'b' : {
+        >>>                 'loc': 3,
+        >>>                 'scale': 5
+        >>>             }
+        >>>              'obs_error': 2
+        >>>         }
+
+        Returns
+        -------
+        model_config : dict
+            A set of default parameters for predictor distributions that allow to save and recreate the model.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_default_sampler_config(self) -> dict:
+        """
+        Returns a class default sampler dict for model builder if no sampler_config is provided on class initialization
+        Useful for understanding structure of required sampler_config to allow its customization by users
+
+        Examples
+        --------
+        >>>     @staticmethod
+        >>>     def default_sampler_config():
+        >>>         return {
+        >>>             'draws': 1_000,
+        >>>             'tune': 1_000,
+        >>>             'chains': 1,
+        >>>             'target_accept': 0.95,
+        >>>         }
+
+        Returns
+        -------
+        sampler_config : dict
+            A set of default settings for used by model in fit process.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _generate_and_preprocess_model_data(
+        self, X: pd.DataFrame | pd.Series, y: pd.Series
+    ) -> None:
+        """
+        Applies preprocessing to the data before fitting the model.
+        if validate is True, it will check if the data is valid for the model.
+        sets self.model_coords based on provided dataset
+
+        In case of optional parameters being passed into the model, this method should implement the conditional
+        logic responsible for correct handling of the optional parameters, and including them into the dataset.
+
+        Parameters
+        ----------
+        X : array, shape (n_obs, n_features)
+        y : array, shape (n_obs,)
+
+        Examples
+        --------
+        >>>     @classmethod
+        >>>     def _generate_and_preprocess_model_data(self, X, y):
+                    coords = {
+                        'x_dim': X.dim_variable,
+                    } #only include if applicable for your model
+        >>>         self.X = X
+        >>>         self.y = y
+
+        Returns
+        -------
+        None
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_model(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+        **kwargs,
+    ) -> None:
+        """
+        Creates an instance of pm.Model based on provided data and model_config, and
+        attaches it to self.
+
+        Parameters
+        ----------
+        X : pd.DataFrame
+            The input data that is going to be used in the model. This should be a DataFrame
+            containing the features (predictors) for the model. For efficiency reasons, it should
+            only contain the necessary data columns, not the entire available dataset, as this
+            will be encoded into the data used to recreate the model.
+
+        y : pd.Series
+            The target data for the model. This should be a Series representing the output
+            or dependent variable for the model.
+
+        kwargs : dict
+            Additional keyword arguments that may be used for model configuration.
+
+        See Also
+        --------
+        default_model_config : returns default model config
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        NotImplementedError
+            This is an abstract method and must be implemented in a subclass.
+        """
+        raise NotImplementedError
+
+    def sample_model(self, **kwargs):
+        """
+        Sample from the PyMC model.
+
+        Parameters
+        ----------
+        **kwargs : dict
+            Additional keyword arguments to pass to the PyMC sampler.
+
+        Returns
+        -------
+        xarray.Dataset
+            The PyMC samples dataset.
+
+        Raises
+        ------
+        RuntimeError
+            If the PyMC model hasn't been built yet.
+
+        Examples
+        --------
+        >>> self.build_model()
+        >>> idata = self.sample_model(draws=100, tune=10)
+        >>> assert isinstance(idata, xr.Dataset)
+        >>> assert "posterior" in idata
+        >>> assert "prior" in idata
+        >>> assert "observed_data" in idata
+        >>> assert "log_likelihood" in idata
+        """
+        if self.model is None:
+            raise RuntimeError(
+                "The model hasn't been built yet, call .build_model() first or call .fit() instead."
+            )
+
+        with self.model:
+            sampler_args = {**self.sampler_config, **kwargs}
+            idata = pm.sample(**sampler_args)
+            idata.extend(pm.sample_prior_predictive(), join="right")
+            idata.extend(pm.sample_posterior_predictive(idata), join="right")
+
+        idata = self.set_idata_attrs(idata)
+        return idata
+
+    def set_idata_attrs(self, idata=None):
+        """
+        Set attributes on an InferenceData object.
+
+        Parameters
+        ----------
+        idata : arviz.InferenceData, optional
+            The InferenceData object to set attributes on.
+
+        Raises
+        ------
+        RuntimeError
+            If no InferenceData object is provided.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> model = MyModel(ModelBuilder)
+        >>> idata = az.InferenceData(your_dataset)
+        >>> model.set_idata_attrs(idata=idata)
+        >>> assert "id" in idata.attrs #this and the following lines are part of doctest, not user manual
+        >>> assert "model_type" in idata.attrs
+        >>> assert "version" in idata.attrs
+        >>> assert "sampler_config" in idata.attrs
+        >>> assert "model_config" in idata.attrs
+        """
+        if idata is None:
+            idata = self.idata
+        if idata is None:
+            raise RuntimeError("No idata provided to set attrs on.")
+        idata.attrs["id"] = self.id
+        idata.attrs["model_type"] = self._model_type
+        idata.attrs["version"] = self.version
+        idata.attrs["sampler_config"] = json.dumps(self.sampler_config)
+        idata.attrs["model_config"] = json.dumps(self._serializable_model_config)
+        # Only classes with non-dataset parameters will implement save_input_params
+        if hasattr(self, "_save_input_params"):
+            self._save_input_params(idata)
+        return idata
+
+    def save(self, fname: str) -> None:
+        """
+        Save the model's inference data to a file.
+
+        Parameters
+        ----------
+        fname : str
+            The name and path of the file to save the inference data with model parameters.
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        RuntimeError
+            If the model hasn't been fit yet (no inference data available).
+
+        Examples
+        --------
+        This method is meant to be overridden and implemented by subclasses.
+        It should not be called directly on the base abstract class or its instances.
+
+        >>> class MyModel(ModelBuilder):
+        >>>     def __init__(self):
+        >>>         super().__init__()
+        >>> model = MyModel()
+        >>> model.fit(data)
+        >>> model.save('model_results.nc')  # This will call the overridden method in MyModel
+        """
+        if self.idata is not None and "posterior" in self.idata:
+            file = Path(str(fname))
+            self.idata.to_netcdf(str(file))
+        else:
+            raise RuntimeError("The model hasn't been fit yet, call .fit() first")
+
+    @classmethod
+    def _model_config_formatting(cls, model_config: dict) -> dict:
+        """
+        Because of json serialization, model_config values that were originally tuples or numpy are being encoded as lists.
+        This function converts them back to tuples and numpy arrays to ensure correct id encoding.
+        """
+        for key in model_config:
+            if isinstance(model_config[key], dict):
+                for sub_key in model_config[key]:
+                    if isinstance(model_config[key][sub_key], list):
+                        # Check if "dims" key to convert it to tuple
+                        if sub_key == "dims":
+                            model_config[key][sub_key] = tuple(model_config[key][sub_key])
+                        # Convert all other lists to numpy arrays
+                        else:
+                            model_config[key][sub_key] = np.array(model_config[key][sub_key])
+        return model_config
+
+    @classmethod
+    def load(cls, fname: str):
+        """
+        Creates a ModelBuilder instance from a file,
+        Loads inference data for the model.
+
+        Parameters
+        ----------
+        fname : string
+            This denotes the name with path from where idata should be loaded from.
+
+        Returns
+        -------
+        Returns an instance of ModelBuilder.
+
+        Raises
+        ------
+        ValueError
+            If the inference data that is loaded doesn't match with the model.
+
+        Examples
+        --------
+        >>> class MyModel(ModelBuilder):
+        >>>     ...
+        >>> name = './mymodel.nc'
+        >>> imported_model = MyModel.load(name)
+        """
+        filepath = Path(str(fname))
+        idata = az.from_netcdf(filepath)
+        # needs to be converted, because json.loads was changing tuple to list
+        model_config = cls._model_config_formatting(json.loads(idata.attrs["model_config"]))
+        model = cls(
+            model_config=model_config,
+            sampler_config=json.loads(idata.attrs["sampler_config"]),
+        )
+        model.idata = idata
+        dataset = idata.fit_data.to_dataframe()
+        X = dataset.drop(columns=[model.output_var])
+        y = dataset[model.output_var]
+        model.build_model(X, y)
+        # All previously used data is in idata.
+
+        if model.id != idata.attrs["id"]:
+            raise ValueError(
+                f"The file '{fname}' does not contain an inference data of the same model or configuration as '{cls._model_type}'"
+            )
+
+        return model
+
+    def fit(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series | None = None,
+        progressbar: bool = True,
+        predictor_names: list[str] | None = None,
+        random_seed: RandomState = None,
+        **kwargs: Any,
+    ) -> az.InferenceData:
+        """
+        Fit a model using the data passed as a parameter.
+        Sets attrs to inference data of the model.
+
+
+        Parameters
+        ----------
+        X : array-like if sklearn is available, otherwise array, shape (n_obs, n_features)
+            The training input samples.
+        y : array-like if sklearn is available, otherwise array, shape (n_obs,)
+            The target values (real numbers).
+        progressbar : bool
+            Specifies whether the fit progressbar should be displayed
+        predictor_names: List[str] = None,
+            Allows for custom naming of predictors given in a form of 2dArray
+            allows for naming of predictors when given in a form of np.ndarray, if not provided the predictors will be named like predictor1, predictor2...
+        random_seed : RandomState
+            Provides sampler with initial random seed for obtaining reproducible samples
+        **kwargs : Any
+            Custom sampler settings can be provided in form of keyword arguments.
+
+        Returns
+        -------
+        self : az.InferenceData
+            returns inference data of the fitted model.
+
+        Examples
+        --------
+        >>> model = MyModel()
+        >>> idata = model.fit(data)
+        Auto-assigning NUTS sampler...
+        Initializing NUTS using jitter+adapt_diag...
+        """
+        if predictor_names is None:
+            predictor_names = []
+        if y is None:
+            y = np.zeros(X.shape[0])
+        y = pd.DataFrame({self.output_var: y})
+        self._generate_and_preprocess_model_data(X, y.values.flatten())
+        self.build_model(self.X, self.y)
+
+        sampler_config = self.sampler_config.copy()
+        sampler_config["progressbar"] = progressbar
+        sampler_config["random_seed"] = random_seed
+        sampler_config.update(**kwargs)
+        self.idata = self.sample_model(**sampler_config)
+
+        X_df = pd.DataFrame(X, columns=X.columns)
+        combined_data = pd.concat([X_df, y], axis=1)
+        assert all(combined_data.columns), "All columns must have non-empty names"
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                category=UserWarning,
+                message="The group fit_data is not defined in the InferenceData scheme",
+            )
+            self.idata.add_groups(fit_data=combined_data.to_xarray())  # type: ignore
+
+        return self.idata  # type: ignore
+
+    def predict(
+        self,
+        X_pred: np.ndarray | pd.DataFrame | pd.Series,
+        extend_idata: bool = True,
+        **kwargs,
+    ) -> np.ndarray:
+        """
+        Uses model to predict on unseen data and return point prediction of all the samples. The point prediction
+        for each input row is the expected output value, computed as the mean of MCMC samples.
+
+        Parameters
+        ----------
+        X_pred : array-like if sklearn is available, otherwise array, shape (n_pred, n_features)
+            The input data used for prediction.
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to True.
+        **kwargs: Additional arguments to pass to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        y_pred : ndarray, shape (n_pred,)
+            Predicted output corresponding to input X_pred.
+
+        Examples
+        --------
+        >>> model = MyModel()
+        >>> idata = model.fit(data)
+        >>> x_pred = []
+        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        >>> pred_mean = model.predict(prediction_data)
+        """
+
+        posterior_predictive_samples = self.sample_posterior_predictive(
+            X_pred, extend_idata, combined=False, **kwargs
+        )
+
+        if self.output_var not in posterior_predictive_samples:
+            raise KeyError(
+                f"Output variable {self.output_var} not found in posterior predictive samples."
+            )
+
+        posterior_means = posterior_predictive_samples[self.output_var].mean(
+            dim=["chain", "draw"], keep_attrs=True
+        )
+        return posterior_means.data
+
+    def sample_prior_predictive(
+        self,
+        X_pred,
+        y_pred=None,
+        samples: int | None = None,
+        extend_idata: bool = False,
+        combined: bool = True,
+        **kwargs,
+    ):
+        """
+        Sample from the model's prior predictive distribution.
+
+        Parameters
+        ----------
+        X_pred : array, shape (n_pred, n_features)
+            The input data used for prediction using prior distribution.
+        samples : int
+            Number of samples from the prior parameter distributions to generate.
+            If not set, uses sampler_config['draws'] if that is available, otherwise defaults to 500.
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to False.
+        combined: Combine chain and draw dims into sample. Won't work if a dim named sample already exists.
+            Defaults to True.
+        **kwargs: Additional arguments to pass to pymc.sample_prior_predictive
+
+        Returns
+        -------
+        prior_predictive_samples : DataArray, shape (n_pred, samples)
+            Prior predictive samples for each input X_pred
+        """
+        if y_pred is None:
+            y_pred = pd.Series(np.zeros(len(X_pred)), name=self.output_var)
+        if samples is None:
+            samples = self.sampler_config.get("draws", 500)
+
+        if self.model is None:
+            self.build_model(X_pred, y_pred)
+        else:
+            self._data_setter(X_pred, y_pred)
+        with self.model:  # sample with new input data
+            prior_pred: az.InferenceData = pm.sample_prior_predictive(samples, **kwargs)
+            self.set_idata_attrs(prior_pred)
+            if extend_idata:
+                if self.idata is not None:
+                    self.idata.extend(prior_pred, join="right")
+                else:
+                    self.idata = prior_pred
+
+        prior_predictive_samples = az.extract(prior_pred, "prior_predictive", combined=combined)
+
+        return prior_predictive_samples
+
+    def sample_posterior_predictive(self, X_pred, extend_idata, combined, **kwargs):
+        """
+        Sample from the model's posterior predictive distribution.
+
+        Parameters
+        ----------
+        X_pred : array, shape (n_pred, n_features)
+            The input data used for prediction using prior distribution..
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to False.
+        combined: Combine chain and draw dims into sample. Won't work if a dim named sample already exists.
+            Defaults to True.
+        **kwargs: Additional arguments to pass to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        posterior_predictive_samples : DataArray, shape (n_pred, samples)
+            Posterior predictive samples for each input X_pred
+        """
+        self._data_setter(X_pred)
+
+        with self.model:  # sample with new input data
+            post_pred = pm.sample_posterior_predictive(self.idata, **kwargs)
+            if extend_idata:
+                self.idata.extend(post_pred, join="right")
+
+        posterior_predictive_samples = az.extract(
+            post_pred, "posterior_predictive", combined=combined
+        )
+
+        return posterior_predictive_samples
+
+    def get_params(self, deep=True):
+        """
+        Get all the model parameters needed to instantiate a copy of the model, not including training data.
+        """
+        return {
+            "model_config": self.model_config,
+            "sampler_config": self.sampler_config,
+        }
+
+    def set_params(self, **params):
+        """
+        Set all the model parameters needed to instantiate the model, not including training data.
+        """
+        self.model_config = params["model_config"]
+        self.sampler_config = params["sampler_config"]
+
+    @property
+    @abstractmethod
+    def _serializable_model_config(self) -> dict[str, int | float | dict]:
+        """
+        Converts non-serializable values from model_config to their serializable reversable equivalent.
+        Data types like pandas DataFrame, Series or datetime aren't JSON serializable,
+        so in order to save the model they need to be formatted.
+
+        Returns
+        -------
+        model_config: dict
+        """
+
+    def predict_proba(
+        self,
+        X_pred: np.ndarray | pd.DataFrame | pd.Series,
+        extend_idata: bool = True,
+        combined: bool = False,
+        **kwargs,
+    ) -> xr.DataArray:
+        """Alias for `predict_posterior`, for consistency with scikit-learn probabilistic estimators."""
+        return self.predict_posterior(X_pred, extend_idata, combined, **kwargs)
+
+    def predict_posterior(
+        self,
+        X_pred: np.ndarray | pd.DataFrame | pd.Series,
+        extend_idata: bool = True,
+        combined: bool = True,
+        **kwargs,
+    ) -> xr.DataArray:
+        """
+        Generate posterior predictive samples on unseen data.
+
+        Parameters
+        ----------
+        X_pred : array-like if sklearn is available, otherwise array, shape (n_pred, n_features)
+            The input data used for prediction.
+        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+            Defaults to True.
+        combined: Combine chain and draw dims into sample. Won't work if a dim named sample already exists.
+            Defaults to True.
+        **kwargs: Additional arguments to pass to pymc.sample_posterior_predictive
+
+        Returns
+        -------
+        y_pred : DataArray, shape (n_pred, chains * draws) if combined is True, otherwise (chains, draws, n_pred)
+            Posterior predictive samples for each input X_pred
+        """
+
+        X_pred = self._validate_data(X_pred)
+        posterior_predictive_samples = self.sample_posterior_predictive(
+            X_pred, extend_idata, combined, **kwargs
+        )
+
+        if self.output_var not in posterior_predictive_samples:
+            raise KeyError(
+                f"Output variable {self.output_var} not found in posterior predictive samples."
+            )
+
+        return posterior_predictive_samples[self.output_var]
+
+    @property
+    def id(self) -> str:
+        """
+        Generate a unique hash value for the model.
+
+        The hash value is created using the last 16 characters of the SHA256 hash encoding, based on the model configuration,
+        version, and model type.
+
+        Returns
+        -------
+        str
+            A string of length 16 characters containing a unique hash of the model.
+
+        Examples
+        --------
+        >>> model = MyModel()
+        >>> model.id
+        '0123456789abcdef'
+        """
+        hasher = hashlib.sha256()
+        hasher.update(str(self.model_config.values()).encode())
+        hasher.update(self.version.encode())
+        hasher.update(self._model_type.encode())
+        return hasher.hexdigest()[:16]