"""Semantic density analysis for PsychiatryNLPKit.
Measures dimensional properties of semantic spaces derived from token embeddings.
Schizophrenia patients can show high or low semantic density, and lower intrinsic
dimensionality indicates more redundant speech.
"""
import logging
import numpy as np
import torch
from skdim.id import MLE
from sklearn.decomposition import PCA
logger = logging.getLogger(__name__)
[docs]
def pca_density_metrics(
token_embedding_vectors: dict[str, torch.Tensor],
sections: list[str] | None = None,
) -> dict[str, dict[str, float]]:
"""PCA-based density metrics per section.
PCA is applied in **token space** (features=tokens, i.e. ``X.T``) so metrics
reflect how many token directions are needed to explain semantic variance.
During statistical analysis, total paragraph length must be controlled.
Notes:
Theoretical basis - While healthy controls tend to have moderately
compressible semantic space, schizophrenia patients can have high or low
semantic density (Palominos et al., 2025).
Args:
token_embedding_vectors: Dict mapping section names to token-level
embedding tensors (from ``TextData.token_embedding_vectors``).
sections: Sections to process. ``None`` processes all sections in the dict.
Returns:
Dict mapping section names to a metric dict with keys
``"Ncomp_90"``, ``"Pcomp_90"``, and ``"ExVar_2"``. Empty sections
receive ``float("nan")`` for all metrics.
References:
Palominos, C., Stein, F., Kircher, T., Ayesa-Arriola, R., Palaniyappan,
L., Homan, P., Sommer, I. E., & Hinzen, W. (2025). Lexical meaning is
lower dimensional in psychosis. Scientific Reports, 16(1), 859.
https://doi.org/10.1038/s41598-025-30443-1
"""
if sections is None:
sections = list(token_embedding_vectors.keys())
results: dict[str, dict[str, float]] = {}
for sec in sections:
if sec not in token_embedding_vectors:
logger.warning("Section %s not found in embeddings, skipping", sec)
continue
vectors = token_embedding_vectors[sec]
n_tokens = vectors.shape[0]
# Handle edge case: fewer than 2 tokens.
if n_tokens < 2:
results[sec] = {
"Ncomp_90": float("nan"),
"Pcomp_90": float("nan"),
"ExVar_2": float("nan"),
}
continue
# Move to CPU for sklearn.
X = vectors.detach().cpu().numpy()
# PCA in **token space**: transpose so tokens are samples, embedding dims are features.
n_components = min(n_tokens, X.shape[1])
pca = PCA(n_components=n_components)
pca.fit(X.T)
explained_variance_ratio = np.asarray(
pca.explained_variance_ratio_, dtype=float
)
cumulative_variance = np.cumsum(explained_variance_ratio)
ncomp_90 = int(np.searchsorted(cumulative_variance, 0.9, side="left") + 1)
pcomp_90 = ncomp_90 / n_tokens
exvar_2 = float(
explained_variance_ratio[: min(2, len(explained_variance_ratio))].sum()
)
results[sec] = {
"Ncomp_90": float(ncomp_90),
"Pcomp_90": float(pcomp_90),
"ExVar_2": exvar_2,
}
return results
[docs]
def intrinsic_dimensionality_density(
token_embedding_vectors: dict[str, torch.Tensor],
sections: list[str] | None = None,
k: int | None = None,
) -> dict[str, dict[str, float]]:
"""Estimate intrinsic dimensionality using MLE (Levina & Bickel, 2004).
Intrinsic dimensionality quantifies the local geometric complexity of the
semantic space formed by token embeddings. Lower values indicate more
redundant speech.
Notes:
Theoretical basis - Lower intrinsic dimensionality indicates more
redundant speech (Palominos et al., 2025).
Args:
token_embedding_vectors: Dict mapping section names to token-level
embedding tensors (from ``TextData.token_embedding_vectors``).
sections: Sections to process. ``None`` processes all sections in the dict.
k: Number of neighbors for MLE estimator. Defaults to
``min(10, n_samples - 1)``.
Returns:
Dict mapping section names to a metric dict with key ``"ID_MLE"``.
Empty sections receive ``float("nan")``.
References:
Palominos, C., Stein, F., Kircher, T., Ayesa-Arriola, R., Palaniyappan,
L., Homan, P., Sommer, I. E., & Hinzen, W. (2025). Lexical meaning is
lower dimensional in psychosis. Scientific Reports, 16(1), 859.
https://doi.org/10.1038/s41598-025-30443-1
"""
if sections is None:
sections = list(token_embedding_vectors.keys())
results: dict[str, dict[str, float]] = {}
for sec in sections:
if sec not in token_embedding_vectors:
logger.warning("Section %s not found in embeddings, skipping", sec)
continue
vectors = token_embedding_vectors[sec]
n_tokens = vectors.shape[0]
# Handle edge case: fewer than 2 tokens.
if n_tokens < 2:
results[sec] = {"ID_MLE": float("nan")}
continue
# Move to CPU for skdim.
X = vectors.detach().cpu().numpy()
k_param = k if k is not None else min(10, n_tokens - 1)
id_estimator = MLE(K=k_param)
intrinsic_dim = float(id_estimator.fit_transform(X))
results[sec] = {"ID_MLE": intrinsic_dim}
return results
__all__: list[str] = [
"pca_density_metrics",
"intrinsic_dimensionality_density",
]