Repository: pymc-devs/mcbackend (language: python)

Issue:
Add an emcee adapter
See https://emcee.readthedocs.io/en/stable/index.html for a minimal example


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 4b02452..559079b 100644
--- a/README.md
+++ b/README.md
@@ -41,7 +41,7 @@ isinstance(backend, mcbackend.Backend)
 ```
 
 ### Part 3: PPL adapters
-Anything that is a `Backend` can be wrapped by an [adapter ](https://en.wikipedia.org/wiki/Adapter_pattern) that makes it compatible with your favorite PPL.
+Anything that is a `Backend` can be wrapped by an [adapter](https://en.wikipedia.org/wiki/Adapter_pattern) that makes it compatible with your favorite PPL.
 
 In the example below, a `ClickHouseBackend` is initialized to store MCMC draws from a PyMC model in a [ClickHouse](http://clickhouse.com/) database.
 See below for [how to run it in Docker](#development).
@@ -56,17 +56,15 @@ ch_client = clickhouse_driver.Client("localhost")
 backend = mcbackend.ClickHouseBackend(ch_client)
 
 with pm.Model():
-    # 3. Create your model
+    # 2. Create your model
     ...
-    # 4. Wrap the PyMC adapter around an `mcbackend.Backend`
-    #    This generates and prints a short `trace.run_id` by which
-    #    this MCMC run is identified in the (database) backend.
-    trace = mcbackend.pymc.TraceBackend(backend)
-
-    # 5. Hit the inference button ™
-    pm.sample(trace=trace)
+    # 3. Hit the inference button ™ while passing the backend!
+    pm.sample(trace=backend)
 ```
 
+In case of PyMC the adapter lives in the PyMC codebase [since version 5.1.1](https://github.com/pymc-devs/pymc/releases/tag/v5.1.1),
+so all you need to do is pass any `mcbackend.Backend` via the `pm.sample(trace=...)` parameter!
+
 Instead of using PyMC's built-in NumPy backend, the MCMC draws now end up in ClickHouse.
 
 ### Retrieving the `draws` & `stats`
@@ -103,7 +101,6 @@ For example:
 * Schema discussion: Which metadata is needed? (related: [PyMC #5160](https://github.com/pymc-devs/pymc/issues/5160))
 * Interface discussion: How should `Backend`/`Run`/`Chain` evolve?
 * Python Backends for disk storage (HDF5, `*.proto`, ...)
-* An `emcee` adapter (#11).
 * C++ `Backend`/`Run`/`Chain` interfaces
 * C++ ClickHouse backend (via [`clickhouse-cpp`](https://github.com/ClickHouse/clickhouse-cpp))
 
diff --git a/mcbackend/__init__.py b/mcbackend/__init__.py
index f8a7f40..47ddd63 100644
--- a/mcbackend/__init__.py
+++ b/mcbackend/__init__.py
@@ -12,4 +12,4 @@
 except ModuleNotFoundError:
     pass
 
-__version__ = "0.4.0"
+__version__ = "0.5.0"
diff --git a/mcbackend/adapters/__init__.py b/mcbackend/adapters/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/mcbackend/adapters/pymc.py b/mcbackend/adapters/pymc.py
deleted file mode 100644
index 8412a08..0000000
--- a/mcbackend/adapters/pymc.py
+++ /dev/null
@@ -1,260 +0,0 @@
-"""
-This module implements an adapter to use any ``mcbackend.Backend`` as a PyMC ``BaseTrace``.
-
-The only PyMC dependency is on the ``BaseTrace`` abstract base class.
-"""
-import base64
-import pickle
-from typing import List, Mapping, Optional, Sequence, Tuple, Union
-
-import hagelkorn
-import numpy
-from pymc.backends.base import BaseTrace
-from pymc.model import Model
-from pytensor.graph.basic import Constant
-from pytensor.tensor.sharedvar import SharedVariable
-
-from mcbackend.meta import Coordinate, DataVariable, Variable
-
-from ..core import Backend, Chain, Run, RunMeta
-from ..npproto.utils import ndarray_from_numpy
-
-
-def find_data(pmodel: Model) -> List[DataVariable]:
-    """Extracts data variables from a model."""
-    observed_rvs = {pmodel.rvs_to_values[rv] for rv in pmodel.observed_RVs}
-    dvars = []
-    # All data containers are named vars!
-    for name, var in pmodel.named_vars.items():
-        dv = DataVariable(name)
-        if isinstance(var, Constant):
-            dv.value = ndarray_from_numpy(var.data)
-        elif isinstance(var, SharedVariable):
-            dv.value = ndarray_from_numpy(var.get_value())
-        else:
-            continue
-        dv.dims = list(pmodel.named_vars_to_dims.get(name, []))
-        dv.is_observed = var in observed_rvs
-        dvars.append(dv)
-    return dvars
-
-
-class ReadOnlyTrace(BaseTrace):
-    """An in-memory read-only PyMC trace.
-
-    This class exists to avoid creating a dependency on
-    the PyMC NDArray that this project is trying to replace.
-    """
-
-    supports_sampler_stats = True
-
-    def __init__(
-        self,
-        *,
-        from_trace: BaseTrace,
-        length: int,
-        draws: Mapping[str, numpy.ndarray],
-        stats: Sequence[Mapping[str, numpy.ndarray]],
-    ):
-        self._length = length
-        self._draws = draws
-        self._stats = stats
-        super().__init__(from_trace.name, from_trace.model, from_trace.vars)
-        self.chain = from_trace.chain
-        self.sampler_vars = from_trace.sampler_vars
-        self.var_shapes = from_trace.var_shapes
-        self.var_dtypes = from_trace.var_dtypes
-        self.varnames = from_trace.varnames
-
-    def __len__(self) -> int:
-        return self._length
-
-    def get_values(self, varname, burn=0, thin=1):
-        return self._draws[varname][burn::thin]
-
-    def _get_sampler_stats(self, stat_name, sampler_idx, burn, thin):
-        if sampler_idx is None:
-            return [sts[stat_name][burn::thin] for sts in self._stats]
-        return self._stats[sampler_idx][stat_name][burn::thin]
-
-    def point(self, idx):
-        return {vn: vals[idx] for vn, vals in self._draws.items()}
-
-    def _slice(self, idx: slice) -> BaseTrace:
-        sliced = ReadOnlyTrace(
-            from_trace=self,
-            length=len(numpy.arange(self._length)[idx]),
-            draws={k: v[idx] for k, v in self._draws.items()},
-            stats=[{k: v[idx] for k, v in sts.items()} for sts in self._stats],
-        )
-        return sliced
-
-
-class TraceBackend(BaseTrace):
-    """Adapter to create a PyMC backend from any McBackend."""
-
-    supports_sampler_stats = True
-
-    def __init__(  # pylint: disable=W0622
-        self,
-        backend: Backend,
-        *,
-        name: Optional[str] = None,
-        model=None,
-        vars=None,
-        test_point=None,
-    ):
-        self.chain: int = -1
-        super().__init__(name, model, vars, test_point)
-        self.run_id = hagelkorn.random(digits=6)
-        print(f"Backend run id: {self.run_id}")
-        self._backend: Backend = backend
-
-        # Sessions created from the underlying backend
-        self._run: Optional[Run] = None
-        self._chain: Optional[Chain] = None
-        self._stat_groups: List[List[Tuple[str, str]]] = []
-        self._length: int = 0
-
-    def __len__(self) -> int:
-        return self._length
-
-    def setup(
-        self,
-        draws: int,
-        chain: int,
-        sampler_vars: Optional[Sequence[Mapping[str, Union[type, numpy.dtype]]]] = None,
-    ) -> None:
-        super().setup(draws, chain, sampler_vars)
-        self.chain = chain
-
-        # Initialize backend sessions
-        free_rv_names = [rv.name for rv in self.model.free_RVs]
-        if not self._run:
-            variables = [
-                Variable(
-                    name,
-                    str(self.var_dtypes[name]),
-                    list(self.var_shapes[name]),
-                    dims=list(self.model.named_vars_to_dims[name])
-                    if name in self.model.named_vars_to_dims
-                    else [],
-                    is_deterministic=(name not in free_rv_names),
-                )
-                for name in self.varnames
-            ]
-
-            self._stat_groups = []
-            sample_stats = [
-                Variable("tune", "bool"),
-            ]
-            if sampler_vars is not None:
-                # In PyMC the sampler stats are grouped by the sampler.
-                # ⚠ PyMC currently does not inform backends about shapes/dims of sampler stats.
-                for s, names_dtypes in enumerate(sampler_vars):
-                    self._stat_groups.append([])
-                    for statname, dtype in names_dtypes.items():
-                        sname = f"sampler_{s}__{statname}"
-                        if statname == "warning":
-                            # SamplerWarnings will be pickled and stored as string!
-                            svar = Variable(sname, "str")
-                        else:
-                            svar = Variable(
-                                name=sname,
-                                dtype=numpy.dtype(dtype).name,
-                                # This 👇 is needed until PyMC provides shapes ahead of time.
-                                undefined_ndim=True,
-                            )
-                        self._stat_groups[s].append((sname, statname))
-                        sample_stats.append(svar)
-
-            coordinates = [
-                Coordinate(dname, ndarray_from_numpy(numpy.array(cvals)))
-                for dname, cvals in self.model.coords.items()
-                if cvals is not None
-            ]
-
-            rmeta = RunMeta(
-                self.run_id,
-                variables=variables,
-                coordinates=coordinates,
-                sample_stats=sample_stats,
-                data=find_data(self.model),
-            )
-            self._run = self._backend.init_run(rmeta)
-        self._chain = self._run.init_chain(chain_number=chain)
-        return
-
-    def record(self, point, sampler_states=None):
-        draw = dict(zip(self.varnames, self.fn(point)))
-        if sampler_states is None:
-            stats = None
-        else:
-            stats = {
-                "tune": False,
-            }
-            # Unpack the stats by sampler to uniquely named stats.
-            for s, sts in enumerate(sampler_states):
-                for statname, sval in sts.items():
-                    sname = f"sampler_{s}__{statname}"
-                    # Automatically pickle SamplerWarnings
-                    if statname == "warning":
-                        sval_bytes = pickle.dumps(sval)
-                        sval = base64.encodebytes(sval_bytes).decode("ascii")
-                    stats[sname] = numpy.asarray(sval)
-                    # Make note whether this is a tuning iteration.
-                    if statname == "tune":
-                        stats["tune"] = sval
-
-        self._chain.append(draw, stats)
-        self._length += 1
-        return
-
-    def get_values(self, varname, burn=0, thin=1) -> numpy.ndarray:
-        if self._chain is None:
-            raise Exception("Trace setup was not completed. Call `.setup()` first.")
-        return self._chain.get_draws(varname)[burn::thin]
-
-    def _get_stats(self, varname, burn=0, thin=1) -> numpy.ndarray:
-        if self._chain is None:
-            raise Exception("Trace setup was not completed. Call `.setup()` first.")
-        values = self._chain.get_stats(varname)[burn::thin]
-        if "warning" in varname:
-            objs = []
-            for v in values:
-                enc = v.encode("ascii")
-                str_ = base64.decodebytes(enc)
-                obj = pickle.loads(str_)
-                objs.append(obj)
-            values = numpy.array(objs, dtype=object)
-        return values
-
-    def _get_sampler_stats(self, stat_name, sampler_idx, burn, thin):
-        if self._chain is None:
-            raise Exception("Trace setup was not completed. Call `.setup()` first.")
-        return self._get_stats(f"sampler_{sampler_idx}__{stat_name}", burn, thin)
-
-    def point(self, idx: int):
-        if self._chain is None:
-            raise Exception("Trace setup was not completed. Call `.setup()` first.")
-        return self._chain.get_draws_at(idx, self.var_names)
-
-    def as_readonly(self) -> ReadOnlyTrace:
-        """Creates a PyMC trace object from this chain."""
-        # Re-organize draws and stats in the PyMC style
-        draws = {varname: self.get_values(varname) for varname in self.varnames}
-        stats = [
-            {statname: self._get_stats(sname) for sname, statname in namemap}
-            for namemap in self._stat_groups
-        ]
-
-        rotrace = ReadOnlyTrace(
-            from_trace=self,
-            length=len(self),
-            draws=draws,
-            stats=stats,
-        )
-        return rotrace
-
-    def _slice(self, idx) -> ReadOnlyTrace:
-        return self.as_readonly()[idx]
diff --git a/requirements-dev.txt b/requirements-dev.txt
index c391556..59a41df 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,7 +1,6 @@
 arviz
 clickhouse-driver
 flake8
-pymc==5.1.2
 pytest
 pytest-cov
 twine


Test patch:
diff --git a/mcbackend/test_adapter_pymc.py b/mcbackend/test_adapter_pymc.py
deleted file mode 100644
index 574a1ce..0000000
--- a/mcbackend/test_adapter_pymc.py
+++ /dev/null
@@ -1,157 +0,0 @@
-import logging
-
-import clickhouse_driver
-import hagelkorn
-import numpy
-import pymc as pm
-import pytest
-
-from mcbackend.meta import RunMeta
-from mcbackend.npproto.utils import ndarray_to_numpy
-
-from .adapters.pymc import TraceBackend, find_data
-from .backends.clickhouse import ClickHouseBackend
-from .test_backend_clickhouse import DB_KWARGS, HAS_REAL_DB
-
-_log = logging.getLogger(__file__)
-
-
-@pytest.fixture
-def simple_model():
-    seconds = numpy.linspace(0, 5)
-    observations = numpy.random.normal(
-        0.5 + numpy.random.uniform(size=3)[:, None] * seconds[None, :]
-    )
-    with pm.Model(
-        coords={
-            "condition": ["A", "B", "C"],
-        }
-    ) as pmodel:
-        x = pm.ConstantData("seconds", seconds, dims="time")
-        a = pm.Normal("scalar")
-        b = pm.Uniform("vector", dims="condition")
-        pm.Deterministic("matrix", a + b[:, None] * x[None, :], dims=("condition", "time"))
-        pm.Bernoulli("integer", p=0.5)
-        obs = pm.MutableData("obs", observations, dims=("condition", "time"))
-        pm.Normal("L", pmodel["matrix"], observed=obs, dims=("condition", "time"))
-    return pmodel
-
-
-def test_find_data(simple_model):
-    dvars = find_data(simple_model)
-    dvars = {dv.name: dv for dv in dvars}
-    assert "seconds" in dvars
-    assert dvars["seconds"].dims == ["time"]
-    assert not dvars["seconds"].is_observed
-    assert "obs" in dvars
-    assert dvars["obs"].dims == ["condition", "time"]
-    assert dvars["obs"].is_observed
-    pass
-
-
-@pytest.mark.skipif(
-    condition=not HAS_REAL_DB,
-    reason="Integration tests need a ClickHouse server on localhost:9000 without authentication.",
-)
-class TestPyMCAdapter:
-    def setup_method(self, method):
-        """Initializes a fresh database just for this test method."""
-        self._db = "testing_" + hagelkorn.random()
-        self._client_main = clickhouse_driver.Client(**DB_KWARGS)
-        self._client_main.execute(f"CREATE DATABASE {self._db};")
-        self._client = clickhouse_driver.Client(**DB_KWARGS, database=self._db)
-        self.backend = ClickHouseBackend(self._client)
-        return
-
-    def teardown_method(self, method):
-        self._client.disconnect()
-        self._client_main.execute(f"DROP DATABASE {self._db};")
-        self._client_main.disconnect()
-        return
-
-    @pytest.mark.parametrize("cores", [1, 3])
-    def test_cores(self, simple_model, cores):
-        backend = ClickHouseBackend(self._client)
-
-        # To extract the run meta that the adapter passes to the backend:
-        args = []
-        original = backend.init_run
-
-        def wrapper(meta: RunMeta):
-            args.append(meta)
-            return original(meta)
-
-        backend.init_run = wrapper
-
-        with simple_model:
-            trace = TraceBackend(backend)
-            idata = pm.sample(
-                trace=trace,
-                tune=30,
-                draws=50,
-                chains=2,
-                cores=cores,
-                discard_tuned_samples=False,
-            )
-        if not len(args) == 1:
-            _log.warning("Run was initialized multiple times.")
-        rmeta = args[0]
-
-        # Chain lenghts after conversion
-        assert idata.posterior.dims["chain"] == 2
-        assert idata.posterior.dims["draw"] == 50
-        assert idata.warmup_posterior.dims["chain"] == 2
-        assert idata.warmup_posterior.dims["draw"] == 30
-
-        # Tracking of named variable dimensions
-        vars = {var.name: var for var in rmeta.variables}
-        assert vars["vector"].dims == ["condition"]
-        assert vars["matrix"].dims == ["condition", "time"]
-
-        # Tracking of coordinates
-        coords = {coord.name: coord for coord in rmeta.coordinates}
-        assert (
-            tuple(ndarray_to_numpy(coords["condition"].values)) == simple_model.coords["condition"]
-        )
-
-        # Meta-information which variables are free vs. deterministic
-        assert not vars["vector"].is_deterministic
-        assert vars["matrix"].is_deterministic
-
-        # Data variables should be included in the RunMeta
-        datavars = {dv.name: dv for dv in rmeta.data}
-        # Constant data variables
-        assert "seconds" in datavars
-        seconds = datavars["seconds"]
-        assert tuple(seconds.dims) == ("time",)
-        assert not seconds.is_observed
-        numpy.testing.assert_array_equal(
-            ndarray_to_numpy(seconds.value), simple_model["seconds"].data
-        )
-        # Observed data variables
-        assert "obs" in datavars
-        obs = datavars["obs"]
-        assert tuple(obs.dims) == ("condition", "time")
-        assert obs.is_observed
-        numpy.testing.assert_array_equal(
-            ndarray_to_numpy(obs.value), simple_model["obs"].get_value()
-        )
-        pass
-
-    def test_uninitialized_exceptions(self, simple_model):
-        backend = ClickHouseBackend(self._client)
-        with simple_model:
-            trace = TraceBackend(backend)
-
-        with pytest.raises(Exception, match="setup was not completed"):
-            trace.get_values("scalar")
-
-        with pytest.raises(Exception, match="setup was not completed"):
-            trace._get_stats("doesntmatter")
-
-        with pytest.raises(Exception, match="setup was not completed"):
-            trace._get_sampler_stats("doesntmatter", 0, 0, 1)
-
-        with pytest.raises(Exception, match="setup was not completed"):
-            trace.point(0)
-        pass