pymc-extras 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

tests/test_model_builder.py

@@ -1,306 +0,0 @@
-#   Copyright 2023 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 hashlib
-import json
-import sys
-import tempfile
-
-import numpy as np
-import pandas as pd
-import pymc as pm
-import pytest
-
-from pymc_extras.model_builder import ModelBuilder
-
-
-@pytest.fixture(scope="module")
-def toy_X():
-    x = np.linspace(start=0, stop=1, num=100)
-    X = pd.DataFrame({"input": x})
-    return X
-
-
-@pytest.fixture(scope="module")
-def toy_y(toy_X):
-    y = 5 * toy_X["input"] + 3
-    y = y + np.random.normal(0, 1, size=len(toy_X))
-    y = pd.Series(y, name="output")
-    return y
-
-
-def get_unfitted_model_instance(X, y):
-    """Creates an unfitted model instance to which idata can be copied in
-    and then used as a fitted model instance. That way a fitted model
-    can be used multiple times without having to run `fit` multiple times."""
-    sampler_config = {
-        "draws": 20,
-        "tune": 10,
-        "chains": 2,
-        "target_accept": 0.95,
-    }
-    model_config = {
-        "a": {"loc": 0, "scale": 10, "dims": ("numbers",)},
-        "b": {"loc": 0, "scale": 10},
-        "obs_error": 2,
-    }
-    model = test_ModelBuilder(
-        model_config=model_config, sampler_config=sampler_config, test_parameter="test_paramter"
-    )
-    # Do the things that `model.fit` does except sample to create idata.
-    model._generate_and_preprocess_model_data(X, y.values.flatten())
-    model.build_model(X, y)
-    return model
-
-
-@pytest.fixture(scope="module")
-def fitted_model_instance_base(toy_X, toy_y):
-    """Because fitting takes a relatively long time, this is intended to
-    be used only once and then have new instances created and fit data patched in
-    for tests that use a fitted model instance. Tests should use
-    `fitted_model_instance` instead of this."""
-    model = get_unfitted_model_instance(toy_X, toy_y)
-    model.fit(toy_X, toy_y)
-    return model
-
-
-@pytest.fixture
-def fitted_model_instance(toy_X, toy_y, fitted_model_instance_base):
-    """Get a fitted model instance. A new instance is created and fit data is
-    patched in, so tests using this fixture can modify the model object without
-    affecting other tests."""
-    model = get_unfitted_model_instance(toy_X, toy_y)
-    model.idata = fitted_model_instance_base.idata.copy()
-    return model
-
-
-class test_ModelBuilder(ModelBuilder):
-    def __init__(self, model_config=None, sampler_config=None, test_parameter=None):
-        self.test_parameter = test_parameter
-        super().__init__(model_config=model_config, sampler_config=sampler_config)
-
-    _model_type = "test_model"
-    version = "0.1"
-
-    def build_model(self, X: pd.DataFrame, y: pd.Series, model_config=None):
-        coords = {"numbers": np.arange(len(X))}
-        self.generate_and_preprocess_model_data(X, y)
-        with pm.Model(coords=coords) as self.model:
-            if model_config is None:
-                model_config = self.model_config
-            x = pm.Data("x", self.X["input"].values)
-            y_data = pm.Data("y_data", self.y)
-
-            # prior parameters
-            a_loc = model_config["a"]["loc"]
-            a_scale = model_config["a"]["scale"]
-            b_loc = model_config["b"]["loc"]
-            b_scale = model_config["b"]["scale"]
-            obs_error = model_config["obs_error"]
-
-            # priors
-            a = pm.Normal("a", a_loc, sigma=a_scale, dims=model_config["a"]["dims"])
-            b = pm.Normal("b", b_loc, sigma=b_scale)
-            obs_error = pm.HalfNormal("σ_model_fmc", obs_error)
-
-            # observed data
-            output = pm.Normal("output", a + b * x, obs_error, shape=x.shape, observed=y_data)
-
-    def _save_input_params(self, idata):
-        idata.attrs["test_paramter"] = json.dumps(self.test_parameter)
-
-    @property
-    def output_var(self):
-        return "output"
-
-    def _data_setter(self, x: pd.Series, y: pd.Series = None):
-        with self.model:
-            pm.set_data({"x": x.values})
-            if y is not None:
-                pm.set_data({"y_data": y.values})
-
-    @property
-    def _serializable_model_config(self):
-        return self.model_config
-
-    def generate_and_preprocess_model_data(self, X: pd.DataFrame, y: pd.Series):
-        self.X = X
-        self.y = y
-
-    @staticmethod
-    def get_default_model_config() -> dict:
-        return {
-            "a": {"loc": 0, "scale": 10, "dims": ("numbers",)},
-            "b": {"loc": 0, "scale": 10},
-            "obs_error": 2,
-        }
-
-    def _generate_and_preprocess_model_data(
-        self, X: pd.DataFrame | pd.Series, y: pd.Series
-    ) -> None:
-        self.X = X
-        self.y = y
-
-    @staticmethod
-    def get_default_sampler_config() -> dict:
-        return {
-            "draws": 10,
-            "tune": 10,
-            "chains": 3,
-            "target_accept": 0.95,
-        }
-
-
-def test_save_input_params(fitted_model_instance):
-    assert fitted_model_instance.idata.attrs["test_paramter"] == '"test_paramter"'
-
-
-@pytest.mark.skipif(
-    sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
-)
-def test_save_load(fitted_model_instance):
-    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
-    fitted_model_instance.save(temp.name)
-    test_builder2 = test_ModelBuilder.load(temp.name)
-    assert fitted_model_instance.idata.groups() == test_builder2.idata.groups()
-    assert fitted_model_instance.id == test_builder2.id
-    x_pred = np.random.uniform(low=0, high=1, size=100)
-    prediction_data = pd.DataFrame({"input": x_pred})
-    pred1 = fitted_model_instance.predict(prediction_data["input"])
-    pred2 = test_builder2.predict(prediction_data["input"])
-    assert pred1.shape == pred2.shape
-    temp.close()
-
-
-def test_initial_build_and_fit(fitted_model_instance, check_idata=True) -> ModelBuilder:
-    if check_idata:
-        assert fitted_model_instance.idata is not None
-        assert "posterior" in fitted_model_instance.idata.groups()
-
-
-def test_save_without_fit_raises_runtime_error():
-    model_builder = test_ModelBuilder()
-    with pytest.raises(RuntimeError):
-        model_builder.save("saved_model")
-
-
-def test_empty_sampler_config_fit(toy_X, toy_y):
-    sampler_config = {}
-    model_builder = test_ModelBuilder(sampler_config=sampler_config)
-    model_builder.idata = model_builder.fit(X=toy_X, y=toy_y)
-    assert model_builder.idata is not None
-    assert "posterior" in model_builder.idata.groups()
-
-
-def test_fit(fitted_model_instance):
-    prediction_data = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
-    pred = fitted_model_instance.predict(prediction_data["input"])
-    post_pred = fitted_model_instance.sample_posterior_predictive(
-        prediction_data["input"], extend_idata=True, combined=True
-    )
-    post_pred[fitted_model_instance.output_var].shape[0] == prediction_data.input.shape
-
-
-def test_fit_no_y(toy_X):
-    model_builder = test_ModelBuilder()
-    model_builder.idata = model_builder.fit(X=toy_X, chains=1, tune=1, draws=1)
-    assert model_builder.model is not None
-    assert model_builder.idata is not None
-    assert "posterior" in model_builder.idata.groups()
-
-
-def test_predict(fitted_model_instance):
-    x_pred = np.random.uniform(low=0, high=1, size=100)
-    prediction_data = pd.DataFrame({"input": x_pred})
-    pred = fitted_model_instance.predict(prediction_data["input"])
-    # Perform elementwise comparison using numpy
-    assert isinstance(pred, np.ndarray)
-    assert len(pred) > 0
-
-
-@pytest.mark.parametrize("combined", [True, False])
-def test_sample_posterior_predictive(fitted_model_instance, combined):
-    n_pred = 100
-    x_pred = np.random.uniform(low=0, high=1, size=n_pred)
-    prediction_data = pd.DataFrame({"input": x_pred})
-    pred = fitted_model_instance.sample_posterior_predictive(
-        prediction_data["input"], combined=combined, extend_idata=True
-    )
-    chains = fitted_model_instance.idata.sample_stats.sizes["chain"]
-    draws = fitted_model_instance.idata.sample_stats.sizes["draw"]
-    expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
-    assert pred[fitted_model_instance.output_var].shape == expected_shape
-    assert np.issubdtype(pred[fitted_model_instance.output_var].dtype, np.floating)
-
-
-@pytest.mark.parametrize("group", ["prior_predictive", "posterior_predictive"])
-@pytest.mark.parametrize("extend_idata", [True, False])
-def test_sample_xxx_extend_idata_param(fitted_model_instance, group, extend_idata):
-    output_var = fitted_model_instance.output_var
-    idata_prev = fitted_model_instance.idata[group][output_var]
-
-    # Since coordinates are provided, the dimension must match
-    n_pred = 100  # Must match toy_x
-    x_pred = np.random.uniform(0, 1, n_pred)
-
-    prediction_data = pd.DataFrame({"input": x_pred})
-    if group == "prior_predictive":
-        prediction_method = fitted_model_instance.sample_prior_predictive
-    else:  # group == "posterior_predictive":
-        prediction_method = fitted_model_instance.sample_posterior_predictive
-
-    pred = prediction_method(prediction_data["input"], combined=False, extend_idata=extend_idata)
-
-    pred_unstacked = pred[output_var].values
-    idata_now = fitted_model_instance.idata[group][output_var].values
-
-    if extend_idata:
-        # After sampling, data in the model should be the same as the predictions
-        np.testing.assert_array_equal(idata_now, pred_unstacked)
-        # Data in the model should NOT be the same as before
-        if idata_now.shape == idata_prev.values.shape:
-            assert np.sum(np.abs(idata_now - idata_prev.values) < 1e-5) <= 2
-    else:
-        # After sampling, data in the model should be the same as it was before
-        np.testing.assert_array_equal(idata_now, idata_prev.values)
-        # Data in the model should NOT be the same as the predictions
-        if idata_now.shape == pred_unstacked.shape:
-            assert np.sum(np.abs(idata_now - pred_unstacked) < 1e-5) <= 2
-
-
-def test_model_config_formatting():
-    model_config = {
-        "a": {
-            "loc": [0, 0],
-            "scale": 10,
-            "dims": [
-                "x",
-            ],
-        },
-    }
-    model_builder = test_ModelBuilder()
-    converted_model_config = model_builder._model_config_formatting(model_config)
-    np.testing.assert_equal(converted_model_config["a"]["dims"], ("x",))
-    np.testing.assert_equal(converted_model_config["a"]["loc"], np.array([0, 0]))
-
-
-def test_id():
-    model_builder = test_ModelBuilder()
-    expected_id = hashlib.sha256(
-        str(model_builder.model_config.values()).encode()
-        + model_builder.version.encode()
-        + model_builder._model_type.encode()
-    ).hexdigest()[:16]
-
-    assert model_builder.id == expected_id

tests/test_pathfinder.py

@@ -1,203 +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 sys
-
-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
-
-
-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])
-
-    with pm.Model() as model:
-        mu = pm.Normal("mu", mu=0.0, sigma=10.0)
-        tau = pm.HalfCauchy("tau", 5.0)
-
-        theta = pm.Normal("theta", mu=0, sigma=1, shape=J)
-        obs = pm.Normal("obs", mu=mu + tau * theta, sigma=sigma, shape=J, observed=y)
-
-    return model
-
-
-@pytest.fixture
-def reference_idata():
-    model = eight_schools_model()
-    with model:
-        idata = pmx.fit(
-            method="pathfinder",
-            num_paths=10,
-            jitter=12.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=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=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)
-    assert idata.posterior["theta"].shape == (1, 1000, 8)
-
-
-@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=10,
-            jitter=12.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=4,
-            jitter=10.0,
-            random_seed=41,
-            inference_backend="pymc",
-        )
-
-        idata_123 = pmx.fit(
-            method="pathfinder",
-            num_paths=4,
-            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)
-
-
-@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)

tests/test_pivoted_cholesky.py

@@ -1,24 +0,0 @@
-# try:
-#     import gpytorch
-#     import torch
-# except ImportError as e:
-#     # print(
-#     #     f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}"
-#     # )
-#     pass
-# import numpy as np
-#
-# import pymc_extras as pmx
-#
-#
-# def test_match_gpytorch_linearcg_output():
-#     N = 10
-#     rank = 5
-#     np.random.seed(1234)  # nans with seed 1234
-#     K = np.random.randn(N, N)
-#     K = K @ K.T + N * np.eye(N)
-#     K_torch = torch.from_numpy(K)
-#
-#     L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3)
-#     L_np, _ = pmx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3)
-#     assert np.allclose(L_gpt, L_np.T)

tests/test_printing.py

@@ -1,98 +0,0 @@
-import numpy as np
-import pymc as pm
-
-from rich.console import Console
-
-from pymc_extras.printing import model_table
-
-
-def get_text(table) -> str:
-    console = Console(width=80)
-    with console.capture() as capture:
-        console.print(table)
-    return capture.get()
-
-
-def test_model_table():
-    with pm.Model(coords={"trial": range(6), "subject": range(20)}) as model:
-        x_data = pm.Data("x_data", np.random.normal(size=(6, 20)), dims=("trial", "subject"))
-        y_data = pm.Data("y_data", np.random.normal(size=(6, 20)), dims=("trial", "subject"))
-
-        mu = pm.Normal("mu", mu=0, sigma=1)
-        sigma = pm.HalfNormal("sigma", sigma=1)
-        global_intercept = pm.Normal("global_intercept", mu=0, sigma=1)
-        intercept_subject = pm.Normal("intercept_subject", mu=0, sigma=1, shape=(20, 1))
-        beta_subject = pm.Normal("beta_subject", mu=mu, sigma=sigma, dims="subject")
-
-        mu_trial = pm.Deterministic(
-            "mu_trial",
-            global_intercept.squeeze() + intercept_subject + beta_subject * x_data,
-            dims=["trial", "subject"],
-        )
-        noise = pm.Exponential("noise", lam=1)
-        y = pm.Normal("y", mu=mu_trial, sigma=noise, observed=y_data, dims=("trial", "subject"))
-
-        pm.Potential("beta_subject_penalty", -pm.math.abs(beta_subject), dims="subject")
-
-    table_txt = get_text(model_table(model))
-    expected = """               Variable  Expression                      Dimensions
-────────────────────────────────────────────────────────────────────────────────
-               x_data =  Data                            trial[6] × subject[20]
-               y_data =  Data                            trial[6] × subject[20]
-
-                   mu ~  Normal(0, 1)
-                sigma ~  HalfNormal(0, 1)
-     global_intercept ~  Normal(0, 1)
-    intercept_subject ~  Normal(0, 1)                    [20, 1]
-         beta_subject ~  Normal(mu, sigma)               subject[20]
-                noise ~  Exponential(f())
-                                                         Parameter count = 44
-
-             mu_trial =  f(intercept_subject,            trial[6] × subject[20]
-                         beta_subject,
-                         global_intercept)
-
- beta_subject_penalty =  Potential(f(beta_subject))      subject[20]
-
-                    y ~  Normal(mu_trial, noise)         trial[6] × subject[20]
-"""
-    assert [s.strip() for s in table_txt.splitlines()] == [s.strip() for s in expected.splitlines()]
-
-    table_txt = get_text(model_table(model, split_groups=False))
-    expected = """               Variable  Expression                      Dimensions
-────────────────────────────────────────────────────────────────────────────────
-               x_data =  Data                            trial[6] × subject[20]
-               y_data =  Data                            trial[6] × subject[20]
-                   mu ~  Normal(0, 1)
-                sigma ~  HalfNormal(0, 1)
-     global_intercept ~  Normal(0, 1)
-    intercept_subject ~  Normal(0, 1)                    [20, 1]
-         beta_subject ~  Normal(mu, sigma)               subject[20]
-             mu_trial =  f(intercept_subject,            trial[6] × subject[20]
-                         beta_subject,
-                         global_intercept)
-                noise ~  Exponential(f())
-                    y ~  Normal(mu_trial, noise)         trial[6] × subject[20]
- beta_subject_penalty =  Potential(f(beta_subject))      subject[20]
-                                                         Parameter count = 44
-"""
-    assert [s.strip() for s in table_txt.splitlines()] == [s.strip() for s in expected.splitlines()]
-
-    table_txt = get_text(
-        model_table(model, split_groups=False, truncate_deterministic=30, parameter_count=False)
-    )
-    expected = """               Variable  Expression                  Dimensions
-────────────────────────────────────────────────────────────────────────────
-               x_data =  Data                        trial[6] × subject[20]
-               y_data =  Data                        trial[6] × subject[20]
-                   mu ~  Normal(0, 1)
-                sigma ~  HalfNormal(0, 1)
-     global_intercept ~  Normal(0, 1)
-    intercept_subject ~  Normal(0, 1)                [20, 1]
-         beta_subject ~  Normal(mu, sigma)           subject[20]
-             mu_trial =  f(intercept_subject, ...)   trial[6] × subject[20]
-                noise ~  Exponential(f())
-                    y ~  Normal(mu_trial, noise)     trial[6] × subject[20]
- beta_subject_penalty =  Potential(f(beta_subject))  subject[20]
-"""
-    assert [s.strip() for s in table_txt.splitlines()] == [s.strip() for s in expected.splitlines()]