pymc-extras 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

tests/test_pathfinder.py

@@ -18,12 +18,12 @@ import numpy as np
 import pymc as pm
 import pytest
 
+pytestmark = pytest.mark.filterwarnings("ignore:compile_pymc was renamed to compile:FutureWarning")
+
 import pymc_extras as pmx
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="JAX not supported on windows.")
-def test_pathfinder():
-    # Data of the Eight Schools Model
+def eight_schools_model() -> pm.Model:
     J = 8
     y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
     sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
@@ -35,11 +35,139 @@ def test_pathfinder():
         theta = pm.Normal("theta", mu=0, sigma=1, shape=J)
         obs = pm.Normal("obs", mu=mu + tau * theta, sigma=sigma, shape=J, observed=y)
 
-        idata = pmx.fit(method="pathfinder", random_seed=41)
+    return model
+
+
+@pytest.fixture
+def reference_idata():
+    model = eight_schools_model()
+    with model:
+        idata = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+        )
+    return idata
+
+
+@pytest.mark.parametrize("inference_backend", ["pymc", "blackjax"])
+def test_pathfinder(inference_backend, reference_idata):
+    if inference_backend == "blackjax" and sys.platform == "win32":
+        pytest.skip("JAX not supported on windows")
+
+    if inference_backend == "blackjax":
+        model = eight_schools_model()
+        with model:
+            idata = pmx.fit(
+                method="pathfinder",
+                num_paths=50,
+                jitter=10.0,
+                random_seed=41,
+                inference_backend=inference_backend,
+            )
+    else:
+        idata = reference_idata
+        np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0, atol=1.6)
+        np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=1.5)
 
     assert idata.posterior["mu"].shape == (1, 1000)
     assert idata.posterior["tau"].shape == (1, 1000)
     assert idata.posterior["theta"].shape == (1, 1000, 8)
-    # FIXME: pathfinder doesn't find a reasonable mean! Fix bug or choose model pathfinder can handle
-    # np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0)
-    np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=0.5)
+
+
+@pytest.mark.parametrize("concurrent", ["thread", "process"])
+def test_concurrent_results(reference_idata, concurrent):
+    model = eight_schools_model()
+    with model:
+        idata_conc = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+            concurrent=concurrent,
+        )
+
+    np.testing.assert_allclose(
+        reference_idata.posterior.mu.data.mean(),
+        idata_conc.posterior.mu.data.mean(),
+        atol=0.4,
+    )
+
+    np.testing.assert_allclose(
+        reference_idata.posterior.tau.data.mean(),
+        idata_conc.posterior.tau.data.mean(),
+        atol=0.4,
+    )
+
+
+def test_seed(reference_idata):
+    model = eight_schools_model()
+    with model:
+        idata_41 = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=41,
+            inference_backend="pymc",
+        )
+
+        idata_123 = pmx.fit(
+            method="pathfinder",
+            num_paths=50,
+            jitter=10.0,
+            random_seed=123,
+            inference_backend="pymc",
+        )
+
+    assert not np.allclose(idata_41.posterior.mu.data.mean(), idata_123.posterior.mu.data.mean())
+
+    assert np.allclose(idata_41.posterior.mu.data.mean(), idata_41.posterior.mu.data.mean())
+
+
+def test_bfgs_sample():
+    import pytensor.tensor as pt
+
+    from pymc_extras.inference.pathfinder.pathfinder import (
+        alpha_recover,
+        bfgs_sample,
+        inverse_hessian_factors,
+    )
+
+    """test BFGS sampling"""
+    Lp1, N = 8, 10
+    L = Lp1 - 1
+    J = 6
+    num_samples = 1000
+
+    # mock data
+    x_data = np.random.randn(Lp1, N)
+    g_data = np.random.randn(Lp1, N)
+
+    # get factors
+    x_full = pt.as_tensor(x_data, dtype="float64")
+    g_full = pt.as_tensor(g_data, dtype="float64")
+    epsilon = 1e-11
+
+    x = x_full[1:]
+    g = g_full[1:]
+    alpha, S, Z, update_mask = alpha_recover(x_full, g_full, epsilon)
+    beta, gamma = inverse_hessian_factors(alpha, S, Z, update_mask, J)
+
+    # sample
+    phi, logq = bfgs_sample(
+        num_samples=num_samples,
+        x=x,
+        g=g,
+        alpha=alpha,
+        beta=beta,
+        gamma=gamma,
+    )
+
+    # check shapes
+    assert beta.eval().shape == (L, N, 2 * J)
+    assert gamma.eval().shape == (L, 2 * J, 2 * J)
+    assert phi.eval().shape == (L, num_samples, N)
+    assert logq.eval().shape == (L, num_samples)

pymc_extras/inference/pathfinder.py

@@ -1,134 +0,0 @@
-#   Copyright 2022 The PyMC Developers
-#
-#   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.
-
-import collections
-import sys
-
-import arviz as az
-import blackjax
-import jax
-import numpy as np
-import pymc as pm
-
-from packaging import version
-from pymc.backends.arviz import coords_and_dims_for_inferencedata
-from pymc.blocking import DictToArrayBijection, RaveledVars
-from pymc.model import modelcontext
-from pymc.sampling.jax import get_jaxified_graph
-from pymc.util import RandomSeed, _get_seeds_per_chain, get_default_varnames
-
-
-def convert_flat_trace_to_idata(
-    samples,
-    include_transformed=False,
-    postprocessing_backend="cpu",
-    model=None,
-):
-    model = modelcontext(model)
-    ip = model.initial_point()
-    ip_point_map_info = pm.blocking.DictToArrayBijection.map(ip).point_map_info
-    trace = collections.defaultdict(list)
-    for sample in samples:
-        raveld_vars = RaveledVars(sample, ip_point_map_info)
-        point = DictToArrayBijection.rmap(raveld_vars, ip)
-        for p, v in point.items():
-            trace[p].append(v.tolist())
-
-    trace = {k: np.asarray(v)[None, ...] for k, v in trace.items()}
-
-    var_names = model.unobserved_value_vars
-    vars_to_sample = list(get_default_varnames(var_names, include_transformed=include_transformed))
-    print("Transforming variables...", file=sys.stdout)
-    jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample)
-    result = jax.vmap(jax.vmap(jax_fn))(
-        *jax.device_put(list(trace.values()), jax.devices(postprocessing_backend)[0])
-    )
-    trace = {v.name: r for v, r in zip(vars_to_sample, result)}
-    coords, dims = coords_and_dims_for_inferencedata(model)
-    idata = az.from_dict(trace, dims=dims, coords=coords)
-
-    return idata
-
-
-def fit_pathfinder(
-    samples=1000,
-    random_seed: RandomSeed | None = None,
-    postprocessing_backend="cpu",
-    model=None,
-    **pathfinder_kwargs,
-):
-    """
-    Fit the pathfinder algorithm as implemented in blackjax
-
-    Requires the JAX backend
-
-    Parameters
-    ----------
-    samples : int
-        Number of samples to draw from the fitted approximation.
-    random_seed : int
-        Random seed to set.
-    postprocessing_backend : str
-        Where to compute transformations of the trace.
-        "cpu" or "gpu".
-    pathfinder_kwargs:
-        kwargs for blackjax.vi.pathfinder.approximate
-
-    Returns
-    -------
-    arviz.InferenceData
-
-    Reference
-    ---------
-    https://arxiv.org/abs/2108.03782
-    """
-    # Temporarily helper
-    if version.parse(blackjax.__version__).major < 1:
-        raise ImportError("fit_pathfinder requires blackjax 1.0 or above")
-
-    model = modelcontext(model)
-
-    ip = model.initial_point()
-    ip_map = DictToArrayBijection.map(ip)
-
-    new_logprob, new_input = pm.pytensorf.join_nonshared_inputs(
-        ip, (model.logp(),), model.value_vars, ()
-    )
-
-    logprob_fn_list = get_jaxified_graph([new_input], new_logprob)
-
-    def logprob_fn(x):
-        return logprob_fn_list(x)[0]
-
-    [pathfinder_seed, sample_seed] = _get_seeds_per_chain(random_seed, 2)
-
-    print("Running pathfinder...", file=sys.stdout)
-    pathfinder_state, _ = blackjax.vi.pathfinder.approximate(
-        rng_key=jax.random.key(pathfinder_seed),
-        logdensity_fn=logprob_fn,
-        initial_position=ip_map.data,
-        **pathfinder_kwargs,
-    )
-    samples, _ = blackjax.vi.pathfinder.sample(
-        rng_key=jax.random.key(sample_seed),
-        state=pathfinder_state,
-        num_samples=samples,
-    )
-
-    idata = convert_flat_trace_to_idata(
-        samples,
-        postprocessing_backend=postprocessing_backend,
-        model=model,
-    )
-    return idata