"""Group lasso over guides grouped by gene, selected or dropped as a set.
This model estimates which guides and genes are associated with a response
while respecting the grouping of guides within genes.
Ordinary lasso penalises each guide on its own, so from a gene's four
correlated guides it keeps whichever one happens to fit best and drops the
rest -- which reads as "one guide works and three do not" when the four
guides are four measurements of the same gene. Group lasso penalises the
gene's guides as a block:
minimise ||y - Xb||^2 / (2n) + lambda * sum_g sqrt(p_g) * ||b_g||_2
The L2 norm inside the sum is not squared, so its subgradient at zero is a
ball rather than a point, and the whole block goes to exactly zero or none of
it does. sqrt(p_g) is the standard correction that stops a gene with six
guides being penalised more than one with three purely for having more.
For a CRISPR screen it is the penalised analogue of the mixed model with
random effects -- "which genes are involved, where each gene is measured by
several guides" -- and it addresses that question with no random effects, no
REML fit, and no `gene_fraction` column. `gene_fraction` is the sum of a
gene's guide fractions, so a design carrying both blocks is singular by
construction. Here the gene enters as a grouping of the guide columns, never
as a column of its own, so the design stays full rank.
It also fits where OLS is undefined. This screen is 610 wells and 823 guides;
the penalty is what makes the answer well posed at p > n. A penalised fit
therefore rests on the assumption that most guides have no effect; the
penalty adds no new information.
The model gives no P value. A penalised coefficient's sampling distribution
is not the OLS one, and a naive P value from it is biased toward
significance. :func:`stability_selection` instead re-fits on many subsamples
and reports the fraction of them in which a gene is selected, an
error-controlled statement (Meinshausen & Buhlmann) that needs no
distributional assumption.
The solver is block coordinate descent with the exact group-soft-threshold
proximal step, in numpy. No new dependency: scikit-learn has no group lasso,
and the two packages that do are unmaintained.
"""
from __future__ import annotations
from typing import List, Tuple
import numpy as np
#: Iterations before block coordinate descent gives up. Reached only on a
#: badly scaled design; the screen-sized problems here converge in tens.
MAX_ITERATIONS = 1000
#: Convergence: the largest coefficient change in a sweep, relative to the
#: largest coefficient.
TOLERANCE = 1e-6
def _groups(labels) -> Tuple[np.ndarray, List[np.ndarray]]:
"""``(unique labels, column indices per label)``, labels in sorted order."""
labels = np.asarray(labels, dtype=object)
unique = np.unique(labels)
return unique, [np.flatnonzero(labels == one) for one in unique]
def _soft_threshold_block(vector: np.ndarray, amount: float) -> np.ndarray:
"""The proximal operator of ``amount * ||.||_2``.
Shrinks the block toward the origin by ``amount`` and sets it to exactly
zero once it is shorter than that -- which is the whole point: this is
what makes a gene drop out as a unit rather than one guide at a time.
"""
norm = float(np.linalg.norm(vector))
if norm <= amount or norm == 0.0:
return np.zeros_like(vector)
return vector * (1.0 - amount / norm)
[docs]
def max_lambda(X, y, labels) -> float:
"""The smallest penalty that zeroes every group.
:param X: design matrix with one column per guide.
:param y: response vector with one value per design row.
:param labels: gene label for each design column.
Useful as the top of a path: any larger value gives the same all-zero fit,
and starting a path above it wastes every iteration spent there.
"""
X = np.asarray(X, dtype=float)
y = np.asarray(y, dtype=float)
_unique, blocks = _groups(labels)
n = X.shape[0]
correlation = X.T @ y / n
return max(
(float(np.linalg.norm(correlation[block]) / np.sqrt(block.size))
for block in blocks), default=0.0)
[docs]
def fit(X, y, labels, *, lam: float = 0.05, max_iterations: int = MAX_ITERATIONS,
tolerance: float = TOLERANCE, fit_intercept: bool = True):
"""Fit the group lasso and return ``(coefficients, intercept, converged)``.
:param X: the design, one column per guide.
:param y: the response, one entry per well.
:param labels: the gene each column belongs to, length ``X.shape[1]``.
:param lam: the penalty. Zero is ordinary least squares and is allowed --
it is how a caller checks the solver against a known answer.
:param fit_intercept: centre both sides, so the intercept is never
penalised. Penalising it would shrink the response's mean toward
zero.
:returns: ``coefficients`` (one per column), ``intercept``, and whether
the sweep converged inside ``max_iterations``.
:raises ValueError: the shapes disagree, or ``lam`` is negative.
"""
X = np.asarray(X, dtype=float)
y = np.asarray(y, dtype=float).ravel()
if X.ndim != 2:
raise ValueError(f"X must be two-dimensional; got shape {X.shape}")
if X.shape[0] != y.size:
raise ValueError(
f"X has {X.shape[0]} row(s) and y has {y.size}; they must match")
if len(labels) != X.shape[1]:
raise ValueError(
f"labels has {len(labels)} entr(ies) and X has {X.shape[1]} "
f"column(s); they must match")
if lam < 0:
raise ValueError(f"lam must not be negative; got {lam!r}")
n = X.shape[0]
if fit_intercept:
x_mean, y_mean = X.mean(axis=0), float(y.mean())
X, y = X - x_mean, y - y_mean
else:
x_mean, y_mean = np.zeros(X.shape[1]), 0.0
_unique, blocks = _groups(labels)
steps = []
for block in blocks:
sub = X[:, block]
largest = float(np.linalg.norm(sub, 2)) ** 2 / n
steps.append(1.0 / largest if largest > 0 else 0.0)
beta = np.zeros(X.shape[1])
residual = y.copy()
converged = False
for _sweep in range(max_iterations):
largest_change = 0.0
for block, step in zip(blocks, steps):
if step == 0.0:
continue
sub = X[:, block]
current = beta[block]
partial = residual + sub @ current
gradient = sub.T @ partial / n
updated = _soft_threshold_block(
current + step * (gradient - sub.T @ (sub @ current) / n),
step * lam * np.sqrt(block.size))
change = float(np.max(np.abs(updated - current))) if block.size \
else 0.0
if change:
residual = partial - sub @ updated
beta[block] = updated
largest_change = max(largest_change, change)
scale = max(1.0, float(np.max(np.abs(beta))) if beta.size else 1.0)
if largest_change <= tolerance * scale:
converged = True
break
intercept = y_mean - float(x_mean @ beta) if fit_intercept else 0.0
return beta, intercept, converged
#: How many penalties a cross-validated path tries, and how far below
#: :func:`max_lambda` it reaches. Both are the conventional choices: a
#: hundred points is more than enough resolution on a log grid, and a
#: thousandth of the ceiling is well past the point where a screen's design
#: stops selecting anything at all.
PATH_POINTS = 40
PATH_DEPTH = 1e-3
#: How many folds :func:`choose_lambda` splits the wells into.
PATH_FOLDS = 5
[docs]
def penalty_path(X, y, labels, *, points: int = PATH_POINTS,
depth: float = PATH_DEPTH):
"""Penalties from :func:`max_lambda` down, log-spaced.
Starting at the ceiling rather than at an absolute number is what makes
the path scale-free: a design of fractions and a design of counts have
ceilings three orders of magnitude apart, and a fixed grid that suits
one is entirely above or entirely below the other.
"""
ceiling = max_lambda(X, y, labels)
if not np.isfinite(ceiling) or ceiling <= 0:
return np.zeros(1)
return np.logspace(np.log10(ceiling), np.log10(ceiling * depth),
int(points))
#: How many further decades :func:`choose_lambda` will reach for when the
#: whole path selects nothing. A screen whose gene blocks only survive four
#: orders of magnitude below the ceiling is unusual and not impossible, and
#: an empty answer is worse than a slow one.
PATH_EXTENSIONS = 3
[docs]
def choose_lambda(X, y, labels, *, folds: int = PATH_FOLDS,
points: int = PATH_POINTS, depth: float = PATH_DEPTH,
required=None, seed: int = 0, **kwargs) -> float:
"""Select a group-lasso penalty by held-out prediction error.
A candidate is eligible only when its fit on the complete dataset selects
at least one column allowed by ``required`` and its held-out error is
finite. If a penalty path contains no eligible candidate, progressively
smaller penalties are evaluated for up to :data:`PATH_EXTENSIONS`
additional ranges. If no eligible candidate is found, the smallest
penalty evaluated is returned.
:param X: two-dimensional design matrix with one row per well and one
column per guide.
:param y: response values aligned to the rows of ``X``.
:param labels: group label for every column of ``X``; guides sharing a
label are selected or dropped as one block.
:param folds: how many held-out splits. Wells, not guides.
:param required: optional boolean mask over columns. A penalty counts as
selecting only when one of these columns is non-zero. Defaults to
every column.
:param seed: fixes the split, so two runs of one screen agree.
:returns: the chosen penalty, as a float.
"""
X = np.asarray(X, dtype=float)
y = np.asarray(y, dtype=float).ravel()
if required is None:
wanted = np.ones(X.shape[1], dtype=bool)
else:
wanted = np.asarray(required, dtype=bool).reshape(-1)
if wanted.size != X.shape[1]:
raise ValueError(
f"required has {wanted.size} entr(ies) and X has "
f"{X.shape[1]} column(s); they must match")
rows = X.shape[0]
folds = int(max(2, min(folds, rows)))
order = np.random.default_rng(seed).permutation(rows)
splits = np.array_split(order, folds)
def score(candidates):
"""(errors, selects) for one path, both aligned to ``candidates``.
The error is held out and the selection is not. Whether a penalty
predicts well is a question about unseen wells; whether it selects
anything is a question about the fit the caller will get, which is
fitted on every well. Reading selection off the folds can choose a
penalty at which most folds keep a gene while the full fit keeps
none (on the tsg101 screen it chose 0.000183 and the full fit was
empty).
"""
errors = np.full(candidates.size, np.inf)
selects = np.zeros(candidates.size, dtype=bool)
for index, lam in enumerate(candidates):
whole, _intercept, _converged = fit(X, y, labels, lam=float(lam),
**kwargs)
selects[index] = bool(np.any(whole[wanted]))
total, counted = 0.0, 0
for held in splits:
keep = np.setdiff1d(order, held, assume_unique=False)
if keep.size < 2 or held.size < 1:
continue
beta, intercept, _converged = fit(X[keep], y[keep], labels,
lam=float(lam), **kwargs)
predicted = X[held] @ beta + intercept
total += float(np.mean((y[held] - predicted) ** 2))
counted += 1
if counted:
errors[index] = total / counted
return errors, selects
floor = depth
smallest = 0.0
for _reach in range(PATH_EXTENSIONS + 1):
candidates = penalty_path(X, y, labels, points=points, depth=floor)
if candidates.size < 2:
return float(candidates[0]) if candidates.size else 0.0
smallest = float(candidates[-1])
errors, selects = score(candidates)
usable = np.where(selects & np.isfinite(errors))[0]
if usable.size:
return float(candidates[usable[int(np.argmin(errors[usable]))]])
floor *= 0.01
return smallest
[docs]
def gene_effects(X, y, labels, **kwargs):
"""One number per gene: the L2 norm of its guide block.
:param X: design matrix with one column per guide.
:param y: response vector with one value per design row.
:param labels: gene label for each design column.
The block is zero or it is not, so the norm is the natural per-gene effect
size. Unlike a coefficient it is never negative: group lasso indicates
that a gene's guides move the response together, not in which direction
each one does.
"""
import pandas as pd
beta, _intercept, converged = fit(X, y, labels, **kwargs)
unique, blocks = _groups(labels)
return pd.DataFrame({
"gene": unique,
"n_guides": [block.size for block in blocks],
"effect": [float(np.linalg.norm(beta[block])) for block in blocks],
"selected": [bool(np.any(beta[block])) for block in blocks],
"converged": converged,
}).sort_values("effect", ascending=False,
kind="stable").reset_index(drop=True)
[docs]
def stability_selection(X, y, labels, *, lam: float = 0.05,
n_boot: int = 100, fraction: float = 0.5,
seed: int = 0, **kwargs):
"""Fraction of subsample re-fits in which each gene is selected.
:param X: two-dimensional design matrix with wells in rows and guides in
columns.
:param y: response values aligned to the rows of ``X``.
:param labels: gene or other group label for every guide column.
:param n_boot: subsamples.
:param fraction: the share of rows in each subsample. Half is
Meinshausen & Buhlmann's choice and the one their error bound is
stated for.
:returns: a DataFrame with a ``selection_frequency`` per gene, sorted.
This replaces a P value, which the group lasso does not provide: the
sampling distribution of a penalised coefficient is not the OLS one, and
a P value computed as though it were is biased toward significance. A
selection frequency states how reproducible the selection is and needs
no distributional assumption.
"""
import pandas as pd
X = np.asarray(X, dtype=float)
y = np.asarray(y, dtype=float).ravel()
if not 0.0 < fraction <= 1.0:
raise ValueError(f"fraction must be in (0, 1]; got {fraction!r}")
if n_boot < 1:
raise ValueError(f"n_boot must be at least 1; got {n_boot!r}")
unique, blocks = _groups(labels)
counts = np.zeros(unique.size)
rng = np.random.default_rng(seed)
size = max(2, int(round(X.shape[0] * fraction)))
for _draw in range(n_boot):
rows = rng.choice(X.shape[0], size=size, replace=False)
beta, _intercept, _converged = fit(X[rows], y[rows], labels, lam=lam,
**kwargs)
for index, block in enumerate(blocks):
if np.any(beta[block]):
counts[index] += 1
return pd.DataFrame({
"gene": unique,
"n_guides": [block.size for block in blocks],
"selection_frequency": counts / n_boot,
}).sort_values("selection_frequency", ascending=False,
kind="stable").reset_index(drop=True)
[docs]
def describe(lam: float = 0.05) -> str:
"""Return the formula and model description shown in the model tab."""
return (
"Group lasso. The gene's guides are penalised as a block:\n\n"
" minimise ||y - Xb||^2 / (2n) + lambda * sum_g sqrt(p_g) * "
"||b_g||_2\n\n"
f"with lambda = {lam:g}. The inner norm is not squared, so a gene's "
"whole block of guides goes to exactly zero or none of it does -- "
"which is the question a screen asks. Ordinary lasso instead keeps "
"whichever one of a gene's correlated guides fits best and drops the "
"rest, reading as 'one guide works' when the truth is 'the gene "
"matters'. sqrt(p_g) stops a gene with six guides being penalised "
"more than one with three for having more.\n\n"
"Recommended for CRISPR screens. It is the penalised analogue of the "
"mixed model's nesting -- which GENES are involved, each measured by "
"several guides -- and the gene enters as a GROUPING of guide "
"columns rather than as a summed `gene_fraction` column, so the "
"design cannot be singular. It fits where OLS is undefined (823 "
"guides, 610 wells).\n\n"
"It gives no P value, and that is deliberate: a penalised "
"coefficient's sampling distribution is not the OLS one. Use "
"stability selection, which reports how often each gene survives a "
"re-fit on half the wells."
)
__all__ = ["MAX_ITERATIONS", "TOLERANCE", "describe", "fit", "gene_effects",
"max_lambda", "stability_selection"]