diff --git a/docs/source/conf.py b/docs/source/conf.py index be8cb5eda..9322a5602 100755 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -164,6 +164,7 @@ # fmt: on numpydoc_xref_aliases = { "TensorVariable": ":class:`~pytensor.tensor.TensorVariable`", + "XTensorVariable": ":class:`~pytensor.xtensor.type.XTensorVariable`", "RandomVariable": ":class:`~pytensor.tensor.random.RandomVariable`", "ndarray": ":class:`~numpy.ndarray`", "DataTree": ":class:`~xarray.DataTree`", diff --git a/pymc_marketing/bass/model.py b/pymc_marketing/bass/model.py index fd2585c63..0a9c43291 100644 --- a/pymc_marketing/bass/model.py +++ b/pymc_marketing/bass/model.py @@ -20,7 +20,9 @@ access to the PyMC model object. The standalone functions :func:`F`, :func:`f`, and :func:`create_bass_model` -are still exposed for direct use. +are still exposed for direct use. :func:`F` and :func:`f` take xtensor +inputs; wrap plain arrays with :func:`pymc.dims.as_xtensor` to call them +outside a model. Adapted from Wiki: https://en.wikipedia.org/wiki/Bass_diffusion_model @@ -134,6 +136,8 @@ """ +from contextlib import contextmanager +from inspect import signature from typing import Any, TypedDict, cast import arviz as az @@ -142,6 +146,7 @@ import numpy.typing as npt import pandas as pd import pymc as pm +import pymc.dims as pmd import pytensor.tensor as pt import xarray as xr from matplotlib.axes import Axes @@ -150,7 +155,14 @@ ) from pymc.model import Model from pymc.util import RandomState -from pymc_extras.prior import Censored, Prior, VariableFactory, create_dim_handler +from pymc_extras.prior import ( + Censored, + MuAlreadyExistsError, + Prior, + UnsupportedDistributionError, + VariableFactory, +) +from pytensor.xtensor.type import XTensorVariable from pymc_marketing.bass import plotting from pymc_marketing.bass.data import to_bass_dataset @@ -159,11 +171,27 @@ from pymc_marketing.version import __version__ +def _check_time(t: object) -> None: + """Reject a ``t`` that ``pymc.dims`` cannot label on its own. + + A scalar tensor converts cleanly, an array or a dim-less vector does not, + and the conversion error it raises does not say what to do about it. + """ + if isinstance(t, XTensorVariable) or ( + isinstance(t, pt.TensorVariable) and t.ndim == 0 + ): + return + raise TypeError( + f"`t` must be an XTensorVariable, got {type(t).__name__}. " + "Wrap plain arrays with `pymc.dims.as_xtensor(t, dims=('T',))`." + ) + + def F( - p: float | pt.TensorVariable, - q: float | pt.TensorVariable, - t: float | pt.TensorVariable, -) -> pt.TensorVariable: + p: float | XTensorVariable, + q: float | XTensorVariable, + t: XTensorVariable | pt.TensorVariable, +) -> XTensorVariable: r"""Installed base fraction (cumulative adoption proportion). This function calculates the cumulative proportion of adopters at time t, @@ -171,16 +199,16 @@ def F( Parameters ---------- - p : float or TensorVariable + p : float or XTensorVariable Coefficient of innovation (external influence) - q : float or TensorVariable + q : float or XTensorVariable Coefficient of imitation (internal influence) - t : array-like or TensorVariable + t : XTensorVariable or scalar TensorVariable Time points Returns ------- - TensorVariable + XTensorVariable The cumulative proportion of adopters at each time point Notes @@ -193,14 +221,15 @@ def F( When :math:`t=0`, :math:`F(t)=0`, and as :math:`t` approaches infinity, :math:`F(t)` approaches 1. """ - return (1 - pt.exp(-(p + q) * t)) / (1 + (q / p) * pt.exp(-(p + q) * t)) + _check_time(t) + return (1 - pmd.math.exp(-(p + q) * t)) / (1 + (q / p) * pmd.math.exp(-(p + q) * t)) def f( - p: float | pt.TensorVariable, - q: float | pt.TensorVariable, - t: float | pt.TensorVariable, -) -> pt.TensorVariable: + p: float | XTensorVariable, + q: float | XTensorVariable, + t: XTensorVariable | pt.TensorVariable, +) -> XTensorVariable: r"""Installed base fraction rate of change (adoption rate). This function calculates the rate of new adoptions at time t as a @@ -209,16 +238,16 @@ def f( Parameters ---------- - p : float or TensorVariable + p : float or XTensorVariable Coefficient of innovation (external influence) - q : float or TensorVariable + q : float or XTensorVariable Coefficient of imitation (internal influence) - t : array-like or TensorVariable + t : XTensorVariable or scalar TensorVariable Time points Returns ------- - TensorVariable + XTensorVariable The adoption rate at each time point as a fraction of potential market Notes @@ -237,9 +266,93 @@ def f( The peak adoption rate occurs at time :math:`t^* = \frac{\ln(q/p)}{p+q}` """ - return (p * pt.square(p + q) * pt.exp(t * (p + q))) / pt.square( - p * pt.exp(t * (p + q)) + q + _check_time(t) + exp_t = pmd.math.exp(t * (p + q)) + return (p * (p + q) ** 2 * exp_t) / (p * exp_t + q) ** 2 + + +def _create_likelihood_variable( + prior: Prior | Censored, + name: str, + mu: XTensorVariable, + observed: XTensorVariable | None, +) -> XTensorVariable: + """Create the outcome variable, observed or not. + + ``create_likelihood_variable`` is for the observed case only: a + likelihood needs data, so pymc_extras refuses ``observed=None`` there + (pymc-devs/pymc-extras#731). Prior predictive still needs the outcome + node, so build it with ``create_variable`` and ``mu`` attached, keeping + the same guards the pymc_extras method applies. + """ + if observed is not None: + return prior.create_likelihood_variable( + name, mu=mu, observed=observed, xdist=True + ) + + # Censored keeps its parameters on the wrapped distribution. + inner = prior.distribution if isinstance(prior, Censored) else prior + if "mu" not in signature(inner.pymc_distribution.dist).parameters: + raise UnsupportedDistributionError( + f"Likelihood distribution {inner.distribution!r} is not supported." + ) + if "mu" in inner.parameters: + raise MuAlreadyExistsError(inner) + + # TODO(pymc-devs/pymc-extras#731): drop this branch once observed=None is + # supported upstream. Rebind rather than mutate the copy's parameters, so + # the model keeps the caller's own tensors instead of deepcopied clones. + unobserved = inner.deepcopy() + unobserved.parameters = {**inner.parameters, "mu": mu} + outcome: Prior | Censored = ( + Censored(unobserved, lower=prior.lower, upper=prior.upper) + if isinstance(prior, Censored) + else unobserved ) + return outcome.create_variable(name, xdist=True) + + +@contextmanager +def _borrow_dims(prior: Prior | Censored, dims: tuple[str, ...]): + """Lend ``dims`` to ``prior`` for the block, leaving it as it was found. + + Setting ``dims`` outright would leave the caller's prior carrying this + model's dims, so a config reused for a second model fails on dims that + model does not have. A copy is not used instead: ``Prior.__deepcopy__`` + also copies ``parameters``, which would hand the model clones of any + tensor the caller passed in. ``Censored.dims`` forwards to the wrapped + distribution, so both types are covered. + """ + original = prior.dims + prior.dims = dims + try: + yield prior + finally: + prior.dims = original + + +def _observed_dims( + observed: Any, model: Model, combined_dims: tuple[str, ...] +) -> tuple[str, ...]: + """Axis labels for ``observed``: its own, else the model's, else positional. + + An ``xr.DataArray`` carries its labels; a registered ``pm.Data`` has them + on the model. Anything else is labelled positionally in ``combined_dims`` + order. + """ + own = getattr(observed, "dims", None) + if own: + return tuple(own) + + # Name equality is not identity: a variable that never reached the model + # can share a name with one that did, and would borrow its dims. + name = getattr(observed, "name", None) + if name is not None and model.named_vars.get(name) is observed: + registered = model.named_vars_to_dims.get(name) + if registered and all(dim is not None for dim in registered): + return tuple(registered) + + return combined_dims class BassPriors(TypedDict): @@ -253,7 +366,7 @@ class BassPriors(TypedDict): def create_bass_model( t: pt.TensorLike, - observed: pt.TensorLike | None, + observed: pt.TensorLike | xr.DataArray | None, priors: BassPriors, coords: dict[str, Any], model: Model | None = None, @@ -279,9 +392,15 @@ def create_bass_model( ---------- t : pt.TensorLike Time points for which the adoption is modeled. - observed : pt.TensorLike | None + observed : pt.TensorLike or xr.DataArray or None Observed adoption data at each time point. If None, only - prior predictive sampling is possible. + prior predictive sampling is possible. Axis labels are read from + the data itself (an ``xr.DataArray``) or from the model (a + ``pm.Data`` registered with dims); anything else, such as a plain + array or a ``pm.Data`` without dims, is labelled positionally in + ``(T, ...)`` order with the extra dims following their first + appearance across the ``p``, ``q``, ``m`` and ``likelihood`` + priors, in that order. priors : BassPriors Dictionary containing priors for: - 'm': Market potential prior @@ -319,48 +438,49 @@ def create_bass_model( """ model = model or pm.Model(coords=coords) with model: - parameter_dims = ( - set(priors["p"].dims or ()) - .union(priors["q"].dims or ()) - .union(priors["m"].dims or ()) + # Declaration order, not set order: `combined_dims` labels the axes of + # `observed` positionally, so an order that varies between processes + # would silently mislabel the data. + declared_dims = ( + *(priors["p"].dims or ()), + *(priors["q"].dims or ()), + *(priors["m"].dims or ()), + *(getattr(priors["likelihood"], "dims", ()) or ()), ) - likelihood_dims = set(getattr(priors["likelihood"], "dims", ()) or ()) - combined_dims = ( "T", - *tuple(parameter_dims.union(likelihood_dims).difference(["T"])), - ) - dim_handler = create_dim_handler(combined_dims) - - m = dim_handler(priors["m"].create_variable("m"), priors["m"].dims) - p = dim_handler(priors["p"].create_variable("p"), priors["p"].dims) - q = dim_handler(priors["q"].create_variable("q"), priors["q"].dims) - - time = dim_handler(t, "T") - - adopters = pm.Deterministic("adopters", m * f(p, q, time), dims=combined_dims) - - pm.Deterministic( - "innovators", - m * p * (1 - F(p, q, time)), - dims=combined_dims, - ) - pm.Deterministic( - "imitators", - m * q * F(p, q, time) * (1 - F(p, q, time)), - dims=combined_dims, + *(dim for dim in dict.fromkeys(declared_dims) if dim != "T"), ) - peak = (pt.log(q) - pt.log(p)) / (p + q) - peak_dims = tuple(parameter_dims) if parameter_dims else None - pm.Deterministic("peak", peak, dims=peak_dims) - - priors["likelihood"].dims = combined_dims - priors["likelihood"].create_likelihood_variable( # type: ignore - "y", - mu=adopters, - observed=observed, + time = pmd.as_xtensor(t, dims=("T",)) + m = priors["m"].create_variable("m", xdist=True) + p = priors["p"].create_variable("p", xdist=True) + q = priors["q"].create_variable("q", xdist=True) + + def deterministic(name: str, value: XTensorVariable) -> XTensorVariable: + """Store ``value`` with the dims it has, in ``combined_dims`` order.""" + order = tuple(dim for dim in combined_dims if dim in value.dims) + return pmd.Deterministic(name, value, dims=order) + + adopters = deterministic("adopters", m * f(p, q, time)) + deterministic("innovators", m * p * (1 - F(p, q, time))) + deterministic("imitators", m * q * F(p, q, time) * (1 - F(p, q, time))) + deterministic("peak", (pmd.math.log(q) - pmd.math.log(p)) / (p + q)) + + observed_xt = ( + None + if observed is None + else pmd.as_xtensor( + observed, dims=_observed_dims(observed, model, combined_dims) + ) ) + with _borrow_dims(priors["likelihood"], combined_dims) as likelihood: + _create_likelihood_variable( + likelihood, + "y", + mu=adopters, + observed=observed_xt, + ) return model diff --git a/pymc_marketing/bass/plotting.py b/pymc_marketing/bass/plotting.py index c95adab88..a93361a7b 100644 --- a/pymc_marketing/bass/plotting.py +++ b/pymc_marketing/bass/plotting.py @@ -15,9 +15,11 @@ Each function takes a fitted :class:`~pymc_marketing.bass.model.BassModel` and returns a ``(Figure, ndarray of Axes)`` tuple, following the convention -of :func:`pymc_marketing.plot.plot_curve`. Multi-series models plot one -subplot per coordinate of the faceting dimension (inferred from the data, -or set with ``dim``); pass ``coord`` to select a single one. +of :func:`pymc_marketing.plot.plot_curve`. Models plot one subplot per +coordinate of the faceting dimension (inferred from the plotted variable, +or set with ``dim``); pass ``coord`` to select a single one. The curves +carry the dims their priors declare, so pooled ``p``, ``q`` and ``m`` give +a single pooled curve even when the data has more dimensions. The functions are also exposed as ``plot_*`` methods on :class:`~pymc_marketing.bass.model.BassModel`. @@ -64,8 +66,10 @@ def _select(da: xr.DataArray, dim: str | None, coord: str | None) -> xr.DataArra return da if dim is None or dim not in da.dims: raise ValueError( - "The model has no faceting dimension. " - "Remove the coord argument for single-series models." + f"The plotted variable has no {'faceting' if dim is None else repr(dim)} " + "dimension to select a coord from. The curves only carry the dims " + "their priors declare, so give `p`, `q` or `m` that dim to get one " + "curve per coordinate." ) return da.sel({dim: coord}) @@ -116,8 +120,8 @@ def plot_adoption_curve( model : BassModel A fitted Bass model. dim : str, optional - Dimension to facet over. Inferred from the fitted data when not - given (the non-time, non-sample dimension). + Dimension to facet over. Inferred from the plotted variable when + not given (its non-time, non-sample dimension). coord : str, optional Plot a single coordinate along ``dim``. Default plots one subplot per coordinate. @@ -195,8 +199,8 @@ def plot_cumulative( model : BassModel A fitted Bass model. dim : str, optional - Dimension to facet over. Inferred from the fitted data when not - given (the non-time, non-sample dimension). + Dimension to facet over. Inferred from the plotted variable when + not given (its non-time, non-sample dimension). coord : str, optional Plot a single coordinate along ``dim``. Default plots one subplot per coordinate. @@ -263,8 +267,8 @@ def plot_decomposition( model : BassModel A fitted Bass model. dim : str, optional - Dimension to facet over. Inferred from the fitted data when not - given (the non-time, non-sample dimension). + Dimension to facet over. Inferred from the plotted variable when + not given (its non-time, non-sample dimension). coord : str, optional Plot a single coordinate along ``dim``. Default plots one subplot per coordinate. @@ -367,8 +371,8 @@ def plot_peak( model : BassModel A fitted Bass model. dim : str, optional - Dimension to facet over. Inferred from the fitted data when not - given (the non-time, non-sample dimension). + Dimension to facet over. Inferred from the plotted variable when + not given (its non-time, non-sample dimension). coord : str, optional Plot a single coordinate along ``dim``. Default plots one subplot per coordinate. diff --git a/pymc_marketing/mlflow.py b/pymc_marketing/mlflow.py index 47afbf304..e672ae1a3 100644 --- a/pymc_marketing/mlflow.py +++ b/pymc_marketing/mlflow.py @@ -561,8 +561,11 @@ def log_model_graph(model: Model, path: str | Path) -> None: def _get_random_variable_name(rv) -> str: - # Taken from new version of pymc/model_graph.py - symbol = rv.owner.op.__class__.__name__ + op = rv.owner.op + # A `pymc.dims` variable wraps the real RV in a generic `XRV`, which would + # otherwise be reported as "X". The distribution is on the wrapped op. + op = getattr(op, "core_op", op) + symbol = op.__class__.__name__ if symbol.endswith("RV"): symbol = symbol[:-2] diff --git a/tests/bass/test_model.py b/tests/bass/test_model.py index 3b3dac089..fb58dfaf8 100644 --- a/tests/bass/test_model.py +++ b/tests/bass/test_model.py @@ -20,12 +20,20 @@ import numpy.typing as npt import pandas as pd import pymc as pm +import pymc.dims as pmd import pytensor import pytensor.tensor as pt import pytest import xarray as xr from pydantic import BaseModel, ConfigDict -from pymc_extras.prior import Censored, Prior, Scaled +from pymc_extras.prior import ( + Censored, + MuAlreadyExistsError, + Prior, + Scaled, + UnsupportedDistributionError, +) +from pytensor.graph import rewrite_graph from pymc_marketing.bass import BassModel from pymc_marketing.bass.model import F, create_bass_model, f @@ -91,7 +99,7 @@ def bass_model_components() -> tuple[ m_true = 1000 # Generate data - adopters_true = m_true * f(p_true, q_true, t).eval() + adopters_true = m_true * f(p_true, q_true, pmd.as_xtensor(t, dims=("T",))).eval() # Add noise rng = np.random.default_rng(42) @@ -125,7 +133,7 @@ def test_f_function(self) -> None: ) # Calculate actual values from the function - actual = f(p, q, t).eval() + actual = f(p, q, pmd.as_xtensor(t, dims=("T",))).eval() np.testing.assert_allclose(actual, expected, rtol=1e-5) @@ -139,7 +147,7 @@ def test_F_function(self) -> None: expected = (1 - np.exp(-(p + q) * t)) / (1 + (q / p) * np.exp(-(p + q) * t)) # Calculate actual values from the function - actual = F(p, q, t).eval() + actual = F(p, q, pmd.as_xtensor(t, dims=("T",))).eval() np.testing.assert_allclose(actual, expected, rtol=1e-5) @@ -152,7 +160,7 @@ def test_f_function_boundary_conditions(self) -> None: # At t=0, f(t) gives approximately 0.00219512 # This is (p+q)/(1+(q/p))^2 = (p+q)*p^2/(p+q)^2 = p^2/(p+q) expected = (p * ((p + q) ** 2)) / (p + q) ** 2 - actual = f(p, q, t).eval() + actual = f(p, q, pmd.as_xtensor(t, dims=("T",))).eval() np.testing.assert_allclose(actual, expected, rtol=1e-5) @@ -164,16 +172,28 @@ def test_F_function_boundary_conditions(self): # At t=0, F(t) should be 0 t = np.array([0]) expected = 0 - actual = F(p, q, t).eval() + actual = F(p, q, pmd.as_xtensor(t, dims=("T",))).eval() np.testing.assert_allclose(actual, expected, rtol=1e-5) # As t approaches infinity, F(t) should approach 1 t = np.array([1000]) # Very large t expected = 1 - actual = F(p, q, t).eval() + actual = F(p, q, pmd.as_xtensor(t, dims=("T",))).eval() np.testing.assert_allclose(actual, expected, rtol=1e-2) +@pytest.mark.parametrize("func", [F, f], ids=["F", "f"]) +@pytest.mark.parametrize( + "t", + [np.array([0.0, 1.0]), [0.0, 1.0], pt.vector("t")], + ids=["ndarray", "list", "vector"], +) +def test_plain_array_time_points_point_at_as_xtensor(func, t) -> None: + """`t` must be labelled; the error has to say how.""" + with pytest.raises(TypeError, match="as_xtensor"): + func(0.03, 0.38, t) + + class TestBassModel: """Test the Bass model creation and behavior.""" @@ -232,10 +252,14 @@ def test_bass_model_deterministics( q_val = 0.38 # Calculate expected values using the formulas directly - expected_adopters = m_val * f(p_val, q_val, t).eval() - expected_innovators = m_val * p_val * (1 - F(p_val, q_val, t)).eval() + t_xt = pmd.as_xtensor(t, dims=("T",)) + expected_adopters = m_val * f(p_val, q_val, t_xt).eval() + expected_innovators = m_val * p_val * (1 - F(p_val, q_val, t_xt)).eval() expected_imitators = ( - m_val * q_val * F(p_val, q_val, t).eval() * (1 - F(p_val, q_val, t).eval()) + m_val + * q_val + * F(p_val, q_val, t_xt).eval() + * (1 - F(p_val, q_val, t_xt).eval()) ) expected_peak = (np.log(q_val) - np.log(p_val)) / (p_val + q_val) @@ -624,13 +648,319 @@ def test_bass_model_simulation(self) -> None: assert "peak" in prior_samples["prior"] +def make_priors(**overrides: Any) -> dict[str, Any]: + """Fresh base priors for create_bass_model; override any key per test.""" + priors: dict[str, Any] = { + "m": Prior("Normal", mu=1000, sigma=200), + "p": Prior("Beta", alpha=1.5, beta=20), + "q": Prior("Beta", alpha=2, beta=5), + "likelihood": Prior("Poisson"), + } + priors.update(overrides) + return priors + + +class TestBassModelLikelihood: + """Guards and dim errors on the outcome variable.""" + + @pytest.fixture + def coords(self) -> dict[str, Any]: + return {"T": np.arange(5), "product": ["A", "B"]} + + def test_censored_likelihood_without_observed(self, coords: dict[str, Any]) -> None: + """A Censored likelihood must still build for prior predictive only.""" + priors = make_priors(likelihood=Censored(Prior("Normal", sigma=1), lower=0)) + + model = create_bass_model( + t=coords["T"], observed=None, priors=priors, coords=coords + ) + + assert "y" in model.named_vars + assert model.named_vars_to_dims["y"] == ("T",) + + def test_mu_on_likelihood_without_observed_raises( + self, coords: dict[str, Any] + ) -> None: + """Setting mu on the likelihood is a misconfiguration, with or without data.""" + priors = make_priors(likelihood=Prior("Normal", mu=5, sigma=1)) + + with pytest.raises(MuAlreadyExistsError): + create_bass_model( + t=coords["T"], observed=None, priors=priors, coords=coords + ) + + def test_mu_less_likelihood_without_observed_raises( + self, coords: dict[str, Any] + ) -> None: + """The unobserved path applies the same guard as create_likelihood_variable.""" + priors = make_priors(likelihood=Prior("Binomial", n=10, p=0.5)) + + with pytest.raises(UnsupportedDistributionError, match="Binomial"): + create_bass_model( + t=coords["T"], observed=None, priors=priors, coords=coords + ) + + def test_likelihood_dim_missing_from_coords_raises(self) -> None: + """A likelihood-only dim without a coord fails with pymc's own error.""" + t = np.arange(5) + priors = make_priors(likelihood=Prior("Poisson", dims=("product",))) + + with pytest.raises(ValueError, match=r"(?i)dims.*are part of the model coords"): + create_bass_model(t=t, observed=None, priors=priors, coords={"T": t}) + + def test_observed_dim_missing_from_coords_raises( + self, coords: dict[str, Any] + ) -> None: + """An observed dim the model does not know still fails, on the dim name.""" + priors = make_priors() + observed = xr.DataArray(np.ones((5, 2)), dims=("T", "geo")) + + with pytest.raises(KeyError, match="geo"): + create_bass_model( + t=coords["T"], observed=observed, priors=priors, coords=coords + ) + + +class TestBassModelDims: + """Dim names and ordering of the deterministics and the likelihood.""" + + def test_deterministics_carry_only_their_own_dims(self) -> None: + """Pooled p, q, m: the curves are pooled too, whatever dims the data has. + + The deterministics are not broadcast up to the likelihood's dims; a + per-product curve is the user's to build from the pooled one. + """ + coords = {"T": np.arange(5), "product": ["A", "B"]} + priors = make_priors(likelihood=Prior("Poisson", dims=("product",))) + observed = np.ones((5, 2)) + + model = create_bass_model( + t=coords["T"], observed=observed, priors=priors, coords=coords + ) + + for name in ["adopters", "innovators", "imitators"]: + assert model.named_vars_to_dims[name] == ("T",) + assert model.named_vars_to_dims["y"] == ("T", "product") + + def test_combined_dims_follow_declaration_order(self) -> None: + """Dim order comes from the priors, not from set iteration order.""" + coords = { + "T": np.arange(4), + "country": ["a", "b"], + "product": ["x", "y", "z"], + } + priors = { + "m": Prior("Normal", mu=1000, sigma=200, dims=("country", "product")), + "p": Prior("Beta", alpha=1.5, beta=20, dims=("country", "product")), + "q": Prior("Beta", alpha=2, beta=5, dims=("country", "product")), + "likelihood": Prior("Poisson"), + } + observed = np.ones((4, 2, 3)) + + model = create_bass_model( + t=coords["T"], observed=observed, priors=priors, coords=coords + ) + + assert model.named_vars_to_dims["adopters"] == ("T", "country", "product") + assert model.named_vars_to_dims["y"] == ("T", "country", "product") + + def test_observed_is_labelled_with_its_own_dims(self) -> None: + """The same data in two layouts must give the same logp. + + ``combined_dims`` need not match the layout of ``observed``, so labelling + the axes with it transposes the data without saying so. + """ + coords = { + "T": np.arange(4), + "country": ["a", "b"], + "product": ["x", "y", "z"], + } + priors = { + "m": Prior("Normal", mu=1000, sigma=200, dims=("product", "country")), + "p": Prior("Beta", alpha=1.5, beta=20, dims=("product", "country")), + "q": Prior("Beta", alpha=2, beta=5, dims=("product", "country")), + "likelihood": Prior("Poisson"), + } + counts = xr.DataArray( + np.arange(24, dtype=float).reshape(4, 2, 3), + dims=("T", "country", "product"), + coords=coords, + ) + + logps = [] + for dims in [("T", "country", "product"), ("T", "product", "country")]: + with pm.Model(coords=coords) as model: + y_obs = pm.Data("y_obs", counts.transpose(*dims).values, dims=dims) + create_bass_model( + t=coords["T"], + observed=y_obs, + priors=priors, + coords=coords, + model=model, + ) + logps.append(model.point_logps()["y"]) + + assert logps[0] == logps[1] + + def test_variable_sharing_a_name_does_not_borrow_its_dims(self) -> None: + """Name equality is not identity. + + A variable that never reached the model can carry the name of one + that did; reading dims off the name alone would mislabel its axes. + """ + coords = {"T": np.arange(4), "product": ["A", "B"]} + priors = make_priors( + m=Prior("Normal", mu=100, sigma=10, dims="product"), + p=Prior("Beta", alpha=1.5, beta=20, dims="product"), + q=Prior("Beta", alpha=2, beta=5, dims="product"), + ) + # never registered on the model, but named after one that is + observed = pt.as_tensor(np.ones((4, 2))) + observed.name = "peak" + + with pm.Model(coords=coords) as model: + create_bass_model( + t=coords["T"], + observed=observed, + priors=priors, + coords=coords, + model=model, + ) + + assert model.named_vars_to_dims["y"] == ("T", "product") + + def test_observed_data_without_dims_is_labelled_positionally(self) -> None: + """A pm.Data registered without dims falls back to combined_dims. + + There is nothing to read the labels from, so the data must already + be laid out in ``(T, ...)`` order. Documented behaviour, same as + before the ``pymc.dims`` migration. + """ + coords = { + "T": np.arange(4), + "product": ["x", "y", "z"], + } + priors = { + "m": Prior("Normal", mu=1000, sigma=200, dims=("product",)), + "p": Prior("Beta", alpha=1.5, beta=20), + "q": Prior("Beta", alpha=2, beta=5), + "likelihood": Prior("Poisson"), + } + counts = np.arange(12, dtype=float).reshape(4, 3) + + logps = [] + for dims in [None, ("T", "product")]: + with pm.Model(coords=coords) as model: + y_obs = pm.Data("y_obs", counts, dims=dims) + create_bass_model( + t=coords["T"], + observed=y_obs, + priors=priors, + coords=coords, + model=model, + ) + assert model.named_vars_to_dims["y"] == ("T", "product") + logps.append(model.point_logps()["y"]) + + assert logps[0] == logps[1] + + def test_observed_dataarray_keeps_its_layout(self) -> None: + """An xr.DataArray carries its own dims; both layouts must agree.""" + coords = { + "T": np.arange(4), + "country": ["a", "b"], + "product": ["x", "y", "z"], + } + priors = { + "m": Prior("Normal", mu=1000, sigma=200, dims=("country", "product")), + "p": Prior("Beta", alpha=1.5, beta=20, dims=("country", "product")), + "q": Prior("Beta", alpha=2, beta=5, dims=("country", "product")), + "likelihood": Prior("Poisson"), + } + counts = xr.DataArray( + np.arange(24, dtype=float).reshape(4, 2, 3), + dims=("T", "country", "product"), + coords=coords, + ) + + logps = [] + for dims in [("T", "country", "product"), ("product", "T", "country")]: + model = create_bass_model( + t=coords["T"], + observed=counts.transpose(*dims), + priors=priors, + coords=coords, + ) + logps.append(model.point_logps()["y"]) + + assert logps[0] == logps[1] + + def test_peak_dims_follow_combined_dims_order(self) -> None: + """peak keeps only the parameter dims, ordered like the other curves.""" + coords = { + "T": np.arange(4), + "product": ["x", "y", "z"], + "country": ["a", "b"], + } + priors = { + "m": Prior("Normal", mu=1000, sigma=200, dims=("product", "country")), + "p": Prior("Beta", alpha=1.5, beta=20, dims=("product",)), + "q": Prior("Beta", alpha=2, beta=5, dims=("country",)), + "likelihood": Prior("Poisson"), + } + + model = create_bass_model( + t=coords["T"], observed=None, priors=priors, coords=coords + ) + + assert model.named_vars_to_dims["adopters"] == ("T", "product", "country") + assert model.named_vars_to_dims["peak"] == ("product", "country") + + def test_dim_known_to_the_model_but_absent_from_coords(self) -> None: + """Sizes come from the model, so a partial coords dict still builds.""" + coords = {"T": np.arange(5), "product": ["A", "B"]} + priors = make_priors(likelihood=Prior("Poisson", dims=("product",))) + + with pm.Model(coords=coords) as model: + create_bass_model( + t=coords["T"], + observed=np.ones((5, 2)), + priors=priors, + # only T, the model already carries product + coords={"T": coords["T"]}, + model=model, + ) + + assert model.named_vars_to_dims["y"] == ("T", "product") + + def test_building_does_not_mutate_the_caller_s_priors(self) -> None: + """The likelihood keeps the dims it declares, not the model's own.""" + priors = make_priors(likelihood=Prior("Poisson", dims=("product",))) + coords = {"T": np.arange(5), "product": ["A", "B"]} + + model = create_bass_model( + t=coords["T"], observed=np.ones((5, 2)), priors=priors, coords=coords + ) + + assert model.named_vars_to_dims["y"] == ("T", "product") + assert priors["likelihood"].dims == ("product",) + + +def lower(var: Any) -> pt.TensorVariable: + """Lower an xtensor graph so ``pytensor.grad`` can walk it. + + Gradients of xtensor ops are pending in pymc-devs/pytensor#2337. + """ + return rewrite_graph(var.values, include=("lower_xtensor",), clone=False) + + def test_derivative() -> None: p = pt.scalar("p") q = pt.scalar("q") t = pt.scalar("t") - F_res = F(p, q, t) + F_res = lower(F(p, q, t)) F_prime_fn = pytensor.function([p, q, t], pytensor.grad(F_res, t)) - f_res = f(p, q, t) + f_res = lower(f(p, q, t)) f_fn = pytensor.function([p, q, t], f_res) for p_, q_, t_ in product([0.01, 0.02, 0.03], [0.3, 0.4, 0.5], [0, 10, 100]): @@ -756,6 +1086,16 @@ def test_user_m_prior_is_untouched(self, y: np.ndarray): assert model.model_config["m"].parameters["sigma"] == 123.0 assert self._graph_m_sigma(model) == pytest.approx(123.0) + def test_likelihood_prior_is_untouched(self, y: np.ndarray): + """Building must not write the model's dims into the config it reads.""" + model = BassModel() + before = model.id + + model.build_model(data=y) + + assert model.model_config["likelihood"].dims is None + assert model.id == before + def test_m_prior_scale_survives_save_load( self, mock_pymc_sample, y: np.ndarray, tmp_path ): diff --git a/tests/bass/test_plotting.py b/tests/bass/test_plotting.py index 538d19187..102445a37 100644 --- a/tests/bass/test_plotting.py +++ b/tests/bass/test_plotting.py @@ -177,3 +177,27 @@ def test_product_on_single_product_raises( ) -> None: with pytest.raises(ValueError, match="no faceting dimension"): getattr(single_product_model, method)(coord="A") + + +@pytest.mark.parametrize("method", PLOT_METHODS) +def test_pooled_priors_on_multi_series_data(mock_pymc_sample, method: str) -> None: + """Pooled priors give one pooled curve, whatever dims the data has. + + The curves carry only the dims the priors declare, so selecting a coord + must say so instead of claiming the model is single-series. + """ + counts = np.random.default_rng(42).poisson(lam=100, size=(20, 3)) + data = xr.Dataset( + {"observed": (("T", "product"), counts)}, + coords={"T": np.arange(20), "product": ["A", "B", "C"]}, + ) + model = BassModel(model_config={"likelihood": Prior("Poisson", dims="product")}) + model.fit(data=data, draws=20, tune=5, chains=1, random_seed=42) + + assert model.idata.posterior["adopters"].dims == ("chain", "draw", "T") + + _, axes = getattr(model, method)() + assert np.asarray(axes).size == 1 + + with pytest.raises(ValueError, match="only carry the dims"): + getattr(model, method)(coord="A") diff --git a/tests/test_mlflow.py b/tests/test_mlflow.py index a67285de0..311581d19 100644 --- a/tests/test_mlflow.py +++ b/tests/test_mlflow.py @@ -824,6 +824,10 @@ def test_autolog_bass(bass_data) -> None: assert params["model_type"] == "BassModel" assert params["version"] == __version__ + # The Bass model builds its variables with pymc.dims, which wraps every RV + # in a generic XRV; the logged name must still be the distribution. + assert params["likelihood"] == "Poisson" + model_config_logged = json.loads(params["model_config"]) assert set(model_config_logged.keys()) == {"m", "p", "q", "likelihood"}