Source code for pyglmGamPoi.offset

"""Low-level matrix API for glmGamPoi offset fitting."""

from __future__ import annotations

from typing import Union

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

from pyglmGamPoi.fitting import fit_glmGamPoi_offset_matrix as _fit_matrix
from pyglmGamPoi.validate import (
    assert_model_pars_contract,
    format_model_pars,
    validate_genes_by_cells_matrix,
)




def fit_glmGamPoi_offset(
    umi: Union[np.ndarray, sparse.spmatrix],
    log10_umi: np.ndarray,
    *,
    gene_names: Union[list[str], np.ndarray, None] = None,
    allow_inf_theta: bool = True,
    apply_shrinkage: bool = True,
) -> pd.DataFrame:
    """
    Fit glmGamPoi offset models (low-level API).

    For trackcell integration, prefer :func:`pyglmGamPoi.fit_offset_model` which
    accepts ``regressor_data`` in the same format as trackcell ``cell_attr``.

    Parameters
    ----------
    umi
        UMI count matrix, genes × cells.
    log10_umi
        Per-cell log10 UMI values (length = n_cells).
    gene_names
        Optional gene index labels.
    allow_inf_theta
        When False, cap theta at mean(mu)/1e-4.
    apply_shrinkage
        When True (default, R ``glm_gp`` default), apply cross-gene dispersion
        trend via ``loc_median_fit`` and refit intercepts.

    Returns
    -------
    DataFrame with columns ``theta``, ``Intercept``, ``log_umi``.
    """
    umi_csr, n_genes, n_cells = validate_genes_by_cells_matrix(umi)
    log10_umi = np.asarray(log10_umi, dtype=np.float64)
    if len(log10_umi) != n_cells:
        raise ValueError("log10_umi length must match number of cells in umi")

    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,
    )

    index = (
        pd.Index(gene_names).astype(str)
        if gene_names is not None
        else pd.Index([str(i) for i in range(n_genes)])
    )
    model_pars = format_model_pars(
        np.asarray(raw["theta"], dtype=np.float64),
        np.asarray(raw["intercept"], dtype=np.float64),
        index,
    )
    return assert_model_pars_contract(model_pars)