test_proba_basic.py

"""Non-suite tests for probability distribution objects."""

# copyright: skpro developers, BSD-3-Clause License (see LICENSE file)
# adapted from sktime

__author__ = ["fkiraly"]

import numpy as np
import pandas as pd
import pytest
from skbase.utils.dependencies import _check_soft_dependencies

from skpro.tests.test_switch import run_test_module_changed


@pytest.mark.skip(reason="Undiagnosed failure. Skipping until resolved. See #918.")
def test_proba_example():
    """Test one subsetting case for BaseDistribution."""
    from skpro.distributions.normal import Normal

    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1)

    assert n.shape == (3, 2)

    one_row = n.loc[[1]]
    assert isinstance(one_row, Normal)
    assert one_row.shape == (1, 2)


@pytest.mark.parametrize("subsetter", ["loc", "iloc"])
def test_proba_subsetters_loc_iloc(subsetter):
    """Test one subsetting case for BaseDistribution."""
    from skpro.distributions.normal import Normal

    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1)

    assert n.shape == (3, 2)

    # should result in 2D array distribution (1, 1)
    nss = getattr(n, subsetter)[1, [1]]
    assert isinstance(nss, Normal)
    assert nss.shape == (1, 1)
    assert nss.mu.shape == (1, 1)

    nss = getattr(n, subsetter)[[1], 1]
    assert isinstance(nss, Normal)
    assert nss.shape == (1, 1)
    assert nss.mu.shape == (1, 1)

    # should result in scalar distribution
    nss = getattr(n, subsetter)[1, 1]
    assert isinstance(nss, Normal)
    assert nss.shape == ()

    nss = getattr(n, subsetter)[1, 1]
    assert isinstance(nss, Normal)
    assert nss.shape == ()


def test_proba_subsetters_at_iat():
    """Test one subsetting case for BaseDistribution."""
    from skpro.distributions.normal import Normal

    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1)

    # should result in scalar distribution
    nss = n.iat[1, 1]
    assert isinstance(nss, Normal)
    assert nss.shape == ()
    assert nss == n.iloc[1, 1]

    nss = n.at[1, 1]
    assert isinstance(nss, Normal)
    assert nss.shape == ()
    assert nss == n.loc[1, 1]


def test_proba_index_coercion():
    """Test index coercion for BaseDistribution."""
    from skpro.distributions.normal import Normal

    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1, columns=["foo", "bar"])

    assert n.shape == (3, 2)
    assert isinstance(n.index, pd.Index)
    assert isinstance(n.columns, pd.Index)
    assert n.index.equals(pd.RangeIndex(3))
    assert n.columns.equals(pd.Index(["foo", "bar"]))

    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1, index=["2", 1, 0])

    assert n.shape == (3, 2)
    assert isinstance(n.index, pd.Index)
    assert isinstance(n.columns, pd.Index)
    assert n.index.equals(pd.Index(["2", 1, 0]))
    assert n.columns.equals(pd.RangeIndex(2))

    # this should coerce to a 2D array of shape (1, 3)
    n = Normal(0, 1, columns=[1, 2, 3])

    assert n.shape == (1, 3)
    assert isinstance(n.index, pd.Index)
    assert isinstance(n.columns, pd.Index)
    assert n.index.equals(pd.RangeIndex(1))
    assert n.columns.equals(pd.Index([1, 2, 3]))


@pytest.mark.skipif(
    not _check_soft_dependencies("matplotlib", severity="none"),
    reason="skip if matplotlib is not available",
)
@pytest.mark.parametrize("fun", ["pdf", "ppf", "cdf"])
def test_proba_plotting(fun):
    """Test that plotting functions do not crash and return ax as expected."""
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure

    from skpro.distributions.normal import Normal

    # default case, 2D distribution with n_columns>1
    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1)
    fig, ax = n.plot(fun=fun)
    assert isinstance(fig, Figure)
    assert isinstance(ax, np.ndarray)
    assert ax.shape == n.shape
    assert all([isinstance(a, Axes) for a in ax.flatten()])
    assert all([a.get_figure() == fig for a in ax.flatten()])

    # 1D case requires special treatment of axes
    n = Normal(mu=[[1], [2], [3]], sigma=1)
    fig, ax = n.plot(fun=fun)
    assert isinstance(fig, Figure)
    assert isinstance(ax, type(ax))
    assert ax.shape == (n.shape[0],)
    assert all([isinstance(a, Axes) for a in ax.flatten()])
    assert all([a.get_figure() == fig for a in ax.flatten()])

    # scalar case
    n = Normal(mu=1, sigma=1)
    ax = n.plot(fun=fun)
    assert isinstance(ax, Axes)


@pytest.mark.skipif(
    not _check_soft_dependencies("matplotlib", severity="none"),
    reason="skip if matplotlib is not available",
)
def test_discrete_pmf_plotting():
    """Test that discrete PMF plotting uses stem plots."""
    from matplotlib.axes import Axes

    from skpro.distributions.binomial import Binomial

    # Test Binomial PMF plotting
    n = Binomial(n=10, p=0.5)
    ax = n.plot(fun="pmf")
    assert isinstance(ax, Axes)

    # Check that stem plot was used (should have containers)
    assert len(ax.containers) > 0, "Stem plot should be used for discrete PMF"

    # For small distributions, check that all support points are plotted.
    # Binomial(n=10) has support [0, 1, 2, ..., 10] = 11 points.
    # A StemContainer always has exactly 3 children (markerline, stemlines,
    # baseline) regardless of the number of data points. The correct way to
    # count plotted points is via stemlines.get_segments().
    stem_container = ax.containers[0]
    assert hasattr(
        stem_container, "stemlines"
    ), "Expected StemContainer from stem plot, check _plot_single uses ax.stem"
    n_plotted = len(stem_container.stemlines.get_segments())
    assert n_plotted > 5, (
        f"Should plot at multiple support points, got {n_plotted}. "
        "Binomial(n=10) has 11 support points (0..10)."
    )


def test_to_df_parametric():
    """Tests coercion to DataFrame via get_params_df and to_df."""
    from skpro.distributions.normal import Normal

    cols = ["foo", "bar"]

    # default case, 2D distribution with n_columns>1
    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1, columns=cols)

    param_names = n.get_param_names()
    params_df = n.get_params_df()
    for k, v in params_df.items():
        assert k in param_names
        assert isinstance(v, pd.DataFrame)
        assert (v.index == n.index).all()
        assert (v.columns == n.columns).all()

    all_params_df = n.to_df()
    assert isinstance(all_params_df, pd.DataFrame)
    assert (all_params_df.index == n.index).all()
    assert isinstance(all_params_df.columns, pd.MultiIndex)

    level0_vals = all_params_df.columns.get_level_values(0).unique()
    level1_vals = all_params_df.columns.get_level_values(1).unique()

    assert (level0_vals == n.columns).all()
    for ix in level1_vals:
        assert ix in param_names
        assert ix not in ["index", "columns"]

    # scalar case
    n = Normal(mu=2, sigma=3)

    param_names = n.get_param_names()
    params_df = n.get_params_df()

    for k, v in params_df.items():
        assert k in param_names
        assert isinstance(v, pd.DataFrame)
        assert (v.index == pd.RangeIndex(1)).all()
        assert (v.columns == pd.RangeIndex(1)).all()

    all_params_df = n.to_df()
    assert isinstance(all_params_df, pd.DataFrame)
    assert (all_params_df.index == pd.RangeIndex(1)).all()
    assert not isinstance(all_params_df.columns, pd.MultiIndex)

    for ix in all_params_df.columns:
        assert ix in param_names
        assert ix not in ["index", "columns"]


def test_head_tail():
    """Test head and tail utility functions."""
    from skpro.distributions.normal import Normal

    cols = ["foo", "bar"]

    # default case, 2D distribution with n_columns>1
    n = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1, columns=cols)

    nh = n.head(2)
    assert isinstance(nh, Normal)
    assert nh.shape == (2, 2)
    assert (nh.columns == n.columns).all()
    assert (nh.index == n.index[:2]).all()

    nh2 = n.head()
    assert isinstance(nh2, Normal)
    assert nh2.shape == (3, 2)
    assert (nh2.columns == n.columns).all()
    assert (nh2.index == n.index).all()

    nt = n.tail(2)
    assert isinstance(nt, Normal)
    assert nt.shape == (2, 2)
    assert (nt.columns == n.columns).all()
    assert (nt.index == n.index[1:]).all()

    nt2 = n.tail()
    assert isinstance(nt2, Normal)
    assert nt2.shape == (3, 2)
    assert (nt2.columns == n.columns).all()
    assert (nt2.index == n.index).all()

    # scalar case
    n = Normal(mu=2, sigma=3)

    nh = n.head()
    assert nh.ndim == 0

    nt = n.tail(42)
    assert nt.ndim == 0


@pytest.mark.skipif(
    not run_test_module_changed("skpro.distributions"),
    reason="run only if skpro.distributions has been changed",
)
def test_multiindex_loc_indexing():
    """Test that loc indexing works with MultiIndex and Index objects."""
    from skpro.distributions.normal import Normal

    index = pd.MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1)])
    columns = pd.Index(["x", "y"])

    dist = Normal(mu=[[0, 1], [2, 3], [4, 5]], sigma=1, index=index, columns=columns)

    x = pd.DataFrame([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], index=index, columns=columns)

    subset = dist.loc[x.index, x.columns]

    assert subset.shape == dist.shape
    assert subset.index.equals(dist.index)
    assert subset.columns.equals(dist.columns)

    result = dist.quantile(0.5)
    assert result.shape == dist.shape


@pytest.mark.skipif(
    not run_test_module_changed("skpro.distributions"),
    reason="run only if skpro.distributions has been changed",
)
def test_pmf_support_method():
    """Test the _pmf_support method for different distribution types."""
    from skpro.distributions.binomial import Binomial
    from skpro.distributions.delta import Delta
    from skpro.distributions.empirical import Empirical
    from skpro.distributions.normal import Normal

    # Test continuous distribution (Normal) - should return empty array
    normal = Normal(mu=0, sigma=1)
    support = normal._pmf_support(-1, 1)
    assert isinstance(support, np.ndarray)
    assert len(support) == 0

    # Test discrete distribution with default integer support (Binomial)
    binomial = Binomial(n=5, p=0.5)
    support = binomial._pmf_support(0, 5)
    assert isinstance(support, np.ndarray)
    assert len(support) > 0
    assert all(isinstance(x, (int, np.integer)) for x in support)
    assert all(x >= 0 for x in support)

    # Test Empirical distribution
    spl = pd.Series([1.5, 2.5, 3.5])
    empirical = Empirical(spl)
    support = empirical._pmf_support(1, 4)
    assert isinstance(support, np.ndarray)
    assert len(support) > 0
    assert 1.5 in support or 2.5 in support or 3.5 in support

    # Test Delta distribution
    delta = Delta(c=2.0)
    support = delta._pmf_support(1, 3)
    assert isinstance(support, np.ndarray)
    assert len(support) == 1
    assert support[0] == 2.0

    # Test Delta outside bounds
    support = delta._pmf_support(3, 4)
    assert isinstance(support, np.ndarray)
    assert len(support) == 0