From 5dfada907ac85d3fa78ea3323a905747618df25d Mon Sep 17 00:00:00 2001 From: Alejandro Schuler Date: Mon, 26 Oct 2026 22:17:39 -0700 Subject: [PATCH 1/3] Store params in NormalFixedVar and NormalFixedMean so they can be sliced Distn.__getitem__ and Distn.__len__ read self._params. NormalFixedVar and NormalFixedMean override __init__ without calling the base constructor, so _params stayed None. On the output of pred_dist, dist[0:2] raised "TypeError: 'NoneType' object is not subscriptable" and len(dist) raised an AttributeError. Both constructors now set self._params, as LogNormal and Exponential do. A parametrized test builds every distribution in ngboost.distns.__all__ and checks len, slicing and integer indexing. Only the two fixed Normal classes failed on master. Integer indexing is marked xfail for Bernoulli, k_categorical and MultivariateNormal, which fail for a different reason (their constructors need two-dimensional params). Found while assessing #168. Co-Authored-By: Claude Sonnet 5.5 --- RELEASE_NOTES.md | 1 + ngboost/distns/normal.py | 2 + tests/test_distns.py | 124 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 126 insertions(+), 1 deletion(-) diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index a7ae2fcf..321d2de6 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -8,6 +8,7 @@ * Replace the deprecated `friedman_mse` criterion with `squared_error` for the default tree learner and distribution tests (issue #408, PR #409) * Support a distinct base learner for each distribution parameter, including scikit-learn nested parameters and feature importances for all-tree configurations (issue #338, PR #411) * Fix the `CRPScore` metric of `Normal`, `NormalFixedMean` and `LogNormal`, which was the correct metric multiplied by the scale. It was too large where the predicted scale is above 1 and too small where it is below 1. Fits that use `CRPScore` with these distributions change wherever the predicted scale is not 1 (issue #421) +* Fix `len(dist)` and slicing such as `dist[0:2]` for `NormalFixedVar` and `NormalFixedMean` predictions, which raised an error because these classes did not store their parameters. New tests check `len` and slicing for every distribution in `ngboost.distns` ## Version 0.5.10 diff --git a/ngboost/distns/normal.py b/ngboost/distns/normal.py index ab9e549f..11080365 100644 --- a/ngboost/distns/normal.py +++ b/ngboost/distns/normal.py @@ -143,6 +143,7 @@ class NormalFixedVar(Normal): # pylint: disable=super-init-not-called def __init__(self, params): + self._params = params self.loc = params[0] self.var = np.ones_like(self.loc) self.scale = np.ones_like(self.loc) @@ -205,6 +206,7 @@ class NormalFixedMean(Normal): # pylint: disable=super-init-not-called def __init__(self, params): + self._params = params self.loc = np.zeros_like(params[0]) self.scale = np.exp(params[0]) self.var = self.scale**2 diff --git a/tests/test_distns.py b/tests/test_distns.py index 3f7fd9c3..2b0b3fa7 100644 --- a/tests/test_distns.py +++ b/tests/test_distns.py @@ -6,10 +6,11 @@ from sklearn.model_selection import train_test_split from sklearn.tree import DecisionTreeRegressor -from ngboost import NGBClassifier, NGBRegressor, NGBSurvival +from ngboost import NGBClassifier, NGBRegressor, NGBSurvival, distns from ngboost.distns import ( Bernoulli, Cauchy, + ClassificationDistn, Distn, Exponential, Gamma, @@ -19,6 +20,7 @@ Normal, NormalFixedMean, NormalFixedVar, + RegressionDistn, T, TFixedDf, TFixedDfFixedVar, @@ -224,3 +226,123 @@ def test_multivariatenormal(k: int, learner): mean = y_dist.mean sample = y_dist.rv() scipy_list = y_dist.scipy_distribution() + + +# Slicing and len. Distn.__getitem__ and Distn.__len__ read self._params, so +# every distribution has to store the params it was built from. + +_N_ROWS = 6 +# Distn.__getitem__ passes a one-dimensional params array to the constructor +# when the key is an integer. Categorical and MultivariateNormal expect a +# two-dimensional array, so dist[0] fails for them. This is separate from +# the missing _params that the NormalFixed classes had. +_NO_INT_INDEX = {"Bernoulli", "k_categorical", "MultivariateNormal"} + + +def _concrete_distns(int_index=False): + """The concrete distributions in ngboost.distns.__all__. + + k_categorical and MultivariateNormal are factories, so they are called + with a small size. Distn, RegressionDistn and ClassificationDistn are + abstract bases and are left out. + """ + bases = (Distn, RegressionDistn, ClassificationDistn) + factories = { + "k_categorical": lambda: k_categorical(3), + "MultivariateNormal": lambda: MultivariateNormal(2), + } + cases = [] + for name in distns.__all__: + obj = getattr(distns, name) + if name in factories: + obj = factories[name]() + if obj in bases: + continue + marks = [] + if int_index and name in _NO_INT_INDEX: + marks.append(pytest.mark.xfail(reason="integer key needs 2D params")) + cases.append(pytest.param(obj, id=name, marks=marks)) + return cases + + +def _build(Dist): + rng = np.random.default_rng(0) + return Dist(rng.normal(scale=0.3, size=(Dist.n_params, _N_ROWS))) + + +def _rows(value, key): + # a few distributions keep a scalar in params, such as the n of BetaBinomial + return value if np.ndim(value) == 0 else value[key] + + +def test_concrete_distns_cover_all_exports(): + names = [case.id for case in _concrete_distns()] + assert len(names) == len(set(names)) + assert {"Normal", "NormalFixedVar", "NormalFixedMean", "Bernoulli"} <= set(names) + + +@pytest.mark.parametrize("Dist", _concrete_distns()) +def test_len(Dist): + assert len(_build(Dist)) == _N_ROWS + + +@pytest.mark.parametrize("Dist", _concrete_distns()) +def test_slice_keeps_class_and_params(Dist): + dist = _build(Dist) + sliced = dist[1:4] + assert sliced.__class__ is Dist + assert len(sliced) == 3 + assert sliced.params.keys() == dist.params.keys() + for name, value in dist.params.items(): + np.testing.assert_allclose(sliced.params[name], _rows(value, slice(1, 4))) + + +@pytest.mark.parametrize("Dist", _concrete_distns(int_index=True)) +def test_integer_index_keeps_class_and_params(Dist): + dist = _build(Dist) + single = dist[2] + assert single.__class__ is Dist + for name, value in dist.params.items(): + np.testing.assert_allclose(single.params[name], _rows(value, 2)) + + +@pytest.mark.parametrize( + "Dist, expected_loc, expected_scale", + [ + (NormalFixedVar, lambda p: p[0], lambda p: np.ones(len(p[0]))), + (NormalFixedMean, lambda p: np.zeros(len(p[0])), lambda p: np.exp(p[0])), + ], + ids=["NormalFixedVar", "NormalFixedMean"], +) +def test_normal_fixed_slice_matches_rows(Dist, expected_loc, expected_scale): + params = np.random.default_rng(1).normal(size=(1, 8)) + dist = Dist(params) + rows = slice(2, 7) + sliced = dist[rows] + assert sliced.__class__ is Dist + assert len(dist) == 8 + assert len(sliced) == 5 + np.testing.assert_allclose(sliced.loc, expected_loc(params)[rows]) + np.testing.assert_allclose(sliced.scale, expected_scale(params)[rows]) + np.testing.assert_allclose(sliced.params["loc"], sliced.loc) + np.testing.assert_allclose(sliced.params["scale"], sliced.scale) + y = np.linspace(-1, 1, 8) + np.testing.assert_allclose( + sliced.dist.logpdf(y[rows]), dist.dist.logpdf(y)[rows], rtol=1e-12 + ) + + +@pytest.mark.parametrize("Dist", [NormalFixedVar, NormalFixedMean]) +def test_normal_fixed_fitted_prediction_supports_len_and_slice(Dist): + rng = np.random.default_rng(2) + X = rng.normal(size=(60, 3)) + Y = X[:, 0] + rng.normal(size=60) + ngb = NGBRegressor(Dist=Dist, n_estimators=5, verbose=False).fit(X, Y) + pred = ngb.pred_dist(X) + assert len(pred) == 60 + sliced = pred[0:2] + assert sliced.__class__ is Dist + assert len(sliced) == 2 + np.testing.assert_allclose(sliced.loc, pred.loc[0:2]) + np.testing.assert_allclose(sliced.scale, pred.scale[0:2]) + np.testing.assert_allclose(pred[0].loc, pred.loc[0]) From 86d7ab5d1e799ace5be349f5eb8d9093fd6a16a9 Mon Sep 17 00:00:00 2001 From: Alejandro Schuler Date: Mon, 26 Oct 2026 22:55:33 -0700 Subject: [PATCH 2/3] Support integer indexing of Categorical and MultivariateNormal predictions Distn.__getitem__ passes params[:, i] to the constructor for an integer key, so the constructor gets a one-dimensional array. The constructors of the k_categorical classes (Bernoulli included) and of MultivariateNormal read params.shape[1], so dist[0] raised an error. Slices worked. The constructors now accept params without the row axis, and the methods that used fixed axis numbers use the last axes instead. dist[i] is then the distribution of row i without the row axis, as for the other distributions and as for a numpy array: class_probs() has shape (K,), predict() gives one class, MultivariateNormal has loc (p,) and cov (p, p), logpdf(y) gives a scalar, and scipy_distribution() gives one scipy object. dist[i:i+1] still gives a batch of one row. ClassificationDistn.predict takes the argmax over the last axis, which is the same as axis 1 for a batch. For batches the outputs are unchanged. The integer-index test no longer marks these three classes as xfail. New tests check class_probs, predict, sample, logpdf, cov, scipy_distribution and iteration on dist[i], and dist[0] of fitted NGBClassifier and NGBRegressor models against predict_proba and predict. Co-Authored-By: Claude Opus 5.5 --- RELEASE_NOTES.md | 1 + ngboost/distns/categorical.py | 8 ++-- ngboost/distns/distn.py | 2 +- ngboost/distns/multivariate_normal.py | 36 ++++++++------ tests/test_distns.py | 69 ++++++++++++++++++++++----- 5 files changed, 85 insertions(+), 31 deletions(-) diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 321d2de6..7244cd66 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -9,6 +9,7 @@ * Support a distinct base learner for each distribution parameter, including scikit-learn nested parameters and feature importances for all-tree configurations (issue #338, PR #411) * Fix the `CRPScore` metric of `Normal`, `NormalFixedMean` and `LogNormal`, which was the correct metric multiplied by the scale. It was too large where the predicted scale is above 1 and too small where it is below 1. Fits that use `CRPScore` with these distributions change wherever the predicted scale is not 1 (issue #421) * Fix `len(dist)` and slicing such as `dist[0:2]` for `NormalFixedVar` and `NormalFixedMean` predictions, which raised an error because these classes did not store their parameters. New tests check `len` and slicing for every distribution in `ngboost.distns` +* Fix integer indexing such as `dist[0]` for `Bernoulli`, `k_categorical` and `MultivariateNormal` predictions, which raised an error because their constructors needed two-dimensional params. `dist[i]` is now the distribution of row i without the row axis, as for the other distributions ## Version 0.5.10 diff --git a/ngboost/distns/categorical.py b/ngboost/distns/categorical.py index 6bce2c93..e9f4a6ed 100644 --- a/ngboost/distns/categorical.py +++ b/ngboost/distns/categorical.py @@ -58,9 +58,9 @@ class Categorical(ClassificationDistn): def __init__(self, params): super().__init__(params) - _, N = params.shape - self.logits = np.zeros((K, N)) - self.logits[1:K, :] = params # default the 0th class logits to 0 + # params is (K - 1) x N, or has only the first axis for one row (dist[0]) + self.logits = np.zeros((K,) + params.shape[1:]) + self.logits[1:K] = params # default the 0th class logits to 0 self.probs = sp.special.softmax(self.logits, axis=0) # self.dist = dist(n=1, p=self.probs) # scipy doesn't allow vectorized multinomial (!?!?) @@ -76,7 +76,7 @@ def fit(Y): def sample1(self): # this is just a helper for sample() cum_p = np.cumsum(self.probs, axis=0)[0:-1] - interval = cum_p < np.random.random((1, len(self))) + interval = cum_p < np.random.random(self.probs.shape[1:]) return np.sum(interval, axis=0) def sample(self, m): diff --git a/ngboost/distns/distn.py b/ngboost/distns/distn.py index 184232f1..7d6abb76 100644 --- a/ngboost/distns/distn.py +++ b/ngboost/distns/distn.py @@ -74,4 +74,4 @@ def predict(self): # predictions for regression are typically conditional means class ClassificationDistn(Distn): def predict(self): # returns class assignments - return np.argmax(self.class_probs(), 1) + return np.argmax(self.class_probs(), -1) diff --git a/ngboost/distns/multivariate_normal.py b/ngboost/distns/multivariate_normal.py index 47fbe650..246a6c98 100644 --- a/ngboost/distns/multivariate_normal.py +++ b/ngboost/distns/multivariate_normal.py @@ -128,15 +128,17 @@ def get_chol_factor(lower_tri_vals): Args: lower_tri_vals: numpy array, shaped as the number of lower triangular - elements, number of observations. + elements, number of observations. The second axis + is absent for one observation. The values ordered according to np.tril_indices(p) where p is the dimension of the multivariate normal distn Returns: Nxpxp numpy array, with the lower triangle filled in. The diagonal is exponentiated. + It is pxp for one observation. """ - lower_size, N = lower_tri_vals.shape + lower_size = lower_tri_vals.shape[0] # solve p(p+3)/2 = lower_size to get the # number of dimensions. @@ -147,13 +149,13 @@ def get_chol_factor(lower_tri_vals): if not isinstance(lower_tri_vals, np.ndarray): lower_tri_vals = np.array(lower_tri_vals) - L = np.zeros((N, p, p)) + L = np.zeros(lower_tri_vals.shape[1:] + (p, p)) for par_ind, (k, l) in enumerate(zip(*np.tril_indices(p))): if k == l: # Add a small number to avoid singular matrices. - L[:, k, l] = np.exp(lower_tri_vals[par_ind, :]) + 1e-6 + L[..., k, l] = np.exp(lower_tri_vals[par_ind]) + 1e-6 else: - L[:, k, l] = lower_tri_vals[par_ind, :] + L[..., k, l] = lower_tri_vals[par_ind] return L @@ -192,7 +194,8 @@ class MVN(RegressionDistn): # pylint: disable=super-init-not-called def __init__(self, params): super().__init__(params) - self.N = params.shape[1] + # params has only the first axis for one row, as in dist[0] + self.N = params.shape[1] if params.ndim > 1 else 1 # Number of MVN dimensions, =k originally self.p = (-3 + (9 + 8 * self.n_params) ** 0.5) / 2 @@ -204,14 +207,14 @@ def __init__(self, params): # rest are the lower triangle of L. # Where Sigma_inverse = L@L.transpose() # Diagonals modelled on the log scale. - self.loc = np.transpose(np.array(params[: self.p, :])) - self.tril_L = np.array(params[self.p :, :]) + self.loc = np.transpose(np.array(params[: self.p])) + self.tril_L = np.array(params[self.p :]) # Returns 3d array, shape (p, p, N) self.L = get_chol_factor(self.tril_L) # The remainder is just for utility. - self.cov_inv = self.L @ self.L.transpose(0, 2, 1) + self.cov_inv = self.L @ np.swapaxes(self.L, -1, -2) # _cov_mat and _Lcov are place holders, relatively expensive to compute # The inverse of self.cov_inv. Access through self.cov @@ -230,20 +233,20 @@ def summaries(self, Y): diff: N x p x1 the residual between the mean and the data eta: N x p which is L@diff """ - diff = np.expand_dims(self.loc - Y, 2) + diff = (self.loc - Y)[..., None] # N x 2 x 2 @ N x p x 1 # -> N x p x 1 we remove the last index - eta = np.squeeze(np.matmul(self.L.transpose(0, 2, 1), diff), axis=2) + eta = np.squeeze(np.matmul(np.swapaxes(self.L, -1, -2), diff), axis=-1) return diff, eta def logpdf(self, Y): _, eta = self.summaries(Y) # the exponentiated part of the pdf: - p1 = -0.5 * np.sum(np.square(eta), axis=1) + p1 = -0.5 * np.sum(np.square(eta), axis=-1) p1 = p1.squeeze() # this is the sqrt(determinant(Sigma)) component of the pdf - p2 = np.sum(np.log(np.diagonal(self.L, axis1=1, axis2=2)), axis=1) + p2 = np.sum(np.log(np.diagonal(self.L, axis1=-2, axis2=-1)), axis=-1) ret = p1 + p2 + self.pdf_constant return ret @@ -263,8 +266,8 @@ def rv(self): if self._Lcov is None: self._Lcov = np.linalg.cholesky(np.linalg.inv(self.cov_inv)) - u = np.random.normal(loc=0, scale=1, size=(self.N, self.p, 1)) - sample = np.expand_dims(self.loc, 2) + self._Lcov @ u + u = np.random.normal(loc=0, scale=1, size=self.loc.shape + (1,)) + sample = self.loc[..., None] + self._Lcov @ u return np.squeeze(sample) def rvs(self, n): @@ -290,8 +293,11 @@ def scipy_distribution(self): """ Returns: List of scipy.stats.multivariate_normal distributions. + One distribution, not a list, for one row such as dist[0]. """ cov_mat = self.cov + if self.loc.ndim == 1: + return sp.stats.multivariate_normal(mean=self.loc, cov=cov_mat) return [ sp.stats.multivariate_normal(mean=self.loc[i, :], cov=cov_mat[i, :, :]) for i in range(self.N) diff --git a/tests/test_distns.py b/tests/test_distns.py index 2b0b3fa7..49d8c852 100644 --- a/tests/test_distns.py +++ b/tests/test_distns.py @@ -232,14 +232,9 @@ def test_multivariatenormal(k: int, learner): # every distribution has to store the params it was built from. _N_ROWS = 6 -# Distn.__getitem__ passes a one-dimensional params array to the constructor -# when the key is an integer. Categorical and MultivariateNormal expect a -# two-dimensional array, so dist[0] fails for them. This is separate from -# the missing _params that the NormalFixed classes had. -_NO_INT_INDEX = {"Bernoulli", "k_categorical", "MultivariateNormal"} -def _concrete_distns(int_index=False): +def _concrete_distns(): """The concrete distributions in ngboost.distns.__all__. k_categorical and MultivariateNormal are factories, so they are called @@ -258,10 +253,7 @@ def _concrete_distns(int_index=False): obj = factories[name]() if obj in bases: continue - marks = [] - if int_index and name in _NO_INT_INDEX: - marks.append(pytest.mark.xfail(reason="integer key needs 2D params")) - cases.append(pytest.param(obj, id=name, marks=marks)) + cases.append(pytest.param(obj, id=name)) return cases @@ -297,7 +289,7 @@ def test_slice_keeps_class_and_params(Dist): np.testing.assert_allclose(sliced.params[name], _rows(value, slice(1, 4))) -@pytest.mark.parametrize("Dist", _concrete_distns(int_index=True)) +@pytest.mark.parametrize("Dist", _concrete_distns()) def test_integer_index_keeps_class_and_params(Dist): dist = _build(Dist) single = dist[2] @@ -346,3 +338,58 @@ def test_normal_fixed_fitted_prediction_supports_len_and_slice(Dist): np.testing.assert_allclose(sliced.loc, pred.loc[0:2]) np.testing.assert_allclose(sliced.scale, pred.scale[0:2]) np.testing.assert_allclose(pred[0].loc, pred.loc[0]) + + +# Distn.__getitem__ passes one-dimensional params to the constructor for an +# integer key, so dist[i] has no row axis, as with numpy. Categorical and +# MultivariateNormal used to need two-dimensional params. + + +@pytest.mark.parametrize( + "Dist", [Bernoulli, k_categorical(3)], ids=["Bernoulli", "k_categorical"] +) +def test_categorical_integer_index_gives_one_row(Dist): + dist = _build(Dist) + for i in [2, -1]: + single = dist[i] + np.testing.assert_allclose(single.class_probs(), dist.class_probs()[i]) + assert single.predict() == dist.predict()[i] + samples = single.sample(50) + assert samples.shape == (50,) + assert set(samples) <= set(range(Dist.K_)) + np.testing.assert_allclose([d.class_probs() for d in dist], dist.class_probs()) + + +@pytest.mark.parametrize("k", [2, 3]) +def test_multivariatenormal_integer_index_gives_one_row(k): + dist = _build(MultivariateNormal(k)) + Y = np.random.default_rng(3).normal(size=(_N_ROWS, k)) + for i in [2, -1]: + single = dist[i] + np.testing.assert_allclose(single.predict(), dist.predict()[i]) + np.testing.assert_allclose(single.cov, dist.cov[i]) + np.testing.assert_allclose(single.logpdf(Y[i]), dist.logpdf(Y)[i]) + np.testing.assert_allclose( + single.scipy_distribution().logpdf(Y[i]), dist.logpdf(Y)[i] + ) + assert np.shape(single.sample(20)) == (20, k) + # one distribution evaluated at several points + np.testing.assert_allclose(single.logpdf(Y), [single.logpdf(y) for y in Y]) + + +def test_integer_index_of_fitted_classifier_and_mvn(): + rng = np.random.default_rng(4) + X = rng.normal(size=(60, 3)) + y = (X[:, 0] > 0).astype(int) + (X[:, 1] > 0).astype(int) + clf = NGBClassifier(Dist=k_categorical(3), n_estimators=5, verbose=False) + clf.fit(X, y) + pred = clf.pred_dist(X) + np.testing.assert_allclose(pred[0].class_probs(), clf.predict_proba(X)[0]) + assert pred[0].predict() == clf.predict(X)[0] + + Y = np.column_stack([X[:, 0], X[:, 1]]) + rng.normal(size=(60, 2)) + reg = NGBRegressor(Dist=MultivariateNormal(2), n_estimators=5, verbose=False) + reg.fit(X, Y) + pred = reg.pred_dist(X) + np.testing.assert_allclose(pred[0].predict(), reg.predict(X)[0]) + np.testing.assert_allclose(pred[0].logpdf(Y[0]), pred.logpdf(Y)[0]) From 48b7183dcc4a36ed9d197d74dbdfc0dc8b56aa97 Mon Sep 17 00:00:00 2001 From: Alejandro Schuler Date: Tue, 10 Nov 2026 10:00:33 -0800 Subject: [PATCH 3/3] Note the one-row shapes in the MultivariateNormal params docstring #438 documented params for a batch of N rows. dist[i] gives loc of shape p and cov of shape p x p. Co-Authored-By: Claude Opus 5.5 --- ngboost/distns/multivariate_normal.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ngboost/distns/multivariate_normal.py b/ngboost/distns/multivariate_normal.py index 61b60213..3fd23ecc 100644 --- a/ngboost/distns/multivariate_normal.py +++ b/ngboost/distns/multivariate_normal.py @@ -323,6 +323,7 @@ def params(self): cov: N x p x p array, the covariance matrices Row i gives scipy.stats.multivariate_normal(mean=loc[i], cov=cov[i]). scipy_distribution() returns these N scipy objects. + For one row, as in dist[0], loc is p and cov is p x p. params['scale'] is a deprecated alias of params['cov'] and gives a FutureWarning. It is the covariance, not a standard deviation.