"""
Diversity analysis functions for AIRR-seq repertoires.
Implements clonality, diversity metrics, and repertoire comparison functions
similar to the R LymphoSeq2 package.
"""
import numpy as np
import polars as pl
import pandas as pd
from typing import Union, Optional, List, Dict, Any
from .statistics import gini_coefficient, shannon_entropy
[docs]
def clonality(
data: Union[pl.DataFrame, pd.DataFrame],
rarefy: bool = False,
min_count: Optional[int] = None,
iterations: int = 100,
group_by: str = "repertoire_id"
) -> Union[pl.DataFrame, pd.DataFrame]:
"""
Calculate clonality and diversity metrics for each repertoire.
This function calculates standard repertoire diversity metrics including
clonality, Gini coefficient, convergence, and unique productive sequences.
Clonality is calculated as 1 - (normalized Shannon entropy), where Shannon
entropy is normalized by dividing by log2(n_unique_sequences). This ensures
clonality is in the range [0, 1], where:
- 0 indicates maximum diversity (all clones equally abundant)
- 1 indicates minimum diversity (single dominant clone)
Args:
data: Input data frame with AIRR-formatted sequences
rarefy: Whether to perform rarefaction analysis
min_count: Minimum sequence count for rarefaction
iterations: Number of rarefaction iterations
group_by: Column to group by (default: repertoire_id)
Returns:
Data frame with diversity metrics for each repertoire:
- clonality: 1 - normalized Shannon entropy (0-1)
- gini_coefficient: Measure of inequality (0-1)
- unique_productive_sequences: Number of unique sequences
- total_count: Total sequence reads
- top_productive_sequence: Frequency of most abundant clone
- convergence: Ratio of total sequences to unique sequences
Examples:
>>> import lymphoseq as ls
>>> data = ls.read_immunoseq("data/")
>>> diversity = ls.clonality(data)
>>> print(diversity)
"""
is_polars = isinstance(data, pl.DataFrame)
if is_polars:
return _clonality_polars(data, rarefy, min_count, iterations, group_by)
else:
return _clonality_pandas(data, rarefy, min_count, iterations, group_by)
def _clonality_polars(
data: pl.DataFrame,
rarefy: bool,
min_count: Optional[int],
iterations: int,
group_by: str
) -> pl.DataFrame:
"""Calculate clonality using Polars."""
# Ensure we have required columns
required_cols = [group_by, "junction_aa", "duplicate_count"]
missing_cols = [col for col in required_cols if col not in data.columns]
if missing_cols:
raise ValueError(f"Missing required columns: {missing_cols}")
# Convert duplicate_count to numeric if it's string
if data[group_by].dtype == pl.Utf8:
data = data.with_columns(
pl.col("duplicate_count").cast(pl.Float64, strict=False).fill_null(1.0)
)
if not rarefy:
# Standard diversity calculation - need to use apply for complex metrics
filtered_data = (
data
.filter(pl.col("productive") == True)
.filter(pl.col("junction_aa") != "")
)
# Calculate using Python for complex metrics
results = []
for repertoire in filtered_data[group_by].unique():
rep_data = filtered_data.filter(pl.col(group_by) == repertoire)
counts = rep_data["duplicate_count"].to_numpy()
frequencies = counts / counts.sum()
# Calculate normalized clonality
n_unique = rep_data["junction_aa"].n_unique()
max_entropy = np.log2(n_unique) if n_unique > 1 else 1.0
normalized_entropy = shannon_entropy(frequencies) / max_entropy if max_entropy > 0 else 0
results.append({
group_by: repertoire,
"total_sequences": len(rep_data),
"unique_productive_sequences": n_unique,
"total_count": int(counts.sum()),
"clonality": 1 - normalized_entropy,
"gini_coefficient": gini_coefficient(counts),
"top_productive_sequence": counts.max() / counts.sum(),
"convergence": len(rep_data) / n_unique
})
results = pl.DataFrame(results)
else:
# Rarefaction analysis
if min_count is None:
min_count = (
data
.group_by(group_by)
.agg(pl.col("duplicate_count").sum().alias("total"))
.select(pl.col("total").min())
.item()
)
results = _perform_rarefaction_polars(data, min_count, iterations, group_by)
return results
def _clonality_pandas(
data: pd.DataFrame,
rarefy: bool,
min_count: Optional[int],
iterations: int,
group_by: str
) -> pd.DataFrame:
"""Calculate clonality using pandas."""
# Ensure we have required columns
required_cols = [group_by, "junction_aa", "duplicate_count"]
missing_cols = [col for col in required_cols if col not in data.columns]
if missing_cols:
raise ValueError(f"Missing required columns: {missing_cols}")
# Filter productive sequences
productive_data = data[
(data["productive"] == True) &
(data["junction_aa"] != "") &
(data["junction_aa"].notna())
].copy()
# Convert duplicate_count to numeric
productive_data["duplicate_count"] = pd.to_numeric(
productive_data["duplicate_count"], errors="coerce"
).fillna(1.0)
if not rarefy:
# Standard diversity calculation
def calculate_metrics(group):
counts = group["duplicate_count"].values
frequencies = counts / counts.sum()
n_unique = group["junction_aa"].nunique()
# Calculate normalized clonality
max_entropy = np.log2(n_unique) if n_unique > 1 else 1.0
normalized_entropy = shannon_entropy(frequencies) / max_entropy if max_entropy > 0 else 0
return pd.Series({
"total_sequences": len(group),
"unique_productive_sequences": n_unique,
"total_count": counts.sum(),
"clonality": 1 - normalized_entropy,
"gini_coefficient": gini_coefficient(counts),
"top_productive_sequence": counts.max() / counts.sum(),
"convergence": len(group) / n_unique
})
results = productive_data.groupby(group_by).apply(calculate_metrics).reset_index()
else:
# Rarefaction analysis
if min_count is None:
min_count = productive_data.groupby(group_by)["duplicate_count"].sum().min()
results = _perform_rarefaction_pandas(productive_data, min_count, iterations, group_by)
return results
def _calculate_clonality_expr():
"""Polars expression for clonality calculation."""
return (
1 - (
-pl.col("duplicate_count").map_elements(
lambda counts: sum(
(c / sum(counts)) * np.log2(c / sum(counts))
for c in counts if c > 0
) if sum(counts) > 0 else 0,
return_dtype=pl.Float64
) / np.log2(pl.len())
)
)
def _calculate_gini_expr():
"""Polars expression for Gini coefficient calculation."""
return pl.col("duplicate_count").map_elements(
lambda counts: gini_coefficient(counts),
return_dtype=pl.Float64
)
def _calculate_convergence_expr():
"""Polars expression for convergence calculation."""
return pl.len() / pl.col("junction_aa").n_unique()
def _perform_rarefaction_polars(
data: pl.DataFrame,
min_count: int,
iterations: int,
group_by: str
) -> pl.DataFrame:
"""Perform rarefaction analysis using Polars."""
results = []
repertoires = data[group_by].unique().to_list()
for repertoire in repertoires:
repertoire_data = data.filter(pl.col(group_by) == repertoire)
total_reads = repertoire_data["duplicate_count"].sum()
# Skip if repertoire doesn't have enough reads
if total_reads < min_count:
continue
# Prepare data for sampling - expand sequences by their counts
sequences = repertoire_data["junction_aa"].to_list()
counts = repertoire_data["duplicate_count"].to_list()
iteration_metrics = []
for _ in range(iterations):
# Sample with replacement proportional to counts
sampled_seqs = _weighted_sample(sequences, counts, min_count)
# Count unique sequences and their frequencies
unique_seqs, seq_counts = np.unique(sampled_seqs, return_counts=True)
frequencies = seq_counts / seq_counts.sum()
# Calculate metrics for this iteration
n_unique = len(unique_seqs)
max_entropy = np.log2(n_unique) if n_unique > 1 else 1.0
normalized_entropy = shannon_entropy(frequencies) / max_entropy if max_entropy > 0 else 0
metrics = {
"total_sequences": len(sampled_seqs),
"unique_productive_sequences": n_unique,
"total_count": min_count,
"clonality": 1 - normalized_entropy,
"gini_coefficient": gini_coefficient(seq_counts),
"top_productive_sequence": seq_counts.max() / seq_counts.sum(),
"convergence": len(sampled_seqs) / n_unique
}
iteration_metrics.append(metrics)
# Average across iterations
avg_metrics = {
group_by: repertoire,
"total_sequences": np.mean([m["total_sequences"] for m in iteration_metrics]),
"unique_productive_sequences": np.mean([m["unique_productive_sequences"] for m in iteration_metrics]),
"total_count": min_count,
"clonality": np.mean([m["clonality"] for m in iteration_metrics]),
"gini_coefficient": np.mean([m["gini_coefficient"] for m in iteration_metrics]),
"top_productive_sequence": np.mean([m["top_productive_sequence"] for m in iteration_metrics]),
"convergence": np.mean([m["convergence"] for m in iteration_metrics])
}
results.append(avg_metrics)
return pl.DataFrame(results)
def _weighted_sample(sequences: list, counts: list, sample_size: int) -> np.ndarray:
"""Sample sequences with replacement, weighted by their counts."""
# Convert counts to probabilities
total = sum(counts)
probabilities = [c / total for c in counts]
# Sample with replacement
sampled_indices = np.random.choice(
len(sequences),
size=sample_size,
replace=True,
p=probabilities
)
return np.array([sequences[i] for i in sampled_indices])
def _perform_rarefaction_pandas(
data: pd.DataFrame,
min_count: int,
iterations: int,
group_by: str
) -> pd.DataFrame:
"""Perform rarefaction analysis using pandas."""
results = []
for repertoire, group in data.groupby(group_by):
total_reads = group["duplicate_count"].sum()
if total_reads >= min_count:
# Prepare sequences and weights
sequences = group["junction_aa"].values
counts = group["duplicate_count"].values
probabilities = counts / counts.sum()
iteration_results = []
for _ in range(iterations):
# Sample sequences proportional to their counts
sampled_seqs = np.random.choice(
sequences,
size=int(min_count),
replace=True,
p=probabilities
)
# Count unique sequences
unique_seqs, seq_counts = np.unique(sampled_seqs, return_counts=True)
frequencies = seq_counts / seq_counts.sum()
# Calculate metrics for this iteration
n_unique = len(unique_seqs)
max_entropy = np.log2(n_unique) if n_unique > 1 else 1.0
normalized_entropy = shannon_entropy(frequencies) / max_entropy if max_entropy > 0 else 0
metrics = {
"total_sequences": len(sampled_seqs),
"unique_productive_sequences": n_unique,
"total_count": min_count,
"clonality": 1 - normalized_entropy,
"gini_coefficient": gini_coefficient(seq_counts),
"top_productive_sequence": seq_counts.max() / seq_counts.sum(),
"convergence": len(sampled_seqs) / n_unique
}
iteration_results.append(metrics)
# Average across iterations
avg_metrics = pd.DataFrame(iteration_results).mean().to_dict()
avg_metrics[group_by] = repertoire
results.append(avg_metrics)
return pd.DataFrame(results)
[docs]
def diversity_metrics(
data: Union[pl.DataFrame, pd.DataFrame],
group_by: str = "repertoire_id"
) -> Union[pl.DataFrame, pd.DataFrame]:
"""
Calculate comprehensive diversity metrics for repertoires.
Args:
data: Input data frame with AIRR-formatted sequences
group_by: Column to group by
Returns:
Data frame with diversity metrics
"""
return clonality(data, rarefy=False, group_by=group_by)
[docs]
def common_sequences(
data: Union[pl.DataFrame, pd.DataFrame],
group_by: str = "repertoire_id",
min_samples: int = 2
) -> Union[pl.DataFrame, pd.DataFrame]:
"""
Find sequences that are common across multiple repertoires.
Args:
data: Input data frame with AIRR-formatted sequences
group_by: Column to group by
min_samples: Minimum number of samples a sequence must appear in
Returns:
Data frame with common sequences
"""
is_polars = isinstance(data, pl.DataFrame)
if is_polars:
# Count how many repertoires each sequence appears in
sequence_counts = (
data
.filter(pl.col("junction_aa") != "")
.group_by("junction_aa")
.agg([
pl.col(group_by).n_unique().alias("repertoire_count"),
pl.col("duplicate_count").sum().alias("total_count")
])
.filter(pl.col("repertoire_count") >= min_samples)
.sort("repertoire_count", descending=True)
)
else:
# Pandas version
sequence_counts = (
data[data["junction_aa"] != ""]
.groupby("junction_aa")
.agg({
group_by: "nunique",
"duplicate_count": "sum"
})
.rename(columns={group_by: "repertoire_count"})
.reset_index()
)
sequence_counts = sequence_counts[
sequence_counts["repertoire_count"] >= min_samples
].sort_values("repertoire_count", ascending=False)
return sequence_counts
[docs]
def rarefaction_curve(
data: Union[pl.DataFrame, pd.DataFrame],
repertoire_ids: Optional[List[str]] = None,
step_size: Optional[int] = None,
iterations: int = 100,
group_by: str = "repertoire_id"
) -> Union[pl.DataFrame, pd.DataFrame]:
"""
Generate rarefaction curves for diversity analysis.
Creates rarefaction curves by subsampling at multiple depths and calculating
diversity metrics. This helps compare repertoires with different sequencing
depths and assess sampling completeness.
Args:
data: Input DataFrame with AIRR-formatted sequences. Must contain
'duplicate_count', 'junction_aa', and group_by columns.
repertoire_ids: List of repertoire IDs to analyze. If None, uses all
repertoires (default: None)
step_size: Number of reads between sampling points. If None, uses
10 steps across the range (default: None)
iterations: Number of iterations for each sampling depth (default: 100)
group_by: Column to group by (default: "repertoire_id")
Returns:
DataFrame with rarefaction curve data containing:
- repertoire_id: Repertoire identifier
- sample_size: Number of reads sampled
- unique_sequences: Mean number of unique sequences
- unique_sequences_sd: Standard deviation across iterations
- clonality: Mean clonality at this depth
- clonality_sd: Standard deviation of clonality
Examples:
>>> import lymphoseq as ls
>>> data = ls.read_immunoseq("data/immunoseq/")
>>> # Generate rarefaction curves for all samples
>>> rarefaction = ls.rarefaction_curve(data, iterations=100)
>>> print(rarefaction.head())
>>> # Generate curves for specific samples
>>> rarefaction = ls.rarefaction_curve(
... data,
... repertoire_ids=["S1", "S2"],
... step_size=1000
... )
>>> # Plot rarefaction curves (if using pandas)
>>> import matplotlib.pyplot as plt
>>> for rep in rarefaction["repertoire_id"].unique():
... subset = rarefaction[rarefaction["repertoire_id"] == rep]
... plt.plot(subset["sample_size"], subset["unique_sequences"], label=rep)
>>> plt.xlabel("Sample size")
>>> plt.ylabel("Unique sequences")
>>> plt.legend()
>>> plt.show()
Notes:
- Rarefaction helps compare diversity across samples with different depths
- Sampling is done with replacement, weighted by sequence abundance
- Curves that haven't plateaued suggest undersampling
- Standard deviations indicate sampling variance
See Also:
- clonality(): Calculate diversity metrics with optional rarefaction
- diversity_metrics(): Comprehensive diversity analysis
"""
is_polars = isinstance(data, pl.DataFrame)
# Filter for specified repertoires
if repertoire_ids is not None:
if is_polars:
filtered = data.filter(pl.col(group_by).is_in(repertoire_ids))
else:
filtered = data[data[group_by].isin(repertoire_ids)].copy()
else:
filtered = data
# Filter productive sequences
if is_polars:
filtered = filtered.filter(
(pl.col("productive") == True) & (pl.col("junction_aa") != "")
)
else:
filtered = filtered[
(filtered["productive"] == True) &
(filtered["junction_aa"] != "") &
(filtered["junction_aa"].notna())
].copy()
results = []
# Process each repertoire
if is_polars:
repertoire_list = filtered[group_by].unique().to_list()
else:
repertoire_list = filtered[group_by].unique()
for repertoire in repertoire_list:
if is_polars:
rep_data = filtered.filter(pl.col(group_by) == repertoire)
sequences = rep_data["junction_aa"].to_list()
counts = rep_data["duplicate_count"].to_list()
total_reads = sum(counts)
else:
rep_data = filtered[filtered[group_by] == repertoire]
sequences = rep_data["junction_aa"].values
counts = rep_data["duplicate_count"].values
total_reads = counts.sum()
# Determine sampling points
if step_size is None:
# Use 10 evenly spaced points
max_size = int(total_reads)
step_size = max(1, max_size // 10)
sample_sizes = list(range(step_size, int(total_reads) + 1, step_size))
if total_reads not in sample_sizes:
sample_sizes.append(int(total_reads))
# Calculate probabilities for sampling
probabilities = [c / total_reads for c in counts]
# For each sampling depth
for sample_size in sample_sizes:
iteration_results = {
"unique_sequences": [],
"clonality": []
}
# Perform multiple iterations
for _ in range(iterations):
# Sample sequences
sampled_seqs = np.random.choice(
sequences,
size=sample_size,
replace=True,
p=probabilities
)
# Count unique sequences
unique_seqs, seq_counts = np.unique(sampled_seqs, return_counts=True)
frequencies = seq_counts / seq_counts.sum()
# Calculate metrics
n_unique = len(unique_seqs)
max_entropy = np.log2(n_unique) if n_unique > 1 else 1.0
normalized_entropy = shannon_entropy(frequencies) / max_entropy if max_entropy > 0 else 0
iteration_results["unique_sequences"].append(n_unique)
iteration_results["clonality"].append(1 - normalized_entropy)
# Calculate mean and SD across iterations
results.append({
group_by: repertoire,
"sample_size": sample_size,
"unique_sequences": np.mean(iteration_results["unique_sequences"]),
"unique_sequences_sd": np.std(iteration_results["unique_sequences"]),
"clonality": np.mean(iteration_results["clonality"]),
"clonality_sd": np.std(iteration_results["clonality"])
})
# Convert to DataFrame
if is_polars:
result_df = pl.DataFrame(results)
result_df = result_df.sort([group_by, "sample_size"])
else:
result_df = pd.DataFrame(results)
result_df = result_df.sort_values([group_by, "sample_size"])
return result_df