Merge branch 'main' into ciguaran_fix_smc_after_bump

Author: ciguaran

.github/workflows/pypi.yml

@@ -45,35 +45,10 @@ jobs:
         with:
           name: artifact
           path: dist/*
-  test:
-    name: upload to test PyPI
-    needs: [build]
-    runs-on: ubuntu-latest
-    if: github.event_name == 'release' && github.event.action == 'published'
-    steps:
-      - uses: actions/download-artifact@v3
-        with:
-          name: artifact
-          path: dist
-      - uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          skip_existing: true
-          user: __token__
-          password: ${{ secrets.TEST_PYPI_API_TOKEN }}
-          repository_url: https://test.pypi.org/legacy/
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-      - name: Test pip install from test.pypi
-        run: |
-          python -m venv venv-test-pypi
-          venv-test-pypi/bin/python -m pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple pymc-experimental
-          echo "Checking import and version number"
-          venv-test-pypi/bin/python -c "import pymc_experimental; assert pymc_experimental.__version__ == '${{ github.ref_name }}'[1:]"
 
   publish:
     name: upload release to PyPI
-    needs: [build, test]
+    needs: [build]
     runs-on: ubuntu-latest
     if: github.event_name == 'release' && github.event.action == 'published'
     steps:

.github/workflows/test.yml

@@ -73,6 +73,7 @@ jobs:
       - uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: conda-envs/windows-environment-test.yml
+          micromamba-version: "1.5.10-0"  # Until https://github.com/mamba-org/mamba/issues/3467 is not fixed
           create-args: >-
             python=${{matrix.python-version}}
           environment-name: pymc-experimental-test

.pre-commit-config.yaml

@@ -1,3 +1,6 @@
+ci:
+  autofix_prs: false
+
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v4.6.0

conda-envs/environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.16.1  # CI was failing to resolve
-  - blackjax>=1.2.3
+  - pymc>=5.17.0  # CI was failing to resolve
+  - blackjax
   - scikit-learn

conda-envs/windows-environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.16.1  # CI was failing to resolve
+  - pymc>=5.17.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

pymc_experimental/__init__.py

@@ -13,9 +13,10 @@
 #   limitations under the License.
 import logging
 
-from pymc_experimental import distributions, gp, statespace, utils
+from pymc_experimental import gp, statespace, utils
+from pymc_experimental.distributions import *
 from pymc_experimental.inference.fit import fit
-from pymc_experimental.model.marginal_model import MarginalModel
+from pymc_experimental.model.marginal.marginal_model import MarginalModel
 from pymc_experimental.model.model_api import as_model
 from pymc_experimental.version import __version__
 
@@ -26,15 +27,3 @@
     if len(_log.handlers) == 0:
         handler = logging.StreamHandler()
         _log.addHandler(handler)
-
-
-__all__ = [
-    "distributions",
-    "gp",
-    "statespace",
-    "utils",
-    "fit",
-    "MarginalModel",
-    "as_model",
-    "__version__",
-]

pymc_experimental/distributions/timeseries.py

@@ -20,7 +20,12 @@
 from pymc.logprob.abstract import _logprob
 from pymc.logprob.basic import logp
 from pymc.pytensorf import constant_fold, intX
-from pymc.util import check_dist_not_registered
+from pymc.step_methods import STEP_METHODS
+from pymc.step_methods.arraystep import ArrayStep
+from pymc.step_methods.compound import Competence
+from pymc.step_methods.metropolis import CategoricalGibbsMetropolis
+from pymc.util import check_dist_not_registered, get_value_vars_from_user_vars
+from pytensor import Mode
 from pytensor.graph.basic import Node
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
@@ -101,10 +106,15 @@ class DiscreteMarkovChain(Distribution):
      Create a Markov Chain of length 100 with 3 states. The number of states is given by the shape of P,
      3 in this case.
 
-    >>> with pm.Model() as markov_chain:
-    >>>     P = pm.Dirichlet("P", a=[1, 1, 1], size=(3,))
-    >>>     init_dist = pm.Categorical.dist(p = np.full(3, 1 / 3))
-    >>>     markov_chain = pm.DiscreteMarkovChain("markov_chain", P=P, init_dist=init_dist, shape=(100,))
+    .. code-block:: python
+
+        import pymc as pm
+        import pymc_experimental as pmx
+
+        with pm.Model() as markov_chain:
+            P = pm.Dirichlet("P", a=[1, 1, 1], size=(3,))
+            init_dist = pm.Categorical.dist(p = np.full(3, 1 / 3))
+            markov_chain = pmx.DiscreteMarkovChain("markov_chain", P=P, init_dist=init_dist, shape=(100,))
 
     """
 
@@ -266,3 +276,69 @@ def discrete_mc_logp(op, values, P, steps, init_dist, state_rng, **kwargs):
         "P must sum to 1 along the last axis, "
         "First dimension of init_dist must be n_lags",
     )
+
+
+class DiscreteMarkovChainGibbsMetropolis(CategoricalGibbsMetropolis):
+    name = "discrete_markov_chain_gibbs_metropolis"
+
+    def __init__(self, vars, proposal="uniform", order="random", model=None):
+        model = pm.modelcontext(model)
+        vars = get_value_vars_from_user_vars(vars, model)
+        initial_point = model.initial_point()
+
+        dimcats = []
+        # The above variable is a list of pairs (aggregate dimension, number
+        # of categories). For example, if vars = [x, y] with x being a 2-D
+        # variable with M categories and y being a 3-D variable with N
+        # categories, we will have dimcats = [(0, M), (1, M), (2, N), (3, N), (4, N)].
+        for v in vars:
+            v_init_val = initial_point[v.name]
+            rv_var = model.values_to_rvs[v]
+            rv_op = rv_var.owner.op
+
+            if not isinstance(rv_op, DiscreteMarkovChainRV):
+                raise TypeError("All variables must be DiscreteMarkovChainRV")
+
+            k_graph = rv_var.owner.inputs[0].shape[-1]
+            (k_graph,) = model.replace_rvs_by_values((k_graph,))
+            k = model.compile_fn(
+                k_graph,
+                inputs=model.value_vars,
+                on_unused_input="ignore",
+                mode=Mode(linker="py", optimizer=None),
+            )(initial_point)
+            start = len(dimcats)
+            dimcats += [(dim, k) for dim in range(start, start + v_init_val.size)]
+
+        if order == "random":
+            self.shuffle_dims = True
+            self.dimcats = dimcats
+        else:
+            if sorted(order) != list(range(len(dimcats))):
+                raise ValueError("Argument 'order' has to be a permutation")
+            self.shuffle_dims = False
+            self.dimcats = [dimcats[j] for j in order]
+
+        if proposal == "uniform":
+            self.astep = self.astep_unif
+        elif proposal == "proportional":
+            # Use the optimized "Metropolized Gibbs Sampler" described in Liu96.
+            self.astep = self.astep_prop
+        else:
+            raise ValueError("Argument 'proposal' should either be 'uniform' or 'proportional'")
+
+        # Doesn't actually tune, but it's required to emit a sampler stat
+        # that indicates whether a draw was done in a tuning phase.
+        self.tune = True
+
+        # We bypass CategoryGibbsMetropolis's __init__ to avoid it's specialiazed initialization logic
+        ArrayStep.__init__(self, vars, [model.compile_logp()])
+
+    @staticmethod
+    def competence(var):
+        if isinstance(var.owner.op, DiscreteMarkovChainRV):
+            return Competence.IDEAL
+        return Competence.INCOMPATIBLE
+
+
+STEP_METHODS.append(DiscreteMarkovChainGibbsMetropolis)