Name the MultivariateNormal covariance cov in params

RELEASE_NOTES.md

@@ -16,6 +16,7 @@
 * Report fitted status to scikit-learn's `check_is_fitted` via `__sklearn_is_fitted__`, fixing `NotFittedError` in `Pipeline.predict` on sklearn >= 1.8 (PR #429)
 * Set scikit-learn classifier metadata on `NGBClassifier` so that scikit-learn correctly identifies it as a classifier (PR #413)
 * Handle array-like sample weights positionally and validate them at fit time; weights with wrong length, shape, NaN, or inf now raise `ValueError` (PR #425)
+* `MultivariateNormal` `params` now gives the covariance matrices under the key `cov`, the name that `scipy.stats.multivariate_normal` uses. The old key `scale` is deprecated: `params['scale']` and `params.get('scale')` still return the covariance, but give a `FutureWarning`. Everywhere else in ngboost `scale` is a standard deviation, so users read the covariance as one. `scale` is no longer in `params.keys()` (issue #168, PR #438)
 * Document saving and loading models with `skops` and test the round trip (PR #427)
 
 ### Bug fixes

ngboost/distns/multivariate_normal.py

@@ -23,6 +23,35 @@ def get_tril_idxs(p):
     return tril_indices, diags, off_diags
 
 
+_SCALE_WARNING = (
+    "params['scale'] of MultivariateNormal is deprecated and will be removed in a "
+    "future release. It is the covariance matrix, not a standard deviation. "
+    "Use params['cov'] instead."
+)
+
+
+class MVNParams(dict):
+    """
+    The dict that MultivariateNormal.params returns, with the keys loc and cov.
+
+    params['scale'] and params.get('scale') still return the covariance, with a
+    FutureWarning. scale is not a key of the dict, so it is not in keys(),
+    items() or the repr, and "scale" in params is False.
+    """
+
+    def __missing__(self, key):
+        if key == "scale":
+            warnings.warn(_SCALE_WARNING, FutureWarning, stacklevel=2)
+            return self["cov"]
+        raise KeyError(key)
+
+    def get(self, key, default=None):
+        if key == "scale" and key not in self:
+            warnings.warn(_SCALE_WARNING, FutureWarning, stacklevel=2)
+            return self["cov"]
+        return super().get(key, default)
+
+
 class MVNLogScore(LogScore):
     def score(self, Y):
         return -self.logpdf(Y)
@@ -284,7 +313,18 @@ def cov(self):
 
         @property
         def params(self):
-            return {"loc": self.loc, "scale": self.cov}
+            """
+            Returns:
+                A dict with two keys:
+                loc: N x p array, the means
+                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.
+
+                params['scale'] is a deprecated alias of params['cov'] and gives
+                a FutureWarning. It is the covariance, not a standard deviation.
+            """
+            return MVNParams(loc=self.loc, cov=self.cov)
 
         def scipy_distribution(self):
             """

tests/test_distns.py

@@ -1,7 +1,9 @@
+import warnings
 from typing import Tuple
 
 import numpy as np
 import pytest
+from scipy.stats import multivariate_normal
 from sklearn.datasets import fetch_california_housing, load_breast_cancer
 from sklearn.model_selection import train_test_split
 from sklearn.tree import DecisionTreeRegressor
@@ -224,3 +226,61 @@ def test_multivariatenormal(k: int, learner):
     mean = y_dist.mean
     sample = y_dist.rv()
     scipy_list = y_dist.scipy_distribution()
+
+
+def _mvn_with_cov(cov, N=4):
+    """N MultivariateNormal rows with means 0, ..., p - 1 and covariance cov"""
+    p = cov.shape[0]
+    # the internal parameters are the means and the lower triangle of L, where
+    # inv(cov) = L @ L.T and the diagonal of L is on the log scale
+    L = np.linalg.cholesky(np.linalg.inv(cov))
+    L[np.diag_indices(p)] = np.log(L[np.diag_indices(p)])
+    theta = np.concatenate([np.arange(p, dtype=float), L[np.tril_indices(p)]])
+    return MultivariateNormal(p)(np.tile(theta[:, None], (1, N)))
+
+
+def test_multivariatenormal_params_cov():
+    cov = np.array([[2.0, 0.6, 0.0], [0.6, 1.0, -0.3], [0.0, -0.3, 0.5]])
+    dist = _mvn_with_cov(cov)
+    params = dist.params
+
+    assert list(params) == ["loc", "cov"]
+    np.testing.assert_allclose(params["loc"], np.tile([0.0, 1.0, 2.0], (4, 1)))
+    # get_chol_factor adds 1e-6 to the diagonal of L, so the match is not exact
+    np.testing.assert_allclose(
+        params["cov"], np.broadcast_to(cov, (4, 3, 3)), rtol=1e-4, atol=1e-5
+    )
+    np.testing.assert_array_equal(params["cov"], dist.cov)
+
+    # the names fit scipy.stats.multivariate_normal row by row
+    Y = np.random.default_rng(0).normal(size=(4, 3))
+    scipy_logpdf = [
+        multivariate_normal(mean=params["loc"][i], cov=params["cov"][i]).logpdf(Y[i])
+        for i in range(4)
+    ]
+    np.testing.assert_allclose(scipy_logpdf, dist.logpdf(Y))
+
+
+def test_multivariatenormal_params_scale_is_deprecated():
+    dist = _mvn_with_cov(np.array([[1.0, 0.3], [0.3, 2.0]]))
+    params = dist.params
+
+    with pytest.warns(FutureWarning, match=r"Use params\['cov'\]") as record:
+        scale = params["scale"]
+    assert record[0].filename == __file__
+    np.testing.assert_array_equal(scale, params["cov"])
+
+    with pytest.warns(FutureWarning, match=r"Use params\['cov'\]") as record:
+        scale = params.get("scale")
+    assert record[0].filename == __file__
+    np.testing.assert_array_equal(scale, params["cov"])
+
+    # the other keys work as in a plain dict and give no warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        assert dict(params.items()).keys() == {"loc", "cov"}
+        np.testing.assert_array_equal(params["loc"], dist.loc)
+        np.testing.assert_array_equal(params.get("cov"), dist.cov)
+        assert params.get("df") is None
+        with pytest.raises(KeyError):
+            params["df"]  # pylint: disable=pointless-statement

tests/test_pickling.py

@@ -59,6 +59,30 @@ def test_model_save(learners_data):
         assert (new_preds == preds).all()
 
 
+def test_mvn_params_after_pickle():
+    """
+    A pickled MultivariateNormal model gives the same params, with the cov key,
+    and the params dict itself keeps the deprecated scale alias through pickle
+    """
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(100, 3))
+    Y = X[:, :2] + rng.normal(size=(100, 2))
+    ngb = NGBRegressor(Dist=MultivariateNormal(2), n_estimators=10, verbose=False)
+    ngb.fit(X, Y)
+    params = ngb.pred_dist(X).params
+
+    model = pickle.loads(pickle.dumps(ngb))
+    new_params = model.pred_dist(X).params
+    assert list(new_params) == ["loc", "cov"]
+    np.testing.assert_array_equal(new_params["loc"], params["loc"])
+    np.testing.assert_array_equal(new_params["cov"], params["cov"])
+
+    params_copy = pickle.loads(pickle.dumps(params))
+    assert list(params_copy) == ["loc", "cov"]
+    with pytest.warns(FutureWarning):
+        np.testing.assert_array_equal(params_copy["scale"], params["cov"])
+
+
 def test_classifier_setstate_restores_missing_label_metadata(breast_cancer_data):
     """Older classifier pickles do not include the sklearn label metadata."""