google-meridian 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

meridian/model/posterior_sampler.py

@@ -0,0 +1,566 @@
+# Copyright 2024 The Meridian Authors.
+#
+# 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.
+
+"""Module for MCMC sampling of posterior distributions in a Meridian model."""
+
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING
+
+import arviz as az
+from meridian import constants
+import numpy as np
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+if TYPE_CHECKING:
+  from meridian.model import model  # pylint: disable=g-bad-import-order,g-import-not-at-top
+
+
+__all__ = [
+    "MCMCSamplingError",
+    "MCMCOOMError",
+    "PosteriorMCMCSampler",
+]
+
+
+class MCMCSamplingError(Exception):
+  """The Markov Chain Monte Carlo (MCMC) sampling failed."""
+
+
+class MCMCOOMError(Exception):
+  """The Markov Chain Monte Carlo (MCMC) sampling exceeds memory limits."""
+
+
+def _get_tau_g(
+    tau_g_excl_baseline: tf.Tensor, baseline_geo_idx: int
+) -> tfp.distributions.Distribution:
+  """Computes `tau_g` from `tau_g_excl_baseline`.
+
+  This function computes `tau_g` by inserting a column of zeros at the
+  `baseline_geo` position in `tau_g_excl_baseline`.
+
+  Args:
+    tau_g_excl_baseline: A tensor of shape `[..., n_geos - 1]` for the
+      user-defined dimensions of the `tau_g` parameter distribution.
+    baseline_geo_idx: The index of the baseline geo to be set to zero.
+
+  Returns:
+    A tensor of shape `[..., n_geos]` with the final distribution of the `tau_g`
+    parameter with zero at position `baseline_geo_idx` and matching
+    `tau_g_excl_baseline` elsewhere.
+  """
+  rank = len(tau_g_excl_baseline.shape)
+  shape = tau_g_excl_baseline.shape[:-1] + [1] if rank != 1 else 1
+  tau_g = tf.concat(
+      [
+          tau_g_excl_baseline[..., :baseline_geo_idx],
+          tf.zeros(shape, dtype=tau_g_excl_baseline.dtype),
+          tau_g_excl_baseline[..., baseline_geo_idx:],
+      ],
+      axis=rank - 1,
+  )
+  return tfp.distributions.Deterministic(tau_g, name="tau_g")
+
+
+@tf.function(autograph=False, jit_compile=True)
+def _xla_windowed_adaptive_nuts(**kwargs):
+  """XLA wrapper for windowed_adaptive_nuts."""
+  return tfp.experimental.mcmc.windowed_adaptive_nuts(**kwargs)
+
+
+class PosteriorMCMCSampler:
+  """A callable that samples from posterior distributions using MCMC."""
+
+  def __init__(self, meridian: "model.Meridian"):
+    self._meridian = meridian
+
+  def _get_joint_dist_unpinned(self) -> tfp.distributions.Distribution:
+    """Returns a `JointDistributionCoroutineAutoBatched` function for MCMC."""
+    mmm = self._meridian
+    mmm.populate_cached_properties()
+
+    # This lists all the derived properties and states of this Meridian object
+    # that are referenced by the joint distribution coroutine.
+    # That is, these are the list of captured parameters.
+    prior_broadcast = mmm.prior_broadcast
+    baseline_geo_idx = mmm.baseline_geo_idx
+    knot_info = mmm.knot_info
+    n_geos = mmm.n_geos
+    n_times = mmm.n_times
+    n_media_channels = mmm.n_media_channels
+    n_rf_channels = mmm.n_rf_channels
+    n_organic_media_channels = mmm.n_organic_media_channels
+    n_organic_rf_channels = mmm.n_organic_rf_channels
+    n_controls = mmm.n_controls
+    n_non_media_channels = mmm.n_non_media_channels
+    holdout_id = mmm.holdout_id
+    media_tensors = mmm.media_tensors
+    rf_tensors = mmm.rf_tensors
+    organic_media_tensors = mmm.organic_media_tensors
+    organic_rf_tensors = mmm.organic_rf_tensors
+    controls_scaled = mmm.controls_scaled
+    non_media_treatments_scaled = mmm.non_media_treatments_scaled
+    media_effects_dist = mmm.media_effects_dist
+    adstock_hill_media_fn = mmm.adstock_hill_media
+    adstock_hill_rf_fn = mmm.adstock_hill_rf
+    get_roi_prior_beta_m_value_fn = (
+        mmm.prior_sampler_callable.get_roi_prior_beta_m_value
+    )
+    get_roi_prior_beta_rf_value_fn = (
+        mmm.prior_sampler_callable.get_roi_prior_beta_rf_value
+    )
+
+    @tfp.distributions.JointDistributionCoroutineAutoBatched
+    def joint_dist_unpinned():
+      # Sample directly from prior.
+      knot_values = yield prior_broadcast.knot_values
+      gamma_c = yield prior_broadcast.gamma_c
+      xi_c = yield prior_broadcast.xi_c
+      sigma = yield prior_broadcast.sigma
+
+      tau_g_excl_baseline = yield tfp.distributions.Sample(
+          prior_broadcast.tau_g_excl_baseline,
+          name=constants.TAU_G_EXCL_BASELINE,
+      )
+      tau_g = yield _get_tau_g(
+          tau_g_excl_baseline=tau_g_excl_baseline,
+          baseline_geo_idx=baseline_geo_idx,
+      )
+      mu_t = yield tfp.distributions.Deterministic(
+          tf.einsum(
+              "k,kt->t",
+              knot_values,
+              tf.convert_to_tensor(knot_info.weights),
+          ),
+          name=constants.MU_T,
+      )
+
+      tau_gt = tau_g[:, tf.newaxis] + mu_t
+      combined_media_transformed = tf.zeros(
+          shape=(n_geos, n_times, 0), dtype=tf.float32
+      )
+      combined_beta = tf.zeros(shape=(n_geos, 0), dtype=tf.float32)
+      if media_tensors.media is not None:
+        alpha_m = yield prior_broadcast.alpha_m
+        ec_m = yield prior_broadcast.ec_m
+        eta_m = yield prior_broadcast.eta_m
+        slope_m = yield prior_broadcast.slope_m
+        beta_gm_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_media_channels],
+            name=constants.BETA_GM_DEV,
+        )
+        media_transformed = adstock_hill_media_fn(
+            media=media_tensors.media_scaled,
+            alpha=alpha_m,
+            ec=ec_m,
+            slope=slope_m,
+        )
+        prior_type = mmm.model_spec.paid_media_prior_type
+        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
+          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
+            roi_or_mroi_m = yield prior_broadcast.roi_m
+          else:
+            roi_or_mroi_m = yield prior_broadcast.mroi_m
+          beta_m_value = get_roi_prior_beta_m_value_fn(
+              alpha_m,
+              beta_gm_dev,
+              ec_m,
+              eta_m,
+              roi_or_mroi_m,
+              slope_m,
+              media_transformed,
+          )
+          beta_m = yield tfp.distributions.Deterministic(
+              beta_m_value, name=constants.BETA_M
+          )
+        else:
+          beta_m = yield prior_broadcast.beta_m
+
+        beta_eta_combined = beta_m + eta_m * beta_gm_dev
+        beta_gm_value = (
+            beta_eta_combined
+            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
+            else tf.math.exp(beta_eta_combined)
+        )
+        beta_gm = yield tfp.distributions.Deterministic(
+            beta_gm_value, name=constants.BETA_GM
+        )
+        combined_media_transformed = tf.concat(
+            [combined_media_transformed, media_transformed], axis=-1
+        )
+        combined_beta = tf.concat([combined_beta, beta_gm], axis=-1)
+
+      if rf_tensors.reach is not None:
+        alpha_rf = yield prior_broadcast.alpha_rf
+        ec_rf = yield prior_broadcast.ec_rf
+        eta_rf = yield prior_broadcast.eta_rf
+        slope_rf = yield prior_broadcast.slope_rf
+        beta_grf_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_rf_channels],
+            name=constants.BETA_GRF_DEV,
+        )
+        rf_transformed = adstock_hill_rf_fn(
+            reach=rf_tensors.reach_scaled,
+            frequency=rf_tensors.frequency,
+            alpha=alpha_rf,
+            ec=ec_rf,
+            slope=slope_rf,
+        )
+
+        prior_type = mmm.model_spec.paid_media_prior_type
+        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
+          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
+            roi_or_mroi_rf = yield prior_broadcast.roi_rf
+          else:
+            roi_or_mroi_rf = yield prior_broadcast.mroi_rf
+          beta_rf_value = get_roi_prior_beta_rf_value_fn(
+              alpha_rf,
+              beta_grf_dev,
+              ec_rf,
+              eta_rf,
+              roi_or_mroi_rf,
+              slope_rf,
+              rf_transformed,
+          )
+          beta_rf = yield tfp.distributions.Deterministic(
+              beta_rf_value,
+              name=constants.BETA_RF,
+          )
+        else:
+          beta_rf = yield prior_broadcast.beta_rf
+
+        beta_eta_combined = beta_rf + eta_rf * beta_grf_dev
+        beta_grf_value = (
+            beta_eta_combined
+            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
+            else tf.math.exp(beta_eta_combined)
+        )
+        beta_grf = yield tfp.distributions.Deterministic(
+            beta_grf_value, name=constants.BETA_GRF
+        )
+        combined_media_transformed = tf.concat(
+            [combined_media_transformed, rf_transformed], axis=-1
+        )
+        combined_beta = tf.concat([combined_beta, beta_grf], axis=-1)
+
+      if organic_media_tensors.organic_media is not None:
+        alpha_om = yield prior_broadcast.alpha_om
+        ec_om = yield prior_broadcast.ec_om
+        eta_om = yield prior_broadcast.eta_om
+        slope_om = yield prior_broadcast.slope_om
+        beta_gom_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_organic_media_channels],
+            name=constants.BETA_GOM_DEV,
+        )
+        organic_media_transformed = adstock_hill_media_fn(
+            media=organic_media_tensors.organic_media_scaled,
+            alpha=alpha_om,
+            ec=ec_om,
+            slope=slope_om,
+        )
+        beta_om = yield prior_broadcast.beta_om
+
+        beta_eta_combined = beta_om + eta_om * beta_gom_dev
+        beta_gom_value = (
+            beta_eta_combined
+            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
+            else tf.math.exp(beta_eta_combined)
+        )
+        beta_gom = yield tfp.distributions.Deterministic(
+            beta_gom_value, name=constants.BETA_GOM
+        )
+        combined_media_transformed = tf.concat(
+            [combined_media_transformed, organic_media_transformed], axis=-1
+        )
+        combined_beta = tf.concat([combined_beta, beta_gom], axis=-1)
+
+      if organic_rf_tensors.organic_reach is not None:
+        alpha_orf = yield prior_broadcast.alpha_orf
+        ec_orf = yield prior_broadcast.ec_orf
+        eta_orf = yield prior_broadcast.eta_orf
+        slope_orf = yield prior_broadcast.slope_orf
+        beta_gorf_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_organic_rf_channels],
+            name=constants.BETA_GORF_DEV,
+        )
+        organic_rf_transformed = adstock_hill_rf_fn(
+            reach=organic_rf_tensors.organic_reach_scaled,
+            frequency=organic_rf_tensors.organic_frequency,
+            alpha=alpha_orf,
+            ec=ec_orf,
+            slope=slope_orf,
+        )
+
+        beta_orf = yield prior_broadcast.beta_orf
+
+        beta_eta_combined = beta_orf + eta_orf * beta_gorf_dev
+        beta_gorf_value = (
+            beta_eta_combined
+            if media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
+            else tf.math.exp(beta_eta_combined)
+        )
+        beta_gorf = yield tfp.distributions.Deterministic(
+            beta_gorf_value, name=constants.BETA_GORF
+        )
+        combined_media_transformed = tf.concat(
+            [combined_media_transformed, organic_rf_transformed], axis=-1
+        )
+        combined_beta = tf.concat([combined_beta, beta_gorf], axis=-1)
+
+      sigma_gt = tf.transpose(tf.broadcast_to(sigma, [n_times, n_geos]))
+      gamma_gc_dev = yield tfp.distributions.Sample(
+          tfp.distributions.Normal(0, 1),
+          [n_geos, n_controls],
+          name=constants.GAMMA_GC_DEV,
+      )
+      gamma_gc = yield tfp.distributions.Deterministic(
+          gamma_c + xi_c * gamma_gc_dev, name=constants.GAMMA_GC
+      )
+      y_pred_combined_media = (
+          tau_gt
+          + tf.einsum("gtm,gm->gt", combined_media_transformed, combined_beta)
+          + tf.einsum("gtc,gc->gt", controls_scaled, gamma_gc)
+      )
+
+      if mmm.non_media_treatments is not None:
+        gamma_n = yield prior_broadcast.gamma_n
+        xi_n = yield prior_broadcast.xi_n
+        gamma_gn_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_non_media_channels],
+            name=constants.GAMMA_GN_DEV,
+        )
+        gamma_gn = yield tfp.distributions.Deterministic(
+            gamma_n + xi_n * gamma_gn_dev, name=constants.GAMMA_GN
+        )
+        y_pred = y_pred_combined_media + tf.einsum(
+            "gtn,gn->gt", non_media_treatments_scaled, gamma_gn
+        )
+      else:
+        y_pred = y_pred_combined_media
+
+      # If there are any holdout observations, the holdout KPI values will
+      # be replaced with zeros using `experimental_pin`. For these
+      # observations, we set the posterior mean equal to zero and standard
+      # deviation to `1/sqrt(2pi)`, so the log-density is 0 regardless of the
+      # sampled posterior parameter values.
+      if holdout_id is not None:
+        y_pred_holdout = tf.where(holdout_id, 0.0, y_pred)
+        test_sd = tf.cast(1.0 / np.sqrt(2.0 * np.pi), tf.float32)
+        sigma_gt_holdout = tf.where(holdout_id, test_sd, sigma_gt)
+        yield tfp.distributions.Normal(
+            y_pred_holdout, sigma_gt_holdout, name="y"
+        )
+      else:
+        yield tfp.distributions.Normal(y_pred, sigma_gt, name="y")
+
+    return joint_dist_unpinned
+
+  def _get_joint_dist(self) -> tfp.distributions.Distribution:
+    mmm = self._meridian
+    y = (
+        tf.where(mmm.holdout_id, 0.0, mmm.kpi_scaled)
+        if mmm.holdout_id is not None
+        else mmm.kpi_scaled
+    )
+    return self._get_joint_dist_unpinned().experimental_pin(y=y)
+
+  def __call__(
+      self,
+      n_chains: Sequence[int] | int,
+      n_adapt: int,
+      n_burnin: int,
+      n_keep: int,
+      current_state: Mapping[str, tf.Tensor] | None = None,
+      init_step_size: int | None = None,
+      dual_averaging_kwargs: Mapping[str, int] | None = None,
+      max_tree_depth: int = 10,
+      max_energy_diff: float = 500.0,
+      unrolled_leapfrog_steps: int = 1,
+      parallel_iterations: int = 10,
+      seed: Sequence[int] | None = None,
+      **pins,
+  ) -> az.InferenceData:
+    """Runs Markov Chain Monte Carlo (MCMC) sampling of posterior distributions.
+
+    For more information about the arguments, see [`windowed_adaptive_nuts`]
+    (https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/mcmc/windowed_adaptive_nuts).
+
+    Args:
+      n_chains: Number of MCMC chains. Given a sequence of integers,
+        `windowed_adaptive_nuts` will be called once for each element. The
+        `n_chains` argument of each `windowed_adaptive_nuts` call will be equal
+        to the respective integer element. Using a list of integers, one can
+        split the chains of a `windowed_adaptive_nuts` call into multiple calls
+        with fewer chains per call. This can reduce memory usage. This might
+        require an increased number of adaptation steps for convergence, as the
+        optimization is occurring across fewer chains per sampling call.
+      n_adapt: Number of adaptation draws per chain.
+      n_burnin: Number of burn-in draws per chain. Burn-in draws occur after
+        adaptation draws and before the kept draws.
+      n_keep: Integer number of draws per chain to keep for inference.
+      current_state: Optional structure of tensors at which to initialize
+        sampling. Use the same shape and structure as
+        `model.experimental_pin(**pins).sample(n_chains)`.
+      init_step_size: Optional integer determining where to initialize the step
+        size for the leapfrog integrator. The structure must broadcast with
+        `current_state`. For example, if the initial state is:  ``` { 'a':
+        tf.zeros(n_chains), 'b': tf.zeros([n_chains, n_features]), } ```  then
+        any of `1.`, `{'a': 1., 'b': 1.}`, or `{'a': tf.ones(n_chains), 'b':
+        tf.ones([n_chains, n_features])}` will work. Defaults to the dimension
+        of the log density to the ¼ power.
+      dual_averaging_kwargs: Optional dict keyword arguments to pass to
+        `tfp.mcmc.DualAveragingStepSizeAdaptation`. By default, a
+        `target_accept_prob` of `0.85` is set, acceptance probabilities across
+        chains are reduced using a harmonic mean, and the class defaults are
+        used otherwise.
+      max_tree_depth: Maximum depth of the tree implicitly built by NUTS. The
+        maximum number of leapfrog steps is bounded by `2**max_tree_depth`, for
+        example, the number of nodes in a binary tree `max_tree_depth` nodes
+        deep. The default setting of `10` takes up to 1024 leapfrog steps.
+      max_energy_diff: Scalar threshold of energy differences at each leapfrog,
+        divergence samples are defined as leapfrog steps that exceed this
+        threshold. Default is `1000`.
+      unrolled_leapfrog_steps: The number of leapfrogs to unroll per tree
+        expansion step. Applies a direct linear multiplier to the maximum
+        trajectory length implied by `max_tree_depth`. Defaults is `1`.
+      parallel_iterations: Number of iterations allowed to run in parallel. Must
+        be a positive integer. For more information, see `tf.while_loop`.
+      seed: Used to set the seed for reproducible results. For more information,
+        see [PRNGS and seeds]
+        (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
+      **pins: These are used to condition the provided joint distribution, and
+        are passed directly to `joint_dist.experimental_pin(**pins)`.
+
+    Returns:
+      An Arviz `InferenceData` object containing posterior samples only.
+
+    Throws:
+      MCMCOOMError: If the model is out of memory. Try reducing `n_keep` or pass
+        a list of integers as `n_chains` to sample chains serially. For more
+        information, see
+        [ResourceExhaustedError when running Meridian.sample_posterior]
+        (https://developers.google.com/meridian/docs/advanced-modeling/model-debugging#gpu-oom-error).
+    """
+    seed = tfp.random.sanitize_seed(seed) if seed else None
+    n_chains_list = [n_chains] if isinstance(n_chains, int) else n_chains
+    total_chains = np.sum(n_chains_list)
+
+    states = []
+    traces = []
+    for n_chains_batch in n_chains_list:
+      try:
+        mcmc = _xla_windowed_adaptive_nuts(
+            n_draws=n_burnin + n_keep,
+            joint_dist=self._get_joint_dist(),
+            n_chains=n_chains_batch,
+            num_adaptation_steps=n_adapt,
+            current_state=current_state,
+            init_step_size=init_step_size,
+            dual_averaging_kwargs=dual_averaging_kwargs,
+            max_tree_depth=max_tree_depth,
+            max_energy_diff=max_energy_diff,
+            unrolled_leapfrog_steps=unrolled_leapfrog_steps,
+            parallel_iterations=parallel_iterations,
+            seed=seed,
+            **pins,
+        )
+      except tf.errors.ResourceExhaustedError as error:
+        raise MCMCOOMError(
+            "ERROR: Out of memory. Try reducing `n_keep` or pass a list of"
+            " integers as `n_chains` to sample chains serially (see"
+            " https://developers.google.com/meridian/docs/advanced-modeling/model-debugging#gpu-oom-error)"
+        ) from error
+      states.append(mcmc.all_states._asdict())
+      traces.append(mcmc.trace)
+
+    mcmc_states = {
+        k: tf.einsum(
+            "ij...->ji...",
+            tf.concat([state[k] for state in states], axis=1)[n_burnin:, ...],
+        )
+        for k in states[0].keys()
+        if k not in constants.UNSAVED_PARAMETERS
+    }
+    # Create Arviz InferenceData for posterior draws.
+    posterior_coords = self._meridian.create_inference_data_coords(
+        total_chains, n_keep
+    )
+    posterior_dims = self._meridian.create_inference_data_dims()
+    infdata_posterior = az.convert_to_inference_data(
+        mcmc_states, coords=posterior_coords, dims=posterior_dims
+    )
+
+    # Save trace metrics in InferenceData.
+    mcmc_trace = {}
+    for k in traces[0].keys():
+      if k not in constants.IGNORED_TRACE_METRICS:
+        mcmc_trace[k] = tf.concat(
+            [
+                tf.broadcast_to(
+                    tf.transpose(trace[k][n_burnin:, ...]),
+                    [n_chains_list[i], n_keep],
+                )
+                for i, trace in enumerate(traces)
+            ],
+            axis=0,
+        )
+
+    trace_coords = {
+        constants.CHAIN: np.arange(total_chains),
+        constants.DRAW: np.arange(n_keep),
+    }
+    trace_dims = {
+        k: [constants.CHAIN, constants.DRAW] for k in mcmc_trace.keys()
+    }
+    infdata_trace = az.convert_to_inference_data(
+        mcmc_trace, coords=trace_coords, dims=trace_dims, group="trace"
+    )
+
+    # Create Arviz InferenceData for divergent transitions and other sampling
+    # statistics. Note that InferenceData has a different naming convention
+    # than Tensorflow, and only certain variables are recongnized.
+    # https://arviz-devs.github.io/arviz/schema/schema.html#sample-stats
+    # The list of values returned by windowed_adaptive_nuts() is the following:
+    # 'step_size', 'tune', 'target_log_prob', 'diverging', 'accept_ratio',
+    # 'variance_scaling', 'n_steps', 'is_accepted'.
+
+    sample_stats = {
+        constants.SAMPLE_STATS_METRICS[k]: v
+        for k, v in mcmc_trace.items()
+        if k in constants.SAMPLE_STATS_METRICS
+    }
+    sample_stats_dims = {
+        constants.SAMPLE_STATS_METRICS[k]: v
+        for k, v in trace_dims.items()
+        if k in constants.SAMPLE_STATS_METRICS
+    }
+    # Tensorflow does not include a "draw" dimension on step size metric if same
+    # step size is used for all chains. Step size must be broadcast to the
+    # correct shape.
+    sample_stats[constants.STEP_SIZE] = tf.broadcast_to(
+        sample_stats[constants.STEP_SIZE], [total_chains, n_keep]
+    )
+    sample_stats_dims[constants.STEP_SIZE] = [constants.CHAIN, constants.DRAW]
+    infdata_sample_stats = az.convert_to_inference_data(
+        sample_stats,
+        coords=trace_coords,
+        dims=sample_stats_dims,
+        group="sample_stats",
+    )
+    return az.concat(infdata_posterior, infdata_trace, infdata_sample_stats)