Source code for arviz.wrappers.base

# pylint: disable=too-many-instance-attributes,too-many-arguments
"""Base class for sampling wrappers."""
from xarray import apply_ufunc

# from import InferenceData
from ..stats import wrap_xarray_ufunc as _wrap_xarray_ufunc

[docs]class SamplingWrapper: """Class wrapping sampling routines for its usage via ArviZ. Using a common class, all inference backends can be supported in ArviZ. Hence, statistical functions requiring refitting like Leave Future Out or Simulation Based Calibration can be performed from ArviZ. For usage examples see user guide pages on :ref:`wrapper_guide`.See other SamplingWrapper classes at :ref:`wrappers api section <wrappers_api>`. Parameters ---------- model The model object used for sampling. idata_orig : InferenceData, optional Original InferenceData object. log_lik_fun : callable, optional For simple cases where the pointwise log likelihood is a Python function, this function will be used to calculate the log likelihood. Otherwise, ``point_log_likelihood`` method must be implemented. It's callback must be ``log_lik_fun(*args, **log_lik_kwargs)`` and will be called using :func:`wrap_xarray_ufunc` or :func:`xarray:xarray.apply_ufunc` depending on the value of `is_ufunc`. For more details on ``args`` or ``log_lik_kwargs`` see the notes and parameters ``posterior_vars`` and ``log_lik_kwargs``. is_ufunc : bool, default True If True, call ``log_lik_fun`` using :func:`xarray:xarray.apply_ufunc` otherwise use :func:`wrap_xarray_ufunc`. posterior_vars : list of str, optional List of variable names to unpack as ``args`` for ``log_lik_fun``. Each string in the list will be used to retrieve a DataArray from the Dataset in the posterior group and passed to ``log_lik_fun``. sample_kwargs : dict, optional Sampling kwargs are stored as class attributes for their usage in the ``sample`` method. idata_kwargs : dict, optional kwargs are stored as class attributes to be used in the ``get_inference_data`` method. log_lik_kwargs : dict, optional Keyword arguments passed to ``log_lik_fun``. apply_ufunc_kwargs : dict, optional Passed to :func:`xarray:xarray.apply_ufunc` or :func:`wrap_xarray_ufunc`. Warnings -------- Sampling wrappers are an experimental feature in a very early stage. Please use them with caution. Notes ----- Example of ``log_like_fun`` usage. """
[docs] def __init__( self, model, idata_orig=None, log_lik_fun=None, is_ufunc=True, posterior_vars=None, sample_kwargs=None, idata_kwargs=None, log_lik_kwargs=None, apply_ufunc_kwargs=None, ): self.model = model # if not isinstance(idata_orig, InferenceData) or idata_orig is not None: # raise TypeError("idata_orig must be of InferenceData type or None") self.idata_orig = idata_orig if log_lik_fun is None or callable(log_lik_fun): self.log_lik_fun = log_lik_fun self.is_ufunc = is_ufunc self.posterior_vars = posterior_vars else: raise TypeError("log_like_fun must be a callable object or None") self.sample_kwargs = {} if sample_kwargs is None else sample_kwargs self.idata_kwargs = {} if idata_kwargs is None else idata_kwargs self.log_lik_kwargs = {} if log_lik_kwargs is None else log_lik_kwargs self.apply_ufunc_kwargs = {} if apply_ufunc_kwargs is None else apply_ufunc_kwargs
[docs] def sel_observations(self, idx): """Select a subset of the observations in idata_orig. **Not implemented**: This method must be implemented by the SamplingWrapper subclasses. It is documented here to show its format and call signature. Parameters ---------- idx Indexes to separate from the rest of the observed data. Returns ------- modified_observed_data Observed data whose index is *not* ``idx`` excluded_observed_data Observed data whose index is ``idx`` """ raise NotImplementedError("sel_observations method must be implemented for each subclass")
[docs] def sample(self, modified_observed_data): """Sample ``self.model`` on the ``modified_observed_data`` subset. **Not implemented**: This method must be implemented by the SamplingWrapper subclasses. It is documented here to show its format and call signature. Parameters ---------- modified_observed_data Data to fit the model on. Returns ------- fitted_model Result of the fit. """ raise NotImplementedError("sample method must be implemented for each subclass")
[docs] def get_inference_data(self, fitted_model): """Convert the ``fitted_model`` to an InferenceData object. **Not implemented**: This method must be implemented by the SamplingWrapper subclasses. It is documented here to show its format and call signature. Parameters ---------- fitted_model Result of the current fit. Returns ------- idata_current: InferenceData InferenceData object containing the samples in ``fitted_model`` """ raise NotImplementedError("get_inference_data method must be implemented for each subclass")
[docs] def log_likelihood__i(self, excluded_obs, idata__i): r"""Get the log likelilhood samples :math:`\log p_{post(-i)}(y_i)`. Calculate the log likelihood of the data contained in excluded_obs using the model fitted with this data excluded, the results of which are stored in ``idata__i``. Parameters ---------- excluded_obs Observations for which to calculate their log likelihood. The second item from the tuple returned by `sel_observations` is passed as this argument. idata__i: InferenceData Inference results of refitting the data excluding some observations. The result of `get_inference_data` is used as this argument. Returns ------- log_likelihood: xr.Dataarray Log likelihood of ``excluded_obs`` evaluated at each of the posterior samples stored in ``idata__i``. """ if self.log_lik_fun is None: raise NotImplementedError( "When `log_like_fun` is not set during class initialization " "log_likelihood__i method must be overwritten" ) posterior = idata__i.posterior arys = (*excluded_obs, *[posterior[var_name] for var_name in self.posterior_vars]) if self.is_ufunc: ufunc_applier = apply_ufunc else: ufunc_applier = _wrap_xarray_ufunc log_lik_idx = ufunc_applier( self.log_lik_fun, *arys, kwargs=self.log_lik_kwargs, **self.apply_ufunc_kwargs, ) return log_lik_idx
def _check_method_is_implemented(self, method, *args): """Check a given method is implemented.""" try: getattr(self, method)(*args) except NotImplementedError: return False except: # pylint: disable=bare-except return True return True
[docs] def check_implemented_methods(self, methods): """Check that all methods listed are implemented. Not all functions that require refitting need to have all the methods implemented in order to work properly. This function shoulg be used before using the SamplingWrapper and its subclasses to get informative error messages. Parameters ---------- methods: list Check all elements in methods are implemented. Returns ------- List with all non implemented methods """ supported_methods_1arg = ( "sel_observations", "sample", "get_inference_data", ) supported_methods_2args = ("log_likelihood__i",) supported_methods = [*supported_methods_1arg, *supported_methods_2args] bad_methods = [method for method in methods if method not in supported_methods] if bad_methods: raise ValueError( "Not all method(s) in {} supported. Supported methods in SamplingWrapper " "subclasses are:{}".format(bad_methods, supported_methods) ) not_implemented = [] for method in methods: if method in supported_methods_1arg: if self._check_method_is_implemented(method, 1): continue not_implemented.append(method) elif method in supported_methods_2args: if self._check_method_is_implemented(method, 1, 1): continue not_implemented.append(method) return not_implemented