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

tests/test_pathfinder.py

@@ -12,14 +12,14 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import re
 import sys
 
 import numpy as np
 import pymc as pm
+import pytensor.tensor as pt
 import pytest
 
-pytestmark = pytest.mark.filterwarnings("ignore:compile_pymc was renamed to compile:FutureWarning")
-
 import pymc_extras as pmx
 
 
@@ -44,15 +44,98 @@ def reference_idata():
     with model:
         idata = pmx.fit(
             method="pathfinder",
-            num_paths=50,
-            jitter=10.0,
+            num_paths=10,
+            jitter=12.0,
             random_seed=41,
             inference_backend="pymc",
         )
     return idata
 
 
+def unstable_lbfgs_update_mask_model() -> pm.Model:
+    # data and model from: https://github.com/pymc-devs/pymc-extras/issues/445
+    # this scenario made LBFGS struggle leading to a lot of rejected iterations, (result.nit being moderate, but only history.count <= 1).
+    # this scenario is used to test that the LBFGS history manager is rejecting iterations as expected and PF can run to completion.
+
+    # fmt: off
+    inp = np.array([0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 2, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 2, 0, 1, 0, 0, 0, 0, 1, 1, 1, 2, 0, 1, 2, 1, 0, 1, 0, 1, 0, 1, 0])
+
+    res = np.array([[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,0,1,0],[1,0,0,0,0],[0,1,0,0,0],[0,0,1,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,1,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,1,0,0,0],[1,0,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,0,1,0,0],[0,1,0,0,0],[0,1,0,0,0],[0,0,0,1,0],[0,0,1,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,1,0,0],[1,0,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,1,0],[1,0,0,0,0],[1,0,0,0,0],[0,1,0,0,0],[1,0,0,0,0],[0,0,1,0,0],[0,0,1,0,0],[1,0,0,0,0],[0,0,0,1,0]])
+    # fmt: on
+
+    n_ordered = res.shape[1]
+    coords = {
+        "obs": np.arange(len(inp)),
+        "inp": np.arange(max(inp) + 1),
+        "outp": np.arange(res.shape[1]),
+    }
+    with pm.Model(coords=coords) as mdl:
+        mu = pm.Normal("intercept", sigma=3.5)[None]
+
+        offset = pm.Normal(
+            "offset", dims=("inp"), transform=pm.distributions.transforms.ZeroSumTransform([0])
+        )
+
+        scale = 3.5 * pm.HalfStudentT("scale", nu=5)
+        mu += (scale * offset)[inp]
+
+        phi_delta = pm.Dirichlet("phi_diffs", [1.0] * (n_ordered - 1))
+        phi = pt.concatenate([[0], pt.cumsum(phi_delta)])
+        s_mu = pm.Normal(
+            "stereotype_intercept",
+            size=n_ordered,
+            transform=pm.distributions.transforms.ZeroSumTransform([-1]),
+        )
+        fprobs = pm.math.softmax(s_mu[None, :] + phi[None, :] * mu[:, None], axis=-1)
+
+        pm.Multinomial("y_res", p=fprobs, n=np.ones(len(inp)), observed=res, dims=("obs", "outp"))
+
+    return mdl
+
+
+@pytest.mark.parametrize("jitter", [12.0, 500.0, 1000.0])
+def test_unstable_lbfgs_update_mask(capsys, jitter):
+    model = unstable_lbfgs_update_mask_model()
+
+    if jitter < 1000:
+        with model:
+            idata = pmx.fit(
+                method="pathfinder",
+                jitter=jitter,
+                random_seed=4,
+            )
+        out, err = capsys.readouterr()
+        status_pattern = [
+            r"INIT_FAILED_LOW_UPDATE_PCT\s+\d+",
+            r"LOW_UPDATE_PCT\s+\d+",
+            r"LBFGS_FAILED\s+\d+",
+            r"SUCCESS\s+\d+",
+        ]
+        for pattern in status_pattern:
+            assert re.search(pattern, out) is not None
+
+    else:
+        with pytest.raises(ValueError, match="All paths failed"):
+            with model:
+                idata = pmx.fit(
+                    method="pathfinder",
+                    jitter=1000,
+                    random_seed=2,
+                    num_paths=4,
+                )
+            out, err = capsys.readouterr()
+
+            status_pattern = [
+                r"INIT_FAILED_LOW_UPDATE_PCT\s+2",
+                r"LOW_UPDATE_PCT\s+2",
+                r"LBFGS_FAILED\s+4",
+            ]
+            for pattern in status_pattern:
+                assert re.search(pattern, out) is not None
+
+
 @pytest.mark.parametrize("inference_backend", ["pymc", "blackjax"])
+@pytest.mark.filterwarnings("ignore:JAXopt is no longer maintained.:DeprecationWarning")
 def test_pathfinder(inference_backend, reference_idata):
     if inference_backend == "blackjax" and sys.platform == "win32":
         pytest.skip("JAX not supported on windows")
@@ -62,15 +145,15 @@ def test_pathfinder(inference_backend, reference_idata):
         with model:
             idata = pmx.fit(
                 method="pathfinder",
-                num_paths=50,
-                jitter=10.0,
+                num_paths=10,
+                jitter=12.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)
+        np.testing.assert_allclose(idata.posterior["mu"].mean(), 5.0, atol=0.95)
+        np.testing.assert_allclose(idata.posterior["tau"].mean(), 4.15, atol=1.35)
 
     assert idata.posterior["mu"].shape == (1, 1000)
     assert idata.posterior["tau"].shape == (1, 1000)
@@ -83,8 +166,8 @@ def test_concurrent_results(reference_idata, concurrent):
     with model:
         idata_conc = pmx.fit(
             method="pathfinder",
-            num_paths=50,
-            jitter=10.0,
+            num_paths=10,
+            jitter=12.0,
             random_seed=41,
             inference_backend="pymc",
             concurrent=concurrent,
@@ -108,7 +191,7 @@ def test_seed(reference_idata):
     with model:
         idata_41 = pmx.fit(
             method="pathfinder",
-            num_paths=50,
+            num_paths=4,
             jitter=10.0,
             random_seed=41,
             inference_backend="pymc",
@@ -116,7 +199,7 @@ def test_seed(reference_idata):
 
         idata_123 = pmx.fit(
             method="pathfinder",
-            num_paths=50,
+            num_paths=4,
             jitter=10.0,
             random_seed=123,
             inference_backend="pymc",
@@ -149,12 +232,11 @@ def test_bfgs_sample():
     # 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)
+    alpha, s, z = alpha_recover(x_full, g_full)
+    beta, gamma = inverse_hessian_factors(alpha, s, z, J)
 
     # sample
     phi, logq = bfgs_sample(
@@ -169,5 +251,47 @@ def test_bfgs_sample():
     # 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)
+    assert all(phi.shape.eval() == (L, num_samples, N))
+    assert all(logq.shape.eval() == (L, num_samples))
+
+
+@pytest.mark.parametrize("importance_sampling", ["psis", "psir", "identity", None])
+def test_pathfinder_importance_sampling(importance_sampling):
+    model = eight_schools_model()
+
+    num_paths = 4
+    num_draws_per_path = 300
+    num_draws = 750
+
+    with model:
+        idata = pmx.fit(
+            method="pathfinder",
+            num_paths=num_paths,
+            num_draws_per_path=num_draws_per_path,
+            num_draws=num_draws,
+            maxiter=5,
+            random_seed=41,
+            inference_backend="pymc",
+            importance_sampling=importance_sampling,
+        )
+
+    if importance_sampling is None:
+        assert idata.posterior["mu"].shape == (num_paths, num_draws_per_path)
+        assert idata.posterior["tau"].shape == (num_paths, num_draws_per_path)
+        assert idata.posterior["theta"].shape == (num_paths, num_draws_per_path, 8)
+    else:
+        assert idata.posterior["mu"].shape == (1, num_draws)
+        assert idata.posterior["tau"].shape == (1, num_draws)
+        assert idata.posterior["theta"].shape == (1, num_draws, 8)
+
+
+def test_pathfinder_initvals():
+    # Run a model with an ordered transform that will fail unless initvals are in place
+    with pm.Model() as mdl:
+        pm.Normal("ordered", size=10, transform=pm.distributions.transforms.ordered)
+        idata = pmx.fit_pathfinder(initvals={"ordered": np.linspace(0, 1, 10)})
+
+    # Check that the samples are ordered to make sure transform was applied
+    assert np.all(
+        idata.posterior["ordered"][..., 1:].values > idata.posterior["ordered"][..., :-1].values
+    )