arviz 0.20.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

arviz/plots/lmplot.py

@@ -300,20 +300,47 @@ def plot_lm(
     # Filter out the required values to generate plotters
     if y_model is not None:
         if kind_model == "lines":
-            y_model = y_model.stack(__sample__=("chain", "draw"))[..., pp_sample_ix]
-
-        y_model = [
-            tup
-            for _, tup in zip(
-                range(len_y),
-                xarray_var_iter(
-                    y_model,
-                    skip_dims=set(y_model.dims),
-                    combined=True,
-                ),
-            )
-        ]
-        y_model = _repeat_flatten_list(y_model, len_x)
+            var_name = y_model.name if y_model.name else "y_model"
+            data = y_model.values
+
+            total_samples = data.shape[0] * data.shape[1]
+            data = data.reshape(total_samples, *data.shape[2:])
+
+            if pp_sample_ix is not None:
+                data = data[pp_sample_ix]
+
+            if plot_dim is not None:
+                # For plot_dim case, transpose to get dimension first
+                data = data.transpose(1, 0, 2)[..., 0]
+
+            # Create plotter tuple(s)
+            if plot_dim is not None:
+                y_model = [(var_name, {}, {}, data) for _ in range(length_plotters)]
+            else:
+                y_model = [(var_name, {}, {}, data)]
+                y_model = _repeat_flatten_list(y_model, len_x)
+
+        elif kind_model == "hdi":
+            var_name = y_model.name if y_model.name else "y_model"
+            data = y_model.values
+
+            if plot_dim is not None:
+                # First transpose to get plot_dim first
+                data = data.transpose(2, 0, 1, 3)
+                # For plot_dim case, we just want HDI for first dimension
+                data = data[..., 0]
+
+                # Reshape to (samples, points)
+                data = data.transpose(1, 2, 0).reshape(-1, data.shape[0])
+                y_model = [(var_name, {}, {}, data) for _ in range(length_plotters)]
+
+            else:
+                data = data.reshape(-1, data.shape[-1])
+                y_model = [(var_name, {}, {}, data)]
+                y_model = _repeat_flatten_list(y_model, len_x)
+
+        if len(y_model) == 1:
+            y_model = _repeat_flatten_list(y_model, len_x)
 
     rows, cols = default_grid(length_plotters)
 

arviz/plots/pairplot.py

@@ -196,9 +196,14 @@ def plot_pair(
             get_coords(dataset, coords), var_names=var_names, skip_dims=combine_dims, combined=True
         )
     )
-    flat_var_names = [
-        labeller.make_label_vert(var_name, sel, isel) for var_name, sel, isel, _ in plotters
-    ]
+    flat_var_names = []
+    flat_ref_slices = []
+    flat_var_labels = []
+    for var_name, sel, isel, _ in plotters:
+        dims = [dim for dim in dataset[var_name].dims if dim not in ["chain", "draw"]]
+        flat_var_names.append(var_name)
+        flat_ref_slices.append(tuple(isel[dim] if dim in isel else slice(None) for dim in dims))
+        flat_var_labels.append(labeller.make_label_vert(var_name, sel, isel))
 
     divergent_data = None
     diverging_mask = None
@@ -253,6 +258,8 @@ def plot_pair(
         diverging_mask=diverging_mask,
         divergences_kwargs=divergences_kwargs,
         flat_var_names=flat_var_names,
+        flat_ref_slices=flat_ref_slices,
+        flat_var_labels=flat_var_labels,
         backend_kwargs=backend_kwargs,
         marginal_kwargs=marginal_kwargs,
         show=show,

arviz/plots/plot_utils.py

@@ -482,16 +482,18 @@ def plot_point_interval(
         if point_estimate:
             point_value = calculate_point_estimate(point_estimate, values)
             if rotated:
-                ax.circle(
+                ax.scatter(
                     x=0,
                     y=point_value,
+                    marker="circle",
                     size=markersize,
                     fill_color=markercolor,
                 )
             else:
-                ax.circle(
+                ax.scatter(
                     x=point_value,
                     y=0,
+                    marker="circle",
                     size=markersize,
                     fill_color=markercolor,
                 )
@@ -534,7 +536,7 @@ def set_bokeh_circular_ticks_labels(ax, hist, labels):
     )
 
     radii_circles = np.linspace(0, np.max(hist) * 1.1, 4)
-    ax.circle(0, 0, radius=radii_circles, fill_color=None, line_color="grey")
+    ax.scatter(0, 0, marker="circle", radius=radii_circles, fill_color=None, line_color="grey")
 
     offset = np.max(hist * 1.05) * 0.15
     ticks_labels_pos_1 = np.max(hist * 1.05)

arviz/preview.py

@@ -1,17 +1,48 @@
-# pylint: disable=unused-import,unused-wildcard-import,wildcard-import
+# pylint: disable=unused-import,unused-wildcard-import,wildcard-import,invalid-name
 """Expose features from arviz-xyz refactored packages inside ``arviz.preview`` namespace."""
+import logging
+
+_log = logging.getLogger(__name__)
+
+info = ""
 
 try:
     from arviz_base import *
+
+    status = "arviz_base available, exposing its functions as part of arviz.preview"
+    _log.info(status)
 except ModuleNotFoundError:
-    pass
+    status = "arviz_base not installed"
+    _log.info(status)
+except ImportError:
+    status = "Unable to import arviz_base"
+    _log.info(status, exc_info=True)
+
+info += status + "\n"
 
 try:
-    import arviz_stats
+    from arviz_stats import *
+
+    status = "arviz_stats available, exposing its functions as part of arviz.preview"
+    _log.info(status)
 except ModuleNotFoundError:
-    pass
+    status = "arviz_stats not installed"
+    _log.info(status)
+except ImportError:
+    status = "Unable to import arviz_stats"
+    _log.info(status, exc_info=True)
+info += status + "\n"
 
 try:
     from arviz_plots import *
+
+    status = "arviz_plots available, exposing its functions as part of arviz.preview"
+    _log.info(status)
 except ModuleNotFoundError:
-    pass
+    status = "arviz_plots not installed"
+    _log.info(status)
+except ImportError:
+    status = "Unable to import arviz_plots"
+    _log.info(status, exc_info=True)
+
+info += status + "\n"

arviz/stats/__init__.py

@@ -9,6 +9,7 @@ from .stats_utils import *
 
 __all__ = [
     "apply_test_function",
+    "bayes_factor",
     "bfmi",
     "compare",
     "hdi",

arviz/stats/density_utils.py

@@ -635,7 +635,7 @@ def _kde_circular(
     cumulative: bool, optional
         Whether return the PDF or the cumulative PDF. Defaults to False.
     grid_len: int, optional
-        The number of intervals used to bin the data pointa i.e. the length of the grid used in the
+        The number of intervals used to bin the data point i.e. the length of the grid used in the
         estimation. Defaults to 512.
     """
     # All values between -pi and pi

arviz/stats/diagnostics.py

@@ -744,8 +744,8 @@ def _ess_sd(ary, relative=False):
     ary = np.asarray(ary)
     if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
         return np.nan
-    ary = _split_chains(ary)
-    return min(_ess(ary, relative=relative), _ess(ary**2, relative=relative))
+    ary = (ary - ary.mean()) ** 2
+    return _ess(_split_chains(ary), relative=relative)
 
 
 def _ess_quantile(ary, prob, relative=False):
@@ -838,13 +838,15 @@ def _mcse_sd(ary):
     ary = np.asarray(ary)
     if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
         return np.nan
-    ess = _ess_sd(ary)
+    sims_c2 = (ary - ary.mean()) ** 2
+    ess = _ess_mean(sims_c2)
+    evar = (sims_c2).mean()
+    varvar = ((sims_c2**2).mean() - evar**2) / ess
+    varsd = varvar / evar / 4
     if _numba_flag:
-        sd = float(_sqrt(svar(np.ravel(ary), ddof=1), np.zeros(1)).item())
+        mcse_sd_value = float(_sqrt(np.ravel(varsd), np.zeros(1)))
     else:
-        sd = np.std(ary, ddof=1)
-    fac_mcse_sd = np.sqrt(np.exp(1) * (1 - 1 / ess) ** (ess - 1) - 1)
-    mcse_sd_value = sd * fac_mcse_sd
+        mcse_sd_value = np.sqrt(varsd)
     return mcse_sd_value
 
 
@@ -973,19 +975,21 @@ def _multichain_statistics(ary, focus="mean"):
         # ess mean
         ess_mean_value = _ess_mean(ary)
 
-        # ess sd
-        ess_sd_value = _ess_sd(ary)
-
         # mcse_mean
-        sd = np.std(ary, ddof=1)
-        mcse_mean_value = sd / np.sqrt(ess_mean_value)
+        sims_c2 = (ary - ary.mean()) ** 2
+        sims_c2_sum = sims_c2.sum()
+        var = sims_c2_sum / (sims_c2.size - 1)
+        mcse_mean_value = np.sqrt(var / ess_mean_value)
 
         # ess bulk
         ess_bulk_value = _ess(z_split)
 
         # mcse_sd
-        fac_mcse_sd = np.sqrt(np.exp(1) * (1 - 1 / ess_sd_value) ** (ess_sd_value - 1) - 1)
-        mcse_sd_value = sd * fac_mcse_sd
+        evar = sims_c2_sum / sims_c2.size
+        ess_mean_sims = _ess_mean(sims_c2)
+        varvar = ((sims_c2**2).mean() - evar**2) / ess_mean_sims
+        varsd = varvar / evar / 4
+        mcse_sd_value = np.sqrt(varsd)
 
         return (
             mcse_mean_value,

arviz/stats/stats.py

@@ -27,6 +27,7 @@ from ..utils import Numba, _numba_var, _var_names, get_coords
 from .density_utils import get_bins as _get_bins
 from .density_utils import histogram as _histogram
 from .density_utils import kde as _kde
+from .density_utils import _kde_linear
 from .diagnostics import _mc_error, _multichain_statistics, ess
 from .stats_utils import ELPDData, _circular_standard_deviation, smooth_data
 from .stats_utils import get_log_likelihood as _get_log_likelihood
@@ -41,6 +42,7 @@ from ..labels import BaseLabeller
 
 __all__ = [
     "apply_test_function",
+    "bayes_factor",
     "compare",
     "hdi",
     "loo",
@@ -867,7 +869,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     )
 
 
-def psislw(log_weights, reff=1.0):
+def psislw(log_weights, reff=1.0, normalize=True):
     """
     Pareto smoothed importance sampling (PSIS).
 
@@ -885,11 +887,13 @@ def psislw(log_weights, reff=1.0):
         Array of size (n_observations, n_samples)
     reff : float, default 1
         relative MCMC efficiency, ``ess / n``
+    normalize : bool, default True
+        return normalized log weights
 
     Returns
     -------
     lw_out : DataArray or (..., N) ndarray
-        Smoothed, truncated and normalized log weights.
+        Smoothed, truncated and possibly normalized log weights.
     kss : DataArray or (...) ndarray
         Estimates of the shape parameter *k* of the generalized Pareto
         distribution.
@@ -934,7 +938,12 @@ def psislw(log_weights, reff=1.0):
     out = np.empty_like(log_weights), np.empty(shape)
 
     # define kwargs
-    func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "out": out}
+    func_kwargs = {
+        "cutoff_ind": cutoff_ind,
+        "cutoffmin": cutoffmin,
+        "out": out,
+        "normalize": normalize,
+    }
     ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
     kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
     log_weights, pareto_shape = _wrap_xarray_ufunc(
@@ -951,7 +960,7 @@ def psislw(log_weights, reff=1.0):
     return log_weights, pareto_shape
 
 
-def _psislw(log_weights, cutoff_ind, cutoffmin):
+def _psislw(log_weights, cutoff_ind, cutoffmin, normalize):
     """
     Pareto smoothed importance sampling (PSIS) for a 1D vector.
 
@@ -961,7 +970,7 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
         Array of length n_observations
     cutoff_ind: int
     cutoffmin: float
-    k_min: float
+    normalize: bool
 
     Returns
     -------
@@ -973,7 +982,8 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
     x = np.asarray(log_weights)
 
     # improve numerical accuracy
-    x -= np.max(x)
+    max_x = np.max(x)
+    x -= max_x
     # sort the array
     x_sort_ind = np.argsort(x)
     # divide log weights into body and right tail
@@ -1005,8 +1015,12 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
             x[tailinds[x_tail_si]] = smoothed_tail
             # truncate smoothed values to the largest raw weight 0
             x[x > 0] = 0
+
     # renormalize weights
-    x -= _logsumexp(x)
+    if normalize:
+        x -= _logsumexp(x)
+    else:
+        x += max_x
 
     return x, k
 
@@ -2337,3 +2351,87 @@ def _cjs_dist(draws, weights):
     bound = cdf_p_int + cdf_q_int
 
     return np.sqrt((cjs_pq + cjs_qp) / bound)
+
+
+def bayes_factor(idata, var_name, ref_val=0, prior=None, return_ref_vals=False):
+    r"""Approximated Bayes Factor for comparing hypothesis of two nested models.
+
+    The Bayes factor is estimated by comparing a model (H1) against a model in which the
+    parameter of interest has been restricted to be a point-null (H0). This computation
+    assumes the models are nested and thus H0 is a special case of H1.
+
+    Notes
+    -----
+    The bayes Factor is approximated as the Savage-Dickey density ratio
+    algorithm presented in [1]_.
+
+    Parameters
+    ----------
+    idata : InferenceData
+        Any object that can be converted to an :class:`arviz.InferenceData` object
+        Refer to documentation of :func:`arviz.convert_to_dataset` for details.
+    var_name : str, optional
+        Name of variable we want to test.
+    ref_val : int, default 0
+        Point-null for Bayes factor estimation.
+    prior : numpy.array, optional
+        In case we want to use different prior, for example for sensitivity analysis.
+    return_ref_vals : bool, optional
+        Whether to return the values of the prior and posterior at the reference value.
+        Used by :func:`arviz.plot_bf` to display the distribution comparison.
+
+
+    Returns
+    -------
+    dict : A dictionary with BF10 (Bayes Factor 10 (H1/H0 ratio), and BF01 (H0/H1 ratio).
+
+    References
+    ----------
+    .. [1] Heck, D., 2019. A caveat on the Savage-Dickey density ratio:
+       The case of computing Bayes factors for regression parameters.
+
+    Examples
+    --------
+    Moderate evidence indicating that the parameter "a" is different from zero.
+
+    .. ipython::
+
+        In [1]: import numpy as np
+           ...: import arviz as az
+           ...: idata = az.from_dict(posterior={"a":np.random.normal(1, 0.5, 5000)},
+           ...: prior={"a":np.random.normal(0, 1, 5000)})
+           ...: az.bayes_factor(idata, var_name="a", ref_val=0)
+
+    """
+
+    posterior = extract(idata, var_names=var_name).values
+
+    if ref_val > posterior.max() or ref_val < posterior.min():
+        _log.warning(
+            "The reference value is outside of the posterior. "
+            "This translate into infinite support for H1, which is most likely an overstatement."
+        )
+
+    if posterior.ndim > 1:
+        _log.warning("Posterior distribution has {posterior.ndim} dimensions")
+
+    if prior is None:
+        prior = extract(idata, var_names=var_name, group="prior").values
+
+    if posterior.dtype.kind == "f":
+        posterior_grid, posterior_pdf, *_ = _kde_linear(posterior)
+        prior_grid, prior_pdf, *_ = _kde_linear(prior)
+        posterior_at_ref_val = np.interp(ref_val, posterior_grid, posterior_pdf)
+        prior_at_ref_val = np.interp(ref_val, prior_grid, prior_pdf)
+
+    elif posterior.dtype.kind == "i":
+        posterior_at_ref_val = (posterior == ref_val).mean()
+        prior_at_ref_val = (prior == ref_val).mean()
+
+    bf_10 = prior_at_ref_val / posterior_at_ref_val
+    bf = {"BF10": bf_10, "BF01": 1 / bf_10}
+
+    if return_ref_vals:
+        return (bf, {"prior": prior_at_ref_val, "posterior": posterior_at_ref_val})
+    else:
+        return bf

arviz/tests/base_tests/test_data.py

@@ -3,6 +3,7 @@
 
 import importlib
 import os
+import warnings
 from collections import namedtuple
 from copy import deepcopy
 from html import escape
@@ -32,7 +33,13 @@ from ... import (
     extract,
 )
 
-from ...data.base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs
+from ...data.base import (
+    dict_to_dataset,
+    generate_dims_coords,
+    infer_stan_dtypes,
+    make_attrs,
+    numpy_to_data_array,
+)
 from ...data.datasets import LOCAL_DATASETS, REMOTE_DATASETS, RemoteFileMetadata
 from ..helpers import (  # pylint: disable=unused-import
     chains,
@@ -231,6 +238,17 @@ def test_dims_coords_skip_event_dims(shape):
     assert "z" not in coords
 
 
+@pytest.mark.parametrize("dims", [None, ["chain", "draw"], ["chain", "draw", None]])
+def test_numpy_to_data_array_with_dims(dims):
+    da = numpy_to_data_array(
+        np.empty((4, 500, 7)),
+        var_name="a",
+        dims=dims,
+        default_dims=["chain", "draw"],
+    )
+    assert list(da.dims) == ["chain", "draw", "a_dim_0"]
+
+
 def test_make_attrs():
     extra_attrs = {"key": "Value"}
     attrs = make_attrs(attrs=extra_attrs)
@@ -921,7 +939,7 @@ class TestInferenceData:  # pylint: disable=too-many-public-methods
         data = np.random.normal(size=(4, 500, 8))
         idata = data_random
         with pytest.warns(UserWarning, match="The group.+not defined in the InferenceData scheme"):
-            idata.add_groups({"new_group": idata.posterior})
+            idata.add_groups({"new_group": idata.posterior}, warn_on_custom_groups=True)
         with pytest.warns(UserWarning, match="the default dims.+will be added automatically"):
             idata.add_groups(constant_data={"a": data[..., 0], "b": data})
         assert idata.new_group.equals(idata.posterior)
@@ -962,8 +980,8 @@ class TestInferenceData:  # pylint: disable=too-many-public-methods
         with pytest.raises(ValueError, match="join must be either"):
             idata.extend(idata2, join="outer")
         idata2.add_groups(new_group=idata2.prior)
-        with pytest.warns(UserWarning):
-            idata.extend(idata2)
+        with pytest.warns(UserWarning, match="new_group"):
+            idata.extend(idata2, warn_on_custom_groups=True)
 
 
 class TestNumpyToDataArray:
@@ -1156,11 +1174,17 @@ def test_bad_inference_data():
         InferenceData(posterior=[1, 2, 3])
 
 
-def test_inference_data_other_groups():
+@pytest.mark.parametrize("warn", [True, False])
+def test_inference_data_other_groups(warn):
     datadict = {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)}
     dataset = convert_to_dataset(datadict, coords={"c": np.arange(10)}, dims={"b": ["c"]})
-    with pytest.warns(UserWarning, match="not.+in.+InferenceData scheme"):
-        idata = InferenceData(other_group=dataset)
+    if warn:
+        with pytest.warns(UserWarning, match="not.+in.+InferenceData scheme"):
+            idata = InferenceData(other_group=dataset, warn_on_custom_groups=True)
+    else:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            idata = InferenceData(other_group=dataset, warn_on_custom_groups=False)
     fails = check_multiple_attrs({"other_group": ["a", "b"]}, idata)
     assert not fails
 
@@ -1477,10 +1501,6 @@ class TestJSON:
         assert not os.path.exists(filepath)
 
 
-@pytest.mark.skipif(
-    not (importlib.util.find_spec("datatree") or "ARVIZ_REQUIRE_ALL_DEPS" in os.environ),
-    reason="test requires xarray-datatree library",
-)
 class TestDataTree:
     def test_datatree(self):
         idata = load_arviz_data("centered_eight")

arviz/tests/base_tests/test_diagnostics.py

@@ -120,10 +120,11 @@ class TestDiagnostics:
         ```
         Reference file:
 
-        Created: 2020-08-31
-        System: Ubuntu 18.04.5 LTS
-        R version 4.0.2 (2020-06-22)
-        posterior 0.1.2
+        Created: 2024-12-20
+        System: Ubuntu 24.04.1 LTS
+        R version 4.4.2 (2024-10-31)
+        posterior version from https://github.com/stan-dev/posterior/pull/388
+            (after release 1.6.0 but before the fixes in the PR were released).
         """
         # download input files
         here = os.path.dirname(os.path.abspath(__file__))

arviz/tests/base_tests/test_plots_bokeh.py

@@ -8,6 +8,7 @@ from pandas import DataFrame  # pylint: disable=wrong-import-position
 from scipy.stats import norm  # pylint: disable=wrong-import-position
 
 from ...data import from_dict, load_arviz_data  # pylint: disable=wrong-import-position
+from ...labels import MapLabeller  # pylint: disable=wrong-import-position
 from ...plots import (  # pylint: disable=wrong-import-position
     plot_autocorr,
     plot_bpv,
@@ -773,7 +774,6 @@ def test_plot_mcse_no_divergences(models):
         {"divergences": True, "var_names": ["theta", "mu"]},
         {"kind": "kde", "var_names": ["theta"]},
         {"kind": "hexbin", "var_names": ["theta"]},
-        {"kind": "hexbin", "var_names": ["theta"]},
         {
             "kind": "hexbin",
             "var_names": ["theta"],
@@ -785,6 +785,21 @@ def test_plot_mcse_no_divergences(models):
             "reference_values": {"mu": 0, "tau": 0},
             "reference_values_kwargs": {"line_color": "blue"},
         },
+        {
+            "var_names": ["mu", "tau"],
+            "reference_values": {"mu": 0, "tau": 0},
+            "labeller": MapLabeller({"mu": r"$\mu$", "theta": r"$\theta"}),
+        },
+        {
+            "var_names": ["theta"],
+            "reference_values": {"theta": [0.0] * 8},
+            "labeller": MapLabeller({"theta": r"$\theta$"}),
+        },
+        {
+            "var_names": ["theta"],
+            "reference_values": {"theta": np.zeros(8)},
+            "labeller": MapLabeller({"theta": r"$\theta$"}),
+        },
     ],
 )
 def test_plot_pair(models, kwargs):
@@ -1201,7 +1216,7 @@ def test_plot_dot_rotated(continuous_model, kwargs):
         },
     ],
 )
-def test_plot_lm(models, kwargs):
+def test_plot_lm_1d(models, kwargs):
     """Test functionality for 1D data."""
     idata = models.model_1
     if "constant_data" not in idata.groups():
@@ -1228,3 +1243,46 @@ def test_plot_lm_list():
     """Test the plots when input data is list or ndarray."""
     y = [1, 2, 3, 4, 5]
     assert plot_lm(y=y, x=np.arange(len(y)), show=False, backend="bokeh")
+
+
+def generate_lm_1d_data():
+    rng = np.random.default_rng()
+    return from_dict(
+        observed_data={"y": rng.normal(size=7)},
+        posterior_predictive={"y": rng.normal(size=(4, 1000, 7)) / 2},
+        posterior={"y_model": rng.normal(size=(4, 1000, 7))},
+        dims={"y": ["dim1"]},
+        coords={"dim1": range(7)},
+    )
+
+
+def generate_lm_2d_data():
+    rng = np.random.default_rng()
+    return from_dict(
+        observed_data={"y": rng.normal(size=(5, 7))},
+        posterior_predictive={"y": rng.normal(size=(4, 1000, 5, 7)) / 2},
+        posterior={"y_model": rng.normal(size=(4, 1000, 5, 7))},
+        dims={"y": ["dim1", "dim2"]},
+        coords={"dim1": range(5), "dim2": range(7)},
+    )
+
+
+@pytest.mark.parametrize("data", ("1d", "2d"))
+@pytest.mark.parametrize("kind", ("lines", "hdi"))
+@pytest.mark.parametrize("use_y_model", (True, False))
+def test_plot_lm(data, kind, use_y_model):
+    if data == "1d":
+        idata = generate_lm_1d_data()
+    else:
+        idata = generate_lm_2d_data()
+
+    kwargs = {"idata": idata, "y": "y", "kind_model": kind, "backend": "bokeh", "show": False}
+    if data == "2d":
+        kwargs["plot_dim"] = "dim1"
+    if use_y_model:
+        kwargs["y_model"] = "y_model"
+    if kind == "lines":
+        kwargs["num_samples"] = 50
+
+    ax = plot_lm(**kwargs)
+    assert ax is not None