[ENH] pandas.Multiindex support for distributions (#580)

Closes #212.

Introdues `MultiIndex` support for distributions, including through
`iloc` and `loc` calls.

Behaves like `pandas`, except that:

* `loc` calls never remove levels
* the dimension after `loc` and `iloc` call always remains 2D, since
there are no 1D distributions

Author: fkiraly

skpro/distributions/base/_base.py

@@ -236,15 +236,52 @@ def _loc(self, rowidx=None, colidx=None):
             colidx = pd.Index([colidx])
 
         if rowidx is not None:
-            row_iloc = self.index.get_indexer_for(rowidx)
+            row_iloc = self._get_indexer_like_pandas(self.index, rowidx)
         else:
             row_iloc = None
         if colidx is not None:
-            col_iloc = self.columns.get_indexer_for(colidx)
+            col_iloc = self._get_indexer_like_pandas(self.columns, colidx)
         else:
             col_iloc = None
         return self._iloc(rowidx=row_iloc, colidx=col_iloc)
 
+    def _get_indexer_like_pandas(self, index, keys):
+        """Return indexer for keys in index.
+
+        A unified helper that mimics pandas' get_indexer_for but supports:
+
+        - scalar key (e.g., "a", ("a", 1))
+        - tuple key (partial or full)
+        - list of keys (partial or full)
+        - works for both Index and MultiIndex
+
+        Returns
+        -------
+        np.ndarray of positions (integers)
+        """
+        # regular index, not multiindex
+        if not isinstance(index, pd.MultiIndex):
+            return index.get_indexer_for(keys)
+
+        # if isinstance(index, pd.MultiIndex):
+
+        if is_scalar_notnone(keys) or isinstance(keys, tuple):
+            keys = [keys]
+
+        # Use get_locs for each key (full or partial)
+        ilocs = []
+        for key in keys:
+            if isinstance(key, slice):
+                ilocs.append(index.slice_indexer(key.start, key.stop, key.step))
+            else:
+                if not isinstance(key, tuple):
+                    key = [key]
+                iloc = index.get_locs(key)
+                if isinstance(iloc, slice):
+                    iloc = np.arange(len(index))[iloc]
+                ilocs.append(iloc)
+        return np.concatenate(ilocs) if ilocs else np.array([], dtype=int)
+
     def _at(self, rowidx=None, colidx=None):
         if rowidx is not None:
             row_iloc = self.index.get_indexer_for([rowidx])[0]
@@ -1762,6 +1799,13 @@ def is_noneslice(obj):
         ref = self.ref
         indexer = getattr(ref, self.method)
 
+        # handle special case of multiindex in loc with single tuple key
+        if isinstance(key, tuple) and not any(isinstance(k, tuple) for k in key):
+            if isinstance(ref.index, pd.MultiIndex) and self.method == "_loc":
+                if type(ref).__name__ != "Empirical":
+                    return indexer(rowidx=key, colidx=None)
+
+        # general case
         if isinstance(key, tuple):
             if not len(key) == 2:
                 raise ValueError(

skpro/distributions/base/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for skpro distribution base class."""

skpro/distributions/base/tests/test_multiindex.py

@@ -0,0 +1,65 @@
+"""Test cases for the MultiIndex functionality of the BaseDistribution.
+
+Uses the Normal distribution, but is intended to trigger the base layer.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from skpro.distributions.normal import Normal
+
+
+@pytest.fixture
+def normal_dist():
+    ix = pd.MultiIndex.from_product([(1, 2), (2, 3)])
+    return Normal(np.array([[1, 2], [2, 3], [4, 5], [6, 7]]), 2, index=ix)
+
+
+def test_loc_partial_level(normal_dist):
+    result = normal_dist.loc[1]
+    expected_index = pd.MultiIndex.from_tuples([(1, 2), (1, 3)])
+    np.testing.assert_array_equal(result.index, expected_index)
+    assert result.mean().shape == (2, 2)
+
+
+def test_loc_full_tuple(normal_dist):
+    result = normal_dist.loc[(2, 2)]
+    expected_index = pd.MultiIndex.from_tuples([(2, 2)])
+    np.testing.assert_array_equal(result.index, expected_index)
+    assert result.mean().shape == (1, 2)
+
+
+def test_loc_list_of_keys(normal_dist):
+    result = normal_dist.loc[[(1, 2), (2, 3)]]
+    expected_index = pd.MultiIndex.from_tuples([(1, 2), (2, 3)])
+    np.testing.assert_array_equal(result.index, expected_index)
+    assert result.mean().shape == (2, 2)
+
+
+def test_iloc_single_row(normal_dist):
+    result = normal_dist.iloc[0]
+    expected_index = pd.MultiIndex.from_tuples([(1, 2)])
+    np.testing.assert_array_equal(result.index, expected_index)
+    assert result.mean().shape == (1, 2)
+
+
+def test_iloc_multiple_rows(normal_dist):
+    result = normal_dist.iloc[[0, 3]]
+    expected_index = pd.MultiIndex.from_tuples([(1, 2), (2, 3)])
+    np.testing.assert_array_equal(result.index, expected_index)
+    assert result.mean().shape == (2, 2)
+
+
+def test_iloc_column_slice(normal_dist):
+    result = normal_dist.iloc[:, 1]
+    expected_index = normal_dist.index
+    assert result.mean().shape == (4, 1)
+    np.testing.assert_array_equal(result.index, expected_index)
+
+
+def test_loc_row_col(normal_dist):
+    result = normal_dist.loc[(1, 2), :]
+    expected_index = pd.MultiIndex.from_tuples([(1, 2)])
+    assert result.mean().shape == (1, 2)
+    np.testing.assert_array_equal(result.index, expected_index)

skpro/distributions/empirical.py

@@ -270,6 +270,24 @@ def _iloc(self, rowidx=None, colidx=None):
             columns=subs_colidx,
         )
 
+    def _loc(self, rowidx=None, colidx=None):
+        if is_scalar_notnone(rowidx) and is_scalar_notnone(colidx):
+            return self._at(rowidx, colidx)
+        if is_scalar_notnone(rowidx):
+            rowidx = pd.Index([rowidx])
+        if is_scalar_notnone(colidx):
+            colidx = pd.Index([colidx])
+
+        if rowidx is not None:
+            row_iloc = self.index.get_indexer_for(rowidx)
+        else:
+            row_iloc = None
+        if colidx is not None:
+            col_iloc = self.columns.get_indexer_for(colidx)
+        else:
+            col_iloc = None
+        return self._iloc(rowidx=row_iloc, colidx=col_iloc)
+
     def _iat(self, rowidx=None, colidx=None):
         if rowidx is None or colidx is None:
             raise ValueError("iat method requires both row and column index")