Source code for acore.drift_correction.cpca_drift_correction

"""
Functions for metabolomics drift correction by
Common Principal Components Analysis (CPCA).
"""

import logging
import math

import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

logger = logging.getLogger(__name__)




def check_missingness(df: pd.DataFrame, rows_to_check: list):
    """
    This function checks for NAs in the data frame inside some user-provided rows.

    Parameters
    ---------
    df              : samples as rows, features as columns
    rows_to_check   : list of rows that should be checked for missingness

    Returns
    --------
    Boolean: True if there is missingness, False if there is no missingness.
    """
    na_counts = df.loc[rows_to_check].isna().sum()
    na_features = df.loc[rows_to_check].isna().any(axis=0).sum()

    logger.debug(
        f"Features (cols) with at least one NA: {na_features} / {len(df.columns)}"
    )
    logger.debug("Samples (rows) with at least one NA: %s", na_counts[na_counts > 0])

    if na_counts.any():
        return True
    else:
        return False





def run_cpca_drift_correction(
    df: pd.DataFrame, sample_rows, qc_rows, n_comps: int = 1
) -> pd.DataFrame:
    """
    Corrects technical drift using Common Principal Components Analysis (CPCA).
    Adapted from https://github.com/m-baralt/metabolomics_incident_diabetes.

    Parameters
    ----------
    df          : samples as rows, features as columns
    sample_rows : list of sample row indices
    qc_rows     : list of QC row indices
    n_comps     : number of common principal components to remove (default 1)

    Returns
    -------
    Full input DataFrame with corrected values applied to the intensity rows
    (sample_rows + qc_rows). Rows outside those arguments are returned unchanged.
    """
    intensity_rows = sample_rows + qc_rows

    X = df.loc[intensity_rows].values.astype(float)  # shape: (n_samples, n_features)

    if np.isnan(X).any():
        raise ValueError("NA values present in dataset. Consider imputation first.")

    sample_idx = list(range(len(sample_rows)))
    qc_idx = list(range(len(sample_rows), len(intensity_rows)))

    scaler = StandardScaler()
    Xs = scaler.fit_transform(X)  # shape: (n_samples, n_features)

    cov_sample = np.cov(Xs[sample_idx], rowvar=False)
    cov_qc = np.cov(Xs[qc_idx], rowvar=False)
    avg_cov = (cov_sample + cov_qc) / 2

    k = max(n_comps, 3)
    eigenvalues, eigenvectors = np.linalg.eigh(avg_cov)
    eigenvectors = eigenvectors[:, np.argsort(eigenvalues)[::-1]]
    cpcs = eigenvectors[:, :k]

    var_projected = np.sum((Xs @ cpcs) ** 2, axis=0)
    var_cpc = np.round(var_projected / np.sum(Xs**2), 3)[:n_comps]
    logger.info(
        "CPC explained variance: %s",
        str(dict(zip([f"CPC{i+1}" for i in range(n_comps)], var_cpc))),
    )

    W = cpcs[:, :n_comps]
    Xs_corrected = Xs - Xs @ W @ W.T
    X_corrected = scaler.inverse_transform(
        Xs_corrected
    )  # shape: (n_samples, n_features)

    df_out = df.astype(float)
    df_out.loc[intensity_rows] = X_corrected
    return df_out





def cpca_centroid(
    df: pd.DataFrame,
    sample_rows,  # list of row indices, OR dict {group_name: [row indices]}
    qc_rows: list,
    log_transform: bool = True,
):
    # Normalise sample_rows to a dict
    if isinstance(sample_rows, list):
        sample_groups = {"Samples": sample_rows}
    else:
        sample_groups = sample_rows

    # Build ordered row + label arrays
    all_rows, labels = [], []
    for group, rows in sample_groups.items():
        for r in rows:
            if r in df.index:
                all_rows.append(r)
                labels.append(group)
    for r in qc_rows:
        if r in df.index:
            all_rows.append(r)
            labels.append("QC")

    # Coerce to numeric, drop features with any NaN
    X = df.loc[all_rows].apply(pd.to_numeric, errors="coerce").dropna(axis=1)

    if log_transform:
        X = np.log1p(X.clip(lower=0))

    X_scaled = StandardScaler().fit_transform(X.values.astype(float))

    pca = PCA(n_components=2)
    coords = pca.fit_transform(X_scaled)

    # define qc points in coords
    for group in list(sample_groups) + ["QC"]:
        mask = [_label == group for _label in labels]
        qc_points = coords[mask]

    # Calculate centroid
    length = qc_points.shape[0]
    sum_x = np.sum(qc_points[:, 0])
    sum_y = np.sum(qc_points[:, 1])
    centroid = sum_x / length, sum_y / length

    distances = [math.dist(point, centroid) for point in qc_points]
    total_distance = sum(distances)
    avg_distance = total_distance / length
    logger.debug(
        f"The average distance of the points to the centroid is {avg_distance:.3f}."
    )

    return avg_distance