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

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()]

tests/test_prior_from_trace.py

@@ -1,172 +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 arviz as az
-import numpy as np
-import pymc as pm
-import pytest
-
-from pymc.distributions import transforms
-
-import pymc_extras as pmx
-
-
-@pytest.mark.parametrize(
-    "case",
-    [
-        (("a", dict(name="b")), dict(name="b", transform=None, dims=None)),
-        (("a", None), dict(name="a", transform=None, dims=None)),
-        (("a", transforms.log), dict(name="a", transform=transforms.log, dims=None)),
-        (
-            ("a", dict(transform=transforms.log)),
-            dict(name="a", transform=transforms.log, dims=None),
-        ),
-        (("a", dict(name="b")), dict(name="b", transform=None, dims=None)),
-        (("a", dict(name="b", dims="test")), dict(name="b", transform=None, dims="test")),
-        (("a", ("test",)), dict(name="a", transform=None, dims=("test",))),
-    ],
-)
-def test_parsing_arguments(case):
-    inp, out = case
-    test = pmx.utils.prior._arg_to_param_cfg(*inp)
-    assert test == out
-
-
-@pytest.fixture
-def coords():
-    return dict(test=range(3), simplex=range(4))
-
-
-@pytest.fixture(
-    params=[
-        [
-            ("t",),
-            dict(
-                a="d",
-                b=dict(transform=transforms.log, dims=("test",)),
-                c=dict(transform=transforms.simplex, dims=("simplex",)),
-            ),
-        ],
-        [("t",), dict()],
-    ]
-)
-def user_param_cfg(request):
-    return request.param
-
-
-@pytest.fixture
-def param_cfg(user_param_cfg):
-    return pmx.utils.prior._parse_args(user_param_cfg[0], **user_param_cfg[1])
-
-
-@pytest.fixture
-def transformed_data(param_cfg, coords):
-    vars = dict()
-    for k, cfg in param_cfg.items():
-        if cfg["dims"] is not None:
-            extra_dims = [len(coords[d]) for d in cfg["dims"]]
-            if cfg["transform"] is not None:
-                t = np.random.randn(*extra_dims)
-                extra_dims = tuple(cfg["transform"].forward(t).shape.eval())
-        else:
-            extra_dims = []
-        orig = np.random.randn(4, 100, *extra_dims)
-        vars[k] = orig
-    return vars
-
-
-@pytest.fixture
-def idata(transformed_data, param_cfg):
-    vars = dict()
-    for k, orig in transformed_data.items():
-        cfg = param_cfg[k]
-        if cfg["transform"] is not None:
-            var = cfg["transform"].backward(orig).eval()
-        else:
-            var = orig
-        assert not np.isnan(var).any()
-        vars[k] = var
-    return az.convert_to_inference_data(vars)
-
-
-def test_idata_for_tests(idata, param_cfg):
-    assert set(idata.posterior.keys()) == set(param_cfg)
-    assert len(idata.posterior.coords["chain"]) == 4
-    assert len(idata.posterior.coords["draw"]) == 100
-
-
-def test_args_compose():
-    cfg = pmx.utils.prior._parse_args(
-        var_names=["a"],
-        b=("test",),
-        c=transforms.log,
-        d="e",
-        f=dict(dims="test"),
-        g=dict(name="h", dims="test", transform=transforms.log),
-    )
-    assert cfg == dict(
-        a=dict(name="a", dims=None, transform=None),
-        b=dict(name="b", dims=("test",), transform=None),
-        c=dict(name="c", dims=None, transform=transforms.log),
-        d=dict(name="e", dims=None, transform=None),
-        f=dict(name="f", dims="test", transform=None),
-        g=dict(name="h", dims="test", transform=transforms.log),
-    )
-
-
-def test_transform_idata(transformed_data, idata, param_cfg):
-    flat_info = pmx.utils.prior._flatten(idata, **param_cfg)
-    expected_shape = 0
-    for v in transformed_data.values():
-        expected_shape += int(np.prod(v.shape[2:]))
-    assert flat_info["data"].shape[1] == expected_shape
-    assert len(flat_info["info"]) == len(param_cfg)
-    assert "sinfo" in flat_info["info"][0]
-    assert "vinfo" in flat_info["info"][0]
-
-
-@pytest.fixture
-def flat_info(idata, param_cfg):
-    return pmx.utils.prior._flatten(idata, **param_cfg)
-
-
-def test_mean_chol(flat_info):
-    mean, chol = pmx.utils.prior._mean_chol(flat_info["data"])
-    assert mean.shape == (flat_info["data"].shape[1],)
-    assert chol.shape == (flat_info["data"].shape[1],) * 2
-
-
-def test_mvn_prior_from_flat_info(flat_info, coords, param_cfg):
-    with pm.Model(coords=coords) as model:
-        priors = pmx.utils.prior._mvn_prior_from_flat_info("trace_prior_", flat_info)
-        test_prior = pm.sample_prior_predictive(1)
-    names = [p["name"] for p in param_cfg.values()]
-    assert set(model.named_vars) == {"trace_prior_", *names}
-
-
-def test_prior_from_idata(idata, user_param_cfg, coords, param_cfg):
-    with pm.Model(coords=coords) as model:
-        priors = pmx.utils.prior.prior_from_idata(
-            idata, var_names=user_param_cfg[0], **user_param_cfg[1]
-        )
-        test_prior = pm.sample_prior_predictive(1)
-    names = [p["name"] for p in param_cfg.values()]
-    assert set(model.named_vars) == {"trace_prior_", *names}
-
-
-def test_empty(idata, coords):
-    with pm.Model(coords=coords):
-        priors = pmx.utils.prior.prior_from_idata(idata)
-        assert not priors

tests/test_splines.py

@@ -1,77 +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 pytensor.tensor as pt
-import pytest
-
-from pytensor.sparse import SparseTensorType
-
-import pymc_extras as pmx
-
-
-@pytest.mark.parametrize("dtype", [np.float32, np.float64])
-@pytest.mark.parametrize("sparse", [True, False])
-def test_spline_construction(dtype, sparse):
-    x = np.linspace(0, 1, 20, dtype=dtype)
-    np_out = pmx.utils.spline.numpy_bspline_basis(x, 10, 3)
-    assert np_out.shape == (20, 10)
-    assert np_out.dtype == dtype
-    spline_op = pmx.utils.spline.BSplineBasis(sparse=sparse)
-    out = spline_op(x, pt.constant(10), pt.constant(3))
-    if not sparse:
-        assert isinstance(out.type, pt.TensorType)
-    else:
-        assert isinstance(out.type, SparseTensorType)
-    B = out.eval()
-    if not sparse:
-        np.testing.assert_allclose(B, np_out)
-    else:
-        np.testing.assert_allclose(B.todense(), np_out)
-    assert B.shape == (20, 10)
-
-
-@pytest.mark.parametrize("shape", [(100,), (100, 5)])
-@pytest.mark.parametrize("sparse", [True, False])
-@pytest.mark.parametrize("points", [dict(n=1001), dict(eval_points=np.linspace(0, 1, 1001))])
-def test_interpolation_api(shape, sparse, points):
-    x = np.random.randn(*shape)
-    yt = pmx.utils.spline.bspline_interpolation(x, **points, sparse=sparse)
-    y = yt.eval()
-    assert y.shape == (1001, *shape[1:])
-
-
-@pytest.mark.parametrize(
-    "params",
-    [
-        (dict(sparse="foo", n=100, degree=1), TypeError, "sparse should be True or False"),
-        (dict(n=100, degree=0.5), TypeError, "degree should be integer"),
-        (
-            dict(n=100, eval_points=np.linspace(0, 1), degree=1),
-            ValueError,
-            "Please provide one of n or eval_points",
-        ),
-        (
-            dict(degree=1),
-            ValueError,
-            "Please provide one of n or eval_points",
-        ),
-    ],
-)
-def test_bad_calls(params):
-    kw, E, err = params
-    x = np.random.randn(10)
-    with pytest.raises(E, match=err):
-        pmx.utils.spline.bspline_interpolation(x, **kw)