From 1363fbade5d11395f707ed60c288336bd5612fb3 Mon Sep 17 00:00:00 2001 From: Mudassir Uddin <145117847+Mudassiruddin7@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:17:50 +0530 Subject: [PATCH] [BUG] fix `iloc` and `iat` for parameters broadcast over rows or columns `BaseDistribution._subset_param` indexed every axis of a parameter by position. A parameter of shape `(n_rows, 1)` or `(1, n_cols)` is broadcast by the constructor, so its length 1 axis has no entry at the selected positions and `iloc` / `iat` raised `IndexError`. Repeat a length 1 axis once per selected row or column instead. Fixes #1202 --- skpro/distributions/base/_base.py | 23 ++++ .../tests/test_subset_broadcast_params.py | 102 ++++++++++++++++++ 2 files changed, 125 insertions(+) create mode 100644 skpro/distributions/tests/test_subset_broadcast_params.py diff --git a/skpro/distributions/base/_base.py b/skpro/distributions/base/_base.py index 4cca5842e..a2594bf97 100644 --- a/skpro/distributions/base/_base.py +++ b/skpro/distributions/base/_base.py @@ -408,11 +408,19 @@ def _subset_param(self, val, rowidx, colidx, coerce_scalar=False): arr = np.array(val) # if len(arr.shape) == 0: # do nothing with arr + # an axis of length 1 is broadcast over the rows or columns of self, + # so it is repeated by the subsetting, rather than indexed by position if len(arr.shape) == 2 and rowidx is not None: + if arr.shape[0] == 1: + rowidx = _broadcast_index(rowidx, self.shape[0]) arr = arr[rowidx, :] if len(arr.shape) == 1 and colidx is not None: + if arr.shape[0] == 1: + colidx = _broadcast_index(colidx, self.shape[1]) arr = arr[colidx] if len(arr.shape) >= 2 and colidx is not None: + if arr.shape[1] == 1: + colidx = _broadcast_index(colidx, self.shape[1]) arr = arr[:, colidx] if np.issubdtype(arr.dtype, np.integer): arr = arr.astype("float") @@ -2233,6 +2241,21 @@ def is_scalar_notnone(obj): return obj is not None and np.isscalar(obj) +def _broadcast_index(idx, size): + """Index into a length 1 axis, picking as many entries as ``idx`` does. + + Parameters + ---------- + idx : scalar, slice, or array-like + index into an axis of length ``size`` + size : int + length of the axis of the distribution that ``idx`` indexes + """ + if np.isscalar(idx): + return 0 + return np.zeros(len(np.arange(size)[idx]), dtype=int) + + def _get_first_argname(fun): """Get the name of the first argument of a function as str.""" from inspect import signature diff --git a/skpro/distributions/tests/test_subset_broadcast_params.py b/skpro/distributions/tests/test_subset_broadcast_params.py new file mode 100644 index 000000000..39a0470a6 --- /dev/null +++ b/skpro/distributions/tests/test_subset_broadcast_params.py @@ -0,0 +1,102 @@ +"""Subsetting distributions whose parameters are broadcast over rows or columns. + +A parameter of shape ``(n_rows, 1)`` or ``(1, n_cols)`` is accepted by the +constructor and broadcast against the other parameters, so ``iloc`` and ``iat`` +have to repeat it for the selected rows and columns, not index it by position. +""" + +import numpy as np +import pandas as pd +import pytest + +from skpro.distributions import Normal, Poisson + +INDEX = pd.Index(["r0", "r1"]) +COLUMNS = pd.Index(["a", "b", "c"]) +MU = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]) + +DISTRIBUTIONS = { + "Normal, sigma per row": lambda: Normal( + mu=MU, sigma=np.array([[1.0], [2.0]]), index=INDEX, columns=COLUMNS + ), + "Normal, sigma per column": lambda: Normal( + mu=MU, sigma=np.array([[1.0, 2.0, 3.0]]), index=INDEX, columns=COLUMNS + ), + "Normal, mu per row and sigma per column": lambda: Normal( + mu=np.array([[0.0], [3.0]]), + sigma=np.array([[1.0, 2.0, 3.0]]), + index=INDEX, + columns=COLUMNS, + ), + "Normal, mu in one row": lambda: Normal( + mu=np.array([[0.0, 1.0, 2.0]]), + sigma=np.array([[1.0], [2.0]]), + index=INDEX, + columns=COLUMNS, + ), + "Poisson, mu per row": lambda: Poisson( + mu=np.array([[1.0], [2.0]]), index=INDEX, columns=COLUMNS + ), +} + +# selection on the distribution, and the equivalent row and column positions +ALL = [0, 1, 2] +SELECTIONS = { + "row": (lambda d: d.iloc[[1]], [1], ALL), + "column": (lambda d: d.iloc[:, [2]], [0, 1], [2]), + "rows and columns": (lambda d: d.iloc[[1, 0], [2, 0]], [1, 0], [2, 0]), + "cell by list iloc": (lambda d: d.iloc[[1], [2]], [1], [2]), + "cell by iat": (lambda d: d.iat[1, 2], [1], [2]), + "cell by scalar iloc": (lambda d: d.iloc[1, 2], [1], [2]), + "scalar row": (lambda d: d.iloc[0], [0], ALL), + "scalar column": (lambda d: d.iloc[:, 1], [0, 1], [1]), +} + + +@pytest.mark.parametrize("make", DISTRIBUTIONS.values(), ids=DISTRIBUTIONS.keys()) +@pytest.mark.parametrize( + "select, rows, cols", SELECTIONS.values(), ids=SELECTIONS.keys() +) +def test_subsetting_keeps_broadcast_parameters_of_the_selected_cells( + make, select, rows, cols +): + distr = make() + # a different point per cell, so that using the wrong cell's parameters shows up + x = pd.DataFrame(np.arange(6.0).reshape(2, 3), index=INDEX, columns=COLUMNS) + expected = distr.cdf(x).iloc[rows, cols] + x_selected = x.iloc[rows, cols] + + sub = select(distr) + + if sub.ndim == 0: + got = sub.cdf(float(x_selected.iloc[0, 0])) + else: + assert sub.shape == x_selected.shape + got = sub.cdf(x_selected) + + np.testing.assert_allclose( + np.asarray(got, dtype=float).reshape(-1), + expected.to_numpy(dtype=float).reshape(-1), + ) + + +@pytest.mark.parametrize("make", DISTRIBUTIONS.values(), ids=DISTRIBUTIONS.keys()) +def test_subsetting_to_no_rows_or_no_columns_keeps_the_shape(make): + distr = make() + + assert distr.iloc[[]].shape == (0, 3) + assert distr.iloc[:, []].shape == (2, 0) + + +@pytest.mark.parametrize("make", DISTRIBUTIONS.values(), ids=DISTRIBUTIONS.keys()) +def test_repeating_a_row_repeats_its_broadcast_parameters(make): + distr = make() + + sub = distr.iloc[[1, 1, 0]] + + assert sub.shape == (3, 3) + for stat in ("mean", "var"): + np.testing.assert_allclose( + getattr(sub, stat)().to_numpy(), + getattr(distr, stat)().iloc[[1, 1, 0]].to_numpy(), + )