from __future__ import annotations
import numpy as np
from anndata import AnnData
from scipy.stats import rankdata
from mantispy._core._corr import CHUNK_BYTES
from mantispy._core._reduce import get_matrix, group_codes
from mantispy._core.logging import get_logger
from mantispy._core.mutation import inplace_or_copy
def _complete_columns(values: np.ndarray) -> np.ndarray:
"""Which columns hold a finite value in every row, reduced one column block at a time."""
n_obs, n_vars = values.shape
out = np.empty(n_vars, dtype=bool)
width = max(int(CHUNK_BYTES / max(n_obs, 1) / values.itemsize), 1)
for start in range(0, n_vars, width):
block = slice(start, start + width)
out[block] = np.isfinite(values[:, block]).all(axis=0)
return out
def _xi(x: np.ndarray, values: np.ndarray, m: int, seed: int) -> np.ndarray:
"""Chatterjee's xi between ``x`` and every column of ``values``, all of which must be finite.
The statistic of :cite:t:`Lin_2022`, centred on its exact mean under independence and scaled by the ``sum(l * (n - l))`` of :cite:t:`Chatterjee_2020`, which is how Chatterjee handles ties in ``y``.
Without ties in ``y`` it is Lin and Han's statistic exactly.
"""
n = x.size
shuffled = np.random.default_rng(seed).permutation(n)
order = shuffled[np.argsort(x[shuffled], kind="stable")]
# As integers: rankdata keeps float32 input float32, whose sums are inexact above 2**24.
ranks = rankdata(values[order], method="max", axis=0).astype(np.int64)
drop = sum(np.maximum(ranks[:-step] - ranks[step:], 0).sum(axis=0) for step in range(1, m + 1))
# The scale is zero only for a constant column, which has no xi.
below = rankdata(values, method="min", axis=0).astype(np.float64) - 1
with np.errstate(invalid="ignore"):
scaled = drop / (below * (n - below)).sum(axis=0) * n * (n - 1)
return -2 + (3 * n * m - scaled) / (n * m + m * (m + 1) / 4)
def chatterjee_xi(x: np.ndarray, y: np.ndarray, m: int = 1, seed: int = 0, min_finite: int = 40) -> np.ndarray:
"""Chatterjee's xi between ``x`` and every column of ``y``.
Args:
x: The variable the others are tested against, such as a group code, a dose or a covariate.
y: One column, or a matrix of them.
Every column is scored against ``x`` in one pass.
m: Right nearest neighbors, as in :cite:t:`Lin_2022`.
Larger ``m`` has the same limit under dependence and a lower noise floor under independence, so a fixed threshold is more reliable.
seed: Seed for breaking ties in ``x``.
min_finite: Fewest finite pairs a column may be scored on, raised to ``m + 2`` when it is below that.
Under independence xi has standard deviation ``sqrt(2 / (5 * n))``, which at 40 pairs equals the 0.1 threshold :func:`feature_select_chatterjee` selects on, so a column measured fewer times than this cannot be told from noise.
Returns:
One xi per column of ``y``, NaN for a constant column or one with fewer than ``min_finite`` finite pairs.
Raises:
ValueError: If ``m`` is below 1.
Notes:
Ties in ``x`` are broken at random, as the coefficient requires.
A group label is almost all ties, and breaking them by row order would score the row order as structure.
Ties in ``y`` are handled as :cite:t:`Chatterjee_2020` handles them, without randomness: tied values share a rank.
A zero-inflated Zernike, Granularity or RadialDistribution column is mostly ties, and breaking them at random scores one that is a step function of ``x`` at 0.51 instead of 0.99.
xi is defined on complete pairs, so each column is scored on the rows where both it and ``x`` are finite, and two columns missing different rows are scored on different subsets.
Ranking a column that still holds NaN sorts the missing rows last, which turns a feature that is merely unmeasured in one group into a step function of the group and scores it as dependence.
How much of a feature is missing is a question for :func:`~mantispy.pp.feature_select` and its ``drop_na_columns`` operation; xi reports on the part that was measured.
"""
if m < 1:
raise ValueError(f"m must be at least 1, got {m}")
values = np.atleast_2d(y.T).T if y.ndim > 1 else y[:, None]
finite_x = np.isfinite(x)
floor = max(min_finite, m + 2)
complete = _complete_columns(values) if finite_x.all() else np.zeros(values.shape[1], dtype=bool)
if complete.all():
return _xi(x, values, m, seed) if x.size >= floor else np.full(values.shape[1], np.nan)
scores = np.full(values.shape[1], np.nan)
if complete.any() and x.size >= floor:
scores[complete] = _xi(x, values[:, complete], m, seed)
for column in np.flatnonzero(~complete):
rows = np.flatnonzero(finite_x & np.isfinite(values[:, column]))
if rows.size >= floor:
scores[column] = _xi(x[rows], values[rows, column][:, None], m, seed)[0]
return scores
[docs]
@inplace_or_copy()
def feature_select_chatterjee(
adata: AnnData,
by: str = "Metadata_Perturbation",
threshold: float = 0.1,
m: int = 1,
seed: int = 0,
min_finite: int = 40,
key_added: str = "selected_chatterjee",
copy: bool = False,
) -> AnnData | None:
"""Keep features whose values depend on the group, monotonically or otherwise.
Args:
adata: Object to select features on.
by: ``obs`` column the features are tested against.
threshold: Keep features scoring above this.
xi is near zero under independence and approaches one when the feature is a deterministic function of the group, so the threshold is comparable across datasets in a way a correlation cutoff is not.
m: Right nearest neighbors :cite:p:`Lin_2022`.
``m=1`` is the coefficient of :cite:t:`Chatterjee_2020` up to a term of order ``1/n``; larger values lower the noise floor without changing what the statistic converges to.
seed: Seed for breaking ties between rows of the same group.
min_finite: Fewest finite values a feature may be scored on.
A feature measured fewer times than this scores NaN, which is above no threshold and so is never selected.
key_added: Name of the boolean ``var`` column written.
copy: Return a modified copy instead of mutating in place.
Returns:
``None``, or the modified copy.
Writes ``var[key_added]`` and the statistic itself to ``var["chatterjee_xi"]``, which is NaN for a feature that is constant or too sparsely measured to score.
Raises:
ValueError: If ``by`` has a single group, so that no feature can depend on it.
Notes:
Run it after :func:`~mantispy.pp.feature_select`, which drops redundant or unmeasurable features; this keeps the features that carry information about the perturbation.
Subset with ``mt.pp.subset_features(adata, key="selected_chatterjee")``.
A feature is scored on the rows where it was measured, so one that is missing in a whole group is scored against the groups that do have it rather than against the pattern of what is missing.
xi reaches one only for a noiseless function of the group, so real values are much lower and the default threshold of ``0.1`` is permissive.
Check the distribution in ``var["chatterjee_xi"]`` before relying on a fixed cutoff.
With shuffled group labels xi collapses toward zero while the real values barely change, so the score reflects dependence on the group rather than noise.
``m=1`` is the default because the threshold was calibrated there.
scmorph uses ``m=5``, and on input free of ties the two implementations agree to 1e-9.
scmorph breaks ties in ``y`` at random, so on tied input the two differ.
"""
codes, keys = group_codes(adata, by)
if len(keys) < 2:
raise ValueError(f"{by!r} has one group, so no feature can depend on it")
scores = chatterjee_xi(codes.astype(float), get_matrix(adata), m=m, seed=seed, min_finite=min_finite)
adata.var["chatterjee_xi"] = scores
# NaN is above no threshold, so a feature that cannot be scored drops out here.
selected = scores > threshold
adata.var[key_added] = selected
get_logger().info(
"chatterjee kept %d of %d features, %d constant or too sparsely measured to score",
int(selected.sum()),
adata.n_vars,
int(np.isnan(scores).sum()),
)
return None
