Support integer indexing of Categorical and MultivariateNormal predictions

RELEASE_NOTES.md

@@ -32,6 +32,8 @@
 * Replace the deprecated `friedman_mse` criterion with `squared_error` for the default tree learner (issue #408, PR #409)
 * Stop the overflow warnings from trial steps that the line search rejects, such as `overflow encountered in exp` in `Normal`. They came most often from `natural_gradient=False` on a target with large units. Fitted models do not change (issue #352, PR #434)
 * `fit` now raises a `ValueError` that says the fit diverged when the gradient has non-finite values, in place of the base learner's `Input y contains NaN` error (PR #434)
+* 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` (PR #441)
+* 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 (PR #442)
 
 ### Documentation and CI
 

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):

ngboost/distns/distn.py

@@ -84,4 +84,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)

ngboost/distns/multivariate_normal.py

@@ -157,15 +157,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.
@@ -176,13 +178,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
 
 
@@ -221,7 +223,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
@@ -233,14 +236,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
@@ -259,20 +262,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
@@ -292,8 +295,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):
@@ -320,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.
@@ -330,8 +334,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)

ngboost/distns/normal.py

@@ -136,6 +136,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)
@@ -198,6 +199,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

tests/test_distns.py

@@ -8,10 +8,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,
@@ -21,6 +22,7 @@
     Normal,
     NormalFixedMean,
     NormalFixedVar,
+    RegressionDistn,
     T,
     TFixedDf,
     TFixedDfFixedVar,
@@ -233,6 +235,173 @@ def test_multivariatenormal(k: int, learner):
     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
+
+
+def _concrete_distns():
+    """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
+        cases.append(pytest.param(obj, id=name))
+    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())
+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])
+
+
+# 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])
+
+
 def _mvn_with_cov(cov, N=4):
     """N MultivariateNormal rows with means 0, ..., p - 1 and covariance cov"""
     p = cov.shape[0]