from __future__ import annotations
import math
import warnings
from collections.abc import Callable
from typing import Any
import numpy as np
from anndata import AnnData
from mantispy._core._numba import group_offsets
from mantispy._core._pca import pca_solver
from mantispy._core._reduce import get_matrix, group_codes, is_lazy, materialize, matrix_dtype, transform_grouped
from mantispy._core.logging import get_logger
from mantispy._core.masks import CONTROL_CLASSES, reference_mask
from mantispy._core.mutation import inplace_or_copy
METHODS = ("ZCA", "ZCA-cor", "PCA", "PCA-cor")
SVD_SOLVERS = ("covariance_eigh", "full")
def _decompose(
reference: np.ndarray, method: str, svd_solver: str | None = None
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
"""Return ``(centre, scale, singular, right, n_obs)``: the epsilon-independent part of :func:`_fit`.
Factored out so a sweep over epsilon pays for the centre, scale and SVD once rather than per candidate.
``svd_solver`` selects the decomposition: ``None`` eigendecomposes the feature covariance when the
reference is tall and narrow and otherwise takes the full SVD, which ``"covariance_eigh"`` and ``"full"``
force directly.
"""
# A missing value would otherwise surface as "LinAlgError: SVD did not converge".
if not np.isfinite(reference).all():
rows = int((~np.isfinite(reference).all(axis=1)).sum())
columns = int((~np.isfinite(reference).all(axis=0)).sum())
raise ValueError(
f"sphering needs a complete reference, but {rows} of {reference.shape[0]} reference row(s) "
f"have a missing or infinite value, across {columns} feature(s). Drop those features "
"(mt.pp.feature_select with drop_na_columns, then mt.pp.subset_features) or those reference wells."
)
centre = reference.mean(axis=0)
centered = reference - centre
if method.endswith("-cor"):
# StandardScaler: population standard deviation.
scale = centered.std(axis=0, ddof=0)
if np.any(scale == 0):
raise ValueError(
"sphering cannot standardize a feature with zero variance; drop constant "
"features first, e.g. mt.pp.filter_features(adata, min_variance=1e-8)"
)
centered = centered / scale
else:
scale = np.ones(reference.shape[1])
n_obs, n_vars = centered.shape
if n_obs <= n_vars:
get_logger().warning(
"sphere: fitted on %d reference rows for %d features, so the covariance is singular",
n_obs,
n_vars,
deep="select fewer features first, or use more controls",
)
# Centering costs one degree of freedom, so a full-rank reference has rank min(n_vars, n_obs - 1).
rank = np.linalg.matrix_rank(centered)
if rank != min(n_vars, n_obs - 1):
raise ValueError(
f"the reference matrix is not full rank: {n_obs} rows, {n_vars} features, rank {rank}. "
"Sphering needs enough independent control profiles to estimate a covariance. "
"Use more control wells, or reduce the feature set with mt.pp.feature_select."
)
if pca_solver(n_obs, n_vars, svd_solver, fallback="full") == "covariance_eigh":
from sklearn.utils.extmath import svd_flip
# The small feature covariance is eigendecomposed in place of the full matrix. The sign of each vector is
# free, so it is fixed by scikit-learn's convention, the one the sklearn-backed sites land on.
eigenvalues, vectors = np.linalg.eigh(centered.T @ centered)
singular = np.sqrt(np.clip(eigenvalues[::-1], 0.0, None))
right = svd_flip(None, vectors[:, ::-1].T, u_based_decision=False)[1]
else:
# Only an underdetermined reference needs the null directions; full_matrices allocates an (n_obs, n_obs) factor.
_, singular, right = np.linalg.svd(centered, full_matrices=n_obs <= n_vars)
if n_obs <= n_vars:
singular = np.concatenate((singular[:rank], np.repeat(singular[rank - 1], n_vars - rank)))
return centre, scale, singular, right, n_obs
def _whiten(singular: np.ndarray, right: np.ndarray, n_obs: int, method: str, epsilon: float) -> np.ndarray:
"""Build the whitening matrix from a decomposition, regularizing the singular values with ``epsilon``."""
singular = singular + epsilon
W = (right / singular[:, np.newaxis]).transpose() * np.sqrt(n_obs - 1)
if method.startswith("ZCA"):
W = W @ right
return W
def _fit(
reference: np.ndarray, method: str, epsilon: float, svd_solver: str | None = None
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Return ``(centre, scale, W)`` for the requested whitening.
Follows pycytominer's ``Spherize``: epsilon is added to the singular values, and when there are no more rows than features the null directions are padded with the smallest non-zero singular value.
"""
centre, scale, singular, right, n_obs = _decompose(reference, method, svd_solver)
return centre, scale, _whiten(singular, right, n_obs, method, epsilon)
# Opportunistic sampling for epsilon="auto": how far past the recipe grid to wander, the step, and when to stop.
_EPSILON_BOUNDS = (1e-8, 1e8)
_EPSILON_STEP = 10.0
_EPSILON_TOLERANCE = 1e-4
_EPSILON_REFINE_ROUNDS = 3
def _extend_epsilon(scored: dict[float, float], evaluate: Callable[[float], float]) -> None:
"""Grow an epsilon sweep past both ends of the grid and refine around the best, scoring in place.
The recipe grid can miss the optimum: it may sit past either end, so walk outward from both the low
and the high end while the score keeps improving (up to a cap); or it may fall between two coarse
points, so bisect towards the best's neighbours. ``scored`` maps each epsilon to its score and gains
an entry for every new candidate.
"""
low_cap, high_cap = _EPSILON_BOUNDS
# Probe both regimes, low and high, so a boundary optimum is never missed.
for direction in (_EPSILON_STEP, 1.0 / _EPSILON_STEP):
candidate = max(scored) if direction > 1.0 else min(scored)
previous = scored[candidate]
while low_cap <= candidate * direction <= high_cap:
candidate *= direction
score = evaluate(candidate)
scored[candidate] = score
if score <= previous + _EPSILON_TOLERANCE:
break
previous = score
# Refine around the best by bisecting towards its neighbours, since the grid is coarse.
for _ in range(_EPSILON_REFINE_ROUNDS):
candidates = sorted(scored)
best = max(candidates, key=lambda candidate: scored[candidate])
index = candidates.index(best)
neighbours = []
if index > 0:
neighbours.append(math.sqrt(candidates[index - 1] * best))
if index < len(candidates) - 1:
neighbours.append(math.sqrt(best * candidates[index + 1]))
fresh = [midpoint for midpoint in neighbours if midpoint not in scored]
if not fresh:
break
improved = False
for midpoint in fresh:
scored[midpoint] = evaluate(midpoint)
if scored[midpoint] > scored[best] + _EPSILON_TOLERANCE:
improved = True
if not improved:
break
def _select_epsilon(
adata: AnnData,
X: np.ndarray,
method: str,
reference: str | None,
label: str,
batch: str,
svd_solver: str | None = None,
) -> tuple[float, np.ndarray]:
"""Pick ``epsilon`` for ``epsilon="auto"``, returning it and the control mask it scored against.
The jump-profiling-recipe's fixed grid gives a reproducible base; each candidate is fitted on the reference and scored by the mean of its activity and replicability mAP. The search then samples past both ends of the grid and between points when the best sits there, so the chosen value can fall outside the recipe's range. The finite argmax, the full set of candidates and their scores are written to ``uns["mantispy"]["sphere_epsilon"]``.
"""
from mantispy._core.frames import as_frame
from mantispy.tl._map import REFERENCE_COLUMN, _resolve_mode, _score_map
reference_column = "Metadata_Control_Type" if reference in CONTROL_CLASSES else reference
missing = []
if label not in adata.obs:
missing.append(f"{label!r} to group replicates by")
if batch not in adata.obs:
missing.append(f"{batch!r} to draw the negatives from")
if reference is None:
missing.append("a negcon reference, but reference=None selects every row as the reference")
elif reference_column not in adata.obs:
missing.append(f"{reference_column!r} to select the negcon reference")
if missing:
raise ValueError(
"sphere(epsilon='auto') scores activity and replicability mAP to choose epsilon, which needs "
+ "; and ".join(missing)
+ ". Pass a float epsilon to skip the search."
)
is_control = reference_mask(adata, reference)
if not is_control.any():
raise ValueError(
f"sphere(epsilon='auto') found no reference rows selected by reference={reference!r}, "
"so there is nothing to score the candidates against."
)
obs = as_frame(adata.obs)
base_meta = obs[[column for column in obs.columns if column.startswith("Metadata_")]].reset_index(drop=True)
activity = _resolve_mode("activity", label=label, batch=batch)
replicability = _resolve_mode("replicability", label=label, batch=batch)
treated = ~is_control
control_positions = np.flatnonzero(is_control)
activity_meta = base_meta.copy()
activity_meta[REFERENCE_COLUMN] = np.where(is_control, np.arange(adata.n_obs), -1)
treated_meta = base_meta[treated].reset_index(drop=True)
# Guarded before _decompose so a no-replicate screen does not pay a throwaway SVD.
if treated_meta.empty or treated_meta[label].value_counts().max() < 2:
raise ValueError(
"sphere(epsilon='auto') needs replicate pairs to score, but no perturbation has two "
"or more wells. Pass a float epsilon."
)
# Everything that does not depend on the candidate is built once, so each score only re-whitens and re-scores.
centre, scale, singular, right, n_obs = _decompose(X[is_control], method, svd_solver)
standardized = (X - centre) / scale
def evaluate(candidate: float) -> float:
"""Score one epsilon by the mean of its activity and replicability mAP; -inf when non-finite."""
features = (standardized @ _whiten(singular, right, n_obs, method, candidate)).astype(np.float32)
# A non-finite profile makes this candidate's mAP non-finite, so it can never be chosen.
if not np.isfinite(features).all():
return -np.inf
activity_table, _ = _score_map(
activity_meta,
features,
activity,
null_size=0,
threshold=1.0,
seed=0,
drop_positions=control_positions,
warn=False,
compute_null=False,
)
replicability_table, _ = _score_map(
treated_meta,
features[treated],
replicability,
null_size=0,
threshold=1.0,
seed=0,
warn=False,
compute_null=False,
)
score = (activity_table["score"].mean() + replicability_table["score"].mean()) / 2
return float(score) if np.isfinite(score) else -np.inf
# The recipe's fixed grid and seed give a reproducible base; opportunistic sampling then grows it.
base = 10.0 ** np.random.default_rng((6, 12, 2022)).uniform(-5.0, 3.0, 25)
scored = {float(candidate): evaluate(float(candidate)) for candidate in base}
_extend_epsilon(scored, evaluate)
if not any(np.isfinite(score) for score in scored.values()):
raise ValueError(
"sphere(epsilon='auto') scored every candidate as non-finite, so none can be chosen. The "
"profiles are not finite after whitening; drop features or rows with missing values first "
"with mt.pp.feature_select (drop_na_columns) or mt.pp.drop_na."
)
grid = np.array(sorted(scored))
scores = np.array([scored[candidate] for candidate in grid])
# Non-finite candidates stay -inf, so argmax never lands on a NaN.
chosen = float(grid[int(np.argmax(scores))])
adata.uns.setdefault("mantispy", {})["sphere_epsilon"] = {"epsilon": chosen, "grid": grid, "scores": scores}
return chosen, is_control
[docs]
@inplace_or_copy()
def sphere(
adata: AnnData,
method: str = "ZCA-cor",
reference: str | None = "negcon",
epsilon: float | str = 1e-6,
by: str | None = None,
key_added: str | None = None,
label: str = "Metadata_Perturbation",
batch: str = "Metadata_Plate",
*,
svd_solver: str | None = None,
copy: bool = False,
) -> AnnData | None:
"""Whiten profiles with a transform fitted on the reference rows.
Args:
adata: Object to sphere.
Usually well-level profiles.
method: ``"ZCA"`` and ``"ZCA-cor"`` rotate back into the original feature basis, so the output columns still correspond to features and ``var`` still describes them.
``"PCA"`` and ``"PCA-cor"`` return principal components, which ``var`` no longer describes, and warn about it unless ``key_added`` is set.
The ``-cor`` variants whiten the correlation instead of the covariance, so high-variance features do not dominate.
reference: Rows to fit on: ``"negcon"`` for the controls, ``None`` for everything, or the name of a boolean ``obs`` column.
epsilon: Regularization added to the singular values, or ``"auto"`` to pick it from the data.
``"auto"`` scores each candidate by the mean of its activity and replicability mAP on the transformed profiles and keeps the argmax, starting from the jump-profiling-recipe grid and then sampling past both ends and between points when the best sits there, so the chosen value can fall outside the recipe's range.
It needs the ``label`` and ``batch`` columns and a negcon reference; the search is global, so ``by`` only changes how the chosen value is applied.
Every candidate, its score and the chosen value are written to ``uns["mantispy"]["sphere_epsilon"]``.
by: Fit and apply separately within each group of this column, e.g. per batch.
key_added: Write to ``layers[key_added]`` instead of overwriting ``X``.
label: The ``obs`` column naming the perturbation, which ``epsilon="auto"`` groups replicates by.
batch: The ``obs`` column naming the plate or batch, which ``epsilon="auto"`` draws negatives from.
svd_solver: How the reference is decomposed. ``None`` eigendecomposes the feature covariance when the reference has many more rows than features and otherwise takes the full SVD; ``"covariance_eigh"`` or ``"full"`` force one. The components match either way, up to the sign of each.
copy: Return a modified copy instead of mutating in place.
Returns:
``None``, or the modified copy.
Writes ``X`` or ``layers[key_added]``.
Raises:
ValueError: If ``method`` is unknown, ``epsilon`` is neither a float nor ``"auto"``, ``svd_solver`` is not ``None``, ``"covariance_eigh"`` or ``"full"``, ``reference`` selects no rows, a group has fewer than two reference rows, the reference holds missing or infinite values, a ``-cor`` method meets a zero-variance feature, the reference matrix is not full rank, ``epsilon="auto"`` is asked without the columns and reference the sweep needs, or every ``epsilon="auto"`` candidate scored non-finite.
"""
if method not in METHODS:
raise ValueError(f"method must be one of {METHODS}, got {method!r}")
if svd_solver is not None and svd_solver not in SVD_SOLVERS:
raise ValueError(f"svd_solver must be None or one of {SVD_SOLVERS}, got {svd_solver!r}")
if method.startswith("PCA") and key_added is None:
warnings.warn(
f"{method} rotates out of the feature basis, so X holds principal components while var still names "
"the features. Pass key_added= or use ZCA.",
UserWarning,
stacklevel=3,
)
X = get_matrix(adata).astype(np.float64)
auto = epsilon == "auto"
if auto:
epsilon, mask = _select_epsilon(adata, X, method, reference, label, batch, svd_solver)
elif isinstance(epsilon, bool) or not isinstance(epsilon, int | float):
# bool is an int subclass, so reject it before the float path turns `singular + True` into a transform.
raise ValueError(f"epsilon must be a float or 'auto', got {epsilon!r}")
else:
mask = reference_mask(adata, reference)
# Keep uns["mantispy"]["sphere_epsilon"] authoritative: a float call leaves no earlier sweep behind.
adata.uns.get("mantispy", {}).pop("sphere_epsilon", None)
if not mask.any():
raise ValueError(
f"no reference rows selected by reference={reference!r}. Check that obs['Metadata_Control_Type'] has "
"'negcon' rows (mt.pp.annotate_controls), or pass reference=None to use all rows."
)
codes, keys = group_codes(adata, by)
out = np.empty_like(X)
for group, key in enumerate(keys):
rows = np.flatnonzero(codes == group)
fit_rows = rows[mask[rows]]
if fit_rows.size < 2:
where = f" in group {key!r}" if by is not None else ""
raise ValueError(
f"sphering needs at least 2 reference rows to estimate a covariance, but "
f"reference={reference!r} selects {fit_rows.size}{where}. Leave by=None to fit across "
"groups, or check that the platemap labels the controls."
)
centre, scale, W = _fit(X[fit_rows], method, epsilon, svd_solver)
out[rows] = ((X[rows] - centre) / scale) @ W
result = out.astype(matrix_dtype(adata))
if key_added is None:
adata.X = result
else:
adata.layers[key_added] = result
return None
def _centre_scale(block: np.ndarray, where: str = "") -> tuple[np.ndarray, np.ndarray]:
"""The mean and spread of the reference rows, which every row is centred and scaled by.
The spread is the population standard deviation, as sklearn's ``StandardScaler`` computes it.
A dimension with no spread among the reference rows is left on its own scale rather than divided by zero, and warns.
"No spread" is read against the resolution of the values rather than as an exact zero, since a dimension an embedding is constant in still carries the rounding of the rotation that built it.
Args:
block: The reference rows.
where: Names the batch in the warning, when the reference rows are one batch's.
Returns:
``(centre, scale)``, with a scale of one for a dimension that has no spread.
"""
centre = block.mean(axis=0)
scale = block.std(axis=0, ddof=0)
# The tolerance scales with the row count because the rounding error in a standard deviation grows with its rows.
degenerate = scale <= block.shape[0] * np.finfo(scale.dtype).eps * np.abs(block).max(axis=0)
if (no_spread := int(degenerate.sum())) > 0:
get_logger().info(
"tvn: %d of %d dimension(s) have no spread among the %d reference row(s)%s and stay unscaled",
no_spread,
scale.size,
block.shape[0],
where,
deep="reduce to fewer components, or use more controls",
)
return centre, np.where(degenerate, 1.0, scale)
def _regularized_covariance(values: np.ndarray, epsilon: float) -> np.ndarray:
"""The covariance of ``values``, with ``epsilon`` added to its diagonal so it can be inverted."""
return np.cov(values, rowvar=False, ddof=1) + epsilon * np.eye(values.shape[1])
def _symmetric_power(matrix: np.ndarray, power: float) -> np.ndarray:
"""Raise a symmetric positive definite matrix to a real power through its eigendecomposition.
Args:
matrix: A symmetric, positive definite matrix.
power: The power to raise it to.
Returns:
The matrix raised to that power.
Raises:
ValueError: The matrix is singular or indefinite, which leaves the power undefined.
"""
eigenvalues, vectors = np.linalg.eigh(matrix)
if eigenvalues.min() <= 0:
raise ValueError(
f"a control covariance has a non-positive eigenvalue ({eigenvalues.min():.3g}), so it cannot be "
"inverted. Raise epsilon, which is added to every covariance diagonal."
)
return (vectors * eigenvalues**power) @ vectors.T
def _fit_tvn(
controls: np.ndarray, batches: np.ndarray, keys: Any, reference: str | None, epsilon: float, svd_solver: str | None
) -> tuple[Callable[[int, np.ndarray], np.ndarray], int]:
"""Fit :func:`tvn` on the control rows alone, which are all its statistics are taken from.
Args:
controls: The control rows, as ``float64``.
batches: The batch of each control row, as a code into ``keys``.
keys: The batches.
reference: How the controls were selected, for the error message.
epsilon: Added to the diagonal of every covariance.
svd_solver: How the control rotation is decomposed.
Returns:
``(align, width)``: ``align(batch, block)`` applies the fitted transform to rows of the batch at that position in ``keys``, and ``width`` is the number of components it returns.
"""
from sklearn.decomposition import PCA
centre, scale = _centre_scale(controls)
values = (controls - centre) / scale
solver = pca_solver(values.shape[0], values.shape[1], svd_solver, fallback="auto")
rotation = PCA(svd_solver=solver).fit(values)
values = rotation.transform(values)
# The two passes cannot be merged: the target below is taken from every control row after the first pass rewrote them.
order, offsets = group_offsets(batches, len(keys))
members = [order[offsets[group] : offsets[group + 1]] for group in range(len(keys))]
rescale = []
thin = []
for key, rows in zip(keys, members, strict=True):
if rows.size < 2:
raise ValueError(
f"batch {key!r} has {rows.size} row(s) selected by reference={reference!r}, "
"and aligning a batch needs at least 2 to estimate its covariance. Drop that batch, or "
"check that the platemap labels its controls."
)
if rows.size <= values.shape[1]:
thin.append(f"{key!r} ({rows.size})")
batch_centre, batch_scale = _centre_scale(values[rows], where=f" of batch {key!r}")
values[rows] = (values[rows] - batch_centre) / batch_scale
rescale.append((batch_centre, batch_scale))
if thin:
get_logger().warning(
"tvn: %d batch(es) have no more reference rows than the %d components (%s)",
len(thin),
values.shape[1],
", ".join(thin[:3]) + (", ..." if len(thin) > 3 else ""),
deep="their empty directions are amplified, so reduce to fewer components or pool the smaller batches",
)
target = _symmetric_power(_regularized_covariance(values, epsilon), 0.5)
whiten = [_symmetric_power(_regularized_covariance(values[rows], epsilon), -0.5) for rows in members]
def align(batch: int, block: np.ndarray) -> np.ndarray:
aligned = rotation.transform((block.astype(np.float64) - centre) / scale)
batch_centre, batch_scale = rescale[batch]
aligned = (aligned - batch_centre) / batch_scale
return aligned @ whiten[batch] @ target
return align, int(rotation.n_components_)
[docs]
@inplace_or_copy()
def tvn(
adata: AnnData,
batch_key: str = "Metadata_Batch",
reference: str | None = "negcon",
use_rep: str | None = "X_pca",
key_added: str = "X_tvn",
epsilon: float = 0.5,
*,
svd_solver: str | None = None,
copy: bool = False,
) -> AnnData | None:
"""Typical variation normalization, then align each batch's controls onto the pooled controls :cite:p:`Celik_2024`.
The controls define what an untreated well looks like, so they are what the transform is fitted on: the profiles are centred and scaled on them, rotated onto the principal components of the controls alone, and centred and scaled on them again within each batch.
The last step is CORAL: each batch is whitened by the covariance of its own controls and recoloured with the covariance of all of them, so a batch whose typical variation points in an unusual direction is brought onto the others rather than merely recentred.
Args:
adata: Object holding the profiles, usually one row per well.
batch_key: ``obs`` column naming the batches to align.
Each needs at least two reference rows.
reference: Rows the transform is fitted on: ``"negcon"`` for the controls, ``None`` for everything, or the name of a boolean ``obs`` column.
use_rep: Embedding to align, as :func:`~mantispy.pp.harmony` takes one, or ``None`` to align ``X`` itself.
Fitting the rotation on the controls of a wide feature matrix is expensive, so the default expects a reduction first, normally :func:`scanpy.pp.pca`.
On a large screen give that call ``svd_solver="covariance_eigh"``: with many more cells than features it decomposes the small feature covariance and keeps the reduction bounded in memory, where the default solver decomposes the whole matrix.
key_added: ``obsm`` key for the result.
epsilon: Added to the diagonal of every covariance before it is inverted.
The profiles are on the controls' own scale by then, so their variances are near one and the reference value of 0.5 is a substantial shrink toward isotropy.
svd_solver: How the control rotation is decomposed. ``None`` eigendecomposes the feature covariance when there are many more controls than components and otherwise uses scikit-learn's default; any scikit-learn solver name forces one. The rotation matches either way.
copy: Return a modified copy instead of writing in place.
Returns:
``None``, or the modified copy.
Writes ``obsm[key_added]``.
Raises:
KeyError: ``obs`` has no column ``batch_key`` or no column named by ``reference``, or ``obsm`` holds nothing under ``use_rep``.
ValueError: ``reference`` selects no rows, or fewer than two in some batch, which leaves that batch's covariance undefined.
ValueError: A control covariance is singular even after ``epsilon``, so it cannot be inverted.
Reachable by passing ``epsilon=0``.
Notes:
The rotation is fitted on the controls, so it keeps ``min(n_controls, n_features)`` components.
With fewer controls than features the result is narrower than the input, which is why this writes ``obsm`` and never ``X``: ``var`` would no longer describe the columns.
Batch correction methods disagree with each other often enough that one metric is not evidence.
Compare this with :func:`~mantispy.pp.harmony` on the same object using :func:`~mantispy.metrics.evaluate_integration`, and on a screen with annotated perturbations also :func:`~mantispy.metrics.known_relationships`, which is the measure :cite:t:`Celik_2024` selects it by.
Measured that way, it tends to trade replicate consistency for relationship recall, where :func:`~mantispy.pp.harmony` trades the other way.
Neither buys the other's gain, so which of the two readouts the screen is for is the question to answer before running either.
A lazy representation, such as the PCA of an object read with :func:`mantispy.io.read` and ``lazy=True``, is computed once and kept in ``obsm[use_rep]``, since it is as narrow as one in memory, and the result is written in memory with the same numbers.
With ``use_rep=None`` on a lazy ``X``, only the control rows are read to fit, and the result stays lazy.
What it needs is controls, per batch and not in total, because the covariance it whitens each batch by is estimated from that batch's controls alone.
A batch with fewer controls than the rotation has components cannot span the space, and the warning that says so is the sign to reduce to fewer components or to pool smaller batches together.
"""
# Checked before the rotation, which is the expensive part, so a mistyped column costs nothing.
if batch_key not in adata.obs:
raise KeyError(f"obs has no column {batch_key!r} naming the batches to align")
if use_rep is not None and use_rep not in adata.obsm:
raise KeyError(f"obsm has no representation {use_rep!r}; run sc.pp.pca first, or pass use_rep=None")
if use_rep is not None and is_lazy(adata.obsm[use_rep]):
# Computed once and kept: the fit and the apply would each run its whole graph, and so would every later step.
adata.obsm[use_rep] = materialize(adata.obsm[use_rep])
controls = reference_mask(adata, reference)
if not controls.any():
raise ValueError(
f"no reference rows selected by reference={reference!r}. Check that obs['Metadata_Control_Type'] has "
"'negcon' rows (mt.pp.annotate_controls), or pass reference=None to use all rows."
)
codes, keys = group_codes(adata, batch_key)
# Only the control rows are read to fit, so a lazy X is never loaded whole.
fitted = get_matrix(adata, rows=np.flatnonzero(controls), use_rep=use_rep).astype(np.float64)
align, width = _fit_tvn(fitted, codes[controls], keys, reference, epsilon, svd_solver)
position = {key: index for index, key in enumerate(keys)}
adata.obsm[key_added] = transform_grouped(
adata,
batch_key,
lambda key, block: align(position[key], block),
use_rep=use_rep,
n_cols=width,
rowwise=True,
grouping=(codes, keys),
)
return None
