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

tests/test_histogram_approximation.py

@@ -1,109 +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 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

@@ -1,265 +0,0 @@
-#   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 GradientBackend, find_MAP
-from pymc_extras.inference.laplace import (
-    fit_laplace,
-    fit_mvn_at_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",
-)
-@pytest.mark.parametrize(
-    "mode, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
-)
-def test_laplace(mode, gradient_backend: GradientBackend):
-    # 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,
-            compile_kwargs={"mode": mode},
-            gradient_backend=gradient_backend,
-        )
-
-    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)
-
-
-@pytest.mark.parametrize(
-    "mode, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
-)
-def test_laplace_only_fit(mode, gradient_backend: GradientBackend):
-    # 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=gradient_backend,
-            compile_kwargs={"mode": mode},
-            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, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
-)
-def test_fit_laplace_coords(rng, transform_samples, mode, gradient_backend: GradientBackend):
-    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=gradient_backend,
-        )
-
-        for value in optimized_point.values():
-            assert value.shape == (3,)
-
-        mu, H_inv = fit_mvn_at_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]",
-    ]
-
-
-@pytest.mark.parametrize(
-    "mode, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
-)
-def test_fit_laplace_ragged_coords(mode, gradient_backend: GradientBackend, 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=gradient_backend,
-            compile_kwargs={"mode": mode},
-        )
-
-    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"],
-)
-@pytest.mark.parametrize(
-    "mode, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
-)
-def test_fit_laplace(fit_in_unconstrained_space, mode, gradient_backend: GradientBackend):
-    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),
-            compile_kwargs={"mode": mode},
-            gradient_backend=gradient_backend,
-        )
-
-        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

@@ -1,208 +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 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)