pymc-extras 0.2.0__py3-none-any.whl

tests/test_histogram_approximation.py

@@ -0,0 +1,109 @@
+#   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 numpy as np
+import pymc as pm
+import pytest
+
+import pymc_extras as pmx
+
+
+@pytest.mark.parametrize("use_dask", [True, False], ids="dask={}".format)
+@pytest.mark.parametrize("zero_inflation", [True, False], ids="ZI={}".format)
+@pytest.mark.parametrize("ndims", [1, 2], ids="ndims={}".format)
+def test_histogram_init_cont(use_dask, zero_inflation, ndims):
+    data = np.random.randn(*(10000, *(2,) * (ndims - 1)))
+    if zero_inflation:
+        data = abs(data)
+        data[:100] = 0
+    if use_dask:
+        dask = pytest.importorskip("dask")
+        dask_df = pytest.importorskip("dask.dataframe")
+        data = dask_df.from_array(data)
+    histogram = pmx.distributions.histogram_utils.quantile_histogram(
+        data, n_quantiles=100, zero_inflation=zero_inflation
+    )
+    if use_dask:
+        (histogram,) = dask.compute(histogram)
+    assert isinstance(histogram, dict)
+    assert isinstance(histogram["mid"], np.ndarray)
+    assert np.issubdtype(histogram["mid"].dtype, np.floating)
+    size = 99 + zero_inflation
+    assert histogram["mid"].shape == (size,) + (1,) * len(data.shape[1:])
+    assert histogram["lower"].shape == (size,) + (1,) * len(data.shape[1:])
+    assert histogram["upper"].shape == (size,) + (1,) * len(data.shape[1:])
+    assert histogram["count"].shape == (size,) + data.shape[1:]
+    assert (histogram["count"].sum(0) == 10000).all()
+    if zero_inflation:
+        (histogram["count"][0] == 100).all()
+
+
+@pytest.mark.parametrize("use_dask", [True, False], ids="dask={}".format)
+@pytest.mark.parametrize("min_count", [None, 5], ids="min_count={}".format)
+@pytest.mark.parametrize("ndims", [1, 2], ids="ndims={}".format)
+def test_histogram_init_discrete(use_dask, min_count, ndims):
+    data = np.random.randint(0, 100, size=(10000,) + (2,) * (ndims - 1))
+    u, c = np.unique(data, return_counts=True)
+    if use_dask:
+        dask = pytest.importorskip("dask")
+        dask_df = pytest.importorskip("dask.dataframe")
+        data = dask_df.from_array(data)
+    histogram = pmx.distributions.histogram_utils.discrete_histogram(data, min_count=min_count)
+    if use_dask:
+        (histogram,) = dask.compute(histogram)
+    assert isinstance(histogram, dict)
+    assert isinstance(histogram["mid"], np.ndarray)
+    assert np.issubdtype(histogram["mid"].dtype, np.integer)
+    if min_count is not None:
+        size = int((c >= min_count).sum())
+    else:
+        size = len(u)
+    assert histogram["mid"].shape == (size,) + (1,) * len(data.shape[1:])
+    assert histogram["count"].shape == (size,) + data.shape[1:]
+    if not min_count:
+        assert (histogram["count"].sum(0) == 10000).all()
+
+
+@pytest.mark.parametrize("use_dask", [True, False], ids="dask={}".format)
+@pytest.mark.parametrize("ndims", [1, 2], ids="ndims={}".format)
+def test_histogram_approx_cont(use_dask, ndims):
+    data = np.random.randn(*(10000, *(2,) * (ndims - 1)))
+    if use_dask:
+        dask = pytest.importorskip("dask")
+        dask_df = pytest.importorskip("dask.dataframe")
+        data = dask_df.from_array(data)
+    with pm.Model():
+        m = pm.Normal("m")
+        s = pm.HalfNormal("s", size=2 if ndims > 1 else 1)
+        pot = pmx.distributions.histogram_utils.histogram_approximation(
+            "histogram_potential", pm.Normal.dist(m, s), observed=data, n_quantiles=1000
+        )
+        trace = pm.sample(10, tune=0)  # very fast
+
+
+@pytest.mark.parametrize("use_dask", [True, False])
+@pytest.mark.parametrize("ndims", [1, 2], ids="ndims={}".format)
+def test_histogram_approx_discrete(use_dask, ndims):
+    data = np.random.randint(0, 100, size=(10000, *(2,) * (ndims - 1)))
+    if use_dask:
+        dask = pytest.importorskip("dask")
+        dask_df = pytest.importorskip("dask.dataframe")
+        data = dask_df.from_array(data)
+    with pm.Model():
+        s = pm.Exponential("s", 1.0, size=2 if ndims > 1 else 1)
+        pot = pmx.distributions.histogram_utils.histogram_approximation(
+            "histogram_potential", pm.Poisson.dist(s), observed=data, min_count=10
+        )
+        trace = pm.sample(10, tune=0)  # very fast

tests/test_laplace.py

@@ -0,0 +1,238 @@
+#   Copyright 2024 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 numpy as np
+import pymc as pm
+import pytest
+
+import pymc_extras as pmx
+
+from pymc_extras.inference.find_map import find_MAP
+from pymc_extras.inference.laplace import (
+    fit_laplace,
+    fit_mvn_to_MAP,
+    sample_laplace_posterior,
+)
+
+
+@pytest.fixture(scope="session")
+def rng():
+    seed = sum(map(ord, "test_laplace"))
+    return np.random.default_rng(seed)
+
+
+@pytest.mark.filterwarnings(
+    "ignore:hessian will stop negating the output in a future version of PyMC.\n"
+    + "To suppress this warning set `negate_output=False`:FutureWarning",
+)
+def test_laplace():
+    # Example originates from Bayesian Data Analyses, 3rd Edition
+    # By Andrew Gelman, John Carlin, Hal Stern, David Dunson,
+    # Aki Vehtari, and Donald Rubin.
+    # See section. 4.1
+
+    y = np.array([2642, 3503, 4358], dtype=np.float64)
+    n = y.size
+    draws = 100000
+
+    with pm.Model() as m:
+        mu = pm.Uniform("mu", -10000, 10000)
+        logsigma = pm.Uniform("logsigma", 1, 100)
+
+        yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+        vars = [mu, logsigma]
+
+        idata = pmx.fit(
+            method="laplace", optimize_method="trust-ncg", draws=draws, random_seed=173300, chains=1
+        )
+
+    assert idata.posterior["mu"].shape == (1, draws)
+    assert idata.posterior["logsigma"].shape == (1, draws)
+    assert idata.observed_data["y"].shape == (n,)
+    assert idata.fit["mean_vector"].shape == (len(vars),)
+    assert idata.fit["covariance_matrix"].shape == (len(vars), len(vars))
+
+    bda_map = [y.mean(), np.log(y.std())]
+    bda_cov = np.array([[y.var() / n, 0], [0, 1 / (2 * n)]])
+
+    np.testing.assert_allclose(idata.fit["mean_vector"].values, bda_map)
+    np.testing.assert_allclose(idata.fit["covariance_matrix"].values, bda_cov, atol=1e-4)
+
+
+def test_laplace_only_fit():
+    # Example originates from Bayesian Data Analyses, 3rd Edition
+    # By Andrew Gelman, John Carlin, Hal Stern, David Dunson,
+    # Aki Vehtari, and Donald Rubin.
+    # See section. 4.1
+
+    y = np.array([2642, 3503, 4358], dtype=np.float64)
+    n = y.size
+
+    with pm.Model() as m:
+        logsigma = pm.Uniform("logsigma", 1, 100)
+        mu = pm.Uniform("mu", -10000, 10000)
+        yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+        vars = [mu, logsigma]
+
+        idata = pmx.fit(
+            method="laplace",
+            optimize_method="BFGS",
+            progressbar=True,
+            gradient_backend="jax",
+            compile_kwargs={"mode": "JAX"},
+            optimizer_kwargs=dict(maxiter=100_000, gtol=1e-100),
+            random_seed=173300,
+        )
+
+    assert idata.fit["mean_vector"].shape == (len(vars),)
+    assert idata.fit["covariance_matrix"].shape == (len(vars), len(vars))
+
+    bda_map = [np.log(y.std()), y.mean()]
+    bda_cov = np.array([[1 / (2 * n), 0], [0, y.var() / n]])
+
+    np.testing.assert_allclose(idata.fit["mean_vector"].values, bda_map)
+    np.testing.assert_allclose(idata.fit["covariance_matrix"].values, bda_cov, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+    "transform_samples",
+    [True, False],
+    ids=["transformed", "untransformed"],
+)
+@pytest.mark.parametrize("mode", ["JAX", None], ids=["jax", "pytensor"])
+def test_fit_laplace_coords(rng, transform_samples, mode):
+    coords = {"city": ["A", "B", "C"], "obs_idx": np.arange(100)}
+    with pm.Model(coords=coords) as model:
+        mu = pm.Normal("mu", mu=3, sigma=0.5, dims=["city"])
+        sigma = pm.Exponential("sigma", 1, dims=["city"])
+        obs = pm.Normal(
+            "obs",
+            mu=mu,
+            sigma=sigma,
+            observed=rng.normal(loc=3, scale=1.5, size=(100, 3)),
+            dims=["obs_idx", "city"],
+        )
+
+        optimized_point = find_MAP(
+            method="trust-ncg",
+            use_grad=True,
+            use_hessp=True,
+            progressbar=False,
+            compile_kwargs=dict(mode=mode),
+            gradient_backend="jax" if mode == "JAX" else "pytensor",
+        )
+
+        for value in optimized_point.values():
+            assert value.shape == (3,)
+
+        mu, H_inv = fit_mvn_to_MAP(
+            optimized_point=optimized_point,
+            model=model,
+            transform_samples=transform_samples,
+        )
+
+        idata = sample_laplace_posterior(
+            mu=mu, H_inv=H_inv, model=model, transform_samples=transform_samples
+        )
+
+    np.testing.assert_allclose(np.mean(idata.posterior.mu, axis=1), np.full((2, 3), 3), atol=0.5)
+    np.testing.assert_allclose(
+        np.mean(idata.posterior.sigma, axis=1), np.full((2, 3), 1.5), atol=0.3
+    )
+
+    suffix = "_log__" if transform_samples else ""
+    assert idata.fit.rows.values.tolist() == [
+        "mu[A]",
+        "mu[B]",
+        "mu[C]",
+        f"sigma{suffix}[A]",
+        f"sigma{suffix}[B]",
+        f"sigma{suffix}[C]",
+    ]
+
+
+def test_fit_laplace_ragged_coords(rng):
+    coords = {"city": ["A", "B", "C"], "feature": [0, 1], "obs_idx": np.arange(100)}
+    with pm.Model(coords=coords) as ragged_dim_model:
+        X = pm.Data("X", np.ones((100, 2)), dims=["obs_idx", "feature"])
+        beta = pm.Normal(
+            "beta", mu=[[-100.0, 100.0], [-100.0, 100.0], [-100.0, 100.0]], dims=["city", "feature"]
+        )
+        mu = pm.Deterministic(
+            "mu", (X[:, None, :] * beta[None]).sum(axis=-1), dims=["obs_idx", "city"]
+        )
+        sigma = pm.Normal("sigma", mu=1.5, sigma=0.5, dims=["city"])
+
+        obs = pm.Normal(
+            "obs",
+            mu=mu,
+            sigma=sigma,
+            observed=rng.normal(loc=3, scale=1.5, size=(100, 3)),
+            dims=["obs_idx", "city"],
+        )
+
+        idata = fit_laplace(
+            optimize_method="Newton-CG",
+            progressbar=False,
+            use_grad=True,
+            use_hessp=True,
+            gradient_backend="jax",
+            compile_kwargs={"mode": "JAX"},
+        )
+
+    assert idata["posterior"].beta.shape[-2:] == (3, 2)
+    assert idata["posterior"].sigma.shape[-1:] == (3,)
+
+    # Check that everything got unraveled correctly -- feature 0 should be strictly negative, feature 1
+    # strictly positive
+    assert (idata["posterior"].beta.sel(feature=0).to_numpy() < 0).all()
+    assert (idata["posterior"].beta.sel(feature=1).to_numpy() > 0).all()
+
+
+@pytest.mark.parametrize(
+    "fit_in_unconstrained_space",
+    [True, False],
+    ids=["transformed", "untransformed"],
+)
+def test_fit_laplace(fit_in_unconstrained_space):
+    with pm.Model() as simp_model:
+        mu = pm.Normal("mu", mu=3, sigma=0.5)
+        sigma = pm.Exponential("sigma", 1)
+        obs = pm.Normal(
+            "obs",
+            mu=mu,
+            sigma=sigma,
+            observed=np.random.default_rng().normal(loc=3, scale=1.5, size=(10000,)),
+        )
+
+        idata = fit_laplace(
+            optimize_method="trust-ncg",
+            use_grad=True,
+            use_hessp=True,
+            fit_in_unconstrained_space=fit_in_unconstrained_space,
+            optimizer_kwargs=dict(maxiter=100_000, tol=1e-100),
+        )
+
+        np.testing.assert_allclose(np.mean(idata.posterior.mu, axis=1), np.full((2,), 3), atol=0.1)
+        np.testing.assert_allclose(
+            np.mean(idata.posterior.sigma, axis=1), np.full((2,), 1.5), atol=0.1
+        )
+
+        if fit_in_unconstrained_space:
+            assert idata.fit.rows.values.tolist() == ["mu", "sigma_log__"]
+            np.testing.assert_allclose(idata.fit.mean_vector.values, np.array([3.0, 0.4]), atol=0.1)
+        else:
+            assert idata.fit.rows.values.tolist() == ["mu", "sigma"]
+            np.testing.assert_allclose(idata.fit.mean_vector.values, np.array([3.0, 1.5]), atol=0.1)

tests/test_linearmodel.py

@@ -0,0 +1,208 @@
+#   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 sys
+import tempfile
+
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+from pymc_extras.linearmodel import LinearModel
+
+try:
+    from sklearn import set_config
+    from sklearn.compose import TransformedTargetRegressor
+    from sklearn.pipeline import Pipeline
+    from sklearn.preprocessing import StandardScaler
+
+    set_config(transform_output="pandas")
+    sklearn_available = True
+except ImportError:
+    sklearn_available = False
+
+
+@pytest.fixture(scope="module")
+def toy_actual_params():
+    return {
+        "intercept": 3,
+        "slope": 5,
+        "obs_error": 0.5,
+    }
+
+
+@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, toy_actual_params):
+    y = toy_actual_params["slope"] * toy_X["input"] + toy_actual_params["intercept"]
+    rng = np.random.default_rng(427)
+    y = y + rng.normal(0, toy_actual_params["obs_error"], size=len(toy_X))
+    y = pd.Series(y, name="output")
+    return y
+
+
+@pytest.fixture(scope="module")
+def fitted_linear_model_instance(toy_X, toy_y):
+    sampler_config = {
+        "draws": 50,
+        "tune": 30,
+        "chains": 2,
+        "target_accept": 0.95,
+    }
+    model = LinearModel(sampler_config=sampler_config)
+    model.fit(toy_X, toy_y, random_seed=312)
+    return model
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
+)
+def test_save_load(fitted_linear_model_instance):
+    model = fitted_linear_model_instance
+    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
+    model.save(temp.name)
+    model2 = LinearModel.load(temp.name)
+    assert model.idata.groups() == model2.idata.groups()
+
+    X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
+    pred1 = model.predict(X_pred, random_seed=423)
+    pred2 = model2.predict(X_pred, random_seed=423)
+    # Predictions should be identical
+    np.testing.assert_array_equal(pred1, pred2)
+    temp.close()
+
+
+def test_save_without_fit_raises_runtime_error(toy_X, toy_y):
+    test_model = LinearModel()
+    with pytest.raises(RuntimeError):
+        test_model.save("saved_model")
+
+
+def test_fit(fitted_linear_model_instance):
+    model = fitted_linear_model_instance
+
+    new_X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
+
+    pred = model.predict(new_X_pred)
+    assert len(new_X_pred) == len(pred)
+    assert isinstance(pred, np.ndarray)
+    post_pred = model.predict_posterior(new_X_pred)
+    assert len(new_X_pred) == len(post_pred)
+    assert isinstance(post_pred, xr.DataArray)
+
+
+def test_parameter_fit(toy_X, toy_y, toy_actual_params):
+    """Check that the fit model recovered the data-generating parameters."""
+    # Fit the model with a sufficient number of samples
+    sampler_config = {
+        "draws": 500,
+        "tune": 300,
+        "chains": 2,
+        "target_accept": 0.95,
+    }
+    model = LinearModel(sampler_config=sampler_config)
+    model.fit(toy_X, toy_y, random_seed=312)
+    fit_params = model.idata.posterior.mean()
+    np.testing.assert_allclose(fit_params["intercept"], toy_actual_params["intercept"], rtol=0.1)
+    np.testing.assert_allclose(fit_params["slope"], toy_actual_params["slope"], rtol=0.1)
+    np.testing.assert_allclose(fit_params["σ_model_fmc"], toy_actual_params["obs_error"], rtol=0.1)
+
+
+def test_predict(fitted_linear_model_instance):
+    model = fitted_linear_model_instance
+    X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
+    pred = model.predict(X_pred)
+    assert len(X_pred) == len(pred)
+    assert np.issubdtype(pred.dtype, np.floating)
+
+
+@pytest.mark.parametrize("combined", [True, False])
+def test_predict_posterior(fitted_linear_model_instance, combined):
+    model = fitted_linear_model_instance
+    n_pred = 150
+    X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=n_pred)})
+    pred = model.predict_posterior(X_pred, combined=combined)
+    chains = model.idata.sample_stats.sizes["chain"]
+    draws = model.idata.sample_stats.sizes["draw"]
+    expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
+    assert pred.shape == expected_shape
+    assert np.issubdtype(pred.dtype, np.floating)
+    # TODO: check that extend_idata has the expected effect
+
+
+@pytest.mark.parametrize("samples", [None, 300])
+@pytest.mark.parametrize("combined", [True, False])
+def test_sample_prior_predictive(samples, combined, toy_X, toy_y):
+    model = LinearModel()
+    prior_pred = model.sample_prior_predictive(toy_X, toy_y, samples, combined=combined)[
+        model.output_var
+    ]
+    draws = model.sampler_config["draws"] if samples is None else samples
+    chains = 1
+    expected_shape = (len(toy_X), chains * draws) if combined else (chains, draws, len(toy_X))
+    assert prior_pred.shape == expected_shape
+    # TODO: check that extend_idata has the expected effect
+
+
+def test_id():
+    model_config = {
+        "intercept": {"loc": 0, "scale": 10},
+        "slope": {"loc": 0, "scale": 10},
+        "obs_error": 2,
+    }
+    sampler_config = {
+        "draws": 1_000,
+        "tune": 1_000,
+        "chains": 3,
+        "target_accept": 0.95,
+    }
+    model = LinearModel(model_config=model_config, sampler_config=sampler_config)
+
+    expected_id = hashlib.sha256(
+        str(model_config.values()).encode() + model.version.encode() + model._model_type.encode()
+    ).hexdigest()[:16]
+
+    assert model.id == expected_id
+
+
+@pytest.mark.skipif(not sklearn_available, reason="scikit-learn package is not available.")
+def test_pipeline_integration(toy_X, toy_y):
+    model_config = {
+        "intercept": {"loc": 0, "scale": 2},
+        "slope": {"loc": 0, "scale": 2},
+        "obs_error": 1,
+        "default_output_var": "y_hat",
+    }
+    model = Pipeline(
+        [
+            ("input_scaling", StandardScaler()),
+            (
+                "linear_model",
+                TransformedTargetRegressor(LinearModel(model_config), transformer=StandardScaler()),
+            ),
+        ]
+    )
+    model.fit(toy_X, toy_y)
+
+    X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
+    model.predict(X_pred)