ngboost 0.5.10.dev0__tar.gz → 0.6.0__tar.gz

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/PKG-INFO

@@ -1,25 +1,25 @@
 Metadata-Version: 2.3
 Name: ngboost
-Version: 0.5.10.dev0
+Version: 0.6.0
 Summary: Library for probabilistic predictions via gradient boosting.
 License: Apache-2.0
 Author: Stanford ML Group
 Author-email: avati@cs.stanford.edu
-Requires-Python: >=3.9,<3.15
+Requires-Python: >=3.10,<3.15
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: joblib (>=1.2)
 Requires-Dist: lifelines (>=0.25)
-Requires-Dist: matplotlib (>=3.0,<3.10) ; python_version < "3.10"
-Requires-Dist: matplotlib (>=3.10.0) ; python_version >= "3.10"
+Requires-Dist: matplotlib (>=3.10.0)
 Requires-Dist: numpy (>=1.21.2) ; python_version < "3.13"
 Requires-Dist: numpy (>=2.1.0) ; python_version >= "3.13"
-Requires-Dist: scikit-learn (>=1.6,<2.0)
+Requires-Dist: scikit-learn (>=1.6,<1.7) ; python_version < "3.14"
+Requires-Dist: scikit-learn (>=1.8.0,<2.0) ; python_version >= "3.14"
 Requires-Dist: scipy (>=1.14.1) ; python_version >= "3.13"
 Requires-Dist: scipy (>=1.7.2) ; python_version < "3.13"
 Requires-Dist: sympy (>=1.12)

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/api.py

@@ -1,6 +1,9 @@
 "The NGBoost library API"
+
 # pylint: disable=too-many-arguments
-from sklearn.base import BaseEstimator
+import numpy as np
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.preprocessing import LabelEncoder
 from sklearn.utils import check_array
 
 from ngboost.distns import (
@@ -31,8 +34,15 @@ class NGBRegressor(NGBoost, BaseEstimator):
                             A distribution from ngboost.distns, e.g. Normal
         Score             : rule to compare probabilistic predictions P̂ to the observed data y.
                             A score from ngboost.scores, e.g. LogScore
-        Base              : base learner to use in the boosting algorithm.
-                            Any instantiated sklearn regressor, e.g. DecisionTreeRegressor()
+        Base              : base learner(s) to use in the boosting algorithm.
+                            Pass a single instantiated sklearn regressor, e.g.
+                            DecisionTreeRegressor(), to use the same learner for
+                            every distribution parameter. To use different learners
+                            per distribution parameter, pass a list/tuple of
+                            instantiated sklearn regressors with length equal to
+                            Dist.n_params. The sequence order matches the distribution
+                            parameter order; for example, Normal uses
+                            [loc_learner, scale_learner].
         natural_gradient  : logical flag indicating whether the natural gradient should be used
         n_estimators      : the number of boosting iterations to fit
         learning_rate     : the learning rate
@@ -49,12 +59,13 @@ class NGBRegressor(NGBoost, BaseEstimator):
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
 
     Output:
         An NGBRegressor object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=Normal,
@@ -71,6 +82,7 @@ class NGBRegressor(NGBoost, BaseEstimator):
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
         assert issubclass(
             Dist, RegressionDistn
@@ -96,6 +108,7 @@ class NGBRegressor(NGBoost, BaseEstimator):
             random_state,
             validation_fraction,
             early_stopping_rounds,
+            n_jobs,
         )
 
         self._estimator_type = "regressor"
@@ -114,7 +127,7 @@ class NGBRegressor(NGBoost, BaseEstimator):
         super().__setstate__(state_dict)
 
 
-class NGBClassifier(NGBoost, BaseEstimator):
+class NGBClassifier(ClassifierMixin, NGBoost, BaseEstimator):
     """
     Constructor for NGBoost classification models.
 
@@ -127,8 +140,15 @@ class NGBClassifier(NGBoost, BaseEstimator):
                             A distribution from ngboost.distns, e.g. Bernoulli
         Score             : rule to compare probabilistic predictions P̂ to the observed data y.
                             A score from ngboost.scores, e.g. LogScore
-        Base              : base learner to use in the boosting algorithm.
-                            Any instantiated sklearn regressor, e.g. DecisionTreeRegressor()
+        Base              : base learner(s) to use in the boosting algorithm.
+                            Pass a single instantiated sklearn regressor, e.g.
+                            DecisionTreeRegressor(), to use the same learner for
+                            every distribution parameter. To use different learners
+                            per distribution parameter, pass a list/tuple of
+                            instantiated sklearn regressors with length equal to
+                            Dist.n_params. The sequence order matches the distribution
+                            parameter order; for example, Normal uses
+                            [loc_learner, scale_learner].
         natural_gradient  : logical flag indicating whether the natural gradient should be used
         n_estimators      : the number of boosting iterations to fit
         learning_rate     : the learning rate
@@ -139,11 +159,18 @@ class NGBClassifier(NGBoost, BaseEstimator):
         tol               : numerical tolerance to be used in optimization
         random_state      : seed for reproducibility. See
                             https://stackoverflow.com/questions/28064634/random-state-pseudo-random-number-in-scikit-learn
+        validation_fraction: Proportion of training data to set
+                             aside as validation data for early stopping.
+        early_stopping_rounds:      The number of consecutive boosting iterations during which the
+                                    loss has to increase before the algorithm stops early.
+                                    Set to None to disable early stopping and validation.
+                                    None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
     Output:
         An NGBClassifier object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=Bernoulli,
@@ -158,6 +185,9 @@ class NGBClassifier(NGBoost, BaseEstimator):
         verbose_eval=100,
         tol=1e-4,
         random_state=None,
+        validation_fraction=0.1,
+        early_stopping_rounds=None,
+        n_jobs=None,
     ):
         assert issubclass(
             Dist, ClassificationDistn
@@ -175,9 +205,108 @@ class NGBClassifier(NGBoost, BaseEstimator):
             verbose_eval,
             tol,
             random_state,
+            validation_fraction,
+            early_stopping_rounds,
+            n_jobs,
         )
         self._estimator_type = "classifier"
 
+    def _fit_label_encoder(self, Y):
+        le = LabelEncoder().fit(Y)
+        self._le = le  # pylint: disable=attribute-defined-outside-init
+        self.classes_ = le.classes_  # pylint: disable=attribute-defined-outside-init
+        n_classes = self.Dist.n_params + 1
+        if len(self.classes_) != n_classes:
+            raise ValueError(
+                "NGBClassifier Dist expects "
+                f"{n_classes} classes, got {len(self.classes_)}."
+            )
+
+    def __setstate__(
+        self, state_dict
+    ):  # pylint: disable=attribute-defined-outside-init
+        super().__setstate__(state_dict)
+        if not hasattr(self, "classes_"):
+            self.classes_ = np.arange(self.Dist.n_params + 1)
+        if not hasattr(self, "_le"):
+            self._le = LabelEncoder()
+            self._le.classes_ = self.classes_
+
+    # fit and partial_fit pass their arguments on to NGBoost unchanged, so
+    # pylint reads the calls as copies of the one in NGBoost.fit.
+    # pylint: disable=too-many-positional-arguments,duplicate-code
+    def fit(
+        self,
+        X,
+        Y,
+        X_val=None,
+        Y_val=None,
+        sample_weight=None,
+        val_sample_weight=None,
+        train_loss_monitor=None,
+        val_loss_monitor=None,
+        early_stopping_rounds=None,
+    ):
+        # Learn the classes on every fit. NGBoost.fit then calls partial_fit,
+        # which encodes Y and Y_val with them.
+        self._fit_label_encoder(Y)
+        super().fit(
+            X,
+            Y,
+            X_val=X_val,
+            Y_val=Y_val,
+            sample_weight=sample_weight,
+            val_sample_weight=val_sample_weight,
+            train_loss_monitor=train_loss_monitor,
+            val_loss_monitor=val_loss_monitor,
+            early_stopping_rounds=early_stopping_rounds,
+        )
+        return self
+
+    # pylint: disable=too-many-positional-arguments
+    def partial_fit(
+        self,
+        X,
+        Y,
+        X_val=None,
+        Y_val=None,
+        sample_weight=None,
+        val_sample_weight=None,
+        train_loss_monitor=None,
+        val_loss_monitor=None,
+        early_stopping_rounds=None,
+    ):
+        if not hasattr(self, "classes_"):
+            self._fit_label_encoder(Y)
+        Y = self._le.transform(Y)
+        if Y_val is not None:
+            Y_val = self._le.transform(Y_val)
+        super().partial_fit(
+            X,
+            Y,
+            X_val=X_val,
+            Y_val=Y_val,
+            sample_weight=sample_weight,
+            val_sample_weight=val_sample_weight,
+            train_loss_monitor=train_loss_monitor,
+            val_loss_monitor=val_loss_monitor,
+            early_stopping_rounds=early_stopping_rounds,
+        )
+        return self
+
+    def predict(self, X, max_iter=None):
+        return self.classes_[super().predict(X, max_iter=max_iter)]
+
+    def staged_predict(self, X, max_iter=None):
+        return [
+            self.classes_[pred] for pred in super().staged_predict(X, max_iter=max_iter)
+        ]
+
+    def score(self, X, Y):  # for sklearn
+        # Call NGBoost.score directly: super() would reach ClassifierMixin.score,
+        # which returns accuracy.
+        return NGBoost.score(self, X, self._le.transform(Y))
+
     def predict_proba(self, X, max_iter=None):
         """
         Probability prediction of Y at the points X=x
@@ -221,8 +350,15 @@ class NGBSurvival(NGBoost, BaseEstimator):
                             A distribution from ngboost.distns, e.g. LogNormal
         Score             : rule to compare probabilistic predictions P̂ to the observed data y.
                             A score from ngboost.scores, e.g. LogScore
-        Base              : base learner to use in the boosting algorithm.
-                            Any instantiated sklearn regressor, e.g. DecisionTreeRegressor()
+        Base              : base learner(s) to use in the boosting algorithm.
+                            Pass a single instantiated sklearn regressor, e.g.
+                            DecisionTreeRegressor(), to use the same learner for
+                            every distribution parameter. To use different learners
+                            per distribution parameter, pass a list/tuple of
+                            instantiated sklearn regressors with length equal to
+                            Dist.n_params. The sequence order matches the distribution
+                            parameter order; for example, Normal uses
+                            [loc_learner, scale_learner].
         natural_gradient  : logical flag indicating whether the natural gradient should be used
         n_estimators      : the number of boosting iterations to fit
         learning_rate     : the learning rate
@@ -233,11 +369,18 @@ class NGBSurvival(NGBoost, BaseEstimator):
         tol               : numerical tolerance to be used in optimization
         random_state      : seed for reproducibility. See
                             https://stackoverflow.com/questions/28064634/random-state-pseudo-random-number-in-scikit-learn
+        validation_fraction: Proportion of training data to set
+                             aside as validation data for early stopping.
+        early_stopping_rounds:      The number of consecutive boosting iterations during which the
+                                    loss has to increase before the algorithm stops early.
+                                    Set to None to disable early stopping and validation.
+                                    None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
     Output:
         An NGBSurvival object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=LogNormal,
@@ -252,6 +395,9 @@ class NGBSurvival(NGBoost, BaseEstimator):
         verbose_eval=100,
         tol=1e-4,
         random_state=None,
+        validation_fraction=0.1,
+        early_stopping_rounds=None,
+        n_jobs=None,
     ):
 
         assert issubclass(
@@ -278,6 +424,9 @@ class NGBSurvival(NGBoost, BaseEstimator):
             verbose_eval,
             tol,
             random_state,
+            validation_fraction,
+            early_stopping_rounds,
+            n_jobs,
         )
 
     def __getstate__(self):

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/distns/exponential.py

@@ -32,16 +32,17 @@ class ExponentialCRPScore(CRPScore):
     def score(self, Y):
         E, T = Y["Event"], Y["Time"]
         score = T + self.scale * (2 * np.exp(-T / self.scale) - 1.5)
-        score[E == 1] -= (
-            0.5 * self.scale[E == 1] * np.exp(-2 * T[E == 1] / self.scale[E == 1])
+        # a censored observation only scores the CDF up to the censoring time
+        score[E == 0] -= (
+            0.5 * self.scale[E == 0] * np.exp(-2 * T[E == 0] / self.scale[E == 0])
         )
         return score
 
     def d_score(self, Y):
         E, T = Y["Event"], Y["Time"]
         deriv = 2 * np.exp(-T / self.scale) * (self.scale + T) - 1.5 * self.scale
-        deriv[E == 1] -= np.exp(-2 * T[E == 1] / self.scale[E == 1]) * (
-            0.5 * self.scale[E == 1] - T[E == 1]
+        deriv[E == 0] -= np.exp(-2 * T[E == 0] / self.scale[E == 0]) * (
+            0.5 * self.scale[E == 0] + T[E == 0]
         )
         return deriv.reshape((-1, 1))
 

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/distns/laplace.py

@@ -37,8 +37,10 @@ class LaplaceCRPScore(CRPScore):
         D[:, 0] = np.sign(self.loc - Y) * (
             1 - np.exp(-np.abs(Y - self.loc) / self.scale)
         )
-        D[:, 1] = np.exp(-np.abs(Y - self.loc) / self.scale) * (
-            self.scale + np.abs(Y - self.loc)
+        D[:, 1] = (
+            np.exp(-np.abs(Y - self.loc) / self.scale)
+            * (self.scale + np.abs(Y - self.loc))
+            - 0.75 * self.scale
         )
         return D
 

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/distns/lognormal.py

@@ -30,11 +30,7 @@ class LogNormalLogScoreCensored(LogScore):
         D_cens[:, 0] = -sp.stats.norm.pdf(lT, loc=self.loc, scale=self.scale) / (
             1 - self.dist.cdf(T) + self.eps
         )
-        D_cens[:, 1] = (
-            -Z
-            * sp.stats.norm.pdf(lT, loc=self.loc, scale=self.scale)
-            / (1 - self.dist.cdf(T) + self.eps)
-        )
+        D_cens[:, 1] = -Z * sp.stats.norm.pdf(Z) / (1 - self.dist.cdf(T) + self.eps)
 
         return (1 - E) * D_cens + E * D_uncens
 
@@ -71,12 +67,9 @@ class LogNormalCRPScoreCensored(CRPScore):
         Z = (lT - self.loc) / self.scale
 
         D = np.zeros((self.loc.shape[0], 2))
+        # the censored CRPS is scale * integral of Phi(z)**2 up to Z
         D[:, 0] = E * -(2 * sp.stats.norm.cdf(Z) - 1) + (1 - E) * -(
             sp.stats.norm.cdf(Z) ** 2
-            + 2 * Z * sp.stats.norm.cdf(Z) * sp.stats.norm.pdf(Z)
-            + 2 * sp.stats.norm.pdf(Z) ** 2
-            - 2 * sp.stats.norm.cdf(Z) * sp.stats.norm.pdf(Z) ** 2
-            - np.sqrt(2 / np.pi) * sp.stats.norm.pdf(np.sqrt(2) * Z)
         )
         D[:, 1] = self.score(Y) + (lT - self.loc) * D[:, 0]
         return D
@@ -85,7 +78,7 @@ class LogNormalCRPScoreCensored(CRPScore):
         I = np.zeros((self.loc.shape[0], 2, 2))
         I[:, 0, 0] = 2
         I[:, 1, 1] = self.scale**2
-        I /= 2 * np.sqrt(np.pi)
+        I /= 2 * np.sqrt(np.pi) * self.scale[:, np.newaxis, np.newaxis]
         return I
 
 

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/distns/normal.py

@@ -49,7 +49,8 @@ class NormalCRPScore(CRPScore):
             self.var,
         ]
         I = I.reshape((self.var.shape[0], 2, 2))
-        I = 1 / (2 * np.sqrt(np.pi)) * I
+        # 2 * integral of grad F grad F^T, which scales like 1 / scale
+        I = 1 / (2 * np.sqrt(np.pi) * self.scale[:, np.newaxis, np.newaxis]) * I
         return I
 
 
@@ -185,7 +186,7 @@ class NormalFixedMeanCRPScore(CRPScore):
         return D
 
     def metric(self):
-        I = np.c_[self.var]
+        I = np.c_[self.scale]
         I = I.reshape((self.var.shape[0], 1, 1))
         I = 1 / (2 * np.sqrt(np.pi)) * I
         return I

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/distns/t.py

@@ -103,7 +103,7 @@ class TFixedDfLogScore(LogScore):
     def metric(self):
         FI = np.zeros((self.var.shape[0], 2, 2))
         FI[:, 0, 0] = (self.df + 1) / ((self.df + 3) * self.var)
-        FI[:, 1, 1] = (self.df) / (2 * (self.df + 3) * self.var)
+        FI[:, 1, 1] = 2 * self.df / (self.df + 3)
         return FI
 
 

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/learners.py

@@ -2,7 +2,7 @@ from sklearn.linear_model import Ridge
 from sklearn.tree import DecisionTreeRegressor
 
 default_tree_learner = DecisionTreeRegressor(
-    criterion="friedman_mse",
+    criterion="squared_error",
     min_samples_split=2,
     min_samples_leaf=1,
     min_weight_fraction_leaf=0.0,

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/ngboost/ngboost.py

@@ -5,6 +5,8 @@
 # pylint: disable=unused-variable,invalid-unary-operand-type,attribute-defined-outside-init
 # pylint: disable=redundant-keyword-arg,protected-access,unnecessary-lambda-assignment
 import numpy as np
+from joblib import Parallel, delayed
+from scipy.sparse import issparse
 from sklearn.base import clone
 from sklearn.model_selection import train_test_split
 from sklearn.tree import DecisionTreeRegressor
@@ -16,6 +18,28 @@ from ngboost.manifold import manifold
 from ngboost.scores import LogScore
 
 
+def _validate_sample_weight(sample_weight, X, input_name="sample_weight"):
+    """Return sample_weight as a 1-D float array with one entry per row of X.
+
+    Weights are matched to rows by position, so a pandas index is ignored.
+    """
+    if sample_weight is None:
+        return None
+    sample_weight = check_array(
+        sample_weight, ensure_2d=False, dtype=np.float64, input_name=input_name
+    )
+    if sample_weight.ndim != 1:
+        raise ValueError(
+            f"{input_name} must be 1-D, got an array of shape {sample_weight.shape}"
+        )
+    if sample_weight.shape[0] != X.shape[0]:
+        raise ValueError(
+            f"{input_name} must have one entry per row: got {sample_weight.shape[0]} "
+            f"entries for {X.shape[0]} rows"
+        )
+    return sample_weight
+
+
 class NGBoost:
     """
     Constructor for all NGBoost models.
@@ -29,8 +53,15 @@ class NGBoost:
                             A distribution from ngboost.distns, e.g. Normal
         Score             : rule to compare probabilistic predictions P̂ to the observed data y.
                             A score from ngboost.scores, e.g. LogScore
-        Base              : base learner to use in the boosting algorithm.
-                            Any instantiated sklearn regressor, e.g. DecisionTreeRegressor()
+        Base              : base learner(s) to use in the boosting algorithm.
+                            Pass a single instantiated sklearn regressor, e.g.
+                            DecisionTreeRegressor(), to use the same learner for
+                            every distribution parameter. To use different learners
+                            per distribution parameter, pass a list/tuple of
+                            instantiated sklearn regressors with length equal to
+                            Dist.n_params. The sequence order matches the distribution
+                            parameter order; for example, Normal uses
+                            [loc_learner, scale_learner].
         natural_gradient  : logical flag indicating whether the natural gradient should be used
         n_estimators      : the number of boosting iterations to fit
         learning_rate     : the learning rate
@@ -46,6 +77,18 @@ class NGBoost:
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : number of joblib threads used to fit the per parameter
+                            base learners in parallel on each boosting round. None
+                            or 1 (the default) fits them serially. A value above 1,
+                            or ``-1`` for all cores, speeds up fitting when the
+                            distribution has many parameters, such as a high
+                            dimensional MultivariateNormal or classification with
+                            many classes; on small problems the thread overhead can
+                            make it slower instead. Threads are used so the shared
+                            data is not copied per worker; when nesting inside an
+                            outer parallel job keep the product of the thread counts
+                            near the core count. The fitted model does not depend on
+                            n_jobs.
 
 
     Output:
@@ -69,6 +112,7 @@ class NGBoost:
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
         self.Dist = Dist
         self.Score = Score
@@ -87,10 +131,11 @@ class NGBoost:
         self.scalings = []
         self.col_idxs = []
         self.tol = tol
-        self.random_state = check_random_state(random_state)
+        self.random_state = random_state
         self.best_val_loss_itr = None
         self.validation_fraction = validation_fraction
         self.early_stopping_rounds = early_stopping_rounds
+        self.n_jobs = n_jobs
 
         if hasattr(self.Dist, "multi_output"):
             self.multi_output = self.Dist.multi_output
@@ -120,6 +165,11 @@ class NGBoost:
         state_dict["Manifold"] = manifold(state_dict["Score"], state_dict["Dist"])
         self.__dict__ = state_dict
 
+    def __sklearn_is_fitted__(self):
+        # check_is_fitted (used by Pipeline.predict) otherwise looks for
+        # attributes that end in "_", and a fitted model has none of those
+        return bool(getattr(self, "base_models", None))
+
     def fit_init_params_to_marginal(self, Y, sample_weight=None, iters=1000):
         self.init_params = self.Manifold.fit(
             Y
@@ -143,18 +193,14 @@ class NGBoost:
 
         if self.minibatch_frac != 1.0:
             sample_size = int(self.minibatch_frac * len(Y))
-            idxs = self.random_state.choice(
-                np.arange(len(Y)), sample_size, replace=False
-            )
+            idxs = self._rng.choice(np.arange(len(Y)), sample_size, replace=False)
 
         if self.col_sample != 1.0:
             if self.col_sample > 0.0:
                 col_size = max(1, int(self.col_sample * X.shape[1]))
             else:
                 col_size = 0
-            col_idx = self.random_state.choice(
-                np.arange(X.shape[1]), col_size, replace=False
-            )
+            col_idx = self._rng.choice(np.arange(X.shape[1]), col_size, replace=False)
 
         weight_batch = None if sample_weight is None else sample_weight[idxs]
 
@@ -167,13 +213,58 @@ class NGBoost:
             params[idxs, :],
         )
 
+    def _base_learners(self):
+        if not isinstance(self.Base, (list, tuple)):
+            return [self.Base] * self.Manifold.n_params
+
+        if len(self.Base) != self.Manifold.n_params:
+            raise ValueError(
+                "Base must be a single estimator or a sequence with one estimator "
+                "per distribution parameter "
+                f"({self.Manifold.n_params} for {self.Dist.__name__}); "
+                f"got {len(self.Base)}."
+            )
+        return self.Base
+
     def fit_base(self, X, grads, sample_weight=None):
-        if sample_weight is None:
-            models = [clone(self.Base).fit(X, g) for g in grads.T]
-        else:
+        base_learners = self._base_learners()
+        # Seed only the base learners that have no random_state of their own,
+        # drawing from self._rng in this single thread so serial and
+        # parallel fits agree. Learners that already fix their own randomness are
+        # left untouched, so their RNG stream is not perturbed (this mirrors how
+        # scikit learn's forests seed their sub estimators).
+        seeds = [
+            (
+                int(self._rng.randint(np.iinfo(np.int32).max))
+                if base.get_params().get("random_state", "set") is None
+                else None
+            )
+            for base in base_learners
+        ]
+
+        def _fit_one(base, g, seed):
+            learner = clone(base)
+            if seed is not None:
+                learner.set_params(random_state=seed)
+            if sample_weight is None:
+                return learner.fit(X, g)
+            return learner.fit(X, g, sample_weight=sample_weight)
+
+        if self.n_jobs in (None, 1):
             models = [
-                clone(self.Base).fit(X, g, sample_weight=sample_weight) for g in grads.T
+                _fit_one(base, g, seed)
+                for base, g, seed in zip(base_learners, grads.T, seeds)
             ]
+        else:
+            if issparse(X) and X.format in ("csr", "csc", "bsr"):
+                # Canonicalize once so the threaded fits do not race on an in
+                # place sort_indices of this shared matrix. Only the compressed
+                # formats have sortable indices; sklearn copies the rest.
+                X.sort_indices()
+            models = Parallel(n_jobs=self.n_jobs, backend="threading")(
+                delayed(_fit_one)(base, g, seed)
+                for base, g, seed in zip(base_learners, grads.T, seeds)
+            )
         fitted = np.array([m.predict(X) for m in models]).T
         self.base_models.append(models)
         return fitted
@@ -235,9 +326,12 @@ class NGBoost:
             Y_val                 : DataFrame object or List or
                                     numpy array of validation-set outcomes in numeric format
             sample_weight         : how much to weigh each example in the training set.
-                                    numpy array of size (n) (defaults to None)
+                                    List, pandas Series or numpy array of size (n)
+                                    (defaults to None). Weights are matched to rows by
+                                    position, so the index of a Series is ignored.
             val_sample_weight     : how much to weigh each example in the validation set.
-                                    (defaults to None)
+                                    List, pandas Series or numpy array of size (n_val),
+                                    matched to rows by position (defaults to None)
             train_loss_monitor    : a custom score or set of scores to track on the training set
                                     during training. Defaults to the score defined in the NGBoost
                                     constructor
@@ -254,6 +348,11 @@ class NGBoost:
         self.base_models = []
         self.scalings = []
         self.col_idxs = []
+        # Start a fresh random stream on every fit so that refitting with an
+        # integer random_state gives the same model
+        self._rng = check_random_state(self.random_state)
+
+        self.best_val_loss_itr = None
 
         return self.partial_fit(
             X,
@@ -298,9 +397,12 @@ class NGBoost:
             Y_val                 : DataFrame object or List or
                                     numpy array of validation-set outcomes in numeric format
             sample_weight         : how much to weigh each example in the training set.
-                                    numpy array of size (n) (defaults to None)
+                                    List, pandas Series or numpy array of size (n)
+                                    (defaults to None). Weights are matched to rows by
+                                    position, so the index of a Series is ignored.
             val_sample_weight     : how much to weigh each example in the validation set.
-                                    (defaults to None)
+                                    List, pandas Series or numpy array of size (n_val),
+                                    matched to rows by position (defaults to None)
             train_loss_monitor    : a custom score or set of scores to track on the training set
                                     during training. Defaults to the score defined in the NGBoost
                                     constructor
@@ -320,6 +422,9 @@ class NGBoost:
             raise RuntimeError(
                 "Base models, scalings, and col_idxs are not the same length"
             )
+        if getattr(self, "_rng", None) is None:
+            self._rng = check_random_state(self.random_state)
+        self._base_learners()
 
         # if early stopping is specified, split X,Y and sample weights (if given) into training and validation sets
         # This will overwrite any X_val and Y_val values passed by the user directly.
@@ -337,7 +442,7 @@ class NGBoost:
                         X,
                         Y,
                         test_size=self.validation_fraction,
-                        random_state=self.random_state,
+                        random_state=self._rng,
                     )
 
                 else:
@@ -356,7 +461,7 @@ class NGBoost:
                         Y,
                         sample_weight,
                         test_size=self.validation_fraction,
-                        random_state=self.random_state,
+                        random_state=self._rng,
                     )
             elif X_val is not None and Y_val is not None:
                 if sample_weight is not None and val_sample_weight is None:
@@ -389,6 +494,7 @@ class NGBoost:
             multi_output=self.multi_output,
             y_numeric=True,
         )
+        sample_weight = _validate_sample_weight(sample_weight, X)
 
         self.n_features = X.shape[1]
 
@@ -406,6 +512,9 @@ class NGBoost:
                 multi_output=self.multi_output,
                 y_numeric=True,
             )
+            val_sample_weight = _validate_sample_weight(
+                val_sample_weight, X_val, input_name="val_sample_weight"
+            )
             val_params = self.pred_param(X_val)
             val_loss_list = []
             best_val_loss = np.inf
@@ -491,16 +600,65 @@ class NGBoost:
 
         return self
 
+    def _set_base_params(self, parameters):
+        if not isinstance(self.Base, (list, tuple)):
+            self.Base.set_params(**parameters)
+            return
+
+        base_learners = list(self.Base)
+        for key, value in parameters.items():
+            index, delim, sub_key = key.partition("__")
+            try:
+                index = int(index)
+            except ValueError as exc:
+                raise ValueError(
+                    "Base sequence parameters must use integer positions, "
+                    f"for example Base__0__max_depth. Got Base__{key}."
+                ) from exc
+            if index < 0 or index >= len(base_learners):
+                raise ValueError(
+                    f"Base sequence index {index} is out of range for "
+                    f"{len(base_learners)} estimators."
+                )
+            if delim:
+                base_learners[index].set_params(**{sub_key: value})
+            else:
+                base_learners[index] = value
+        self.Base = type(self.Base)(base_learners)
+
     def set_params(self, **parameters):
-        for parameter, value in parameters.items():
-            setattr(self, parameter, value)
+        if not parameters:
+            return self
+
+        valid_params = self.get_params(deep=True)
+        nested_params = {}
+        for key, value in parameters.items():
+            key, delim, sub_key = key.partition("__")
+            if key not in valid_params:
+                local_valid_params = self.get_params(deep=False).keys()
+                raise ValueError(
+                    f"Invalid parameter {key!r} for estimator {self}. "
+                    f"Valid parameters are: {sorted(local_valid_params)!r}."
+                )
+            if delim:
+                nested_params.setdefault(key, {})[sub_key] = value
+            else:
+                setattr(self, key, value)
+                valid_params[key] = value
+
+        for key, sub_params in nested_params.items():
+            if key == "Base":
+                self._set_base_params(sub_params)
+            else:
+                valid_params[key].set_params(**sub_params)
         return self
 
     def get_params(self, deep=True):
         """
         Parameters
         ----------
-        deep : Ignored. (for compatibility with sklearn)
+        deep : bool, default=True
+            If True, return nested parameters for the base learner or learners.
         Returns
         ----------
         params : returns an dictionary of parameters.
@@ -515,15 +673,45 @@ class NGBoost:
             "minibatch_frac": self.minibatch_frac,
             "col_sample": self.col_sample,
             "verbose": self.verbose,
+            "verbose_eval": self.verbose_eval,
+            "tol": self.tol,
             "random_state": self.random_state,
             "validation_fraction": self.validation_fraction,
             "early_stopping_rounds": self.early_stopping_rounds,
+            "n_jobs": self.n_jobs,
         }
 
+        if not deep:
+            return params
+
+        if isinstance(self.Base, (list, tuple)):
+            for i, base in enumerate(self.Base):
+                params[f"Base__{i}"] = base
+                if hasattr(base, "get_params"):
+                    for key, value in base.get_params(deep=True).items():
+                        params[f"Base__{i}__{key}"] = value
+        elif hasattr(self.Base, "get_params"):
+            for key, value in self.Base.get_params(deep=True).items():
+                params[f"Base__{key}"] = value
+
         return params
 
     def score(self, X, Y):  # for sklearn
-        return self.Manifold(self.pred_dist(X)._params).total_score(Y)
+        """
+        Score the predicted distributions at the points X=x against the outcomes Y.
+        Higher is better, as scikit-learn expects, so this is the negative of the
+        average scoring rule: the mean log-likelihood for LogScore and the negative
+        mean CRPS for CRPScore.
+
+        Parameters:
+            X         : DataFrame object or List or
+                        numpy array of predictors (n x p) in numeric format.
+            Y         : numpy array of outcomes (n)
+
+        Output:
+            A float, the negative of the average score
+        """
+        return -self.Manifold(self.pred_dist(X)._params).total_score(Y)
 
     def pred_dist(self, X, max_iter=None):
         """
@@ -616,8 +804,12 @@ class NGBoost:
         # Check whether the model is fitted
         if not self.base_models:
             return None
-        # Check whether the base model is DecisionTreeRegressor
-        if not isinstance(self.base_models[0][0], DecisionTreeRegressor):
+        # Check whether all base models are DecisionTreeRegressor
+        if any(
+            not isinstance(model, DecisionTreeRegressor)
+            for models in self.base_models
+            for model in models
+        ):
             return None
         # Reshape the base_models
         params_trees = zip(*self.base_models)

{ngboost-0.5.10.dev0 → ngboost-0.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ngboost"
-version = "0.5.10dev"
+version = "0.6.0"
 description = "Library for probabilistic predictions via gradient boosting."
 authors = ["Stanford ML Group <avati@cs.stanford.edu>"]
 readme = "README.md"
@@ -13,8 +13,11 @@ classifiers = [
 license = "Apache License 2.0"
 
 [tool.poetry.dependencies]
-python = ">=3.9, <3.15"
-scikit-learn = "^1.6"
+python = ">=3.10, <3.15"
+scikit-learn = [
+    {version = ">=1.6,<1.7", markers = "python_version < '3.14'"},
+    {version = ">=1.8.0,<2.0", markers = "python_version >= '3.14'"}
+]
 numpy = [
     {version = ">=1.21.2", markers = "python_version < '3.13'"},
     {version = ">=2.1.0", markers = "python_version >= '3.13'"}
@@ -26,18 +29,18 @@ scipy = [
 tqdm = ">=4.3"
 lifelines = ">=0.25"
 sympy = ">=1.12"
-matplotlib = [
-    {version = ">=3.0,<3.10", markers = "python_version < '3.10'"},
-    {version = ">=3.10.0", markers = "python_version >= '3.10'"}
-]
+joblib = ">=1.2"
+matplotlib = ">=3.10.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"
+pandas = ">=1.0"
 black = "^24.0.0"
 pre-commit = "^4.0.0"
 isort = "^5.13.0"
 pylint = "^3.2.0"
 flake8 = "^7.1.0"
+skops = ">=0.16"
 
 
 [build-system]