Merge branch 'main' into add-zerosumnormal

Author: AlexAndorra

pymc/distributions/discrete.py

@@ -18,6 +18,7 @@
 
 from aesara.tensor.random.basic import (
     RandomVariable,
+    ScipyRandomVariable,
     bernoulli,
     betabinom,
     binomial,
@@ -1117,15 +1118,15 @@ def logcdf(value, good, bad, n):
         )
 
 
-class DiscreteUniformRV(RandomVariable):
+class DiscreteUniformRV(ScipyRandomVariable):
     name = "discrete_uniform"
     ndim_supp = 0
     ndims_params = [0, 0]
     dtype = "int64"
     _print_name = ("DiscreteUniform", "\\operatorname{DiscreteUniform}")
 
     @classmethod
-    def rng_fn(cls, rng, lower, upper, size=None):
+    def rng_fn_scipy(cls, rng, lower, upper, size=None):
         return stats.randint.rvs(lower, upper + 1, size=size, random_state=rng)
 
 

pymc/distributions/timeseries.py

@@ -19,7 +19,6 @@
 import aesara.tensor as at
 import numpy as np
 
-from aeppl.abstract import _get_measurable_outputs
 from aeppl.logprob import _logprob
 from aesara.graph.basic import Node, clone_replace
 from aesara.tensor import TensorVariable
@@ -144,12 +143,6 @@ def rv_op(cls, init_dist, innovation_dist, steps, size=None):
         )(init_dist, innovation_dist, steps)
 
 
-@_get_measurable_outputs.register(RandomWalkRV)
-def _get_measurable_outputs_random_walk(op, node):
-    # Ignore steps output
-    return [node.default_output()]
-
-
 @_change_dist_size.register(RandomWalkRV)
 def change_random_walk_size(op, dist, new_size, expand):
     init_dist, innovation_dist, steps = dist.owner.inputs

pymc/model.py

@@ -1149,7 +1149,7 @@ def add_coord(
             length = len(values)
         if not isinstance(length, Variable):
             if mutable:
-                length = aesara.shared(length)
+                length = aesara.shared(length, name=name)
             else:
                 length = aesara.tensor.constant(length)
         self._dim_lengths[name] = length

pymc/sampling.py

@@ -1621,6 +1621,8 @@ def compile_forward_sampling_function(
     vars_in_trace: List[Variable],
     basic_rvs: Optional[List[Variable]] = None,
     givens_dict: Optional[Dict[Variable, Any]] = None,
+    constant_data: Optional[Dict[str, np.ndarray]] = None,
+    constant_coords: Optional[Set[str]] = None,
     **kwargs,
 ) -> Tuple[Callable[..., Union[np.ndarray, List[np.ndarray]]], Set[Variable]]:
     """Compile a function to draw samples, conditioned on the values of some variables.
@@ -1634,18 +1636,18 @@ def compile_forward_sampling_function(
     compiled function or after inference has been run. These variables are:
 
     - Variables in the outputs list
-    - ``SharedVariable`` instances that are not ``RandomStateSharedVariable`` or ``RandomGeneratorSharedVariable``
+    - ``SharedVariable`` instances that are not ``RandomStateSharedVariable`` or ``RandomGeneratorSharedVariable``, and whose values changed with respect to what they were at inference time
     - Variables that are in the `basic_rvs` list but not in the ``vars_in_trace`` list
     - Variables that are keys in the ``givens_dict``
     - Variables that have volatile inputs
 
     Concretely, this function can be used to compile a function to sample from the
     posterior predictive distribution of a model that has variables that are conditioned
-    on ``MutableData`` instances. The variables that depend on the mutable data will be
-    considered volatile, and as such, they wont be included as inputs into the compiled function.
-    This means that if they have values stored in the posterior, these values will be ignored
-    and new values will be computed (in the case of deterministics and potentials) or sampled
-    (in the case of random variables).
+    on ``MutableData`` instances. The variables that depend on the mutable data that have changed
+    will be considered volatile, and as such, they wont be included as inputs into the compiled
+    function. This means that if they have values stored in the posterior, these values will be
+    ignored and new values will be computed (in the case of deterministics and potentials) or
+    sampled (in the case of random variables).
 
     This function also enables a way to impute values for any variable in the computational
     graph that produces the desired outputs: the ``givens_dict``. This dictionary can be used
@@ -1672,6 +1674,25 @@ def compile_forward_sampling_function(
         A dictionary that maps tensor variables to the values that should be used to replace them
         in the compiled function. The types of the key and value should match or an error will be
         raised during compilation.
+    constant_data : Optional[Dict[str, numpy.ndarray]]
+        A dictionary that maps the names of ``MutableData`` or ``ConstantData`` instances to their
+        corresponding values at inference time. If a model was created with ``MutableData``, these
+        are stored as ``SharedVariable`` with the name of the data variable and a value equal to
+        the initial data. At inference time, this information is stored in ``InferenceData``
+        objects under the ``constant_data`` group, which allows us to check whether a
+        ``SharedVariable`` instance changed its values after inference or not. If the values have
+        changed, then the ``SharedVariable`` is assumed to be volatile. If it has not changed, then
+        the ``SharedVariable`` is assumed to not be volatile. If a ``SharedVariable`` is not found
+        in either ``constant_data`` or ``constant_coords``, then it is assumed to be volatile.
+        Setting ``constant_data`` to ``None`` is equivalent to passing an empty dictionary.
+    constant_coords : Optional[Set[str]]
+        A set with the names of the mutable coordinates that have not changed their shape after
+        inference. If a model was created with mutable coordinates, these are stored as
+        ``SharedVariable`` with the name of the coordinate and a value equal to the length of said
+        coordinate. This set let's us check if a ``SharedVariable`` is a mutated coordinate, in
+        which case, it is considered volatile. If a ``SharedVariable`` is not found
+        in either ``constant_data`` or ``constant_coords``, then it is assumed to be volatile.
+        Setting ``constant_coords`` to ``None`` is equivalent to passing an empty set.
 
     Returns
     -------
@@ -1687,6 +1708,20 @@ def compile_forward_sampling_function(
     if basic_rvs is None:
         basic_rvs = []
 
+    if constant_data is None:
+        constant_data = {}
+    if constant_coords is None:
+        constant_coords = set()
+
+    # We define a helper function to check if shared values match to an array
+    def shared_value_matches(var):
+        try:
+            old_array_value = constant_data[var.name]
+        except KeyError:
+            return var.name in constant_coords
+        current_shared_value = var.get_value(borrow=True)
+        return np.array_equal(old_array_value, current_shared_value)
+
     # We need a function graph to walk the clients and propagate the volatile property
     fg = FunctionGraph(outputs=outputs, clone=False)
 
@@ -1702,6 +1737,7 @@ def compile_forward_sampling_function(
             or (  # SharedVariables, except RandomState/Generators
                 isinstance(node, SharedVariable)
                 and not isinstance(node, (RandomStateSharedVariable, RandomGeneratorSharedVariable))
+                and not shared_value_matches(node)
             )
             or (  # Basic RVs that are not in the trace
                 node in basic_rvs and node not in vars_in_trace
@@ -1835,16 +1871,24 @@ def sample_posterior_predictive(
         idata_kwargs = {}
     else:
         idata_kwargs = idata_kwargs.copy()
+    constant_data: Dict[str, np.ndarray] = {}
+    trace_coords: Dict[str, np.ndarray] = {}
     if "coords" not in idata_kwargs:
         idata_kwargs["coords"] = {}
     if isinstance(trace, InferenceData):
         idata_kwargs["coords"].setdefault("draw", trace["posterior"]["draw"])
         idata_kwargs["coords"].setdefault("chain", trace["posterior"]["chain"])
+        _constant_data = getattr(trace, "constant_data", None)
+        if _constant_data is not None:
+            trace_coords.update({str(k): v.data for k, v in _constant_data.coords.items()})
+            constant_data.update({str(k): v.data for k, v in _constant_data.items()})
+        trace_coords.update({str(k): v.data for k, v in trace["posterior"].coords.items()})
         _trace = dataset_to_point_list(trace["posterior"])
         nchain, len_trace = chains_and_samples(trace)
     elif isinstance(trace, xarray.Dataset):
         idata_kwargs["coords"].setdefault("draw", trace["draw"])
         idata_kwargs["coords"].setdefault("chain", trace["chain"])
+        trace_coords.update({str(k): v.data for k, v in trace.coords.items()})
         _trace = dataset_to_point_list(trace)
         nchain, len_trace = chains_and_samples(trace)
     elif isinstance(trace, MultiTrace):
@@ -1901,6 +1945,16 @@ def sample_posterior_predictive(
             stacklevel=2,
         )
 
+    constant_coords = set()
+    for dim, coord in trace_coords.items():
+        current_coord = model.coords.get(dim, None)
+        if (
+            current_coord is not None
+            and len(coord) == len(current_coord)
+            and np.all(coord == current_coord)
+        ):
+            constant_coords.add(dim)
+
     if var_names is not None:
         vars_ = [model[x] for x in var_names]
     else:
@@ -1935,6 +1989,8 @@ def sample_posterior_predictive(
         basic_rvs=model.basic_RVs,
         givens_dict=None,
         random_seed=random_seed,
+        constant_data=constant_data,
+        constant_coords=constant_coords,
         **compile_kwargs,
     )
     sampler_fn = point_wrapper(_sampler_fn)

pymc/smc/smc.py

@@ -274,9 +274,7 @@ def update_beta_and_weights(self):
 
     def resample(self):
         """Resample particles based on importance weights"""
-        self.resampling_indexes = self.rng.choice(
-            np.arange(self.draws), size=self.draws, p=self.weights
-        )
+        self.resampling_indexes = systematic_resampling(self.weights, self.rng)
 
         self.tempered_posterior = self.tempered_posterior[self.resampling_indexes]
         self.prior_logp = self.prior_logp[self.resampling_indexes]
@@ -546,6 +544,36 @@ def sample_settings(self):
         return stats
 
 
+def systematic_resampling(weights, rng):
+    """
+    Systematic resampling.
+
+    Parameters
+    ----------
+    weights :
+        The weights should be probabilities and the total sum should be 1.
+
+    Returns
+    -------
+    new_indices: array
+        A vector of indices in the interval 0, ..., len(normalized_weights)
+    """
+    lnw = len(weights)
+    arange = np.arange(lnw)
+    uniform = (rng.random(1) + arange) / lnw
+
+    idx = 0
+    weight_accu = weights[0]
+    new_indices = np.empty(lnw, dtype=int)
+    for i in arange:
+        while uniform[i] > weight_accu:
+            idx += 1
+            weight_accu += weights[idx]
+        new_indices[i] = idx
+
+    return new_indices
+
+
 def _logp_forw(point, out_vars, in_vars, shared):
     """Compile Aesara function of the model and the input and output variables.
 

pymc/tests/distributions/test_discrete.py

@@ -1042,6 +1042,10 @@ def discrete_uniform_rng_fn(self, size, lower, upper, rng):
         "check_rv_size",
     ]
 
+    def test_implied_degenerate_shape(self):
+        x = pm.DiscreteUniform.dist(0, [1])
+        assert x.eval().shape == (1,)
+
 
 class TestDiracDelta(BaseTestDistributionRandom):
     def diracdelta_rng_fn(self, size, c):

pymc/tests/distributions/test_distribution.py

@@ -339,11 +339,25 @@ class TestInlinedSymbolicRV(SymbolicRandomVariable):
         x_inline = TestInlinedSymbolicRV([], [Flat.dist()], ndim_supp=0)()
         assert np.isclose(logp(x_inline, 0).eval(), 0)
 
-    def test_measurable_outputs(self):
+    def test_measurable_outputs_rng_ignored(self):
+        """Test that any RandomType outputs are ignored as a measurable_outputs"""
+
         class TestSymbolicRV(SymbolicRandomVariable):
             pass
 
         next_rng_, dirac_delta_ = DiracDelta.dist(5).owner.outputs
         next_rng, dirac_delta = TestSymbolicRV([], [next_rng_, dirac_delta_], ndim_supp=0)()
         node = dirac_delta.owner
         assert get_measurable_outputs(node.op, node) == [dirac_delta]
+
+    @pytest.mark.parametrize("default_output_idx", (0, 1))
+    def test_measurable_outputs_default_output(self, default_output_idx):
+        """Test that if provided, a default output is considered the only measurable_output"""
+
+        class TestSymbolicRV(SymbolicRandomVariable):
+            default_output = default_output_idx
+
+        dirac_delta_1_ = DiracDelta.dist(5)
+        dirac_delta_2_ = DiracDelta.dist(10)
+        node = TestSymbolicRV([], [dirac_delta_1_, dirac_delta_2_], ndim_supp=0)().owner
+        assert get_measurable_outputs(node.op, node) == [node.outputs[default_output_idx]]