Merge branch 'pymc-devs:main' into MvStudentT

Author: patel-zeel

pymc/backends/arviz.py

@@ -42,6 +42,26 @@
 Var = Any  # pylint: disable=invalid-name
 
 
+def find_observations(model: Optional["Model"]) -> Optional[Dict[str, Var]]:
+    """If there are observations available, return them as a dictionary."""
+    if model is None:
+        return None
+
+    observations = {}
+    for obs in model.observed_RVs:
+        aux_obs = getattr(obs.tag, "observations", None)
+        if aux_obs is not None:
+            try:
+                obs_data = extract_obs_data(aux_obs)
+                observations[obs.name] = obs_data
+            except TypeError:
+                warnings.warn(f"Could not extract data from symbolic observation {obs}")
+        else:
+            warnings.warn(f"No data for observation {obs}")
+
+    return observations
+
+
 class _DefaultTrace:
     """
     Utility for collecting samples into a dictionary.
@@ -196,25 +216,7 @@ def arbitrary_element(dct: Dict[Any, np.ndarray]) -> np.ndarray:
             self.dims = {**model_dims, **self.dims}
 
         self.density_dist_obs = density_dist_obs
-        self.observations = self.find_observations()
-
-    def find_observations(self) -> Optional[Dict[str, Var]]:
-        """If there are observations available, return them as a dictionary."""
-        if self.model is None:
-            return None
-        observations = {}
-        for obs in self.model.observed_RVs:
-            aux_obs = getattr(obs.tag, "observations", None)
-            if aux_obs is not None:
-                try:
-                    obs_data = extract_obs_data(aux_obs)
-                    observations[obs.name] = obs_data
-                except TypeError:
-                    warnings.warn(f"Could not extract data from symbolic observation {obs}")
-            else:
-                warnings.warn(f"No data for observation {obs}")
-
-        return observations
+        self.observations = find_observations(self.model)
 
     def split_trace(self) -> Tuple[Union[None, "MultiTrace"], Union[None, "MultiTrace"]]:
         """Split MultiTrace object into posterior and warmup.

pymc/bart/bart.py

@@ -19,7 +19,7 @@
 from aesara.tensor.random.op import RandomVariable, default_shape_from_params
 from pandas import DataFrame, Series
 
-from pymc.distributions.distribution import NoDistribution
+from pymc.distributions.distribution import NoDistribution, _get_moment
 
 __all__ = ["BART"]
 
@@ -110,6 +110,10 @@ def __new__(
 
         NoDistribution.register(BARTRV)
 
+        @_get_moment.register(BARTRV)
+        def get_moment(rv, size, *rv_inputs):
+            return cls.get_moment(rv, size, *rv_inputs)
+
         cls.rv_op = bart_op
         params = [X, Y, m, alpha, k]
         return super().__new__(cls, name, *params, **kwargs)
@@ -132,6 +136,11 @@ def logp(x, *inputs):
         """
         return at.zeros_like(x)
 
+    @classmethod
+    def get_moment(cls, rv, size, *rv_inputs):
+        mean = at.fill(size, rv.Y.mean())
+        return mean
+
 
 def preprocess_XY(X, Y):
     if isinstance(Y, (Series, DataFrame)):

pymc/distributions/continuous.py

@@ -2334,23 +2334,19 @@ def dist(cls, alpha=None, beta=None, mu=None, sigma=None, sd=None, *args, **kwar
         alpha = at.as_tensor_variable(floatX(alpha))
         beta = at.as_tensor_variable(floatX(beta))
 
-        # m = beta / (alpha - 1.0)
-        # try:
-        #     mean = (alpha > 1) * m or np.inf
-        # except ValueError:  # alpha is an array
-        #     m[alpha <= 1] = np.inf
-        #     mean = m
-
-        # mode = beta / (alpha + 1.0)
-        # variance = at.switch(
-        #     at.gt(alpha, 2), (beta ** 2) / ((alpha - 2) * (alpha - 1.0) ** 2), np.inf
-        # )
-
         assert_negative_support(alpha, "alpha", "InverseGamma")
         assert_negative_support(beta, "beta", "InverseGamma")
 
         return super().dist([alpha, beta], **kwargs)
 
+    def get_moment(rv, size, alpha, beta):
+        mean = beta / (alpha - 1.0)
+        mode = beta / (alpha + 1.0)
+        moment = at.switch(alpha > 1, mean, mode)
+        if not rv_size_is_none(size):
+            moment = at.full(size, moment)
+        return moment
+
     @classmethod
     def _get_alpha_beta(cls, alpha, beta, mu, sigma):
         if alpha is not None:

pymc/distributions/discrete.py

@@ -1462,7 +1462,7 @@ class ZeroInflatedBinomial(Discrete):
 
     ========  ==========================
     Support   :math:`x \in \mathbb{N}_0`
-    Mean      :math:`(1 - \psi) n p`
+    Mean      :math:`\psi n p`
     Variance  :math:`(1-\psi) n p [1 - p(1 - \psi n)].`
     ========  ==========================
 
@@ -1487,7 +1487,7 @@ def dist(cls, psi, n, p, *args, **kwargs):
         return super().dist([psi, n, p], *args, **kwargs)
 
     def get_moment(rv, size, psi, n, p):
-        mean = at.round((1 - psi) * n * p)
+        mean = at.round(psi * n * p)
         if not rv_size_is_none(size):
             mean = at.full(size, mean)
         return mean
@@ -1650,6 +1650,12 @@ def dist(cls, psi, mu, alpha, *args, **kwargs):
         p = at.as_tensor_variable(floatX(p))
         return super().dist([psi, n, p], *args, **kwargs)
 
+    def get_moment(rv, size, psi, n, p):
+        mean = at.floor(psi * n * (1 - p) / p)
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logp(value, psi, n, p):
         r"""
         Calculate log-probability of ZeroInflatedNegativeBinomial distribution at specified value.

pymc/distributions/simulator.py

@@ -25,7 +25,7 @@
 from scipy.spatial import cKDTree
 
 from pymc.aesaraf import floatX
-from pymc.distributions.distribution import NoDistribution
+from pymc.distributions.distribution import NoDistribution, _get_moment
 
 __all__ = ["Simulator"]
 
@@ -223,13 +223,23 @@ def logp(op, value_var_list, *dist_params, **kwargs):
             value_var = value_var_list[0]
             return cls.logp(value_var, op, dist_params)
 
+        @_get_moment.register(SimulatorRV)
+        def get_moment(op, rv, size, *rv_inputs):
+            return cls.get_moment(rv, size, *rv_inputs)
+
         cls.rv_op = sim_op
         return super().__new__(cls, name, *params, **kwargs)
 
     @classmethod
     def dist(cls, *params, **kwargs):
         return super().dist(params, **kwargs)
 
+    @classmethod
+    def get_moment(cls, rv, size, *sim_inputs):
+        # Take the mean of 10 draws
+        multiple_sim = rv.owner.op(*sim_inputs, size=at.concatenate([[10], rv.shape]))
+        return at.mean(multiple_sim, axis=0)
+
     @classmethod
     def logp(cls, value, sim_op, sim_inputs):
         # Use a new rng to avoid non-randomness in parallel sampling

pymc/sampling_jax.py

@@ -26,7 +26,9 @@
 from aesara.link.jax.dispatch import jax_funcify
 
 from pymc import Model, modelcontext
-from pymc.aesaraf import compile_rv_inplace, inputvars
+from pymc.aesaraf import compile_rv_inplace
+from pymc.backends.arviz import find_observations
+from pymc.distributions import logpt
 from pymc.util import get_default_varnames
 
 warnings.warn("This module is experimental.")
@@ -95,6 +97,40 @@ def logp_fn_wrap(x):
     return logp_fn_wrap
 
 
+# Adopted from arviz numpyro extractor
+def _sample_stats_to_xarray(posterior):
+    """Extract sample_stats from NumPyro posterior."""
+    rename_key = {
+        "potential_energy": "lp",
+        "adapt_state.step_size": "step_size",
+        "num_steps": "n_steps",
+        "accept_prob": "acceptance_rate",
+    }
+    data = {}
+    for stat, value in posterior.get_extra_fields(group_by_chain=True).items():
+        if isinstance(value, (dict, tuple)):
+            continue
+        name = rename_key.get(stat, stat)
+        value = value.copy()
+        data[name] = value
+        if stat == "num_steps":
+            data["tree_depth"] = np.log2(value).astype(int) + 1
+    return data
+
+
+def _get_log_likelihood(model, samples):
+    "Compute log-likelihood for all observations"
+    data = {}
+    for v in model.observed_RVs:
+        logp_v = replace_shared_variables([logpt(v)])
+        fgraph = FunctionGraph(model.value_vars, logp_v, clone=False)
+        optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
+        jax_fn = jax_funcify(fgraph)
+        result = jax.jit(jax.vmap(jax.vmap(jax_fn)))(*samples)[0]
+        data[v.name] = result
+    return data
+
+
 def sample_numpyro_nuts(
     draws=1000,
     tune=1000,
@@ -115,6 +151,20 @@ def sample_numpyro_nuts(
 
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
+    coords = {
+        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
+        for cname, cvals in model.coords.items()
+        if cvals is not None
+    }
+
+    if hasattr(model, "RV_dims"):
+        dims = {
+            var_name: [dim for dim in dims if dim is not None]
+            for var_name, dims in model.RV_dims.items()
+        }
+    else:
+        dims = {}
+
     tic1 = pd.Timestamp.now()
     print("Compiling...", file=sys.stdout)
 
@@ -151,9 +201,23 @@ def sample_numpyro_nuts(
     map_seed = jax.random.split(seed, chains)
 
     if chains == 1:
-        pmap_numpyro.run(seed, init_params=init_state, extra_fields=("num_steps",))
+        init_params = init_state
+        map_seed = seed
     else:
-        pmap_numpyro.run(map_seed, init_params=init_state_batched, extra_fields=("num_steps",))
+        init_params = init_state_batched
+
+    pmap_numpyro.run(
+        map_seed,
+        init_params=init_params,
+        extra_fields=(
+            "num_steps",
+            "potential_energy",
+            "energy",
+            "adapt_state.step_size",
+            "accept_prob",
+            "diverging",
+        ),
+    )
 
     raw_mcmc_samples = pmap_numpyro.get_samples(group_by_chain=True)
 
@@ -164,6 +228,7 @@ def sample_numpyro_nuts(
     mcmc_samples = {}
     for v in vars_to_sample:
         fgraph = FunctionGraph(model.value_vars, [v], clone=False)
+        optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
         jax_fn = jax_funcify(fgraph)
         result = jax.vmap(jax.vmap(jax_fn))(*raw_mcmc_samples)[0]
         mcmc_samples[v.name] = result
@@ -172,6 +237,13 @@ def sample_numpyro_nuts(
     print("Transformation time = ", tic4 - tic3, file=sys.stdout)
 
     posterior = mcmc_samples
-    az_trace = az.from_dict(posterior=posterior)
+    az_trace = az.from_dict(
+        posterior=posterior,
+        log_likelihood=_get_log_likelihood(model, raw_mcmc_samples),
+        observed_data=find_observations(model),
+        sample_stats=_sample_stats_to_xarray(pmap_numpyro),
+        coords=coords,
+        dims=dims,
+    )
 
     return az_trace

pymc/tests/test_bart.py

@@ -6,6 +6,8 @@
 
 import pymc as pm
 
+from pymc.tests.test_distributions_moments import assert_moment_is_expected
+
 
 def test_split_node():
     split_node = pm.bart.tree.SplitNode(index=5, idx_split_variable=2, split_value=3.0)
@@ -97,3 +99,17 @@ def test_predict(self):
     )
     def test_pdp(self, kwargs):
         pm.bart.utils.plot_dependence(self.idata, X=self.X, Y=self.Y, **kwargs)
+
+
+@pytest.mark.parametrize(
+    "size, expected",
+    [
+        (None, np.zeros(50)),
+    ],
+)
+def test_bart_moment(size, expected):
+    X = np.zeros((50, 2))
+    Y = np.zeros(50)
+    with pm.Model() as model:
+        pm.BART("x", X=X, Y=Y, size=size)
+    assert_moment_is_expected(model, expected)