API Reference

Estimators

PretrainedLasso

ptlasso.PretrainedLasso

Bases: RegressorMixin, BasePretrainedLasso

Pretrained Lasso estimator.

Two-step training: 1. Fit an overall Lasso on all samples (lambda selected by internal CV). 2. For each group, fit a group-specific Lasso with offset (1 - alpha) * eta_overall, where eta_overall is the overall linear predictor (before the link function). Features not selected by the overall model receive a stronger penalty of 1 / alpha.

Parameters:

Name	Type	Description	Default
alpha	float in [0, 1]	Pretraining strength. 0 = overall model with fine-tuning (maximum pretraining); 1 = individual per-group Lasso (no pretraining). Matches the R ptLasso convention.	0.5
family	(gaussian, binomial, multinomial)	Response distribution.	"gaussian"
overall_lambda	('lambda.1se', 'lambda.min')	Lambda selection rule for the stage-1 overall model. "lambda.1se" (default, matching R) gives a sparser offset; "lambda.min" uses the CV minimum.	"lambda.1se"
fit_intercept	bool	Whether to fit an intercept in every sub-model.	True
lmda_path_size	int	Number of lambdas in the regularisation path.	100
min_ratio	float	Ratio of the smallest to largest lambda on the path.	0.0001
verbose	bool	Whether to display fitting progress and a summary after training. Adelie's internal output is always suppressed regardless of this setting.	True
n_threads	int	Number of threads passed to adelie's solver. Set to a higher value to parallelise the coordinate descent within each model fit. -1 uses all available CPU cores (os.cpu_count()).	-1

Attributes:

Name	Type	Description
overall_model_	adelie state	Fitted overall Lasso (stage 1).
overall_coef_	ndarray of shape (n_features,) or (n_features, K)	Coefficients from the overall model at the selected lambda. Shape is (n_features, K) for multinomial.
overall_intercept_	float	Intercept from the overall model. Not set for multinomial.
overall_lmda_idx_	int	Index into overall_model_.lmdas for the selected lambda.
pretrain_models_	dict {group -> adelie state}	Per-group fitted Lasso models (stage 2, with pretraining offset).
pretrain_lmda_idx_	dict {group -> int}	CV-selected lambda index for each pretrain group model (lambda.min).
individual_models_	dict {group -> adelie state}	Per-group fitted Lasso models without any pretraining offset.
individual_lmda_idx_	dict {group -> int}	CV-selected lambda index for each individual group model (lambda.min).
groups_	ndarray	Unique group labels seen during fit.
n_features_in_	int	Number of features seen during fit.
feature_names_in_	ndarray of str or None	Feature names, if provided.
n_classes_	int or None	Number of classes for multinomial; None otherwise.

References

Craig, E., Pilanci, M., Le Menestrel, T., Narasimhan, B., Rivas, M. A., Gullaksen, S. E., & Tibshirani, R. (2025). Pretraining and the lasso. Journal of the Royal Statistical Society Series B, qkaf050.

Source code in src/ptlasso/_ptlasso.py

class PretrainedLasso(RegressorMixin, BasePretrainedLasso):
    """Pretrained Lasso estimator.

    Two-step training:
    1. Fit an overall Lasso on all samples (lambda selected by internal CV).
    2. For each group, fit a group-specific Lasso with offset
       ``(1 - alpha) * eta_overall``, where ``eta_overall`` is the overall
       linear predictor (before the link function).  Features not selected
       by the overall model receive a stronger penalty of ``1 / alpha``.

    Parameters
    ----------
    alpha : float in [0, 1], default=0.5
        Pretraining strength.  ``0`` = overall model with fine-tuning
        (maximum pretraining); ``1`` = individual per-group Lasso
        (no pretraining).  Matches the R ptLasso convention.
    family : {"gaussian", "binomial", "multinomial"}, default="gaussian"
        Response distribution.
    overall_lambda : {"lambda.1se", "lambda.min"}, default="lambda.1se"
        Lambda selection rule for the stage-1 overall model.
        ``"lambda.1se"`` (default, matching R) gives a sparser offset;
        ``"lambda.min"`` uses the CV minimum.
    fit_intercept : bool, default=True
        Whether to fit an intercept in every sub-model.
    lmda_path_size : int, default=100
        Number of lambdas in the regularisation path.
    min_ratio : float, default=0.0001
        Ratio of the smallest to largest lambda on the path.
    verbose : bool, default=True
        Whether to display fitting progress and a summary after training.
        Adelie's internal output is always suppressed regardless of this setting.
    n_threads : int, default=-1
        Number of threads passed to adelie's solver.  Set to a higher value
        to parallelise the coordinate descent within each model fit.
        ``-1`` uses all available CPU cores (``os.cpu_count()``).

    Attributes
    ----------
    overall_model_ : adelie state
        Fitted overall Lasso (stage 1).
    overall_coef_ : ndarray of shape (n_features,) or (n_features, K)
        Coefficients from the overall model at the selected lambda.
        Shape is ``(n_features, K)`` for multinomial.
    overall_intercept_ : float
        Intercept from the overall model.  Not set for multinomial.
    overall_lmda_idx_ : int
        Index into ``overall_model_.lmdas`` for the selected lambda.
    pretrain_models_ : dict {group -> adelie state}
        Per-group fitted Lasso models (stage 2, with pretraining offset).
    pretrain_lmda_idx_ : dict {group -> int}
        CV-selected lambda index for each pretrain group model (``lambda.min``).
    individual_models_ : dict {group -> adelie state}
        Per-group fitted Lasso models without any pretraining offset.
    individual_lmda_idx_ : dict {group -> int}
        CV-selected lambda index for each individual group model (``lambda.min``).
    groups_ : ndarray
        Unique group labels seen during fit.
    n_features_in_ : int
        Number of features seen during fit.
    feature_names_in_ : ndarray of str or None
        Feature names, if provided.
    n_classes_ : int or None
        Number of classes for multinomial; ``None`` otherwise.

    References
    ----------
    Craig, E., Pilanci, M., Le Menestrel, T., Narasimhan, B., Rivas, M. A.,
    Gullaksen, S. E., & Tibshirani, R. (2025). Pretraining and the lasso.
    *Journal of the Royal Statistical Society Series B*, qkaf050.
    """

    def __init__(
        self,
        alpha=0.5,
        family="gaussian",
        overall_lambda="lambda.1se",
        fit_intercept=True,
        lmda_path_size=100,
        min_ratio=0.0001,
        verbose=True,
        n_threads=-1,
        standardize=True,
        n_folds=10,
    ):
        self.alpha = alpha
        self.family = family
        self.overall_lambda = overall_lambda
        self.fit_intercept = fit_intercept
        self.lmda_path_size = lmda_path_size
        self.min_ratio = min_ratio
        self.verbose = verbose
        self.n_threads = n_threads
        self.standardize = standardize
        self.n_folds = n_folds

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _validate_params(self):
        if not isinstance(self.alpha, (int, float)):
            raise TypeError(f"alpha must be a number, got {type(self.alpha).__name__}")
        if not (0.0 <= self.alpha <= 1.0):
            raise ValueError(f"alpha must be in [0, 1], got {self.alpha}")
        if self.family not in FAMILIES:
            raise ValueError(f"family must be one of {FAMILIES}, got '{self.family}'")
        if self.overall_lambda not in LMDA_MODES:
            raise ValueError(
                f"overall_lambda must be one of {LMDA_MODES}, got '{self.overall_lambda}'"
            )
        if not isinstance(self.lmda_path_size, int) or self.lmda_path_size < 1:
            raise ValueError(
                f"lmda_path_size must be a positive integer, got {self.lmda_path_size}"
            )
        if not isinstance(self.min_ratio, (int, float)):
            raise TypeError(f"min_ratio must be a number, got {type(self.min_ratio).__name__}")
        if not (0.0 < self.min_ratio < 1.0):
            raise ValueError(f"min_ratio must be in (0, 1), got {self.min_ratio}")

    def _label(self, g):
        return self.group_labels_.get(g, g)

    def _names_or_indices(self, indices):
        if self.feature_names_in_ is not None:
            return self.feature_names_in_[indices]
        return indices

    def _make_onehot(self, groups):
        """Build (n, k-1) group indicator matrix.

        The first group in ``self.groups_`` is the reference category (all
        zeros), matching R's ``model.matrix(~groups - 1)[, 2:k]`` convention.
        These columns are prepended to X in the overall model with zero
        penalty so they act as free group-specific intercepts (R's
        ``group.intercepts = TRUE``).
        """
        k = len(self.groups_)
        n = len(groups)
        if k <= 1:
            return np.empty((n, 0), dtype=np.float64, order="F")
        onehot = np.zeros((n, k - 1), dtype=np.float64, order="F")
        for col, g in enumerate(self.groups_[1:]):  # groups_[0] is reference
            onehot[groups == g, col] = 1.0
        return onehot

    def _grpnet_kwargs(self):
        n_threads = os.cpu_count() if self.n_threads == -1 else self.n_threads
        return dict(
            alpha=1,  # pure lasso (no group penalty mixing)
            intercept=self.fit_intercept,
            lmda_path_size=self.lmda_path_size,
            min_ratio=self.min_ratio,
            progress_bar=False,  # replaced by ptlasso's own tqdm output
            n_threads=n_threads,
        )

    def _cv_grpnet_kwargs(self):
        return {**self._grpnet_kwargs(), "n_folds": self.n_folds}

    def _wrap_matrix(self, X):
        # adelie incurs ~20x slowdown when matrix and solver use different
        # n_threads (OpenMP switching cost). Wrap X so both use the same value.
        n_threads = os.cpu_count() if self.n_threads == -1 else self.n_threads
        return ad.matrix.dense(X, method="naive", n_threads=n_threads)

    def _overall_eta(self, X, groups):
        """Overall linear predictor at the selected lambda.

        The overall model was trained on ``[onehot | X]``, so we reconstruct
        the augmented matrix before computing the linear predictor.  This
        matches R's predict behaviour where group-intercept columns are always
        included.
        """
        onehot = self._make_onehot(groups)
        X_aug = (
            np.asfortranarray(np.hstack([onehot, X]))
            if onehot.shape[1] > 0
            else np.asfortranarray(X)
        )
        return _eta_from_state(
            self.overall_model_, X_aug, self.overall_lmda_idx_, self.family, self.n_classes_
        )

    def _compute_oof_eta(self, X_overall, y, overall_pf, groups):
        """Compute out-of-fold linear predictors for the overall model.

        Mirrors R's ``cv.glmnet(..., keep=TRUE)`` which stores prevalidated
        (out-of-fold) predictions.  Each sample's prediction comes from a
        model that never saw that sample during training, so the offset used
        in stage-2 fitting is free of in-sample optimism.

        Parameters
        ----------
        X_overall : (n, k-1+p) ndarray
            Augmented feature matrix including group-indicator columns.
        y : (n,) ndarray
            Target values.
        overall_pf : (k-1+p,) ndarray
            Penalty factors (0 for group-dummy columns, 1 for features).
        groups : (n,) ndarray
            Group membership for each sample, used to create stratified folds.

        Returns
        -------
        oof_eta : (n,) or (n, K) ndarray
            Out-of-fold linear predictors in the same shape as the overall
            model's ``_eta_from_state`` output.
        """
        n = X_overall.shape[0]
        lamhat = float(self.overall_model_.lmdas[self.overall_lmda_idx_])

        if self.family == "multinomial":
            oof_eta = np.zeros((n, self.n_classes_))
        else:
            oof_eta = np.zeros(n)

        n_folds = self._n_folds_oof_

        splitter = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
        _show_oof = getattr(self, "_show_overall_progress", False)
        if _show_oof:
            _t_oof = time.time()
            _logger.info("    OOF predictions:")
        for _oof_i, (train_idx, test_idx) in enumerate(splitter.split(X_overall, groups)):
            X_tr = np.asfortranarray(X_overall[train_idx])
            y_tr = y[train_idx]
            X_te = np.asfortranarray(X_overall[test_idx])

            glm_fold = _make_glm(self.family, y_tr)
            with _silence():
                fold_state = ad.grpnet(
                    self._wrap_matrix(X_tr), glm_fold, penalty=overall_pf, **self._grpnet_kwargs()
                )

            # Find the lambda in this fold's path closest to the full-data lamhat.
            lmdas = np.asarray(fold_state.lmdas)
            lmda_idx = int(np.argmin(np.abs(lmdas - lamhat))) if len(lmdas) > 0 else 0

            oof_eta[test_idx] = _eta_from_state(
                fold_state, X_te, lmda_idx, self.family, self.n_classes_
            )
            if _show_oof:
                _logger.info(f"    OOF fold {_oof_i + 1}/{n_folds} done")
        if _show_oof:
            _logger.info(f"    OOF done ({time.time() - _t_oof:.1f}s)")

        return oof_eta

    def _compute_group_oof_eta(self, X_g, y_g, g_offset_oof, pf, lamhat_g):
        """OOF linear predictor for a group model (group contribution only, no offset).

        Mirrors R's ``cv.glmnet(..., keep=TRUE)`` for stage-2 group models.
        Each sample is predicted by a fold model that never saw it during
        training.  The caller adds back ``(1-alpha)*overall_oof_eta[mask]`` to
        obtain the full combined pretrain predictor.

        Parameters
        ----------
        X_g : (n_g, p) ndarray
            Standardized group feature matrix (already passed through scaler).
        y_g : (n_g,) ndarray
            Group targets.
        g_offset_oof : (n_g,) ndarray or None
            OOF offset for training fold models: ``(1-alpha)*overall_oof_eta``
            for this group.  Pass ``None`` for individual (no-offset) models.
        pf : (p,) ndarray or None
            Stage-2 penalty factors.  ``None`` means uniform (no weighting).
        lamhat_g : float
            Selected lambda from the full-data group model; each fold model's
            nearest lambda is used for prediction (matches glmnet keep=TRUE).

        Returns
        -------
        oof_eta : (n_g,) or (n_g, K) ndarray
            Group model OOF linear predictor WITHOUT the overall offset.
        """
        n_g = X_g.shape[0]

        n_folds = self._n_folds_oof_
        if self.family in ("binomial", "multinomial"):
            y_int = y_g.astype(int)
            splitter = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
            fold_iter = list(splitter.split(X_g, y_int))
        else:
            fold_iter = list(KFold(n_splits=n_folds, shuffle=True, random_state=42).split(X_g))

        if self.family == "multinomial":
            oof_eta = np.zeros((n_g, self.n_classes_))
        else:
            oof_eta = np.zeros(n_g)

        for train_idx, test_idx in fold_iter:
            X_tr = np.asfortranarray(X_g[train_idx])
            y_tr = y_g[train_idx]
            X_te = np.asfortranarray(X_g[test_idx])

            glm_fold = _make_glm(self.family, y_tr)
            kw = self._grpnet_kwargs()
            if pf is not None:
                kw["penalty"] = pf
            if g_offset_oof is not None:
                kw["offsets"] = np.asfortranarray(g_offset_oof[train_idx])
            with _silence():
                fold_state = ad.grpnet(self._wrap_matrix(X_tr), glm_fold, **kw)

            lmdas = np.asarray(fold_state.lmdas)
            lmda_idx = int(np.argmin(np.abs(lmdas - lamhat_g))) if len(lmdas) > 0 else 0
            oof_eta[test_idx] = _eta_from_state(
                fold_state, X_te, lmda_idx, self.family, self.n_classes_
            )

        return oof_eta

    # ------------------------------------------------------------------
    # Progress / display helpers
    # ------------------------------------------------------------------

    def _support_size(self, state, lmda_idx):
        """Number of nonzero features in a fitted state at a given lambda index."""
        c = _coef_at(state, lmda_idx)
        if self.n_classes_ is not None:
            c = c.reshape(self.n_features_in_, self.n_classes_, order="F")
            return int(np.sum(np.any(c != 0, axis=1)))
        return int(np.sum(c != 0))

    def _print_fit_summary(self, elapsed):
        SEP = "─" * 54
        if self.family == "multinomial":
            n_ov = int(np.sum(np.any(self.overall_coef_ != 0, axis=1)))
        else:
            n_ov = int(np.sum(self.overall_coef_ != 0))
        pre_parts = "   ".join(
            f"{self._label(g)}: |S|={self._support_size(self.pretrain_models_[g], self.pretrain_lmda_idx_[g])}"
            for g in self.groups_
        )
        ind_parts = "   ".join(
            f"{self._label(g)}: |S|={self._support_size(self.individual_models_[g], self.individual_lmda_idx_[g])}"
            for g in self.groups_
        )
        _logger.info(SEP)
        _logger.info(f"  {'overall':<13}|S| = {n_ov}")
        _logger.info(f"  {'pretrain':<13}{pre_parts}")
        _logger.info(f"  {'individual':<13}{ind_parts}")
        _logger.info(SEP)
        _logger.info(f"  Fitted in {elapsed:.1f}s")

    # ------------------------------------------------------------------
    # fit
    # ------------------------------------------------------------------

    def fit(self, X, y, groups, group_labels=None, feature_names=None):
        """Fit the pretrained Lasso.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values.  For ``family="binomial"``, must contain only 0
            and 1.  For ``family="multinomial"``, must contain non-negative
            integer class labels 0..K-1.
        groups : array-like of shape (n_samples,)
            Group membership for each sample.  Must contain at least two
            distinct values.
        group_labels : dict or None, default=None
            Optional mapping from group values to display names used in
            ``__repr__`` and ``get_coef()``.
        feature_names : array-like of str or None, default=None
            Feature names.  Inferred from ``X.columns`` when ``X`` is a
            DataFrame and this argument is ``None``.

        Returns
        -------
        self : PretrainedLasso
            Fitted estimator.
        """
        t0 = time.time()
        self._validate_params()

        if feature_names is None and hasattr(X, "columns"):
            feature_names = list(X.columns)

        X, y = check_X_y(X, y, dtype=np.float64, order="F")
        self.n_features_in_ = X.shape[1]

        # Standardize features to zero mean / unit std, matching glmnet's
        # default (standardize=TRUE).  The one-hot group-indicator columns
        # added later are built after this step and left un-standardized.
        self.scaler_ = StandardScaler(with_mean=True, with_std=self.standardize)
        X = self.scaler_.fit_transform(X)

        if feature_names is not None:
            feature_names = np.asarray(feature_names)
            if len(feature_names) != self.n_features_in_:
                raise ValueError(
                    f"feature_names has {len(feature_names)} entries "
                    f"but X has {self.n_features_in_} features"
                )
        self.feature_names_in_ = feature_names

        if self.family == "binomial":
            unique_y = np.unique(y)
            if not np.all(np.isin(unique_y, [0.0, 1.0])):
                raise ValueError(
                    f"For family='binomial', y must contain only 0 and 1, "
                    f"got unique values {unique_y}"
                )
        if self.family == "multinomial":
            y_int = np.asarray(y)
            if not (np.all(y_int == np.round(y_int)) and np.all(y_int >= 0)):
                raise ValueError(
                    "For family='multinomial', y must contain non-negative integer class labels"
                )

        groups = np.asarray(groups)
        if len(groups) != X.shape[0]:
            raise ValueError(f"groups has {len(groups)} elements but X has {X.shape[0]} samples")
        self.groups_ = np.unique(groups)
        if len(self.groups_) < 2:
            raise ValueError("groups must contain at least 2 unique values")

        # Single OOF fold count used by both _compute_oof_eta and
        # _compute_group_oof_eta, matching R's uniform nfolds across all models.
        _min_group = min(int(np.sum(groups == g)) for g in self.groups_)
        if self.family in ("binomial", "multinomial"):
            _min_class = min(
                int(np.bincount(np.asarray(y)[groups == g].astype(int)).min()) for g in self.groups_
            )
            self._n_folds_oof_ = max(2, min(self.n_folds, _min_group, _min_class))
        else:
            self._n_folds_oof_ = max(2, min(self.n_folds, _min_group))

        if group_labels is not None:
            if not isinstance(group_labels, dict):
                raise TypeError(f"group_labels must be a dict, got {type(group_labels).__name__}")
            unknown = set(group_labels) - set(self.groups_)
            if unknown:
                raise ValueError(f"group_labels contains unknown group keys: {unknown}")
        self.group_labels_ = group_labels or {}

        self.n_classes_ = (
            int(np.asarray(y, dtype=int).max()) + 1 if self.family == "multinomial" else None
        )

        if self.verbose:
            _enable_verbose_logging()
            k = len(self.groups_)
            glabels = [str(self._label(g)) for g in self.groups_]
            gstr = ", ".join(glabels[:4]) + (f" +{k - 4} more" if k > 4 else "")
            _logger.info(
                f"\nPretrainedLasso  {self.family}  ·  {k} groups ({gstr})"
                f"  ·  {self.n_features_in_} features  ·  α={self.alpha}"
            )

        # ----------------------------------------------------------
        # Build augmented X for the overall model (group intercepts).
        # Matches R's group.intercepts = TRUE: one-hot group dummies (k-1
        # columns, first group = reference) are prepended to X with zero
        # penalty so they act as free intercepts per group.
        # ----------------------------------------------------------
        onehot = self._make_onehot(groups)  # (n, k-1)
        n_onehot = onehot.shape[1]  # = k-1 >= 1
        self._n_onehot_ = n_onehot
        X_overall = (
            np.asfortranarray(np.hstack([onehot, X])) if n_onehot > 0 else np.asfortranarray(X)
        )
        # Zero penalty for group-dummy columns; unit penalty for X features.
        overall_pf = np.concatenate([np.zeros(n_onehot), np.ones(self.n_features_in_)])

        # ----------------------------------------------------------
        # Step 1: overall model (lambda selected by CV)
        # ----------------------------------------------------------
        glm_all = _make_glm(self.family, y)
        _show_progress = self.verbose or getattr(self, "_show_overall_progress", False)

        # cv_grpnet runs K internal CV folds with OpenMP — progress_bar=True would
        # spawn K bars simultaneously, so we always silence it and show a timer.
        if _show_progress:
            _t_cv = time.time()
        with _silence():
            cv_overall = ad.cv_grpnet(
                self._wrap_matrix(X_overall),
                glm_all,
                penalty=overall_pf,
                **self._cv_grpnet_kwargs(),
            )
        self.overall_lmda_idx_ = (
            cv_overall.best_idx
            if self.overall_lambda == "lambda.min"
            else _lmda_1se_idx(cv_overall)
        )
        if _show_progress:
            _logger.info(f"    CV λ-selection done ({time.time() - _t_cv:.1f}s)")
            _logger.info("    Full-data refit (λ-path):")
            self.overall_model_ = cv_overall.fit(
                self._wrap_matrix(X_overall),
                glm_all,
                penalty=overall_pf,
                **{**self._grpnet_kwargs(), "progress_bar": True},
            )
        else:
            with _silence():
                self.overall_model_ = cv_overall.fit(
                    self._wrap_matrix(X_overall),
                    glm_all,
                    penalty=overall_pf,
                    **self._grpnet_kwargs(),
                )

        # Extract X-feature coefficients only (skip the k-1 onehot columns).
        # overall_coef_ has shape (p,) or (p, K) — identical to the no-group-
        # intercepts case from the caller's perspective.
        if self.family == "multinomial":
            flat = _coef_at(self.overall_model_, self.overall_lmda_idx_)
            p_aug = X_overall.shape[1]
            coef_mat = flat.reshape(p_aug, self.n_classes_, order="F")
            self.overall_coef_ = coef_mat[n_onehot:, :]  # (p, K)
        else:
            coef_full = _coef_at(self.overall_model_, self.overall_lmda_idx_)
            self.overall_coef_ = coef_full[n_onehot:]  # (p,)
            self.overall_intercept_ = float(self.overall_model_.intercepts[self.overall_lmda_idx_])

        # ----------------------------------------------------------
        # Compute OOF overall linear predictor (matches R's keep=TRUE).
        # Using in-sample predictions as the offset introduces optimism:
        # the overall model has already seen the training samples, so its
        # predictions are sharper than true held-out predictions.  OOF
        # predictions avoid this leakage.
        # ----------------------------------------------------------
        preval_offset = self._compute_oof_eta(X_overall, y, overall_pf, groups)
        self._preval_offset = preval_offset  # cached for PretrainedLassoCV reuse

        # ----------------------------------------------------------
        # Step 2: per-group models
        # ----------------------------------------------------------
        # Penalty factor: features NOT in the overall X-support get penalty
        # 1/alpha, steering the group model to prefer features already
        # selected overall.  Matches R's fac = rep(1/alpha, p); fac[supall] = 1.
        if self.family == "multinomial":
            overall_support = np.where(np.any(self.overall_coef_ != 0, axis=1))[0]
        else:
            overall_support = np.where(self.overall_coef_ != 0)[0]

        _alpha_pf = self.alpha if self.alpha > 0 else 1e-9
        pf = np.full(self.n_features_in_, 1.0 / _alpha_pf)
        pf[overall_support] = 1.0

        if self.verbose:
            if self.family == "multinomial":
                n_ov = int(np.sum(np.any(self.overall_coef_ != 0, axis=1)))
            else:
                n_ov = int(np.sum(self.overall_coef_ != 0))
            _logger.info(f"  Overall model done  |S|={n_ov}")

        self.pretrain_models_ = {}
        self.pretrain_lmda_idx_ = {}
        self.individual_models_ = {}
        self.individual_lmda_idx_ = {}

        if self.verbose:
            try:
                from tqdm import tqdm as _tqdm

                group_iter = _tqdm(
                    self.groups_, desc="  [2/2] Group models", unit="group", leave=True
                )
            except ImportError:
                _logger.info("  [2/2] Group models")
                group_iter = self.groups_
        else:
            group_iter = self.groups_

        for g in group_iter:
            mask = groups == g
            X_g = np.asfortranarray(X[mask])
            glm_g = _make_glm(self.family, y[mask])
            # OOF offset: each sample's overall prediction came from a model
            # trained without that sample — mirrors R's use of fit.preval.
            offset = np.asfortranarray((1 - self.alpha) * preval_offset[mask])

            # Mirrors R's cv.glmnet for per-group models: select lambda by CV
            # (lambda.min), then refit on the full group data at that lambda.
            with _silence():
                cv_pre = ad.cv_grpnet(
                    self._wrap_matrix(X_g),
                    glm_g,
                    offsets=offset,
                    penalty=pf,
                    **self._cv_grpnet_kwargs(),
                )
                self.pretrain_lmda_idx_[g] = cv_pre.best_idx
                self.pretrain_models_[g] = cv_pre.fit(
                    self._wrap_matrix(X_g),
                    glm_g,
                    offsets=offset,
                    penalty=pf,
                    **self._grpnet_kwargs(),
                )

                cv_ind = ad.cv_grpnet(self._wrap_matrix(X_g), glm_g, **self._cv_grpnet_kwargs())
                self.individual_lmda_idx_[g] = cv_ind.best_idx
                self.individual_models_[g] = cv_ind.fit(
                    self._wrap_matrix(X_g), glm_g, **self._grpnet_kwargs()
                )

        if self.verbose:
            self._print_fit_summary(time.time() - t0)

        return self

    def _fit_groups_only(self, X, y, groups, preval_offset):
        """Fit stage-2 group models, reusing the stage-1 state already on self.

        Called by PretrainedLassoCV to avoid re-running the expensive overall
        model + OOF computation for every alpha within the same CV fold.
        Requires that all stage-1 attributes (overall_model_, overall_coef_,
        overall_lmda_idx_, n_features_in_, groups_, _n_onehot_, etc.) are
        already set (typically copied from a template estimator).
        """
        X = np.asfortranarray(self.scaler_.transform(np.asarray(X, dtype=np.float64)))
        y = np.asarray(y, dtype=np.float64)
        groups = np.asarray(groups)

        if self.family == "multinomial":
            overall_support = np.where(np.any(self.overall_coef_ != 0, axis=1))[0]
        else:
            overall_support = np.where(self.overall_coef_ != 0)[0]

        _alpha_pf = self.alpha if self.alpha > 0 else 1e-9
        pf = np.full(self.n_features_in_, 1.0 / _alpha_pf)
        pf[overall_support] = 1.0

        self.pretrain_models_ = {}
        self.pretrain_lmda_idx_ = {}
        self.individual_models_ = {}
        self.individual_lmda_idx_ = {}

        for g in self.groups_:
            mask = groups == g
            X_g = np.asfortranarray(X[mask])
            glm_g = _make_glm(self.family, y[mask])
            offset = np.asfortranarray((1 - self.alpha) * preval_offset[mask])

            with _silence():
                cv_pre = ad.cv_grpnet(
                    self._wrap_matrix(X_g),
                    glm_g,
                    offsets=offset,
                    penalty=pf,
                    **self._cv_grpnet_kwargs(),
                )
                self.pretrain_lmda_idx_[g] = cv_pre.best_idx
                self.pretrain_models_[g] = cv_pre.fit(
                    self._wrap_matrix(X_g),
                    glm_g,
                    offsets=offset,
                    penalty=pf,
                    **self._grpnet_kwargs(),
                )

                cv_ind = ad.cv_grpnet(self._wrap_matrix(X_g), glm_g, **self._cv_grpnet_kwargs())
                self.individual_lmda_idx_[g] = cv_ind.best_idx
                self.individual_models_[g] = cv_ind.fit(
                    self._wrap_matrix(X_g), glm_g, **self._grpnet_kwargs()
                )

    # ------------------------------------------------------------------
    # predict / score / evaluate
    # ------------------------------------------------------------------

    def predict(self, X, groups, model="pretrain", type="response", lmda_idx=None):
        """Predict target values.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        groups : array-like of shape (n_samples,)
        model : {"pretrain", "individual", "overall"}, default="pretrain"
        type : {"response", "link", "class"}, default="response"
            Scale of the returned predictions.

            ``"response"`` — fitted values on the response scale: probabilities
            for binomial/multinomial, fitted values for gaussian.

            ``"link"`` — raw linear predictor before the link function
            (``η = Xβ + intercept``).  For gaussian this is identical to
            ``"response"``.

            ``"class"`` — predicted class label: ``0`` or ``1`` for binomial
            (threshold 0.5), integer argmax for multinomial.  Not valid for
            gaussian.
        lmda_idx : int or None
            Lambda index for group models.  ``None`` uses the CV-selected
            lambda for each group (``lambda.min``, matching R).

        Returns
        -------
        y_pred : ndarray
            Shape ``(n,)`` for gaussian/binomial and for multinomial
            ``"class"``.  Shape ``(n, K)`` for multinomial ``"response"``
            or ``"link"``.
        """
        check_is_fitted(self)
        X = check_array(X, dtype=np.float64, order="F")
        X = np.asfortranarray(self.scaler_.transform(X))
        groups = np.asarray(groups)

        if model not in PREDICT_MODELS:
            raise ValueError(f"model must be one of {PREDICT_MODELS}, got '{model}'")
        if type not in PREDICT_TYPES:
            raise ValueError(f"type must be one of {PREDICT_TYPES}, got '{type}'")
        if type == "class" and self.family == "gaussian":
            raise ValueError("type='class' is not valid for family='gaussian'")
        if len(groups) != X.shape[0]:
            raise ValueError(f"groups has {len(groups)} elements but X has {X.shape[0]} samples")
        unknown = set(np.unique(groups)) - set(self.groups_)
        if unknown:
            raise ValueError(f"predict received groups not seen during fit: {unknown}")

        if model == "overall":
            return _eta_to_output(self._overall_eta(X, groups), self.family, type)

        n_out = (X.shape[0], self.n_classes_) if self.family == "multinomial" else (X.shape[0],)
        eta_out = np.empty(n_out)

        # Overall eta is only needed when combining with a group model (pretrain).
        eta_ov = self._overall_eta(X, groups) if model == "pretrain" else None

        for g in self.groups_:
            mask = groups == g
            X_g = X[mask]

            if model == "pretrain":
                g_idx = self.pretrain_lmda_idx_.get(g, -1) if lmda_idx is None else lmda_idx
                eta_group = _eta_from_state(
                    self.pretrain_models_[g], X_g, g_idx, self.family, self.n_classes_
                )
                eta_out[mask] = (1 - self.alpha) * eta_ov[mask] + eta_group
            else:
                g_idx = self.individual_lmda_idx_.get(g, -1) if lmda_idx is None else lmda_idx
                eta_out[mask] = _eta_from_state(
                    self.individual_models_[g], X_g, g_idx, self.family, self.n_classes_
                )

        return _eta_to_output(eta_out, self.family, type)

    def score(self, X, y, groups):
        """Return a scalar performance metric using the pretrained model.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        y : array-like of shape (n_samples,)
        groups : array-like of shape (n_samples,)

        Returns
        -------
        score : float
            R² for gaussian; classification accuracy for binomial/multinomial.
        """
        return _model_score(y, self.predict(X, groups), self.family)

    def evaluate(self, X, y, groups):
        """Predict and score with all three sub-models.

        Convenience method matching R's ``predict(fit, xtest, ytest=ytest)``.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        y : array-like of shape (n_samples,)
        groups : array-like of shape (n_samples,)

        Returns
        -------
        result : dict
            Keys are ``"pretrain"``, ``"individual"``, ``"overall"``.
            Each value is a dict with entries:

            ``"predictions"`` : ndarray
                Predicted values, shape ``(n,)`` or ``(n, K)`` for multinomial.
            ``"score"`` : float
                R² for gaussian; accuracy for binomial/multinomial.
        """
        check_is_fitted(self)
        result = {}
        for m in ("pretrain", "individual", "overall"):
            preds = self.predict(X, groups, model=m)
            result[m] = {"predictions": preds, "score": _model_score(y, preds, self.family)}
        return result

    # ------------------------------------------------------------------
    # repr / get_coef
    # ------------------------------------------------------------------

    def __repr__(self):
        header = (
            f"PretrainedLasso(alpha={self.alpha}, family='{self.family}', "
            f"overall_lambda='{self.overall_lambda}', "
            f"fit_intercept={self.fit_intercept}, lmda_path_size={self.lmda_path_size})"
        )
        if not hasattr(self, "overall_model_"):
            return header + "\n  [not fitted]"

        coef = self.overall_coef_
        nz_idx = (
            np.where(np.any(coef != 0, axis=1))[0]
            if self.family == "multinomial"
            else np.where(coef != 0)[0]
        )
        names = self._names_or_indices(nz_idx)
        preview = list(names[:5]) + (["..."] if len(nz_idx) > 5 else [])
        ov_str = (
            f"|Ŝ| = {len(nz_idx)} / {self.n_features_in_}  [{', '.join(str(v) for v in preview)}]"
        )

        group_sizes = {}
        for g in self.groups_:
            c = _coef_at(self.pretrain_models_[g], self.pretrain_lmda_idx_.get(g, -1))
            if self.family == "multinomial":
                c = c.reshape(self.n_features_in_, self.n_classes_, order="F")
                group_sizes[self._label(g)] = int(np.sum(np.any(c != 0, axis=1)))
            else:
                group_sizes[self._label(g)] = int(np.sum(c != 0))
        group_str = ", ".join(f"{lbl}: |Ŝ|={v}" for lbl, v in group_sizes.items())

        return (
            f"{header}\n"
            f"  family       : {self.family}\n"
            f"  n_features   : {self.n_features_in_}\n"
            f"  n_groups     : {len(self.groups_)}\n"
            f"  overall |Ŝ|  : {ov_str}\n"
            f"  pretrain |Ŝ| : {group_str}\n"
            f"  overall λ    : {self.overall_lambda}  (idx {self.overall_lmda_idx_})"
        )

    def get_coef(self, model="all", lmda_idx=None):
        """Return fitted coefficients as a nested dict.

        Parameters
        ----------
        model : {"all", "overall", "pretrain", "individual"}, default="all"
            Which sub-model(s) to return.
        lmda_idx : int or None
            Lambda index for group models.  ``None`` uses the last lambda (``-1``).

        Returns
        -------
        coefs : dict
            When ``model="all"``, keys are ``"overall"``, ``"pretrain"``,
            ``"individual"``.  For ``"overall"``, value is
            ``{"coef": ndarray, "intercept": ndarray}``.  For ``"pretrain"``
            and ``"individual"``, value is a dict keyed by group label, each
            containing ``{"coef": ndarray, "intercept": ndarray}``.
            When a specific model is requested, only that sub-dict is returned.
        """
        if model not in COEF_MODELS:
            raise ValueError(f"model must be one of {COEF_MODELS}, got '{model}'")
        check_is_fitted(self)

        def _group_coefs(models, lmda_idxs):
            result = {}
            for g, state in models.items():
                use_idx = lmda_idxs.get(g, -1) if lmda_idx is None else lmda_idx
                result[self._label(g)] = {
                    "coef": _coef_at(state, use_idx),
                    "intercept": np.asarray(state.intercepts[use_idx]).ravel(),
                }
            return result

        result = {}
        if model in ("all", "overall"):
            result["overall"] = {
                "coef": self.overall_coef_,
                "intercept": np.asarray(
                    self.overall_model_.intercepts[self.overall_lmda_idx_]
                ).ravel(),
            }
        if model in ("all", "pretrain"):
            result["pretrain"] = _group_coefs(self.pretrain_models_, self.pretrain_lmda_idx_)
        if model in ("all", "individual"):
            result["individual"] = _group_coefs(self.individual_models_, self.individual_lmda_idx_)
        return result if model == "all" else result[model]

    # ------------------------------------------------------------------
    # Serialisation
    # ------------------------------------------------------------------

    def __getstate__(self):
        d = self.__dict__.copy()
        _proxify_models(d)
        return d

fit

fit(X, y, groups, group_labels=None, feature_names=None)

Fit the pretrained Lasso.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)	Training data.	required
y	array-like of shape (n_samples,)	Target values. For family="binomial", must contain only 0 and 1. For family="multinomial", must contain non-negative integer class labels 0..K-1.	required
groups	array-like of shape (n_samples,)	Group membership for each sample. Must contain at least two distinct values.	required
group_labels	dict or None	Optional mapping from group values to display names used in __repr__ and get_coef().	None
feature_names	array-like of str or None	Feature names. Inferred from X.columns when X is a DataFrame and this argument is None.	None

Returns:

Name	Type	Description
self	PretrainedLasso	Fitted estimator.

Source code in src/ptlasso/_ptlasso.py

def fit(self, X, y, groups, group_labels=None, feature_names=None):
    """Fit the pretrained Lasso.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
        Training data.
    y : array-like of shape (n_samples,)
        Target values.  For ``family="binomial"``, must contain only 0
        and 1.  For ``family="multinomial"``, must contain non-negative
        integer class labels 0..K-1.
    groups : array-like of shape (n_samples,)
        Group membership for each sample.  Must contain at least two
        distinct values.
    group_labels : dict or None, default=None
        Optional mapping from group values to display names used in
        ``__repr__`` and ``get_coef()``.
    feature_names : array-like of str or None, default=None
        Feature names.  Inferred from ``X.columns`` when ``X`` is a
        DataFrame and this argument is ``None``.

    Returns
    -------
    self : PretrainedLasso
        Fitted estimator.
    """
    t0 = time.time()
    self._validate_params()

    if feature_names is None and hasattr(X, "columns"):
        feature_names = list(X.columns)

    X, y = check_X_y(X, y, dtype=np.float64, order="F")
    self.n_features_in_ = X.shape[1]

    # Standardize features to zero mean / unit std, matching glmnet's
    # default (standardize=TRUE).  The one-hot group-indicator columns
    # added later are built after this step and left un-standardized.
    self.scaler_ = StandardScaler(with_mean=True, with_std=self.standardize)
    X = self.scaler_.fit_transform(X)

    if feature_names is not None:
        feature_names = np.asarray(feature_names)
        if len(feature_names) != self.n_features_in_:
            raise ValueError(
                f"feature_names has {len(feature_names)} entries "
                f"but X has {self.n_features_in_} features"
            )
    self.feature_names_in_ = feature_names

    if self.family == "binomial":
        unique_y = np.unique(y)
        if not np.all(np.isin(unique_y, [0.0, 1.0])):
            raise ValueError(
                f"For family='binomial', y must contain only 0 and 1, "
                f"got unique values {unique_y}"
            )
    if self.family == "multinomial":
        y_int = np.asarray(y)
        if not (np.all(y_int == np.round(y_int)) and np.all(y_int >= 0)):
            raise ValueError(
                "For family='multinomial', y must contain non-negative integer class labels"
            )

    groups = np.asarray(groups)
    if len(groups) != X.shape[0]:
        raise ValueError(f"groups has {len(groups)} elements but X has {X.shape[0]} samples")
    self.groups_ = np.unique(groups)
    if len(self.groups_) < 2:
        raise ValueError("groups must contain at least 2 unique values")

    # Single OOF fold count used by both _compute_oof_eta and
    # _compute_group_oof_eta, matching R's uniform nfolds across all models.
    _min_group = min(int(np.sum(groups == g)) for g in self.groups_)
    if self.family in ("binomial", "multinomial"):
        _min_class = min(
            int(np.bincount(np.asarray(y)[groups == g].astype(int)).min()) for g in self.groups_
        )
        self._n_folds_oof_ = max(2, min(self.n_folds, _min_group, _min_class))
    else:
        self._n_folds_oof_ = max(2, min(self.n_folds, _min_group))

    if group_labels is not None:
        if not isinstance(group_labels, dict):
            raise TypeError(f"group_labels must be a dict, got {type(group_labels).__name__}")
        unknown = set(group_labels) - set(self.groups_)
        if unknown:
            raise ValueError(f"group_labels contains unknown group keys: {unknown}")
    self.group_labels_ = group_labels or {}

    self.n_classes_ = (
        int(np.asarray(y, dtype=int).max()) + 1 if self.family == "multinomial" else None
    )

    if self.verbose:
        _enable_verbose_logging()
        k = len(self.groups_)
        glabels = [str(self._label(g)) for g in self.groups_]
        gstr = ", ".join(glabels[:4]) + (f" +{k - 4} more" if k > 4 else "")
        _logger.info(
            f"\nPretrainedLasso  {self.family}  ·  {k} groups ({gstr})"
            f"  ·  {self.n_features_in_} features  ·  α={self.alpha}"
        )

    # ----------------------------------------------------------
    # Build augmented X for the overall model (group intercepts).
    # Matches R's group.intercepts = TRUE: one-hot group dummies (k-1
    # columns, first group = reference) are prepended to X with zero
    # penalty so they act as free intercepts per group.
    # ----------------------------------------------------------
    onehot = self._make_onehot(groups)  # (n, k-1)
    n_onehot = onehot.shape[1]  # = k-1 >= 1
    self._n_onehot_ = n_onehot
    X_overall = (
        np.asfortranarray(np.hstack([onehot, X])) if n_onehot > 0 else np.asfortranarray(X)
    )
    # Zero penalty for group-dummy columns; unit penalty for X features.
    overall_pf = np.concatenate([np.zeros(n_onehot), np.ones(self.n_features_in_)])

    # ----------------------------------------------------------
    # Step 1: overall model (lambda selected by CV)
    # ----------------------------------------------------------
    glm_all = _make_glm(self.family, y)
    _show_progress = self.verbose or getattr(self, "_show_overall_progress", False)

    # cv_grpnet runs K internal CV folds with OpenMP — progress_bar=True would
    # spawn K bars simultaneously, so we always silence it and show a timer.
    if _show_progress:
        _t_cv = time.time()
    with _silence():
        cv_overall = ad.cv_grpnet(
            self._wrap_matrix(X_overall),
            glm_all,
            penalty=overall_pf,
            **self._cv_grpnet_kwargs(),
        )
    self.overall_lmda_idx_ = (
        cv_overall.best_idx
        if self.overall_lambda == "lambda.min"
        else _lmda_1se_idx(cv_overall)
    )
    if _show_progress:
        _logger.info(f"    CV λ-selection done ({time.time() - _t_cv:.1f}s)")
        _logger.info("    Full-data refit (λ-path):")
        self.overall_model_ = cv_overall.fit(
            self._wrap_matrix(X_overall),
            glm_all,
            penalty=overall_pf,
            **{**self._grpnet_kwargs(), "progress_bar": True},
        )
    else:
        with _silence():
            self.overall_model_ = cv_overall.fit(
                self._wrap_matrix(X_overall),
                glm_all,
                penalty=overall_pf,
                **self._grpnet_kwargs(),
            )

    # Extract X-feature coefficients only (skip the k-1 onehot columns).
    # overall_coef_ has shape (p,) or (p, K) — identical to the no-group-
    # intercepts case from the caller's perspective.
    if self.family == "multinomial":
        flat = _coef_at(self.overall_model_, self.overall_lmda_idx_)
        p_aug = X_overall.shape[1]
        coef_mat = flat.reshape(p_aug, self.n_classes_, order="F")
        self.overall_coef_ = coef_mat[n_onehot:, :]  # (p, K)
    else:
        coef_full = _coef_at(self.overall_model_, self.overall_lmda_idx_)
        self.overall_coef_ = coef_full[n_onehot:]  # (p,)
        self.overall_intercept_ = float(self.overall_model_.intercepts[self.overall_lmda_idx_])

    # ----------------------------------------------------------
    # Compute OOF overall linear predictor (matches R's keep=TRUE).
    # Using in-sample predictions as the offset introduces optimism:
    # the overall model has already seen the training samples, so its
    # predictions are sharper than true held-out predictions.  OOF
    # predictions avoid this leakage.
    # ----------------------------------------------------------
    preval_offset = self._compute_oof_eta(X_overall, y, overall_pf, groups)
    self._preval_offset = preval_offset  # cached for PretrainedLassoCV reuse

    # ----------------------------------------------------------
    # Step 2: per-group models
    # ----------------------------------------------------------
    # Penalty factor: features NOT in the overall X-support get penalty
    # 1/alpha, steering the group model to prefer features already
    # selected overall.  Matches R's fac = rep(1/alpha, p); fac[supall] = 1.
    if self.family == "multinomial":
        overall_support = np.where(np.any(self.overall_coef_ != 0, axis=1))[0]
    else:
        overall_support = np.where(self.overall_coef_ != 0)[0]

    _alpha_pf = self.alpha if self.alpha > 0 else 1e-9
    pf = np.full(self.n_features_in_, 1.0 / _alpha_pf)
    pf[overall_support] = 1.0

    if self.verbose:
        if self.family == "multinomial":
            n_ov = int(np.sum(np.any(self.overall_coef_ != 0, axis=1)))
        else:
            n_ov = int(np.sum(self.overall_coef_ != 0))
        _logger.info(f"  Overall model done  |S|={n_ov}")

    self.pretrain_models_ = {}
    self.pretrain_lmda_idx_ = {}
    self.individual_models_ = {}
    self.individual_lmda_idx_ = {}

    if self.verbose:
        try:
            from tqdm import tqdm as _tqdm

            group_iter = _tqdm(
                self.groups_, desc="  [2/2] Group models", unit="group", leave=True
            )
        except ImportError:
            _logger.info("  [2/2] Group models")
            group_iter = self.groups_
    else:
        group_iter = self.groups_

    for g in group_iter:
        mask = groups == g
        X_g = np.asfortranarray(X[mask])
        glm_g = _make_glm(self.family, y[mask])
        # OOF offset: each sample's overall prediction came from a model
        # trained without that sample — mirrors R's use of fit.preval.
        offset = np.asfortranarray((1 - self.alpha) * preval_offset[mask])

        # Mirrors R's cv.glmnet for per-group models: select lambda by CV
        # (lambda.min), then refit on the full group data at that lambda.
        with _silence():
            cv_pre = ad.cv_grpnet(
                self._wrap_matrix(X_g),
                glm_g,
                offsets=offset,
                penalty=pf,
                **self._cv_grpnet_kwargs(),
            )
            self.pretrain_lmda_idx_[g] = cv_pre.best_idx
            self.pretrain_models_[g] = cv_pre.fit(
                self._wrap_matrix(X_g),
                glm_g,
                offsets=offset,
                penalty=pf,
                **self._grpnet_kwargs(),
            )

            cv_ind = ad.cv_grpnet(self._wrap_matrix(X_g), glm_g, **self._cv_grpnet_kwargs())
            self.individual_lmda_idx_[g] = cv_ind.best_idx
            self.individual_models_[g] = cv_ind.fit(
                self._wrap_matrix(X_g), glm_g, **self._grpnet_kwargs()
            )

    if self.verbose:
        self._print_fit_summary(time.time() - t0)

    return self

predict

predict(X, groups, model='pretrain', type='response', lmda_idx=None)

Predict target values.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
groups	array-like of shape (n_samples,)		required
model	(pretrain, individual, overall)		"pretrain"
type	(response, link, 'class')	Scale of the returned predictions. "response" — fitted values on the response scale: probabilities for binomial/multinomial, fitted values for gaussian. "link" — raw linear predictor before the link function (η = Xβ + intercept). For gaussian this is identical to "response". "class" — predicted class label: 0 or 1 for binomial (threshold 0.5), integer argmax for multinomial. Not valid for gaussian.	"response"
lmda_idx	int or None	Lambda index for group models. None uses the CV-selected lambda for each group (lambda.min, matching R).	None

Returns:

Name	Type	Description
y_pred	ndarray	Shape (n,) for gaussian/binomial and for multinomial "class". Shape (n, K) for multinomial "response" or "link".

Source code in src/ptlasso/_ptlasso.py

def predict(self, X, groups, model="pretrain", type="response", lmda_idx=None):
    """Predict target values.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    groups : array-like of shape (n_samples,)
    model : {"pretrain", "individual", "overall"}, default="pretrain"
    type : {"response", "link", "class"}, default="response"
        Scale of the returned predictions.

        ``"response"`` — fitted values on the response scale: probabilities
        for binomial/multinomial, fitted values for gaussian.

        ``"link"`` — raw linear predictor before the link function
        (``η = Xβ + intercept``).  For gaussian this is identical to
        ``"response"``.

        ``"class"`` — predicted class label: ``0`` or ``1`` for binomial
        (threshold 0.5), integer argmax for multinomial.  Not valid for
        gaussian.
    lmda_idx : int or None
        Lambda index for group models.  ``None`` uses the CV-selected
        lambda for each group (``lambda.min``, matching R).

    Returns
    -------
    y_pred : ndarray
        Shape ``(n,)`` for gaussian/binomial and for multinomial
        ``"class"``.  Shape ``(n, K)`` for multinomial ``"response"``
        or ``"link"``.
    """
    check_is_fitted(self)
    X = check_array(X, dtype=np.float64, order="F")
    X = np.asfortranarray(self.scaler_.transform(X))
    groups = np.asarray(groups)

    if model not in PREDICT_MODELS:
        raise ValueError(f"model must be one of {PREDICT_MODELS}, got '{model}'")
    if type not in PREDICT_TYPES:
        raise ValueError(f"type must be one of {PREDICT_TYPES}, got '{type}'")
    if type == "class" and self.family == "gaussian":
        raise ValueError("type='class' is not valid for family='gaussian'")
    if len(groups) != X.shape[0]:
        raise ValueError(f"groups has {len(groups)} elements but X has {X.shape[0]} samples")
    unknown = set(np.unique(groups)) - set(self.groups_)
    if unknown:
        raise ValueError(f"predict received groups not seen during fit: {unknown}")

    if model == "overall":
        return _eta_to_output(self._overall_eta(X, groups), self.family, type)

    n_out = (X.shape[0], self.n_classes_) if self.family == "multinomial" else (X.shape[0],)
    eta_out = np.empty(n_out)

    # Overall eta is only needed when combining with a group model (pretrain).
    eta_ov = self._overall_eta(X, groups) if model == "pretrain" else None

    for g in self.groups_:
        mask = groups == g
        X_g = X[mask]

        if model == "pretrain":
            g_idx = self.pretrain_lmda_idx_.get(g, -1) if lmda_idx is None else lmda_idx
            eta_group = _eta_from_state(
                self.pretrain_models_[g], X_g, g_idx, self.family, self.n_classes_
            )
            eta_out[mask] = (1 - self.alpha) * eta_ov[mask] + eta_group
        else:
            g_idx = self.individual_lmda_idx_.get(g, -1) if lmda_idx is None else lmda_idx
            eta_out[mask] = _eta_from_state(
                self.individual_models_[g], X_g, g_idx, self.family, self.n_classes_
            )

    return _eta_to_output(eta_out, self.family, type)

score

score(X, y, groups)

Return a scalar performance metric using the pretrained model.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
y	array-like of shape (n_samples,)		required
groups	array-like of shape (n_samples,)		required

Returns:

Name	Type	Description
score	float	R² for gaussian; classification accuracy for binomial/multinomial.

Source code in src/ptlasso/_ptlasso.py

def score(self, X, y, groups):
    """Return a scalar performance metric using the pretrained model.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    y : array-like of shape (n_samples,)
    groups : array-like of shape (n_samples,)

    Returns
    -------
    score : float
        R² for gaussian; classification accuracy for binomial/multinomial.
    """
    return _model_score(y, self.predict(X, groups), self.family)

evaluate

evaluate(X, y, groups)

Predict and score with all three sub-models.

Convenience method matching R's predict(fit, xtest, ytest=ytest).

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
y	array-like of shape (n_samples,)		required
groups	array-like of shape (n_samples,)		required

Returns:

Name	Type	Description
result	dict	Keys are "pretrain", "individual", "overall". Each value is a dict with entries: "predictions" : ndarray Predicted values, shape (n,) or (n, K) for multinomial. "score" : float R² for gaussian; accuracy for binomial/multinomial.

Source code in src/ptlasso/_ptlasso.py

def evaluate(self, X, y, groups):
    """Predict and score with all three sub-models.

    Convenience method matching R's ``predict(fit, xtest, ytest=ytest)``.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    y : array-like of shape (n_samples,)
    groups : array-like of shape (n_samples,)

    Returns
    -------
    result : dict
        Keys are ``"pretrain"``, ``"individual"``, ``"overall"``.
        Each value is a dict with entries:

        ``"predictions"`` : ndarray
            Predicted values, shape ``(n,)`` or ``(n, K)`` for multinomial.
        ``"score"`` : float
            R² for gaussian; accuracy for binomial/multinomial.
    """
    check_is_fitted(self)
    result = {}
    for m in ("pretrain", "individual", "overall"):
        preds = self.predict(X, groups, model=m)
        result[m] = {"predictions": preds, "score": _model_score(y, preds, self.family)}
    return result

get_coef

get_coef(model='all', lmda_idx=None)

Return fitted coefficients as a nested dict.

Parameters:

Name	Type	Description	Default
model	(all, overall, pretrain, individual)	Which sub-model(s) to return.	"all"
lmda_idx	int or None	Lambda index for group models. None uses the last lambda (-1).	None

Returns:

Name	Type	Description
coefs	dict	When model="all", keys are "overall", "pretrain", "individual". For "overall", value is {"coef": ndarray, "intercept": ndarray}. For "pretrain" and "individual", value is a dict keyed by group label, each containing {"coef": ndarray, "intercept": ndarray}. When a specific model is requested, only that sub-dict is returned.

Source code in src/ptlasso/_ptlasso.py

def get_coef(self, model="all", lmda_idx=None):
    """Return fitted coefficients as a nested dict.

    Parameters
    ----------
    model : {"all", "overall", "pretrain", "individual"}, default="all"
        Which sub-model(s) to return.
    lmda_idx : int or None
        Lambda index for group models.  ``None`` uses the last lambda (``-1``).

    Returns
    -------
    coefs : dict
        When ``model="all"``, keys are ``"overall"``, ``"pretrain"``,
        ``"individual"``.  For ``"overall"``, value is
        ``{"coef": ndarray, "intercept": ndarray}``.  For ``"pretrain"``
        and ``"individual"``, value is a dict keyed by group label, each
        containing ``{"coef": ndarray, "intercept": ndarray}``.
        When a specific model is requested, only that sub-dict is returned.
    """
    if model not in COEF_MODELS:
        raise ValueError(f"model must be one of {COEF_MODELS}, got '{model}'")
    check_is_fitted(self)

    def _group_coefs(models, lmda_idxs):
        result = {}
        for g, state in models.items():
            use_idx = lmda_idxs.get(g, -1) if lmda_idx is None else lmda_idx
            result[self._label(g)] = {
                "coef": _coef_at(state, use_idx),
                "intercept": np.asarray(state.intercepts[use_idx]).ravel(),
            }
        return result

    result = {}
    if model in ("all", "overall"):
        result["overall"] = {
            "coef": self.overall_coef_,
            "intercept": np.asarray(
                self.overall_model_.intercepts[self.overall_lmda_idx_]
            ).ravel(),
        }
    if model in ("all", "pretrain"):
        result["pretrain"] = _group_coefs(self.pretrain_models_, self.pretrain_lmda_idx_)
    if model in ("all", "individual"):
        result["individual"] = _group_coefs(self.individual_models_, self.individual_lmda_idx_)
    return result if model == "all" else result[model]

---

PretrainedLassoCV

ptlasso.PretrainedLassoCV

Bases: RegressorMixin, BasePretrainedLasso

Pretrained Lasso with cross-validation over alpha.

Parameters:

Name	Type	Description	Default
alphas	array - like or None	Candidate pretraining strengths. Defaults to [0, 0.25, 0.5, 0.75, 1.0]. 0 = maximum pretraining; 1 = no pretraining (individual models).	None
cv	int	Number of CV folds.	5
alphahat_choice	(overall, mean)	"overall" minimises the global CV error; "mean" minimises the unweighted mean of per-group CV errors.	"overall"
family	(gaussian, binomial, multinomial)	Response distribution.	"gaussian"
overall_lambda	('lambda.1se', 'lambda.min')	Lambda selection rule for the stage-1 overall model.	"lambda.1se"
fit_intercept	bool	Whether to fit an intercept in every sub-model.	True
lmda_path_size	int	Number of lambdas in the regularisation path.	100
min_ratio	float	Ratio of the smallest to largest lambda on the path.	0.0001
verbose	bool	Whether to display fitting progress and a summary after training. Adelie's internal output is always suppressed regardless of this setting.	True
foldid	array-like of int or None	Fold assignments, one integer per sample. When provided, overrides the internal StratifiedKFold splitter.	None
n_threads	int	Number of threads passed to adelie's solver. Set to a higher value to parallelise the coordinate descent within each model fit. -1 uses all available CPU cores (os.cpu_count()).	-1
scoring	str, callable, or None	CV criterion used to select the best alpha. All values follow the higher = better convention (matching sklearn). None — family-based loss: MSE for gaussian, log-loss for binomial/multinomial. Matches the current default behaviour. "roc_auc" — area under the ROC curve (binomial only). Matches R's type.measure="auc". "neg_log_loss" — negative log-loss (binomial/multinomial). "accuracy" — classification accuracy (binomial/multinomial). "neg_mean_squared_error" — negative MSE (gaussian). "r2" — coefficient of determination (gaussian). callable — any function (y_true, y_pred) -> float where higher is better, e.g. a custom sklearn scorer. cv_results_ always stores values as losses (lower = better), so AUC will appear as -AUC.	None

Attributes:

Name	Type	Description
alpha_	float	Best alpha selected by CV (based on alphahat_choice).
varying_alphahat_	dict {group -> float}	Per-group best alpha.
cv_results_	dict {alpha -> float}	Global mean CV loss per alpha.
cv_results_se_	dict {alpha -> float}	Standard error of global CV loss.
cv_results_per_group_	dict {alpha -> {group -> float}}	Mean CV loss per alpha per group.
cv_results_mean_	dict {alpha -> float}	Unweighted mean of per-group CV losses.
cv_results_wtd_mean_	dict {alpha -> float}	Group-size-weighted mean of per-group CV losses.
cv_results_individual_	float	Global CV loss for the individual (no-pretraining) model.
cv_results_overall_	float	Global CV loss for the overall model.
best_estimator_	PretrainedLasso	Full-data refit at the globally selected alpha_.
all_estimators_	dict {alpha -> PretrainedLasso}	Full-data refits for each unique alpha needed by varying_alphahat_.

References

Craig, E., Pilanci, M., Le Menestrel, T., Narasimhan, B., Rivas, M. A., Gullaksen, S. E., & Tibshirani, R. (2025). Pretraining and the lasso. Journal of the Royal Statistical Society Series B, qkaf050.

Source code in src/ptlasso/_ptlasso.py

class PretrainedLassoCV(RegressorMixin, BasePretrainedLasso):
    """Pretrained Lasso with cross-validation over alpha.

    Parameters
    ----------
    alphas : array-like or None, default=None
        Candidate pretraining strengths.  Defaults to [0, 0.25, 0.5, 0.75, 1.0].
        ``0`` = maximum pretraining; ``1`` = no pretraining (individual models).
    cv : int, default=5
        Number of CV folds.
    alphahat_choice : {"overall", "mean"}, default="overall"
        ``"overall"`` minimises the global CV error; ``"mean"`` minimises the
        unweighted mean of per-group CV errors.
    family : {"gaussian", "binomial", "multinomial"}, default="gaussian"
        Response distribution.
    overall_lambda : {"lambda.1se", "lambda.min"}, default="lambda.1se"
        Lambda selection rule for the stage-1 overall model.
    fit_intercept : bool, default=True
        Whether to fit an intercept in every sub-model.
    lmda_path_size : int, default=100
        Number of lambdas in the regularisation path.
    min_ratio : float, default=0.0001
        Ratio of the smallest to largest lambda on the path.
    verbose : bool, default=True
        Whether to display fitting progress and a summary after training.
        Adelie's internal output is always suppressed regardless of this setting.
    foldid : array-like of int or None, default=None
        Fold assignments, one integer per sample.  When provided, overrides
        the internal ``StratifiedKFold`` splitter.
    n_threads : int, default=-1
        Number of threads passed to adelie's solver.  Set to a higher value
        to parallelise the coordinate descent within each model fit.
        ``-1`` uses all available CPU cores (``os.cpu_count()``).
    scoring : str, callable, or None, default=None
        CV criterion used to select the best alpha.  All values follow the
        **higher = better** convention (matching sklearn).

        - ``None`` — family-based loss: MSE for gaussian, log-loss for
          binomial/multinomial.  Matches the current default behaviour.
        - ``"roc_auc"`` — area under the ROC curve (binomial only).
          Matches R's ``type.measure="auc"``.
        - ``"neg_log_loss"`` — negative log-loss (binomial/multinomial).
        - ``"accuracy"`` — classification accuracy (binomial/multinomial).
        - ``"neg_mean_squared_error"`` — negative MSE (gaussian).
        - ``"r2"`` — coefficient of determination (gaussian).
        - callable — any function ``(y_true, y_pred) -> float`` where
          **higher is better**, e.g. a custom sklearn scorer.

        ``cv_results_`` always stores values as *losses* (lower = better),
        so AUC will appear as ``-AUC``.

    Attributes
    ----------
    alpha_ : float
        Best alpha selected by CV (based on ``alphahat_choice``).
    varying_alphahat_ : dict {group -> float}
        Per-group best alpha.
    cv_results_ : dict {alpha -> float}
        Global mean CV loss per alpha.
    cv_results_se_ : dict {alpha -> float}
        Standard error of global CV loss.
    cv_results_per_group_ : dict {alpha -> {group -> float}}
        Mean CV loss per alpha per group.
    cv_results_mean_ : dict {alpha -> float}
        Unweighted mean of per-group CV losses.
    cv_results_wtd_mean_ : dict {alpha -> float}
        Group-size-weighted mean of per-group CV losses.
    cv_results_individual_ : float
        Global CV loss for the individual (no-pretraining) model.
    cv_results_overall_ : float
        Global CV loss for the overall model.
    best_estimator_ : PretrainedLasso
        Full-data refit at the globally selected ``alpha_``.
    all_estimators_ : dict {alpha -> PretrainedLasso}
        Full-data refits for each unique alpha needed by ``varying_alphahat_``.

    References
    ----------
    Craig, E., Pilanci, M., Le Menestrel, T., Narasimhan, B., Rivas, M. A.,
    Gullaksen, S. E., & Tibshirani, R. (2025). Pretraining and the lasso.
    *Journal of the Royal Statistical Society Series B*, qkaf050.
    """

    def __init__(
        self,
        alphas=DEFAULT_ALPHAS,
        cv=5,
        alphahat_choice="overall",
        family="gaussian",
        overall_lambda="lambda.1se",
        fit_intercept=True,
        lmda_path_size=100,
        min_ratio=0.0001,
        verbose=True,
        foldid=None,
        scoring=None,
        n_threads=-1,
        standardize=True,
        n_folds=10,
    ):
        self.alphas = alphas
        self.cv = cv
        self.alphahat_choice = alphahat_choice
        self.family = family
        self.overall_lambda = overall_lambda
        self.fit_intercept = fit_intercept
        self.lmda_path_size = lmda_path_size
        self.min_ratio = min_ratio
        self.verbose = verbose
        self.foldid = foldid
        self.scoring = scoring
        self.n_threads = n_threads
        self.standardize = standardize
        self.n_folds = n_folds

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _validate_params(self):
        alphas = self.alphas
        bad = [a for a in alphas if not (0.0 <= a <= 1.0)]
        if bad:
            raise ValueError(f"All alphas must be in [0, 1], got {bad}")
        if len(set(alphas)) != len(alphas):
            raise ValueError(f"alphas must not contain duplicates, got {alphas}")
        if not isinstance(self.cv, int) or self.cv < 2:
            raise ValueError(f"cv must be an integer >= 2, got {self.cv}")
        if self.alphahat_choice not in ("overall", "mean"):
            raise ValueError(
                f"alphahat_choice must be 'overall' or 'mean', got '{self.alphahat_choice}'"
            )
        if self.family not in FAMILIES:
            raise ValueError(f"family must be one of {FAMILIES}, got '{self.family}'")
        if self.overall_lambda not in LMDA_MODES:
            raise ValueError(f"overall_lambda must be one of {LMDA_MODES}")
        if not isinstance(self.lmda_path_size, int) or self.lmda_path_size < 1:
            raise ValueError("lmda_path_size must be a positive integer")
        if not (0.0 < self.min_ratio < 1.0):
            raise ValueError("min_ratio must be in (0, 1)")
        if self.scoring is not None and not callable(self.scoring):
            if self.scoring not in _VALID_SCORERS:
                raise ValueError(
                    f"scoring must be None, a callable, or one of {_VALID_SCORERS}, "
                    f"got '{self.scoring}'"
                )

    def _get_alphas(self):
        """Return the list of alpha candidates."""
        return list(self.alphas if self.alphas is not None else DEFAULT_ALPHAS)

    def _base_estimator(self, alpha):
        """Return a configured PretrainedLasso for a given alpha."""
        return PretrainedLasso(
            alpha=alpha,
            family=self.family,
            overall_lambda=self.overall_lambda,
            fit_intercept=self.fit_intercept,
            lmda_path_size=self.lmda_path_size,
            min_ratio=self.min_ratio,
            verbose=False,  # CV sub-fits are silent; PretrainedLassoCV owns the progress bar
            n_threads=self.n_threads,
            standardize=self.standardize,
            n_folds=self.n_folds,
        )

    def _fold_iter(self, X, groups):
        """Yield (train_idx, test_idx) pairs for CV."""
        if self.foldid is not None:
            foldid = np.asarray(self.foldid)
            if len(foldid) != len(groups):
                raise ValueError(
                    f"foldid has {len(foldid)} elements but X has {len(groups)} samples"
                )
            if len(np.unique(foldid)) < 2:
                raise ValueError("foldid must contain at least 2 distinct fold values")
            for f in np.unique(foldid):
                yield np.where(foldid != f)[0], np.where(foldid == f)[0]
        else:
            splitter = StratifiedKFold(n_splits=self.cv, shuffle=True, random_state=0)
            yield from splitter.split(X, groups)

    # ------------------------------------------------------------------
    # Progress / display helpers
    # ------------------------------------------------------------------

    def _print_cv_summary(self, elapsed):
        SEP = "─" * 54
        rows = [SEP, f"    {'α':<8} {'CV loss':<12} {'±SE'}", f"  {'─' * 34}"]
        for a in self.alphalist_:
            marker = "►" if a == self.alpha_ else " "
            rows.append(
                f" {marker}  {a:<8.2f} {self.cv_results_[a]:<12.4f} {self.cv_results_se_[a]:.4f}"
            )
        rows += [
            f"  {'─' * 34}",
            f"  {'individual':<13}{self.cv_results_individual_:.4f}",
            f"  {'overall':<13}{self.cv_results_overall_:.4f}",
            SEP,
        ]
        _logger.info("\n".join(rows))
        best = self.best_estimator_
        stage2 = set()
        for g in self.groups_:
            c = _coef_at(best.pretrain_models_[g], best.pretrain_lmda_idx_[g])
            if best.n_classes_ is not None:
                c = c.reshape(best.n_features_in_, best.n_classes_, order="F")
                stage2 |= set(int(i) for i in np.where(np.any(c != 0, axis=1))[0])
            else:
                stage2 |= set(int(i) for i in np.where(c != 0)[0])
        _logger.info(
            f"  Best α = {self.alpha_:.2f}   |S| = {len(stage2)}"
            f"   Fitted in {elapsed:.1f}s"
            f"  ({len(self.alphalist_)} alphas · OOF-based)"
        )

    # ------------------------------------------------------------------
    # fit
    # ------------------------------------------------------------------

    def fit(self, X, y, groups, group_labels=None, feature_names=None):
        """Fit PretrainedLassoCV.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target values.  For ``family="binomial"``, must contain only 0
            and 1.  For ``family="multinomial"``, must contain non-negative
            integer class labels 0..K-1.
        groups : array-like of shape (n_samples,)
            Group membership for each sample.  Must contain at least two
            distinct values.
        group_labels : dict or None, default=None
            Optional mapping from group values to display names used in
            ``__repr__`` and ``get_coef()``.
        feature_names : array-like of str or None, default=None
            Feature names.  Inferred from ``X.columns`` when ``X`` is a
            DataFrame and this argument is ``None``.

        Returns
        -------
        self : PretrainedLassoCV
            Fitted estimator.
        """
        t0 = time.time()
        self._validate_params()

        if feature_names is None and hasattr(X, "columns"):
            feature_names = list(X.columns)

        X, y = check_X_y(X, y, dtype=np.float64, order="F")
        groups = np.asarray(groups)

        self.n_features_in_ = X.shape[1]
        self.group_labels_ = group_labels or {}

        alphas = self._get_alphas()
        self.alphalist_ = np.asarray(alphas)

        unique_groups = np.unique(groups)
        group_sizes = {g: int(np.sum(groups == g)) for g in unique_groups}

        scorer_fn = _resolve_scorer(self.scoring)

        if self.verbose:
            _enable_verbose_logging()
            _logger.info(
                f"\nPretrainedLassoCV  {self.family}  ·  {len(unique_groups)} groups"
                f"  ·  {self.n_features_in_} features"
                f"  ·  {len(alphas)} alphas · OOF-based (all data)"
            )
            from tqdm import tqdm as _tqdm

        # --- R-style architecture: train all models on full data, evaluate
        #     via within-model OOF predictions (mirrors cv.glmnet keep=TRUE).
        #     No data is withheld from model fitting. ---

        if self.verbose:
            _logger.info("\n── Full-data fit (all samples) " + "─" * 26)
            _logger.info("  Overall model (λ-path):")

        # Step 1: fit stage-1 (overall model + OOF) once on ALL data.
        _stage1 = self._base_estimator(alphas[0])
        _stage1._show_overall_progress = self.verbose
        _stage1.fit(X, y, groups, group_labels=group_labels, feature_names=feature_names)
        _preval_offset = _stage1._preval_offset  # (n,) OOF overall linear predictor

        if self.verbose:
            _n_support = (
                int(np.sum(_stage1.overall_coef_ != 0))
                if self.family != "multinomial"
                else int(np.sum(np.any(_stage1.overall_coef_ != 0, axis=1)))
            )
            _logger.info(f"  Overall model done  |S|={_n_support}")

        # Standardized X used by all group models.
        X_std = np.asfortranarray(_stage1.scaler_.transform(np.asarray(X, dtype=np.float64)))

        # Per-alpha OOF accumulators (all n samples, no fold averaging).
        oof_losses = {}
        oof_losses_grp = {a: {} for a in alphas}
        all_estimators_full = {}

        # Pre-compute individual OOF once — doesn't depend on alpha.
        if self.family == "multinomial":
            _oof_ind_eta = np.zeros((len(y), _stage1.n_classes_))
        else:
            _oof_ind_eta = np.zeros(len(y))
        for _g in unique_groups:
            _mask = groups == _g
            _lhat_ind = float(
                _stage1.individual_models_[_g].lmdas[_stage1.individual_lmda_idx_[_g]]
            )
            _oof_ind_eta[_mask] = _stage1._compute_group_oof_eta(
                np.asfortranarray(X_std[_mask]), y[_mask], None, None, _lhat_ind
            )
        self.cv_results_individual_ = _cv_loss(
            scorer_fn, y, _apply_link(_oof_ind_eta, self.family), self.family
        )

        # Overall OOF score comes directly from _preval_offset.
        self.cv_results_overall_ = _cv_loss(
            scorer_fn, y, _apply_link(_preval_offset, self.family), self.family
        )

        _pbar = None
        if self.verbose:
            _pbar = _tqdm(total=len(alphas), desc="  Group fits", unit="α", leave=True)
            _pbar.refresh()

        for i, a in enumerate(alphas):
            # Step 2: fit stage-2 group models for this alpha on ALL data.
            if i == 0:
                est = _stage1  # already fully fitted
            else:
                est = self._base_estimator(a)
                est.n_features_in_ = _stage1.n_features_in_
                est.groups_ = _stage1.groups_
                est.group_labels_ = _stage1.group_labels_
                est.n_classes_ = _stage1.n_classes_
                est.feature_names_in_ = _stage1.feature_names_in_
                est._n_onehot_ = _stage1._n_onehot_
                est.scaler_ = _stage1.scaler_
                est.overall_model_ = _stage1.overall_model_
                est.overall_lmda_idx_ = _stage1.overall_lmda_idx_
                est.overall_coef_ = _stage1.overall_coef_
                est._n_folds_oof_ = _stage1._n_folds_oof_
                if self.family != "multinomial":
                    est.overall_intercept_ = _stage1.overall_intercept_
                est._fit_groups_only(X, y, groups, _preval_offset)
            all_estimators_full[a] = est

            # Step 3: compute OOF pretrain predictions on all n samples.
            # Penalty factor matches _fit_groups_only.
            if self.family == "multinomial":
                _ovsupp = np.where(np.any(est.overall_coef_ != 0, axis=1))[0]
            else:
                _ovsupp = np.where(est.overall_coef_ != 0)[0]
            _apf = a if a > 0 else 1e-9
            _pf = np.full(est.n_features_in_, 1.0 / _apf)
            _pf[_ovsupp] = 1.0

            if self.family == "multinomial":
                oof_pre_eta = np.zeros((len(y), est.n_classes_))
            else:
                oof_pre_eta = np.zeros(len(y))

            for g in unique_groups:
                mask = groups == g
                X_g = np.asfortranarray(X_std[mask])
                y_g = y[mask]
                g_offset = np.asfortranarray((1.0 - a) * _preval_offset[mask])
                lamhat_g = float(est.pretrain_models_[g].lmdas[est.pretrain_lmda_idx_[g]])
                # Group model OOF (without offset) + overall offset = full OOF pretrain eta.
                oof_g = est._compute_group_oof_eta(X_g, y_g, g_offset, _pf, lamhat_g)
                oof_pre_eta[mask] = oof_g + (1.0 - a) * _preval_offset[mask]

                y_pred_g = _apply_link(oof_pre_eta[mask], self.family)
                oof_losses_grp[a][g] = _cv_loss(scorer_fn, y[mask], y_pred_g, self.family)

            y_pred_pre = _apply_link(oof_pre_eta, self.family)
            oof_losses[a] = _cv_loss(scorer_fn, y, y_pred_pre, self.family)

            if _pbar is not None:
                _pbar.update(1)
                _pbar.set_postfix(alpha=f"{a:.2f}", refresh=False)

        if _pbar is not None:
            _pbar.close()

        # Aggregate OOF results.
        self.cv_results_ = oof_losses
        self.cv_results_se_ = {a: 0.0 for a in alphas}  # no SE without fold averaging
        self.cv_results_per_group_ = oof_losses_grp
        self.cv_results_mean_ = {
            a: float(np.nanmean(list(oof_losses_grp[a].values()))) for a in alphas
        }
        self.cv_results_wtd_mean_ = {
            a: float(
                np.nansum(
                    [
                        oof_losses_grp[a][g] * group_sizes[g] / len(y)
                        for g in unique_groups
                        if not np.isnan(oof_losses_grp[a].get(g, float("nan")))
                    ]
                )
            )
            for a in alphas
        }

        # Select global best alpha.
        criterion = self.cv_results_mean_ if self.alphahat_choice == "mean" else self.cv_results_
        self.alpha_ = min(alphas, key=criterion.__getitem__)

        # Per-group best alpha.
        self.varying_alphahat_ = {
            g: min(alphas, key=lambda a, _g=g: self.cv_results_per_group_[a].get(_g, np.inf))
            for g in unique_groups
        }

        # All models were already fitted on full data above — no refit needed.
        unique_alphas = set(self.varying_alphahat_.values()) | {self.alpha_}
        fit_kwargs = dict(group_labels=group_labels, feature_names=feature_names)
        self.all_estimators_ = {}
        for a in unique_alphas:
            if a in all_estimators_full:
                self.all_estimators_[a] = all_estimators_full[a]
            else:
                # Shouldn't happen: all alphalist_ alphas were fitted above.
                self.all_estimators_[a] = self._base_estimator(a).fit(X, y, groups, **fit_kwargs)
        self.best_estimator_ = self.all_estimators_[self.alpha_]

        # Mirror fitted attributes from the best estimator for a uniform interface
        for attr in (
            "overall_model_",
            "overall_lmda_idx_",
            "overall_coef_",
            "pretrain_models_",
            "pretrain_lmda_idx_",
            "individual_models_",
            "individual_lmda_idx_",
            "groups_",
            "feature_names_in_",
            "n_classes_",
            "_n_onehot_",
        ):
            setattr(self, attr, getattr(self.best_estimator_, attr))
        if self.family != "multinomial":
            self.overall_intercept_ = self.best_estimator_.overall_intercept_

        if self.verbose:
            self._print_cv_summary(time.time() - t0)

        return self

    # ------------------------------------------------------------------
    # predict / score / evaluate
    # ------------------------------------------------------------------

    @property
    def alpha(self):
        """Selected alpha — mirrors PretrainedLasso.alpha for a uniform interface."""
        check_is_fitted(self, ["alpha_"])
        return self.alpha_

    def predict(
        self, X, groups, model="pretrain", type="response", alphatype="best", lmda_idx=None
    ):
        """Predict target values.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        groups : array-like of shape (n_samples,)
        model : {"pretrain", "individual", "overall"}, default="pretrain"
        type : {"response", "link", "class"}, default="response"
            Scale of the returned predictions.  See :meth:`PretrainedLasso.predict`
            for full documentation.
        alphatype : {"best", "varying"}, default="best"
            ``"best"`` uses the globally selected ``alpha_``; ``"varying"``
            uses each group's own optimal alpha from ``varying_alphahat_``.
        lmda_idx : int or None

        Returns
        -------
        y_pred : ndarray of shape (n_samples,) or (n_samples, K)
            Shape is ``(n_samples, K)`` for multinomial ``"response"`` or
            ``"link"``.
        """
        check_is_fitted(self)
        groups = np.asarray(groups)

        if model not in PREDICT_MODELS:
            raise ValueError(f"model must be one of {PREDICT_MODELS}, got '{model}'")
        if type not in PREDICT_TYPES:
            raise ValueError(f"type must be one of {PREDICT_TYPES}, got '{type}'")
        if alphatype not in ALPHATYPES:
            raise ValueError(f"alphatype must be one of {ALPHATYPES}, got '{alphatype}'")

        if alphatype == "varying":
            X = check_array(X, dtype=np.float64, order="F")
            n_out = (X.shape[0], self.n_classes_) if self.family == "multinomial" else (X.shape[0],)
            y_pred = np.empty(n_out)
            for g in np.unique(groups):
                mask = groups == g
                a = self.varying_alphahat_.get(g, self.alpha_)
                y_pred[mask] = self.all_estimators_[a].predict(
                    X[mask], groups[mask], model=model, type=type, lmda_idx=lmda_idx
                )
            return y_pred

        return self.best_estimator_.predict(
            X=X, groups=groups, model=model, type=type, lmda_idx=lmda_idx
        )

    def evaluate(self, X, y, groups, alphatype="best"):
        """Predict and score with all three sub-models.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        y : array-like of shape (n_samples,)
        groups : array-like of shape (n_samples,)
        alphatype : {"best", "varying"}, default="best"
            ``"best"`` uses ``alpha_``; ``"varying"`` uses each group's own
            optimal alpha from ``varying_alphahat_``.

        Returns
        -------
        result : dict
            Keys are ``"pretrain"``, ``"individual"``, ``"overall"``.
            Each value is a dict with entries:

            ``"predictions"`` : ndarray
                Predicted values.
            ``"score"`` : float
                R² for gaussian; accuracy for binomial/multinomial.
        """
        check_is_fitted(self)
        result = {}
        for m in ("pretrain", "individual", "overall"):
            preds = self.predict(X, groups, model=m, alphatype=alphatype)
            result[m] = {"predictions": preds, "score": _model_score(y, preds, self.family)}
        return result

    def score(self, X, y, groups):
        """Return a scalar performance metric using the best estimator.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        y : array-like of shape (n_samples,)
        groups : array-like of shape (n_samples,)

        Returns
        -------
        score : float
            R² for gaussian; classification accuracy for binomial/multinomial.
        """
        check_is_fitted(self)
        return self.best_estimator_.score(X, y, groups)

    # ------------------------------------------------------------------
    # repr / get_coef
    # ------------------------------------------------------------------

    def __repr__(self):
        alphas = self._get_alphas()
        header = (
            f"PretrainedLassoCV(alphas={alphas}, cv={self.cv}, "
            f"alphahat_choice='{self.alphahat_choice}', "
            f"family='{self.family}', overall_lambda='{self.overall_lambda}')"
        )
        if not hasattr(self, "best_estimator_"):
            return header + "\n  [not fitted]"

        rows = [
            f"  alpha={a:.2f}  overall={self.cv_results_[a]:.4f}"
            f"  mean={self.cv_results_mean_[a]:.4f}"
            f"  wtd_mean={self.cv_results_wtd_mean_[a]:.4f}"
            for a in self.alphalist_
        ]
        return (
            f"{header}\n"
            f"  family     : {self.family}\n"
            f"  n_features : {self.n_features_in_}\n"
            f"  n_groups   : {len(self.groups_)}\n"
            f"  alpha_     : {self.alpha_}  ({self.alphahat_choice})\n"
            f"  individual : {self.cv_results_individual_:.4f}\n"
            f"  overall    : {self.cv_results_overall_:.4f}\n"
            f"  CV results :\n" + "\n".join(rows)
        )

    # ------------------------------------------------------------------
    # Serialisation
    # ------------------------------------------------------------------

    def __getstate__(self):
        d = self.__dict__.copy()
        # Convert the adelie state refs mirrored directly onto this object.
        # all_estimators_ / best_estimator_ are PretrainedLasso instances and
        # go through PretrainedLasso.__getstate__ automatically.
        _proxify_models(d)
        return d

    # ------------------------------------------------------------------
    # get_coef
    # ------------------------------------------------------------------

    def get_coef(self, model="all", **kwargs):
        """Return fitted coefficients from the best estimator.

        Delegates to :meth:`PretrainedLasso.get_coef`.

        Parameters
        ----------
        model : {"all", "overall", "pretrain", "individual"}, default="all"
        **kwargs
            Forwarded to :meth:`PretrainedLasso.get_coef`.

        Returns
        -------
        coefs : dict
            See :meth:`PretrainedLasso.get_coef` for the structure.
        """
        check_is_fitted(self)
        return self.best_estimator_.get_coef(model=model, **kwargs)

alpha property

alpha

Selected alpha — mirrors PretrainedLasso.alpha for a uniform interface.

fit

fit(X, y, groups, group_labels=None, feature_names=None)

Fit PretrainedLassoCV.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)	Training data.	required
y	array-like of shape (n_samples,)	Target values. For family="binomial", must contain only 0 and 1. For family="multinomial", must contain non-negative integer class labels 0..K-1.	required
groups	array-like of shape (n_samples,)	Group membership for each sample. Must contain at least two distinct values.	required
group_labels	dict or None	Optional mapping from group values to display names used in __repr__ and get_coef().	None
feature_names	array-like of str or None	Feature names. Inferred from X.columns when X is a DataFrame and this argument is None.	None

Returns:

Name	Type	Description
self	PretrainedLassoCV	Fitted estimator.

Source code in src/ptlasso/_ptlasso.py

def fit(self, X, y, groups, group_labels=None, feature_names=None):
    """Fit PretrainedLassoCV.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
        Training data.
    y : array-like of shape (n_samples,)
        Target values.  For ``family="binomial"``, must contain only 0
        and 1.  For ``family="multinomial"``, must contain non-negative
        integer class labels 0..K-1.
    groups : array-like of shape (n_samples,)
        Group membership for each sample.  Must contain at least two
        distinct values.
    group_labels : dict or None, default=None
        Optional mapping from group values to display names used in
        ``__repr__`` and ``get_coef()``.
    feature_names : array-like of str or None, default=None
        Feature names.  Inferred from ``X.columns`` when ``X`` is a
        DataFrame and this argument is ``None``.

    Returns
    -------
    self : PretrainedLassoCV
        Fitted estimator.
    """
    t0 = time.time()
    self._validate_params()

    if feature_names is None and hasattr(X, "columns"):
        feature_names = list(X.columns)

    X, y = check_X_y(X, y, dtype=np.float64, order="F")
    groups = np.asarray(groups)

    self.n_features_in_ = X.shape[1]
    self.group_labels_ = group_labels or {}

    alphas = self._get_alphas()
    self.alphalist_ = np.asarray(alphas)

    unique_groups = np.unique(groups)
    group_sizes = {g: int(np.sum(groups == g)) for g in unique_groups}

    scorer_fn = _resolve_scorer(self.scoring)

    if self.verbose:
        _enable_verbose_logging()
        _logger.info(
            f"\nPretrainedLassoCV  {self.family}  ·  {len(unique_groups)} groups"
            f"  ·  {self.n_features_in_} features"
            f"  ·  {len(alphas)} alphas · OOF-based (all data)"
        )
        from tqdm import tqdm as _tqdm

    # --- R-style architecture: train all models on full data, evaluate
    #     via within-model OOF predictions (mirrors cv.glmnet keep=TRUE).
    #     No data is withheld from model fitting. ---

    if self.verbose:
        _logger.info("\n── Full-data fit (all samples) " + "─" * 26)
        _logger.info("  Overall model (λ-path):")

    # Step 1: fit stage-1 (overall model + OOF) once on ALL data.
    _stage1 = self._base_estimator(alphas[0])
    _stage1._show_overall_progress = self.verbose
    _stage1.fit(X, y, groups, group_labels=group_labels, feature_names=feature_names)
    _preval_offset = _stage1._preval_offset  # (n,) OOF overall linear predictor

    if self.verbose:
        _n_support = (
            int(np.sum(_stage1.overall_coef_ != 0))
            if self.family != "multinomial"
            else int(np.sum(np.any(_stage1.overall_coef_ != 0, axis=1)))
        )
        _logger.info(f"  Overall model done  |S|={_n_support}")

    # Standardized X used by all group models.
    X_std = np.asfortranarray(_stage1.scaler_.transform(np.asarray(X, dtype=np.float64)))

    # Per-alpha OOF accumulators (all n samples, no fold averaging).
    oof_losses = {}
    oof_losses_grp = {a: {} for a in alphas}
    all_estimators_full = {}

    # Pre-compute individual OOF once — doesn't depend on alpha.
    if self.family == "multinomial":
        _oof_ind_eta = np.zeros((len(y), _stage1.n_classes_))
    else:
        _oof_ind_eta = np.zeros(len(y))
    for _g in unique_groups:
        _mask = groups == _g
        _lhat_ind = float(
            _stage1.individual_models_[_g].lmdas[_stage1.individual_lmda_idx_[_g]]
        )
        _oof_ind_eta[_mask] = _stage1._compute_group_oof_eta(
            np.asfortranarray(X_std[_mask]), y[_mask], None, None, _lhat_ind
        )
    self.cv_results_individual_ = _cv_loss(
        scorer_fn, y, _apply_link(_oof_ind_eta, self.family), self.family
    )

    # Overall OOF score comes directly from _preval_offset.
    self.cv_results_overall_ = _cv_loss(
        scorer_fn, y, _apply_link(_preval_offset, self.family), self.family
    )

    _pbar = None
    if self.verbose:
        _pbar = _tqdm(total=len(alphas), desc="  Group fits", unit="α", leave=True)
        _pbar.refresh()

    for i, a in enumerate(alphas):
        # Step 2: fit stage-2 group models for this alpha on ALL data.
        if i == 0:
            est = _stage1  # already fully fitted
        else:
            est = self._base_estimator(a)
            est.n_features_in_ = _stage1.n_features_in_
            est.groups_ = _stage1.groups_
            est.group_labels_ = _stage1.group_labels_
            est.n_classes_ = _stage1.n_classes_
            est.feature_names_in_ = _stage1.feature_names_in_
            est._n_onehot_ = _stage1._n_onehot_
            est.scaler_ = _stage1.scaler_
            est.overall_model_ = _stage1.overall_model_
            est.overall_lmda_idx_ = _stage1.overall_lmda_idx_
            est.overall_coef_ = _stage1.overall_coef_
            est._n_folds_oof_ = _stage1._n_folds_oof_
            if self.family != "multinomial":
                est.overall_intercept_ = _stage1.overall_intercept_
            est._fit_groups_only(X, y, groups, _preval_offset)
        all_estimators_full[a] = est

        # Step 3: compute OOF pretrain predictions on all n samples.
        # Penalty factor matches _fit_groups_only.
        if self.family == "multinomial":
            _ovsupp = np.where(np.any(est.overall_coef_ != 0, axis=1))[0]
        else:
            _ovsupp = np.where(est.overall_coef_ != 0)[0]
        _apf = a if a > 0 else 1e-9
        _pf = np.full(est.n_features_in_, 1.0 / _apf)
        _pf[_ovsupp] = 1.0

        if self.family == "multinomial":
            oof_pre_eta = np.zeros((len(y), est.n_classes_))
        else:
            oof_pre_eta = np.zeros(len(y))

        for g in unique_groups:
            mask = groups == g
            X_g = np.asfortranarray(X_std[mask])
            y_g = y[mask]
            g_offset = np.asfortranarray((1.0 - a) * _preval_offset[mask])
            lamhat_g = float(est.pretrain_models_[g].lmdas[est.pretrain_lmda_idx_[g]])
            # Group model OOF (without offset) + overall offset = full OOF pretrain eta.
            oof_g = est._compute_group_oof_eta(X_g, y_g, g_offset, _pf, lamhat_g)
            oof_pre_eta[mask] = oof_g + (1.0 - a) * _preval_offset[mask]

            y_pred_g = _apply_link(oof_pre_eta[mask], self.family)
            oof_losses_grp[a][g] = _cv_loss(scorer_fn, y[mask], y_pred_g, self.family)

        y_pred_pre = _apply_link(oof_pre_eta, self.family)
        oof_losses[a] = _cv_loss(scorer_fn, y, y_pred_pre, self.family)

        if _pbar is not None:
            _pbar.update(1)
            _pbar.set_postfix(alpha=f"{a:.2f}", refresh=False)

    if _pbar is not None:
        _pbar.close()

    # Aggregate OOF results.
    self.cv_results_ = oof_losses
    self.cv_results_se_ = {a: 0.0 for a in alphas}  # no SE without fold averaging
    self.cv_results_per_group_ = oof_losses_grp
    self.cv_results_mean_ = {
        a: float(np.nanmean(list(oof_losses_grp[a].values()))) for a in alphas
    }
    self.cv_results_wtd_mean_ = {
        a: float(
            np.nansum(
                [
                    oof_losses_grp[a][g] * group_sizes[g] / len(y)
                    for g in unique_groups
                    if not np.isnan(oof_losses_grp[a].get(g, float("nan")))
                ]
            )
        )
        for a in alphas
    }

    # Select global best alpha.
    criterion = self.cv_results_mean_ if self.alphahat_choice == "mean" else self.cv_results_
    self.alpha_ = min(alphas, key=criterion.__getitem__)

    # Per-group best alpha.
    self.varying_alphahat_ = {
        g: min(alphas, key=lambda a, _g=g: self.cv_results_per_group_[a].get(_g, np.inf))
        for g in unique_groups
    }

    # All models were already fitted on full data above — no refit needed.
    unique_alphas = set(self.varying_alphahat_.values()) | {self.alpha_}
    fit_kwargs = dict(group_labels=group_labels, feature_names=feature_names)
    self.all_estimators_ = {}
    for a in unique_alphas:
        if a in all_estimators_full:
            self.all_estimators_[a] = all_estimators_full[a]
        else:
            # Shouldn't happen: all alphalist_ alphas were fitted above.
            self.all_estimators_[a] = self._base_estimator(a).fit(X, y, groups, **fit_kwargs)
    self.best_estimator_ = self.all_estimators_[self.alpha_]

    # Mirror fitted attributes from the best estimator for a uniform interface
    for attr in (
        "overall_model_",
        "overall_lmda_idx_",
        "overall_coef_",
        "pretrain_models_",
        "pretrain_lmda_idx_",
        "individual_models_",
        "individual_lmda_idx_",
        "groups_",
        "feature_names_in_",
        "n_classes_",
        "_n_onehot_",
    ):
        setattr(self, attr, getattr(self.best_estimator_, attr))
    if self.family != "multinomial":
        self.overall_intercept_ = self.best_estimator_.overall_intercept_

    if self.verbose:
        self._print_cv_summary(time.time() - t0)

    return self

predict

predict(X, groups, model='pretrain', type='response', alphatype='best', lmda_idx=None)

Predict target values.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
groups	array-like of shape (n_samples,)		required
model	(pretrain, individual, overall)		"pretrain"
type	(response, link, 'class')	Scale of the returned predictions. See :meth:PretrainedLasso.predict for full documentation.	"response"
alphatype	(best, varying)	"best" uses the globally selected alpha_; "varying" uses each group's own optimal alpha from varying_alphahat_.	"best"
lmda_idx	int or None		None

Returns:

Name	Type	Description
y_pred	ndarray of shape (n_samples,) or (n_samples, K)	Shape is (n_samples, K) for multinomial "response" or "link".

Source code in src/ptlasso/_ptlasso.py

def predict(
    self, X, groups, model="pretrain", type="response", alphatype="best", lmda_idx=None
):
    """Predict target values.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    groups : array-like of shape (n_samples,)
    model : {"pretrain", "individual", "overall"}, default="pretrain"
    type : {"response", "link", "class"}, default="response"
        Scale of the returned predictions.  See :meth:`PretrainedLasso.predict`
        for full documentation.
    alphatype : {"best", "varying"}, default="best"
        ``"best"`` uses the globally selected ``alpha_``; ``"varying"``
        uses each group's own optimal alpha from ``varying_alphahat_``.
    lmda_idx : int or None

    Returns
    -------
    y_pred : ndarray of shape (n_samples,) or (n_samples, K)
        Shape is ``(n_samples, K)`` for multinomial ``"response"`` or
        ``"link"``.
    """
    check_is_fitted(self)
    groups = np.asarray(groups)

    if model not in PREDICT_MODELS:
        raise ValueError(f"model must be one of {PREDICT_MODELS}, got '{model}'")
    if type not in PREDICT_TYPES:
        raise ValueError(f"type must be one of {PREDICT_TYPES}, got '{type}'")
    if alphatype not in ALPHATYPES:
        raise ValueError(f"alphatype must be one of {ALPHATYPES}, got '{alphatype}'")

    if alphatype == "varying":
        X = check_array(X, dtype=np.float64, order="F")
        n_out = (X.shape[0], self.n_classes_) if self.family == "multinomial" else (X.shape[0],)
        y_pred = np.empty(n_out)
        for g in np.unique(groups):
            mask = groups == g
            a = self.varying_alphahat_.get(g, self.alpha_)
            y_pred[mask] = self.all_estimators_[a].predict(
                X[mask], groups[mask], model=model, type=type, lmda_idx=lmda_idx
            )
        return y_pred

    return self.best_estimator_.predict(
        X=X, groups=groups, model=model, type=type, lmda_idx=lmda_idx
    )

evaluate

evaluate(X, y, groups, alphatype='best')

Predict and score with all three sub-models.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
y	array-like of shape (n_samples,)		required
groups	array-like of shape (n_samples,)		required
alphatype	(best, varying)	"best" uses alpha_; "varying" uses each group's own optimal alpha from varying_alphahat_.	"best"

Returns:

Name	Type	Description
result	dict	Keys are "pretrain", "individual", "overall". Each value is a dict with entries: "predictions" : ndarray Predicted values. "score" : float R² for gaussian; accuracy for binomial/multinomial.

Source code in src/ptlasso/_ptlasso.py

def evaluate(self, X, y, groups, alphatype="best"):
    """Predict and score with all three sub-models.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    y : array-like of shape (n_samples,)
    groups : array-like of shape (n_samples,)
    alphatype : {"best", "varying"}, default="best"
        ``"best"`` uses ``alpha_``; ``"varying"`` uses each group's own
        optimal alpha from ``varying_alphahat_``.

    Returns
    -------
    result : dict
        Keys are ``"pretrain"``, ``"individual"``, ``"overall"``.
        Each value is a dict with entries:

        ``"predictions"`` : ndarray
            Predicted values.
        ``"score"`` : float
            R² for gaussian; accuracy for binomial/multinomial.
    """
    check_is_fitted(self)
    result = {}
    for m in ("pretrain", "individual", "overall"):
        preds = self.predict(X, groups, model=m, alphatype=alphatype)
        result[m] = {"predictions": preds, "score": _model_score(y, preds, self.family)}
    return result

score

score(X, y, groups)

Return a scalar performance metric using the best estimator.

Parameters:

Name	Type	Description	Default
X	array-like of shape (n_samples, n_features)		required
y	array-like of shape (n_samples,)		required
groups	array-like of shape (n_samples,)		required

Returns:

Name	Type	Description
score	float	R² for gaussian; classification accuracy for binomial/multinomial.

Source code in src/ptlasso/_ptlasso.py

def score(self, X, y, groups):
    """Return a scalar performance metric using the best estimator.

    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
    y : array-like of shape (n_samples,)
    groups : array-like of shape (n_samples,)

    Returns
    -------
    score : float
        R² for gaussian; classification accuracy for binomial/multinomial.
    """
    check_is_fitted(self)
    return self.best_estimator_.score(X, y, groups)

get_coef

get_coef(model='all', **kwargs)

Return fitted coefficients from the best estimator.

Delegates to :meth:PretrainedLasso.get_coef.

Parameters:

Name	Type	Description	Default
model	(all, overall, pretrain, individual)		"all"
**kwargs		Forwarded to :meth:PretrainedLasso.get_coef.	{}

Returns:

Name	Type	Description
coefs	dict	See :meth:PretrainedLasso.get_coef for the structure.

Source code in src/ptlasso/_ptlasso.py

def get_coef(self, model="all", **kwargs):
    """Return fitted coefficients from the best estimator.

    Delegates to :meth:`PretrainedLasso.get_coef`.

    Parameters
    ----------
    model : {"all", "overall", "pretrain", "individual"}, default="all"
    **kwargs
        Forwarded to :meth:`PretrainedLasso.get_coef`.

    Returns
    -------
    coefs : dict
        See :meth:`PretrainedLasso.get_coef` for the structure.
    """
    check_is_fitted(self)
    return self.best_estimator_.get_coef(model=model, **kwargs)

---

Support utilities

ptlasso.get_overall_support

get_overall_support(fit, lmda_idx=None)

Nonzero features from the overall model.

Parameters:

Name	Type	Description	Default
fit	PretrainedLasso or PretrainedLassoCV		required
lmda_idx	int or None	Index into the overall model's lambda path. Defaults to overall_lmda_idx_ (the CV-selected lambda).	None

Returns:

Name	Type	Description
support	ndarray of int or str

Source code in src/ptlasso/_support.py

def get_overall_support(fit, lmda_idx=None):
    """Nonzero features from the overall model.

    Parameters
    ----------
    fit : PretrainedLasso or PretrainedLassoCV
    lmda_idx : int or None
        Index into the overall model's lambda path.
        Defaults to ``overall_lmda_idx_`` (the CV-selected lambda).

    Returns
    -------
    support : ndarray of int or str
    """
    check_is_fitted(fit)
    idx = lmda_idx if lmda_idx is not None else fit.overall_lmda_idx_
    return _resolve(fit, _nonzero_overall(fit, idx))

ptlasso.get_pretrain_support

get_pretrain_support(fit, lmda_idx=None, groups=None, include_overall=True, common_only=False)

Nonzero features from the per-group pretrained models.

Parameters:

Name	Type	Description	Default
fit	PretrainedLasso or PretrainedLassoCV		required
lmda_idx	int or None	Lambda index for the pretrained group models. When None (default), each group uses its own CV-selected index (pretrain_lmda_idx_), matching the lambda used by predict.	None
groups	array - like or None	Subset of group labels to consider. Default is all groups.	None
include_overall	bool	Union the per-group support with the overall model support. Ignored when alpha=1 (R convention).	True
common_only	bool	If True, return only features selected by more than half the groups.	False

Returns:

Name	Type	Description
support	ndarray of int or str

Source code in src/ptlasso/_support.py

def get_pretrain_support(fit, lmda_idx=None, groups=None, include_overall=True, common_only=False):
    """Nonzero features from the per-group pretrained models.

    Parameters
    ----------
    fit : PretrainedLasso or PretrainedLassoCV
    lmda_idx : int or None
        Lambda index for the pretrained group models.  When ``None`` (default),
        each group uses its own CV-selected index (``pretrain_lmda_idx_``),
        matching the lambda used by ``predict``.
    groups : array-like or None
        Subset of group labels to consider.  Default is all groups.
    include_overall : bool, default=True
        Union the per-group support with the overall model support.
        Ignored when ``alpha=1`` (R convention).
    common_only : bool, default=False
        If True, return only features selected by more than half the groups.

    Returns
    -------
    support : ndarray of int or str
    """
    check_is_fitted(fit)
    groups = fit.groups_ if groups is None else np.asarray(groups)

    base = (
        _nonzero_overall(fit, fit.overall_lmda_idx_)
        if include_overall and fit.alpha < 1
        else np.array([], dtype=int)
    )
    # Use each group's CV-selected lambda when lmda_idx is not specified.
    per_group_idxs = getattr(fit, "pretrain_lmda_idx_", {})
    per_group = [
        _nonzero(
            fit,
            fit.pretrain_models_[g],
            lmda_idx if lmda_idx is not None else per_group_idxs.get(g, -1),
        )
        for g in groups
    ]

    return _resolve(fit, _combine(per_group, base, common_only, len(groups)))

ptlasso.get_pretrain_support_split

get_pretrain_support_split(fit, lmda_idx=None, groups=None)

Split pretrain support into stage-1 ("common") and stage-2 ("individual") parts.

Mirrors the R package's suppre.common / suppre.individual convention:

• common : features selected by the overall model (stage 1).
• individual : features selected by the per-group models (stage 2) that are not in the common support.

Parameters:

Name	Type	Description	Default
fit	PretrainedLasso or PretrainedLassoCV		required
lmda_idx	int or None	Lambda index for the group models. When None (default), each group uses its own CV-selected index (pretrain_lmda_idx_).	None
groups	array - like or None	Subset of group labels. Default is all groups.	None

Returns:

Name	Type	Description
common	ndarray of int or str
individual	ndarray of int or str

Source code in src/ptlasso/_support.py

def get_pretrain_support_split(fit, lmda_idx=None, groups=None):
    """Split pretrain support into stage-1 ("common") and stage-2 ("individual") parts.

    Mirrors the R package's ``suppre.common`` / ``suppre.individual`` convention:

    - **common**     : features selected by the overall model (stage 1).
    - **individual** : features selected by the per-group models (stage 2)
                       that are *not* in the common support.

    Parameters
    ----------
    fit : PretrainedLasso or PretrainedLassoCV
    lmda_idx : int or None
        Lambda index for the group models.  When ``None`` (default), each
        group uses its own CV-selected index (``pretrain_lmda_idx_``).
    groups : array-like or None
        Subset of group labels.  Default is all groups.

    Returns
    -------
    common : ndarray of int or str
    individual : ndarray of int or str
    """
    check_is_fitted(fit)
    groups = fit.groups_ if groups is None else np.asarray(groups)

    per_group_idxs = getattr(fit, "pretrain_lmda_idx_", {})
    overall_idx = set(_nonzero_overall(fit, fit.overall_lmda_idx_).tolist())
    stage2_idx = set()
    for g in groups:
        idx = lmda_idx if lmda_idx is not None else per_group_idxs.get(g, -1)
        stage2_idx |= set(_nonzero(fit, fit.pretrain_models_[g], idx).tolist())

    common_idx = np.sort(np.array(list(overall_idx), dtype=int))
    individual_idx = np.sort(np.array(list(stage2_idx - overall_idx), dtype=int))

    return _resolve(fit, common_idx), _resolve(fit, individual_idx)

ptlasso.get_individual_support

get_individual_support(fit, lmda_idx=None, groups=None, common_only=False)

Nonzero features from the per-group individual (no-pretraining) models.

Parameters:

Name	Type	Description	Default
fit	PretrainedLasso or PretrainedLassoCV		required
lmda_idx	int or None	Lambda index for the individual group models. When None (default), each group uses its own CV-selected index (individual_lmda_idx_), matching the lambda used by predict.	None
groups	array - like or None	Subset of group labels to consider. Default is all groups.	None
common_only	bool	If True, return only features selected by more than half the groups.	False

Returns:

Name	Type	Description
support	ndarray of int or str

Source code in src/ptlasso/_support.py

def get_individual_support(fit, lmda_idx=None, groups=None, common_only=False):
    """Nonzero features from the per-group individual (no-pretraining) models.

    Parameters
    ----------
    fit : PretrainedLasso or PretrainedLassoCV
    lmda_idx : int or None
        Lambda index for the individual group models.  When ``None`` (default),
        each group uses its own CV-selected index (``individual_lmda_idx_``),
        matching the lambda used by ``predict``.
    groups : array-like or None
        Subset of group labels to consider.  Default is all groups.
    common_only : bool, default=False
        If True, return only features selected by more than half the groups.

    Returns
    -------
    support : ndarray of int or str
    """
    check_is_fitted(fit)
    groups = fit.groups_ if groups is None else np.asarray(groups)

    per_group_idxs = getattr(fit, "individual_lmda_idx_", {})
    per_group = [
        _nonzero(
            fit,
            fit.individual_models_[g],
            lmda_idx if lmda_idx is not None else per_group_idxs.get(g, -1),
        )
        for g in groups
    ]

    return _resolve(fit, _combine(per_group, np.array([], dtype=int), common_only, len(groups)))

---

Plotting

ptlasso.plot_cv

plot_cv(fit, ax=None, plot_alphahat=True, column='single', save=None, colors=None, figure_widths=None)

Plot the cross-validation curve for a :class:PretrainedLassoCV.

Draws the mean CV loss ±1 SE band over alpha, with horizontal reference lines for the individual and overall baselines.

Parameters:

Name	Type	Description	Default
fit	PretrainedLassoCV	A fitted CV estimator.	required
ax	Axes or None	Axes to draw on. A new figure is created when None.	None
plot_alphahat	bool	Whether to draw a vertical line at the selected alpha_.	True
column	(single, double)	Target figure width — "single" ≈ 3.5 in, "double" ≈ 7 in.	"single"
save	str or None	File path to save the figure (300 dpi). No file is written when None.	None
colors	dict or None	Override plot colours for "overall", "pretrain", and/or "individual". Missing keys fall back to :data:ptlasso.COLORS.	None
figure_widths	dict or None	Override figure widths for "single" and/or "double". Missing keys fall back to :data:ptlasso.FIGURE_WIDTHS.	None

Returns:

Name	Type	Description
fig	Figure
ax	Axes

Source code in src/ptlasso/_plot.py

def plot_cv(
    fit, ax=None, plot_alphahat=True, column="single", save=None, colors=None, figure_widths=None
):
    """Plot the cross-validation curve for a :class:`PretrainedLassoCV`.

    Draws the mean CV loss ±1 SE band over alpha, with horizontal reference
    lines for the individual and overall baselines.

    Parameters
    ----------
    fit : PretrainedLassoCV
        A fitted CV estimator.
    ax : matplotlib.axes.Axes or None, default=None
        Axes to draw on.  A new figure is created when ``None``.
    plot_alphahat : bool, default=True
        Whether to draw a vertical line at the selected ``alpha_``.
    column : {"single", "double"}, default="single"
        Target figure width — ``"single"`` ≈ 3.5 in, ``"double"`` ≈ 7 in.
    save : str or None, default=None
        File path to save the figure (300 dpi).  No file is written when ``None``.
    colors : dict or None, default=None
        Override plot colours for ``"overall"``, ``"pretrain"``, and/or
        ``"individual"``.  Missing keys fall back to :data:`ptlasso.COLORS`.
    figure_widths : dict or None, default=None
        Override figure widths for ``"single"`` and/or ``"double"``.
        Missing keys fall back to :data:`ptlasso.FIGURE_WIDTHS`.

    Returns
    -------
    fig : matplotlib.figure.Figure
    ax : matplotlib.axes.Axes
    """
    check_is_fitted(fit)

    c = _resolve_colors(colors)
    w = _resolve_widths(figure_widths).get(column, 3.5)

    if ax is None:
        fig, ax = plt.subplots(figsize=(w, w * 0.8))
    else:
        fig = ax.get_figure()

    alphas = list(fit.alphalist_)
    cv_mean = np.array([fit.cv_results_[a] for a in alphas])
    cv_se = np.array([fit.cv_results_se_[a] for a in alphas])

    # ±1 SE band + pretrain curve
    ax.fill_between(
        alphas, cv_mean - cv_se, cv_mean + cv_se, color=c["pretrain"], alpha=0.15, linewidth=0
    )
    ax.plot(
        alphas,
        cv_mean,
        color=c["pretrain"],
        marker="o",
        markersize=5,
        linewidth=2,
        label="Pretrain",
        zorder=3,
    )

    # Individual and overall baselines
    ax.axhline(fit.cv_results_individual_, color=c["individual"], linestyle="--", linewidth=1.8)
    ax.axhline(fit.cv_results_overall_, color=c["overall"], linestyle="--", linewidth=1.8)

    # Selected alpha
    if plot_alphahat:
        ax.axvline(
            fit.alpha_,
            color="#555555",
            linestyle=":",
            linewidth=1.5,
            label=f"$\\hat{{\\alpha}}={fit.alpha_}$",
        )

    # Support sizes at right edge
    n_pre = len(get_pretrain_support(fit))
    n_ind = len(get_individual_support(fit))
    n_ov = len(get_overall_support(fit))
    xmax = max(alphas)
    ax.text(
        xmax + 0.01,
        fit.cv_results_individual_,
        f"Individual  $|\\hat{{S}}|={n_ind}$",
        va="center",
        fontsize=8,
        color=c["individual"],
        clip_on=False,
    )
    ax.text(
        xmax + 0.01,
        fit.cv_results_overall_,
        f"Overall  $|\\hat{{S}}|={n_ov}$",
        va="center",
        fontsize=8,
        color=c["overall"],
        clip_on=False,
    )

    # Top axis: support size at selected alpha
    ax2 = ax.twiny()
    ax2.set_xlim(ax.get_xlim())
    ax2.set_xticks([fit.alpha_])
    ax2.set_xticklabels([f"$|\\hat{{S}}|={n_pre}$"], fontsize=8, color=c["pretrain"])
    ax2.tick_params(length=0)
    ax2.spines["top"].set_visible(False)
    ax2.spines["right"].set_visible(False)

    _despine(ax)
    ax.set_xlabel("$\\alpha$", fontsize=11)
    ax.set_ylabel(_cv_ylabel(fit.family), fontsize=11)
    ax.set_title("Cross-validation over $\\alpha$", fontsize=12, pad=14)
    ax.set_xticks(alphas)
    ax.tick_params(labelsize=9)
    ax.legend(frameon=False, fontsize=9, loc="upper right")

    fig.tight_layout()
    if save:
        fig.savefig(save, dpi=300, bbox_inches="tight")
    return fig, ax

ptlasso.plot_paths

plot_paths(fit, column='double', save=None, colors=None, figure_widths=None)

Plot regularisation paths for all sub-models in a :class:PretrainedLasso.

Produces a 3-row grid: overall model (full width), per-group pretrained models, and per-group individual models. Features in the final support are coloured; inactive features are shown in grey.

Parameters:

Name	Type	Description	Default
fit	PretrainedLasso or PretrainedLassoCV	A fitted estimator.	required
column	(single, double)	Target figure width — "single" ≈ 3.5 in, "double" ≈ 7 in.	"single"
save	str or None	File path to save the figure (300 dpi). No file is written when None.	None
colors	dict or None	Override plot colours for "overall", "pretrain", and/or "individual". Missing keys fall back to :data:ptlasso.COLORS.	None
figure_widths	dict or None	Override figure widths for "single" and/or "double". Missing keys fall back to :data:ptlasso.FIGURE_WIDTHS.	None

Returns:

Name	Type	Description
fig	Figure

Source code in src/ptlasso/_plot.py

def plot_paths(fit, column="double", save=None, colors=None, figure_widths=None):
    """Plot regularisation paths for all sub-models in a :class:`PretrainedLasso`.

    Produces a 3-row grid: overall model (full width), per-group pretrained
    models, and per-group individual models.  Features in the final support are
    coloured; inactive features are shown in grey.

    Parameters
    ----------
    fit : PretrainedLasso or PretrainedLassoCV
        A fitted estimator.
    column : {"single", "double"}, default="double"
        Target figure width — ``"single"`` ≈ 3.5 in, ``"double"`` ≈ 7 in.
    save : str or None, default=None
        File path to save the figure (300 dpi).  No file is written when ``None``.
    colors : dict or None, default=None
        Override plot colours for ``"overall"``, ``"pretrain"``, and/or
        ``"individual"``.  Missing keys fall back to :data:`ptlasso.COLORS`.
    figure_widths : dict or None, default=None
        Override figure widths for ``"single"`` and/or ``"double"``.
        Missing keys fall back to :data:`ptlasso.FIGURE_WIDTHS`.

    Returns
    -------
    fig : matplotlib.figure.Figure
    """
    check_is_fitted(fit)

    c = _resolve_colors(colors)
    k = len(fit.groups_)
    w = _resolve_widths(figure_widths).get(column, 7.0)
    labels = _label_map(fit)
    feature_colors = _feature_color_map(fit)

    all_lmdas = np.log(np.asarray(fit.overall_model_.lmdas))
    xlim = (all_lmdas[0], all_lmdas[-1] + 0.18 * (all_lmdas[-1] - all_lmdas[0]))
    overall_sup = set(
        np.where(np.reshape(fit.overall_coef_, (fit.n_features_in_, -1)).any(axis=1))[0]
    )

    # The overall model was trained on [onehot | X]; strip the onehot columns
    # from betas so that _draw_paths sees only the p X-feature coefficients.
    n_onehot = getattr(fit, "_n_onehot_", 0)

    class _StrippedOverallState:
        """Thin view of the overall model state with onehot columns removed."""

        def __init__(self, state, n_skip):
            raw = _betas_dense(state)  # (L, k-1+p)
            self.betas = raw[:, n_skip:]  # (L, p)
            self.lmdas = state.lmdas
            self.intercepts = state.intercepts

    overall_state_for_plot = (
        _StrippedOverallState(fit.overall_model_, n_onehot) if n_onehot > 0 else fit.overall_model_
    )

    fig = plt.figure(figsize=(w, w * 0.42 * 3))
    gs = gridspec.GridSpec(3, k, figure=fig, hspace=0.7, wspace=0.45)

    # Row 0: overall model (spans all columns)
    _draw_paths(
        fig.add_subplot(gs[0, :]),
        overall_state_for_plot,
        "Overall",
        c["overall"],
        feature_colors,
        overall_sup,
        labels,
        xlim,
    )

    # Rows 1 & 2: per-group pretrain and individual
    pre_idxs = getattr(fit, "pretrain_lmda_idx_", {})
    ind_idxs = getattr(fit, "individual_lmda_idx_", {})
    for col, g in enumerate(fit.groups_):
        lbl = fit._label(g)
        pre_i = pre_idxs.get(g, -1)
        ind_i = ind_idxs.get(g, -1)
        pre_sup = set(np.where(_betas_dense(fit.pretrain_models_[g])[pre_i] != 0)[0])
        ind_sup = set(np.where(_betas_dense(fit.individual_models_[g])[ind_i] != 0)[0])

        _draw_paths(
            fig.add_subplot(gs[1, col]),
            fit.pretrain_models_[g],
            lbl,
            c["pretrain"],
            feature_colors,
            pre_sup,
            labels,
            xlim,
        )
        _draw_paths(
            fig.add_subplot(gs[2, col]),
            fit.individual_models_[g],
            lbl,
            c["individual"],
            feature_colors,
            ind_sup,
            labels,
            xlim,
        )

    # Row labels in left margin
    for row, (text, color) in enumerate(
        [
            ("Overall", c["overall"]),
            ("Pretrain", c["pretrain"]),
            ("Individual", c["individual"]),
        ]
    ):
        fig.text(
            0.01,
            1 - (row + 0.5) / 3,
            text,
            va="center",
            ha="center",
            fontsize=10,
            fontweight="bold",
            color=color,
            rotation=90,
            transform=fig.transFigure,
        )

    fig.suptitle("Regularisation paths", fontsize=13, y=1.01)
    if save:
        fig.savefig(save, dpi=300, bbox_inches="tight")
    return fig

---

Simulation

ptlasso.make_data

make_data(k, class_sizes, s_common, s_indiv, beta_common, beta_indiv, intercepts=None, sigma=1.0, family='gaussian', seed=None)

Generate synthetic grouped data for ptlasso.

Feature layout

• Columns 0 … s_common-1 : shared features (active in all groups)
• Columns s_common … s_common+s_indiv[0]-1 : features specific to group 0
• ... and so on for each group
• Remaining columns (if any) : pure noise

Parameters:

Name	Type	Description	Default
k	int	Number of groups.	required
class_sizes	array-like of length k	Number of observations per group.	required
s_common	int	Number of shared (common) features.	required
s_indiv	int or array-like of length k	Number of group-specific features per group.	required
beta_common	float or array - like	Coefficients for common features. Scalar → same value for all s_common features in every group. 1-D array of length s_common → per-feature coefficient (same across groups). List of k arrays → per-group, per-feature coefficients.	required
beta_indiv	float or array - like	Coefficients for group-specific features, same shapes as beta_common.	required
intercepts	array-like of length k or None	Per-group intercepts. Default is 0.	None
sigma	float	Gaussian noise std (only used for family="gaussian").	1.0
family	(gaussian, binomial)		"gaussian"
seed	int or None	Random seed passed to numpy.random.default_rng.	None

Returns:

Type	Description
dict with keys	X : ndarray (n, p) y : ndarray (n,) groups : ndarray (n,) of int

Source code in src/ptlasso/_simulate.py

def make_data(
    k,
    class_sizes,
    s_common,
    s_indiv,
    beta_common,
    beta_indiv,
    intercepts=None,
    sigma=1.0,
    family="gaussian",
    seed=None,
):
    """Generate synthetic grouped data for ptlasso.

    Feature layout
    --------------
    - Columns 0 … s_common-1             : shared features (active in all groups)
    - Columns s_common … s_common+s_indiv[0]-1 : features specific to group 0
    - ... and so on for each group
    - Remaining columns (if any)          : pure noise

    Parameters
    ----------
    k : int
        Number of groups.
    class_sizes : array-like of length k
        Number of observations per group.
    s_common : int
        Number of shared (common) features.
    s_indiv : int or array-like of length k
        Number of group-specific features per group.
    beta_common : float or array-like
        Coefficients for common features.  Scalar → same value for all s_common
        features in every group.  1-D array of length s_common → per-feature
        coefficient (same across groups).  List of k arrays → per-group,
        per-feature coefficients.
    beta_indiv : float or array-like
        Coefficients for group-specific features, same shapes as beta_common.
    intercepts : array-like of length k or None
        Per-group intercepts.  Default is 0.
    sigma : float, default=1.0
        Gaussian noise std (only used for ``family="gaussian"``).
    family : {"gaussian", "binomial"}, default="gaussian"
    seed : int or None
        Random seed passed to ``numpy.random.default_rng``.

    Returns
    -------
    dict with keys
        ``X``      : ndarray (n, p)
        ``y``      : ndarray (n,)
        ``groups`` : ndarray (n,) of int
    """
    rng = np.random.default_rng(seed)

    class_sizes = np.asarray(class_sizes, dtype=int)
    s_indiv = np.broadcast_to(s_indiv, (k,)).copy()
    n = int(class_sizes.sum())
    p = s_common + int(s_indiv.sum())

    X = rng.standard_normal((n, p))
    y = np.empty(n)
    groups = np.empty(n, dtype=int)

    row_start = 0
    feat_start = s_common

    if intercepts is None:
        intercepts = np.zeros(k)

    for kk in range(k):
        row_end = row_start + class_sizes[kk]
        feat_end = feat_start + s_indiv[kk]

        beta = np.zeros(p)
        beta[:s_common] = _expand_coef(beta_common, kk, s_common)
        beta[feat_start:feat_end] = _expand_coef(beta_indiv, kk, s_indiv[kk])

        mu = X[row_start:row_end] @ beta + intercepts[kk]

        if family == "gaussian":
            y[row_start:row_end] = mu + sigma * rng.standard_normal(class_sizes[kk])
        elif family == "binomial":
            prob = 1.0 / (1.0 + np.exp(-mu))
            y[row_start:row_end] = rng.binomial(1, prob).astype(float)
        else:
            raise ValueError(f"family must be 'gaussian' or 'binomial', got '{family}'")

        groups[row_start:row_end] = kk
        row_start = row_end
        feat_start = feat_end

    return {"X": X, "y": y, "groups": groups}

ptlasso.gaussian_example_data

gaussian_example_data(k=2, class_sizes=None, s_common=5, s_indiv=5, beta_common=1.0, beta_indiv=0.5, sigma=1.0, seed=None)

Convenience wrapper for Gaussian grouped data.

Parameters:

Name	Type	Description	Default
k	int		2
class_sizes	array - like or None	Defaults to 50 observations per group.	None
s_common	int		5
s_indiv	int or array - like		5
beta_common	float		1.0
beta_indiv	float		0.5
sigma	float		1.0
seed	int or None		None

Returns:

Type	Description
dict with keys ``X``, ``y``, ``groups``

Source code in src/ptlasso/_simulate.py

def gaussian_example_data(
    k=2,
    class_sizes=None,
    s_common=5,
    s_indiv=5,
    beta_common=1.0,
    beta_indiv=0.5,
    sigma=1.0,
    seed=None,
):
    """Convenience wrapper for Gaussian grouped data.

    Parameters
    ----------
    k : int, default=2
    class_sizes : array-like or None
        Defaults to 50 observations per group.
    s_common : int, default=5
    s_indiv : int or array-like, default=5
    beta_common : float, default=1.0
    beta_indiv : float, default=0.5
    sigma : float, default=1.0
    seed : int or None

    Returns
    -------
    dict with keys ``X``, ``y``, ``groups``
    """
    if class_sizes is None:
        class_sizes = [50] * k
    return make_data(
        k=k,
        class_sizes=class_sizes,
        s_common=s_common,
        s_indiv=s_indiv,
        beta_common=beta_common,
        beta_indiv=beta_indiv,
        sigma=sigma,
        family="gaussian",
        seed=seed,
    )

ptlasso.binomial_example_data

binomial_example_data(k=2, class_sizes=None, s_common=5, s_indiv=5, beta_common=0.5, beta_indiv=0.3, seed=None)

Convenience wrapper for binary grouped data.

Parameters:

Name	Type	Description	Default
k	int		2
class_sizes	array - like or None	Defaults to 100 observations per group.	None
s_common	int		5
s_indiv	int or array - like		5
beta_common	float		0.5
beta_indiv	float		0.3
seed	int or None		None

Returns:

Type	Description
dict with keys ``X``, ``y``, ``groups``

Source code in src/ptlasso/_simulate.py

def binomial_example_data(
    k=2,
    class_sizes=None,
    s_common=5,
    s_indiv=5,
    beta_common=0.5,
    beta_indiv=0.3,
    seed=None,
):
    """Convenience wrapper for binary grouped data.

    Parameters
    ----------
    k : int, default=2
    class_sizes : array-like or None
        Defaults to 100 observations per group.
    s_common : int, default=5
    s_indiv : int or array-like, default=5
    beta_common : float, default=0.5
    beta_indiv : float, default=0.3
    seed : int or None

    Returns
    -------
    dict with keys ``X``, ``y``, ``groups``
    """
    if class_sizes is None:
        class_sizes = [100] * k
    return make_data(
        k=k,
        class_sizes=class_sizes,
        s_common=s_common,
        s_indiv=s_indiv,
        beta_common=beta_common,
        beta_indiv=beta_indiv,
        family="binomial",
        seed=seed,
    )