Source code for pyglmGamPoi.trackcell_compat

"""
Trackcell-compatible offset model fitting API.

Primary integration surface — see ``docs/INTEGRATION.md``.
"""

from __future__ import annotations

from typing import Union

import numpy as np
import pandas as pd
from scipy import sparse

from pyglmGamPoi.constants import METHOD_GLMGAMPOI_OFFSET
from pyglmGamPoi.fitting import fit_glmGamPoi_offset_matrix as _fit_matrix
from pyglmGamPoi.validate import (
    assert_model_pars_contract,
    format_model_pars,
    normalize_model_pars_columns,
    validate_genes_by_cells_matrix,
    validate_gene_index,
    validate_regressor_data,
)




def is_available() -> bool:
    """Return True when the native C++ extension is importable."""
    try:
        from pyglmGamPoi import _core as _c  # noqa: F401

        return True
    except ImportError:
        return False





def fit_offset_model(
    umi: Union[np.ndarray, sparse.spmatrix],
    regressor_data: pd.DataFrame,
    gene_index: pd.Index,
    *,
    method: str = METHOD_GLMGAMPOI_OFFSET,
    allow_inf_theta: bool = True,
    apply_shrinkage: bool = True,
) -> pd.DataFrame:
    """
    Fit glmGamPoi offset models (sctransform ``fit_glmGamPoi_offset``).

    Parameters
    ----------
    umi
        UMI count matrix, **genes × cells** (CSR sparse or dense).
    regressor_data
        Cell-level covariates, one row per cell. Must contain ``log_umi`` in
        **log10** scale, indexed by cell barcodes in the same order as umi columns.
    gene_index
        Gene identifiers, length ``umi.shape[0]``.
    method
        Must be ``"glmGamPoi_offset"``.
    allow_inf_theta
        When False, cap theta at ``mean(mu)/1e-4`` (Seurat default).
    apply_shrinkage
        When True (default), apply ``loc_median_fit`` dispersion trend and refit
        intercepts (R ``glm_gp(overdispersion_shrinkage=TRUE)``).

    Returns
    -------
    DataFrame with columns ``theta``, ``Intercept``, ``log_umi``.
    """
    if method != METHOD_GLMGAMPOI_OFFSET:
        raise ValueError(
            f"Unsupported method '{method}'. pyglmGamPoi only implements "
            f"'{METHOD_GLMGAMPOI_OFFSET}'."
        )
    umi_csr, n_genes, n_cells = validate_genes_by_cells_matrix(umi)
    gene_index = validate_gene_index(gene_index, n_genes)
    validate_regressor_data(regressor_data, n_cells)

    log10_umi = regressor_data["log_umi"].to_numpy(dtype=np.float64)
    offset = np.log(np.power(10.0, log10_umi))

    raw = _fit_matrix(
        umi_csr,
        offset,
        allow_inf_theta=allow_inf_theta,
        apply_shrinkage=apply_shrinkage,
        n_cells=n_cells,
    )

    model_pars = format_model_pars(
        np.asarray(raw["theta"], dtype=np.float64),
        np.asarray(raw["intercept"], dtype=np.float64),
        gene_index,
    )
    return assert_model_pars_contract(model_pars)



__all__ = [
    "fit_offset_model",
    "is_available",
    "normalize_model_pars_columns",
    "METHOD_GLMGAMPOI_OFFSET",
]