return_inferencedata option for pm.sample

pymc3/backends/ndarray.py

@@ -55,7 +55,7 @@ def save_trace(trace: MultiTrace, directory: Optional[str]=None, overwrite=False
     """
     warnings.warn(
         'The `save_trace` function will soon be removed.'
-        'Instead, use ArviZ to save/load traces.',
+        'Instead, use `arviz.to_netcdf` to save traces.',
         DeprecationWarning,
     )
 
@@ -98,7 +98,7 @@ def load_trace(directory: str, model=None) -> MultiTrace:
     """
     warnings.warn(
         'The `load_trace` function will soon be removed.'
-        'Instead, use ArviZ to save/load traces.',
+        'Instead, use `arviz.from_netcdf` to load traces.',
         DeprecationWarning,
     )
     straces = []

pymc3/backends/report.py

@@ -18,6 +18,7 @@
 import typing
 from ..util import is_transformed_name, get_untransformed_name
 
+import arviz
 
 logger = logging.getLogger('pymc3')
 
@@ -98,31 +99,27 @@ def raise_ok(self, level='error'):
         if errors:
             raise ValueError('Serious convergence issues during sampling.')
 
-    def _run_convergence_checks(self, trace, model):
-        if trace.nchains == 1:
+    def _run_convergence_checks(self, idata:arviz.InferenceData, model):
+        if idata.posterior.sizes['chain'] == 1:
             msg = ("Only one chain was sampled, this makes it impossible to "
                    "run some convergence checks")
             warn = SamplerWarning(WarningType.BAD_PARAMS, msg, 'info',
                                   None, None, None)
             self._add_warnings([warn])
             return
 
-        from pymc3 import rhat, ess
-        from arviz import from_pymc3
-
         valid_name = [rv.name for rv in model.free_RVs + model.deterministics]
         varnames = []
         for rv in model.free_RVs:
             rv_name = rv.name
             if is_transformed_name(rv_name):
                 rv_name2 = get_untransformed_name(rv_name)
                 rv_name = rv_name2 if rv_name2 in valid_name else rv_name
-            if rv_name in trace.varnames:
+            if rv_name in idata.posterior:
                 varnames.append(rv_name)
 
-        idata = from_pymc3(trace, log_likelihood=False)
-        self._ess = ess = ess(idata, var_names=varnames)
-        self._rhat = rhat = rhat(idata, var_names=varnames)
+        self._ess = ess = arviz.ess(idata, var_names=varnames)
+        self._rhat = rhat = arviz.rhat(idata, var_names=varnames)
 
         warnings = []
         rhat_max = max(val.max() for val in rhat.values())
@@ -147,7 +144,7 @@ def _run_convergence_checks(self, trace, model):
             warnings.append(warn)
 
         eff_min = min(val.min() for val in ess.values())
-        n_samples = len(trace) * trace.nchains
+        n_samples = idata.posterior.sizes['chain'] * idata.posterior.sizes['draw']
         if eff_min < 200 and n_samples >= 500:
             msg = ("The estimated number of effective samples is smaller than "
                    "200 for some parameters.")

pymc3/backends/text.py

@@ -194,7 +194,7 @@ def load(name, model=None):
     """
     warnings.warn(
         'The `load` function will soon be removed. '
-        'Please use ArviZ to save traces. '
+        'Please use `arviz.from_netcdf` to load traces. '
         'If you have good reasons for using the `load` function, file an issue and tell us about them. ',
         DeprecationWarning,
     )
@@ -239,7 +239,7 @@ def dump(name, trace, chains=None):
     """
     warnings.warn(
         'The `dump` function will soon be removed. '
-        'Please use ArviZ to save traces. '
+        'Please use `arviz.to_netcdf` to save traces. '
         'If you have good reasons for using the `dump` function, file an issue and tell us about them. ',
         DeprecationWarning,
     )