API Reference

parafac2

Main exports.

CompressedData dataclass

Compressed representation of a multi-condition single-cell dataset.

Stores the small dense per-condition cores and the orthonormal bases for the gene mode (Q) and cell mode (Q_k), enabling fast rank sweeps without touching the raw data again.

Source code in parafac2/compress.py

@dataclass
class CompressedData:
    """Compressed representation of a multi-condition single-cell dataset.

    Stores the small dense per-condition cores and the orthonormal bases for
    the gene mode (``Q``) and cell mode (``Q_k``), enabling fast rank sweeps
    without touching the raw data again.
    """

    cores: list[np.ndarray]
    """List of length ``n_cond`` with dense cores ``Y_k`` of shape ``(L_c_k, L_g)``."""

    Q: np.ndarray
    """Gene projection basis of shape ``(n_genes, L_g)`` with orthonormal columns."""

    Q_k: list[np.ndarray | None] | None
    """Per-condition cell projection bases ``(n_cells_k, L_c_k)`` or ``None``."""

    condition_unique_idxs: np.ndarray
    """Integer condition indices for each cell."""

    norm_tensor: float
    """Total squared Frobenius norm of the original mean-centered dataset."""

    lost_var: float
    """Variance discarded by the compression projectors."""

    total_cells: int
    """Total number of cells across all conditions."""

    n_genes: int
    """Number of genes in the uncompressed dataset."""

    n_cond: int
    """Number of conditions."""

    slice_weights: np.ndarray | None = None
    """Optional per-condition slice weights for normalized ALS."""

    means: np.ndarray | None = None
    """Per-gene means subtracted during compression."""

    adata: anndata.AnnData | None = None
    """Reference to original AnnData object if available."""

    @property
    def L_g(self) -> int:
        """Gene compression dimension."""
        return self.Q.shape[1]

    @property
    def max_cell_dim(self) -> int:
        """Maximum cell dimension across cores."""
        return max(c.shape[0] for c in self.cores)

L_g property

Gene compression dimension.

Q instance-attribute

Gene projection basis of shape (n_genes, L_g) with orthonormal columns.

Q_k instance-attribute

Per-condition cell projection bases (n_cells_k, L_c_k) or None.

adata = None class-attribute instance-attribute

Reference to original AnnData object if available.

condition_unique_idxs instance-attribute

Integer condition indices for each cell.

cores instance-attribute

List of length n_cond with dense cores Y_k of shape (L_c_k, L_g).

lost_var instance-attribute

Variance discarded by the compression projectors.

max_cell_dim property

Maximum cell dimension across cores.

means = None class-attribute instance-attribute

Per-gene means subtracted during compression.

n_cond instance-attribute

Number of conditions.

n_genes instance-attribute

Number of genes in the uncompressed dataset.

norm_tensor instance-attribute

Total squared Frobenius norm of the original mean-centered dataset.

slice_weights = None class-attribute instance-attribute

Optional per-condition slice weights for normalized ALS.

total_cells instance-attribute

Total number of cells across all conditions.

GPUMatrix

Wrapper for a single matrix (csr_array or np.ndarray) stored on GPU memory (CuPy or MLX) or CPU. Evaluates matrix products on the device and returns results as NumPy ndarrays.

Parameters:

Name	Type	Description	Default
mat	ndarray | csr_array	The matrix to wrap and transfer to the selected device.	required
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:get_backend).	None

Source code in parafac2/backend.py

class GPUMatrix:
    """
    Wrapper for a single matrix (csr_array or np.ndarray) stored on GPU memory
    (CuPy or MLX) or CPU. Evaluates matrix products on the device and returns
    results as NumPy ndarrays.

    Parameters
    ----------
    mat : np.ndarray | csr_array
        The matrix to wrap and transfer to the selected device.
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is auto-detected (see :func:`get_backend`).
    """

    __array_priority__ = 1000

    def __init__(self, mat: np.ndarray | csr_array, backend: str | None = None) -> None:
        """Transfer ``mat`` to the resolved backend's device memory."""
        self.backend = get_backend(backend)
        self.shape = mat.shape
        self.dtype = mat.dtype
        self.is_sparse = issparse(mat)

        if self.backend == "cupy":
            self.device_mat = _to_cupy_matrix(mat)
        elif self.backend == "mlx":
            self.device_mat = _to_mlx_matrix(mat)
        else:
            self.device_mat = mat

    def matmul(self, rhs: np.ndarray) -> np.ndarray:
        """Compute ``self @ rhs`` on the wrapped device.

        Parameters
        ----------
        rhs : np.ndarray
            The right-hand operand.

        Returns
        -------
        np.ndarray
            The product, as a NumPy array.
        """
        if self.backend == "cupy":
            return _matmul_cupy(self.device_mat, rhs)
        elif self.backend == "mlx":
            return _matmul_mlx(
                self.device_mat, rhs, is_sparse=self.is_sparse, shape=self.shape
            )
        return self.device_mat @ rhs

    def rmatmul(self, lhs: np.ndarray) -> np.ndarray:
        """Compute ``lhs @ self`` on the wrapped device.

        Parameters
        ----------
        lhs : np.ndarray
            The left-hand operand.

        Returns
        -------
        np.ndarray
            The product, as a NumPy array.
        """
        if self.backend == "cupy":
            return _rmatmul_cupy(lhs, self.device_mat)
        elif self.backend == "mlx":
            return _rmatmul_mlx(
                lhs, self.device_mat, is_sparse=self.is_sparse, shape=self.shape
            )
        return lhs @ self.device_mat

    def __matmul__(self, rhs: np.ndarray) -> np.ndarray:
        """Operator form of :meth:`matmul`, enabling ``gpu_matrix @ rhs``."""
        return self.matmul(rhs)

    def __rmatmul__(self, lhs: np.ndarray) -> np.ndarray:
        """Operator form of :meth:`rmatmul`, enabling ``lhs @ gpu_matrix``."""
        return self.rmatmul(lhs)

__init__(mat, backend=None)

Transfer mat to the resolved backend's device memory.

Source code in parafac2/backend.py

def __init__(self, mat: np.ndarray | csr_array, backend: str | None = None) -> None:
    """Transfer ``mat`` to the resolved backend's device memory."""
    self.backend = get_backend(backend)
    self.shape = mat.shape
    self.dtype = mat.dtype
    self.is_sparse = issparse(mat)

    if self.backend == "cupy":
        self.device_mat = _to_cupy_matrix(mat)
    elif self.backend == "mlx":
        self.device_mat = _to_mlx_matrix(mat)
    else:
        self.device_mat = mat

__matmul__(rhs)

Operator form of :meth:matmul, enabling gpu_matrix @ rhs.

Source code in parafac2/backend.py

def __matmul__(self, rhs: np.ndarray) -> np.ndarray:
    """Operator form of :meth:`matmul`, enabling ``gpu_matrix @ rhs``."""
    return self.matmul(rhs)

__rmatmul__(lhs)

Operator form of :meth:rmatmul, enabling lhs @ gpu_matrix.

Source code in parafac2/backend.py

def __rmatmul__(self, lhs: np.ndarray) -> np.ndarray:
    """Operator form of :meth:`rmatmul`, enabling ``lhs @ gpu_matrix``."""
    return self.rmatmul(lhs)

matmul(rhs)

Compute self @ rhs on the wrapped device.

Parameters:

Name	Type	Description	Default
rhs	ndarray	The right-hand operand.	required

Returns:

Type	Description
ndarray	The product, as a NumPy array.

Source code in parafac2/backend.py

def matmul(self, rhs: np.ndarray) -> np.ndarray:
    """Compute ``self @ rhs`` on the wrapped device.

    Parameters
    ----------
    rhs : np.ndarray
        The right-hand operand.

    Returns
    -------
    np.ndarray
        The product, as a NumPy array.
    """
    if self.backend == "cupy":
        return _matmul_cupy(self.device_mat, rhs)
    elif self.backend == "mlx":
        return _matmul_mlx(
            self.device_mat, rhs, is_sparse=self.is_sparse, shape=self.shape
        )
    return self.device_mat @ rhs

rmatmul(lhs)

Compute lhs @ self on the wrapped device.

Parameters:

Name	Type	Description	Default
lhs	ndarray	The left-hand operand.	required

Returns:

Type	Description
ndarray	The product, as a NumPy array.

Source code in parafac2/backend.py

def rmatmul(self, lhs: np.ndarray) -> np.ndarray:
    """Compute ``lhs @ self`` on the wrapped device.

    Parameters
    ----------
    lhs : np.ndarray
        The left-hand operand.

    Returns
    -------
    np.ndarray
        The product, as a NumPy array.
    """
    if self.backend == "cupy":
        return _rmatmul_cupy(lhs, self.device_mat)
    elif self.backend == "mlx":
        return _rmatmul_mlx(
            lhs, self.device_mat, is_sparse=self.is_sparse, shape=self.shape
        )
    return lhs @ self.device_mat

compress_dataset(X_in, L='auto', rank=None, n_power_iter=2, random_state=None, normalize_slices=False, backend=None)

Compress an AnnData dataset in gene and cell modes.

Parameters:

Name	Type	Description	Default
X_in	AnnData	Input dataset with data in X_in.X, condition indices in X_in.obs["condition_unique_idxs"], and optional means in X_in.var["means"].	required
L	int | tuple[int, int | None] | str	Compression dimension(s). If "auto", picks dimensions based on rank (or default rank 30 if rank is None). If an int, sets both L_g = L and L_c = L. If a tuple (L_g, L_c), sets gene and cell dimensions individually (pass L_c=None for gene-only compression).	"auto"
rank	int | None	Expected maximum rank to fit on the compressed data. Used when L="auto".	None
n_power_iter	int	Number of power iterations for randomized SVD.	2
random_state	int | Generator | None	Random seed or generator.	None
normalize_slices	bool	Whether to precalculate slice weights for normalized ALS.	False
backend	str | None	Compute backend for raw matrix products.	None

Returns:

Type	Description
CompressedData	The compressed dataset ready for fast PARAFAC2 fitting.

Source code in parafac2/compress.py

def compress_dataset(
    X_in: anndata.AnnData,
    L: int | tuple[int, int | None] | str = "auto",
    rank: int | None = None,
    n_power_iter: int = 2,
    random_state: int | np.random.Generator | None = None,
    normalize_slices: bool = False,
    backend: str | None = None,
) -> CompressedData:
    """Compress an AnnData dataset in gene and cell modes.

    Parameters
    ----------
    X_in : anndata.AnnData
        Input dataset with data in ``X_in.X``, condition indices in
        ``X_in.obs["condition_unique_idxs"]``, and optional means in
        ``X_in.var["means"]``.
    L : int | tuple[int, int | None] | str, default "auto"
        Compression dimension(s). If ``"auto"``, picks dimensions based on
        ``rank`` (or default rank 30 if ``rank`` is None). If an int, sets
        both ``L_g = L`` and ``L_c = L``. If a tuple ``(L_g, L_c)``, sets
        gene and cell dimensions individually (pass ``L_c=None`` for
        gene-only compression).
    rank : int | None, default None
        Expected maximum rank to fit on the compressed data. Used when
        ``L="auto"``.
    n_power_iter : int, default 2
        Number of power iterations for randomized SVD.
    random_state : int | np.random.Generator | None, default None
        Random seed or generator.
    normalize_slices : bool, default False
        Whether to precalculate slice weights for normalized ALS.
    backend : str | None, default None
        Compute backend for raw matrix products.

    Returns
    -------
    CompressedData
        The compressed dataset ready for fast PARAFAC2 fitting.
    """
    (
        X_mat,
        condition_unique_idxs,
        means,
        norm_tensor,
        slice_weights,
    ) = extract_dataset_info(X_in, normalize_slices=normalize_slices)
    total_cells, n_genes = X_mat.shape
    n_cond = int(np.amax(condition_unique_idxs)) + 1

    # Determine L_g and L_c
    target_rank = rank if rank is not None else 30
    if isinstance(L, str) and L == "auto":
        L_g_val = min(n_genes, max(4 * target_rank, target_rank + 20))
        L_c_val: int | None = max(4 * target_rank, target_rank + 20)
    elif isinstance(L, tuple):
        L_g_val, L_c_val = L
        L_g_val = min(n_genes, L_g_val)
    elif isinstance(L, (int, np.integer)):
        L_g_val = min(n_genes, int(L))
        L_c_val = int(L)
    else:
        raise ValueError(f"Invalid compression parameter L: {L}")

    X_raw = to_gpu(X_mat, backend=backend)
    X_c, Q, _norm_Xc_sq = compress_genes(
        X_raw,
        means,
        L_g=L_g_val,
        n_power_iter=n_power_iter,
        random_state=random_state,
    )

    cores, Q_k, norm_cores_sq = compress_cells(
        X_c,
        condition_unique_idxs,
        L_c=L_c_val,
    )

    lost_var = float(np.maximum(0.0, norm_tensor - norm_cores_sq))

    return CompressedData(
        cores=cores,
        Q=Q,
        Q_k=Q_k,
        condition_unique_idxs=condition_unique_idxs,
        norm_tensor=norm_tensor,
        lost_var=lost_var,
        total_cells=total_cells,
        n_genes=n_genes,
        n_cond=n_cond,
        slice_weights=slice_weights,
        means=means,
        adata=X_in,
    )

get_backend(backend=None)

Return the requested backend, or auto-detect the first available one.

Parameters:

Name	Type	Description	Default
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is chosen by attempting to import cupy then mlx.core, falling back to 'cpu' if neither is installed.	None

Returns:

Type	Description
str	The resolved backend name: 'mlx', 'cupy', or 'cpu'.

Raises:

Type	Description
ValueError	If backend is given but is not one of the supported names.

Source code in parafac2/backend.py

def get_backend(backend: str | None = None) -> str:
    """Return the requested backend, or auto-detect the first available one.

    Parameters
    ----------
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is chosen by attempting to import ``cupy``
        then ``mlx.core``, falling back to ``'cpu'`` if neither is
        installed.

    Returns
    -------
    str
        The resolved backend name: ``'mlx'``, ``'cupy'``, or ``'cpu'``.

    Raises
    ------
    ValueError
        If ``backend`` is given but is not one of the supported names.
    """
    if backend is not None:
        backend_lower = backend.lower()
        if backend_lower in ("mlx", "cupy", "cpu"):
            return backend_lower
        raise ValueError(
            f"Unknown backend '{backend}'. Supported backends: 'mlx', 'cupy', 'cpu'."
        )

    try:
        import cupy  # noqa: F401  # ty: ignore[unresolved-import]

        return "cupy"
    except ImportError:
        pass

    try:
        import mlx.core  # noqa: F401  # ty: ignore[unresolved-import]

        return "mlx"
    except ImportError:
        pass

    return "cpu"

parafac2_init(X, condition_unique_idxs, rank=3, means=None, random_state=None, n_oversamples=10, n_iter=2, norm_tensor=None)

Compute initial factors using randomized SVD.

Parameters:

Name	Type	Description	Default
X	Any	The (optionally sparse or GPU-backed) data matrix, stacked across all conditions, with shape (total_cells, n_genes).	required
condition_unique_idxs	ndarray	Integer array of length total_cells giving each row's condition index.	required
rank	int	The number of components to compute.	3
means	ndarray | None	Per-gene means to mean-center X by, or None to skip centering.	None
random_state	int | Generator | None	Seed or generator controlling the random projection used by the randomized SVD.	None
n_oversamples	int	Extra dimensions added to rank when forming the random projection, to improve the accuracy of the randomized SVD.	10
n_iter	int	Number of power iterations used to refine the random projection.	2
norm_tensor	float | None	Precomputed squared Frobenius norm of the mean-centered X. If None, it is computed via :func:~parafac2.utils.calc_norm_sq.	None

Returns:

Type	Description
tuple[list[ndarray], float]	The initial [A, B, C] factor matrices (with A all-ones, B the identity, and C the top right-singular vectors of the mean-centered X), and the squared Frobenius norm norm_tensor.

Source code in parafac2/parafac2.py

def parafac2_init(
    X: Any,
    condition_unique_idxs: np.ndarray,
    rank: int = 3,
    means: np.ndarray | None = None,
    random_state: int | np.random.Generator | None = None,
    n_oversamples: int = 10,
    n_iter: int = 2,
    norm_tensor: float | None = None,
) -> tuple[list[np.ndarray], float]:
    """Compute initial factors using randomized SVD.

    Parameters
    ----------
    X : Any
        The (optionally sparse or GPU-backed) data matrix, stacked across
        all conditions, with shape ``(total_cells, n_genes)``.
    condition_unique_idxs : np.ndarray
        Integer array of length ``total_cells`` giving each row's condition
        index.
    rank : int, default 3
        The number of components to compute.
    means : np.ndarray | None, default None
        Per-gene means to mean-center ``X`` by, or ``None`` to skip
        centering.
    random_state : int | np.random.Generator | None, default None
        Seed or generator controlling the random projection used by the
        randomized SVD.
    n_oversamples : int, default 10
        Extra dimensions added to ``rank`` when forming the random
        projection, to improve the accuracy of the randomized SVD.
    n_iter : int, default 2
        Number of power iterations used to refine the random projection.
    norm_tensor : float | None, default None
        Precomputed squared Frobenius norm of the mean-centered ``X``. If
        ``None``, it is computed via :func:`~parafac2.utils.calc_norm_sq`.

    Returns
    -------
    tuple[list[np.ndarray], float]
        The initial ``[A, B, C]`` factor matrices (with ``A`` all-ones,
        ``B`` the identity, and ``C`` the top right-singular vectors of the
        mean-centered ``X``), and the squared Frobenius norm ``norm_tensor``.
    """
    n_cond = int(np.amax(condition_unique_idxs)) + 1
    if norm_tensor is None:
        norm_tensor = calc_norm_sq(X, means)

    C = randomized_svd_right(
        X,
        means,
        n_components=rank,
        n_oversamples=n_oversamples,
        n_power_iter=n_iter,
        random_state=random_state,
    )

    factors = [
        np.ones((n_cond, rank), dtype=np.float64),
        np.eye(rank, dtype=np.float64),
        C,
    ]
    return factors, norm_tensor

parafac2_nd(X_in, rank, n_iter_max=100, tol=1e-06, random_state=None, callback=None, backend=None, normalize_slices=False, n_inner=1, compress=None)

The same interface as regular PARAFAC2 with optional CANDELINC compression.

If compress is specified (or if X_in is already a :class:~parafac2.compress.CompressedData), PARAFAC2 is fit in the compressed subspace (Bro's "compress-then-fit"), eliminating per-sweep raw data passes.

If normalize_slices is True, each condition's contribution to the factor updates is rescaled by the inverse of its (mean-centered) Frobenius norm. This prevents conditions with many more cells (or much higher variance) from dominating the shared factors, e.g. the A matrix. The weighting is computed from small per-condition summary statistics and applied to per-condition intermediates only, so X is never copied or modified. The reported error/R2X are unaffected, since they are still computed from the unweighted fit.

Parameters:

Name	Type	Description	Default
X_in	AnnData | CompressedData	Input dataset with the (optionally sparse) data matrix in X_in.X, condition labels in X_in.obs["condition_unique_idxs"], and optionally per-gene means in X_in.var["means"] (defaults to zero, i.e. no centering, if absent). Alternatively, a pre-compressed :class:~parafac2.compress.CompressedData object.	required
rank	int	The number of components to fit.	required
n_iter_max	int	Maximum number of ALS iterations.	100
tol	float	Convergence tolerance: iteration stops once the (non-negative) decrease in relative error between successive iterations drops below this value.	1e-6
random_state	int | None	Seed controlling the randomized SVD initialization.	None
callback	Callable[[int, float, list[ndarray]], None] | None	Optional callback invoked after each iteration with the iteration index, the relative error, and the current factor matrices.	None
backend	str | None	Compute backend to run matrix products on: one of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:~parafac2.backend.get_backend).	None
normalize_slices	bool	Whether to rescale each condition's contribution to the factor updates by the inverse of its Frobenius norm, as described above.	False
n_inner	int	Number of (projection, A, B) sub-iterations per sweep. These read the data only through the cached W, costing O(n_cells * rank^2) against the O(nnz * rank) of a raw-data product, so they are close to free on sparse inputs. Raising n_inner trades that cheap compute for fewer sweeps, and so for fewer raw-data passes: on structured test data, n_inner=2 cut the sweeps needed to reach a fixed error by ~20%, with little further gain beyond 3. Whether that is a net win depends on how strongly the raw-data products dominate, so it is worth benchmarking per dataset. The default of 1 reproduces the classic one-update-per-mode ALS sweep.	1
compress	int | tuple[int, int | None] | str | bool | None	Compression mode. If None or False (default), exact ALS is used. If "auto" or True, sets compression dimensions automatically based on rank. If an integer, sets both gene and cell compression dimensions to that value. If a tuple (L_g, L_c), sets dimensions separately (pass L_c=None for gene-only compression). Ignored if X_in is already a :class:~parafac2.compress.CompressedData.	None

Returns:

Type	Description
tuple[tuple[ndarray, list[ndarray], list[ndarray]], float]	A ((weights, factors, projections), R2X) tuple: the standardized weights and [A, B, C] factor matrices, the per-condition projection matrices P_k, and the final fraction of variance explained.

Source code in parafac2/parafac2.py

def parafac2_nd(
    X_in: anndata.AnnData | CompressedData,
    rank: int,
    n_iter_max: int = 100,
    tol: float = 1e-6,
    random_state: int | None = None,
    callback: Callable[[int, float, list[np.ndarray]], None] | None = None,
    backend: str | None = None,
    normalize_slices: bool = False,
    n_inner: int = 1,
    compress: int | tuple[int, int | None] | str | bool | None = None,
) -> tuple[tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], float]:
    r"""The same interface as regular PARAFAC2 with optional CANDELINC compression.

    If ``compress`` is specified (or if ``X_in`` is already a
    :class:`~parafac2.compress.CompressedData`), PARAFAC2 is fit in the
    compressed subspace (Bro's "compress-then-fit"), eliminating per-sweep raw
    data passes.

    If ``normalize_slices`` is True, each condition's contribution to the
    factor updates is rescaled by the inverse of its (mean-centered)
    Frobenius norm. This prevents conditions with many more cells (or much
    higher variance) from dominating the shared factors, e.g. the ``A``
    matrix. The weighting is computed from small per-condition summary
    statistics and applied to per-condition intermediates only, so ``X`` is
    never copied or modified. The reported error/R2X are unaffected, since
    they are still computed from the unweighted fit.

    Parameters
    ----------
    X_in : anndata.AnnData | CompressedData
        Input dataset with the (optionally sparse) data matrix in ``X_in.X``,
        condition labels in ``X_in.obs["condition_unique_idxs"]``, and
        optionally per-gene means in ``X_in.var["means"]`` (defaults to zero,
        i.e. no centering, if absent). Alternatively, a pre-compressed
        :class:`~parafac2.compress.CompressedData` object.
    rank : int
        The number of components to fit.
    n_iter_max : int, default 100
        Maximum number of ALS iterations.
    tol : float, default 1e-6
        Convergence tolerance: iteration stops once the (non-negative)
        decrease in relative error between successive iterations drops
        below this value.
    random_state : int | None, default None
        Seed controlling the randomized SVD initialization.
    callback : Callable[[int, float, list[np.ndarray]], None] | None, default None
        Optional callback invoked after each iteration with the iteration
        index, the relative error, and the current factor matrices.
    backend : str | None, default None
        Compute backend to run matrix products on: one of ``'mlx'``,
        ``'cupy'``, or ``'cpu'``. If ``None``, the first available
        accelerator is auto-detected (see
        :func:`~parafac2.backend.get_backend`).
    normalize_slices : bool, default False
        Whether to rescale each condition's contribution to the factor
        updates by the inverse of its Frobenius norm, as described above.
    n_inner : int, default 1
        Number of ``(projection, A, B)`` sub-iterations per sweep. These read
        the data only through the cached ``W``, costing ``O(n_cells *
        rank^2)`` against the ``O(nnz * rank)`` of a raw-data product, so
        they are close to free on sparse inputs. Raising ``n_inner`` trades
        that cheap compute for fewer sweeps, and so for fewer raw-data
        passes: on structured test data, ``n_inner=2`` cut the sweeps needed
        to reach a fixed error by ~20%, with little further gain beyond 3.
        Whether that is a net win depends on how strongly the raw-data
        products dominate, so it is worth benchmarking per dataset. The
        default of 1 reproduces the classic one-update-per-mode ALS sweep.
    compress : int | tuple[int, int | None] | str | bool | None, default None
        Compression mode. If ``None`` or ``False`` (default), exact ALS is
        used. If ``"auto"`` or ``True``, sets compression dimensions
        automatically based on ``rank``. If an integer, sets both gene and cell
        compression dimensions to that value. If a tuple ``(L_g, L_c)``, sets
        dimensions separately (pass ``L_c=None`` for gene-only compression).
        Ignored if ``X_in`` is already a
        :class:`~parafac2.compress.CompressedData`.

    Returns
    -------
    tuple[tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], float]
        A ``((weights, factors, projections), R2X)`` tuple: the standardized
        weights and ``[A, B, C]`` factor matrices, the per-condition
        projection matrices ``P_k``, and the final fraction of variance
        explained.
    """
    # Verbose if this is not an automated build
    verbose = "CI" not in os.environ

    if isinstance(X_in, CompressedData):
        return _fit_parafac2_compressed(
            X_in,
            rank=rank,
            n_iter_max=n_iter_max,
            tol=tol,
            random_state=random_state,
            callback=callback,
            verbose=verbose,
        )

    if compress is not None and compress is not False:
        compressed = compress_dataset(
            X_in,
            L=compress,
            rank=rank,
            random_state=random_state,
            normalize_slices=normalize_slices,
            backend=backend,
        )
        return _fit_parafac2_compressed(
            compressed,
            rank=rank,
            n_iter_max=n_iter_max,
            tol=tol,
            random_state=random_state,
            callback=callback,
            verbose=verbose,
        )

    (
        X_mat,
        condition_unique_idxs,
        means,
        norm_tensor,
        slice_weights,
    ) = extract_dataset_info(X_in, normalize_slices=normalize_slices)

    X_raw = to_gpu(X_mat, backend=backend)

    factors, _ = parafac2_init(
        X_raw,
        condition_unique_idxs,
        rank=rank,
        means=means,
        random_state=random_state,
        norm_tensor=norm_tensor,
    )

    cond_slices = condition_slices(
        condition_unique_idxs, int(np.amax(condition_unique_idxs)) + 1
    )

    # W depends only on C, so it stays valid across the A and B updates and is
    # recomputed only once C changes. Each sweep therefore costs exactly two
    # raw-data products: this one and the X^T @ H inside the mode-2 update.
    W = calc_W(X_raw, means, factors[2])
    projections, S = project_data(W, factors, cond_slices)
    errs = [calc_err(S, factors, norm_tensor) / norm_tensor]

    tq = tqdm(range(n_iter_max), disable=(not verbose), delay=0.5)
    for iteration in tq:
        # The (P, A, B) block reads the data only through the cached W, so
        # extra inner passes buy convergence at no raw-data cost.
        for _ in range(n_inner):
            factors = parafac_update(factors, 0, S, slice_weights=slice_weights)
            projections, S = project_data(W, factors, cond_slices)
            factors = parafac_update(factors, 1, S, slice_weights=slice_weights)
            projections, S = project_data(W, factors, cond_slices)

        factors = parafac_update(
            factors,
            2,
            S,
            projections,
            X=X_raw,
            means=means,
            cond_slices=cond_slices,
            slice_weights=slice_weights,
        )

        # C changed, so refresh W; this also yields the projections and error
        # for the factors as they stand at the end of this sweep.
        W = calc_W(X_raw, means, factors[2])
        projections, S = project_data(W, factors, cond_slices)
        errs.append(calc_err(S, factors, norm_tensor) / norm_tensor)

        delta = errs[-2] - errs[-1]
        tq.set_postfix(error=errs[-1], R2X=1.0 - errs[-1], Δ=delta, refresh=False)
        if callback is not None:
            callback(iteration, errs[-1], factors)

        if 0 <= delta < tol:
            break

    R2X = 1 - errs[-1]

    # Standardize the results and return
    return standardize_pf2(factors, projections), R2X

prepare_dataset(X, condition_name, geneThreshold)

Preprocess and normalize an AnnData dataset for PARAFAC2 factorization.

Performs quality control filtering of low-count cells and low-expression genes, normalizes total cell counts and gene sums, applies a log10 transformation, and computes metadata required by PARAFAC2 (condition indices and gene means).

Parameters:

Name	Type	Description	Default
X	AnnData	Input single-cell dataset with raw count matrix stored in X.X (must be a sparse matrix with non-negative values).	required
condition_name	str	Column name in X.obs identifying the sample or experimental condition grouping for each cell.	required
geneThreshold	float	Minimum threshold fraction for gene inclusion. Genes with total counts less than geneThreshold * total_cells are filtered out.	required

Returns:

Type	Description
AnnData	A filtered and normalized copy of the AnnData object. Contains the log-transformed normalized counts in X.X, integer condition codes in X.obs["condition_unique_idxs"], and per-gene mean expression values in X.var["means"].

Source code in parafac2/normalize.py

def prepare_dataset(
    X: anndata.AnnData, condition_name: str, geneThreshold: float
) -> anndata.AnnData:
    """Preprocess and normalize an AnnData dataset for PARAFAC2 factorization.

    Performs quality control filtering of low-count cells and low-expression
    genes, normalizes total cell counts and gene sums, applies a log10
    transformation, and computes metadata required by PARAFAC2 (condition
    indices and gene means).

    Parameters
    ----------
    X : anndata.AnnData
        Input single-cell dataset with raw count matrix stored in ``X.X``
        (must be a sparse matrix with non-negative values).
    condition_name : str
        Column name in ``X.obs`` identifying the sample or experimental
        condition grouping for each cell.
    geneThreshold : float
        Minimum threshold fraction for gene inclusion. Genes with total counts
        less than ``geneThreshold * total_cells`` are filtered out.

    Returns
    -------
    anndata.AnnData
        A filtered and normalized copy of the AnnData object. Contains the
        log-transformed normalized counts in ``X.X``, integer condition
        codes in ``X.obs["condition_unique_idxs"]``, and per-gene mean
        expression values in ``X.var["means"]``.
    """
    assert issparse(X.X)
    X_X_raw = cast("csr_array", X.X)
    assert np.amin(X_X_raw.data) >= 0.0

    # Filter out genes with too few reads, and cells with fewer than 10 counts
    cell_mask = np.ravel(X_X_raw.sum(axis=1)) > 10
    gene_mask = np.ravel(X_X_raw.sum(axis=0)) > (geneThreshold * X_X_raw.shape[0])

    # Subset and materialize actual AnnData object before modifying X.X
    if cell_mask.all() and gene_mask.all():
        X = X.copy()
    else:
        X = X[cell_mask, gene_mask].copy()

    # Convert subset to csr_array and float32 data
    X.X = csr_array(X.X)
    X_X = cast("csr_array", X.X)

    if X_X.dtype != np.float32:
        X_X.data = X_X.data.astype(np.float32)

    ## Normalize total counts per cell
    # Keep the counts on a reasonable scale to avoid accuracy issues
    counts_per_cell = np.ravel(X_X.sum(axis=1)).astype(np.float32, copy=False)
    counts_per_cell /= np.median(counts_per_cell)
    # In-place CSR row scaling
    X_X.data /= np.repeat(counts_per_cell, np.diff(X_X.indptr))

    # Scale genes by sum, in-place CSR column scaling
    gene_sums = np.ravel(X_X.sum(axis=0)).astype(np.float32, copy=False)
    X_X.data /= gene_sums[X_X.indices]

    # Transform values in-place to avoid nnz-sized temporaries
    X_X.data *= np.float32(1000.0)
    X_X.data += np.float32(1.0)
    np.log10(X_X.data, out=X_X.data)

    # Get the indices for subsetting the data
    X.obs["condition_unique_idxs"] = pd.Categorical(X.obs[condition_name]).codes

    # Pre-calculate gene means
    X.var["means"] = np.ravel(X_X.mean(axis=0))

    return X

store_pf2(X, parafac2_output)

Store the Pf2 results into the anndata object.

Parameters:

Name	Type	Description	Default
X	AnnData | CompressedData	The dataset the factorization was fit on. Must have X.obs["condition_unique_idxs"] set (as produced by :func:~parafac2.normalize.prepare_dataset or equivalent). If a :class:~parafac2.compress.CompressedData is provided, factors are written to its underlying .adata object.	required
parafac2_output	tuple[ndarray, list[ndarray], list[ndarray]]	The (weights, factors, projections) output of :func:parafac2_nd, where factors is the [A, B, C] factor matrices and projections is the per-condition projection matrices P_k.	required

Returns:

Type	Description
AnnData	The target AnnData object, mutated in place, with the weights in X.uns["Pf2_weights"], factors in X.uns["Pf2_A"]/X.uns["Pf2_B"]/X.varm["Pf2_C"], and per-cell projections in X.obsm["projections"] and X.obsm["weighted_projections"] (projections composed with B).

Source code in parafac2/parafac2.py

def store_pf2(
    X: anndata.AnnData | CompressedData,
    parafac2_output: tuple[np.ndarray, list[np.ndarray], list[np.ndarray]],
) -> anndata.AnnData:
    """Store the Pf2 results into the anndata object.

    Parameters
    ----------
    X : anndata.AnnData | CompressedData
        The dataset the factorization was fit on. Must have
        ``X.obs["condition_unique_idxs"]`` set (as produced by
        :func:`~parafac2.normalize.prepare_dataset` or equivalent). If a
        :class:`~parafac2.compress.CompressedData` is provided, factors are
        written to its underlying ``.adata`` object.
    parafac2_output : tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]
        The ``(weights, factors, projections)`` output of :func:`parafac2_nd`,
        where ``factors`` is the ``[A, B, C]`` factor matrices and
        ``projections`` is the per-condition projection matrices ``P_k``.

    Returns
    -------
    anndata.AnnData
        The target AnnData object, mutated in place, with the weights in
        ``X.uns["Pf2_weights"]``, factors in
        ``X.uns["Pf2_A"]``/``X.uns["Pf2_B"]``/``X.varm["Pf2_C"]``, and
        per-cell projections in ``X.obsm["projections"]`` and
        ``X.obsm["weighted_projections"]`` (projections composed with ``B``).
    """
    if isinstance(X, CompressedData):
        if X.adata is None:
            raise ValueError("CompressedData has no associated AnnData object.")
        target_adata = X.adata
        condition_unique_idxs = X.condition_unique_idxs
    else:
        target_adata = X
        condition_unique_idxs = target_adata.obs["condition_unique_idxs"]

    target_adata.uns["Pf2_weights"] = parafac2_output[0]
    target_adata.uns["Pf2_A"], target_adata.uns["Pf2_B"], target_adata.varm["Pf2_C"] = (
        parafac2_output[1]
    )

    target_adata.obsm["projections"] = np.zeros(
        (target_adata.shape[0], len(target_adata.uns["Pf2_weights"])), dtype=np.float32
    )
    for i, p in enumerate(parafac2_output[2]):
        target_adata.obsm["projections"][condition_unique_idxs == i, :] = p

    target_adata.obsm["weighted_projections"] = (
        target_adata.obsm["projections"] @ target_adata.uns["Pf2_B"]
    ).astype(np.float32, copy=False)

    return target_adata

to_gpu(mat, backend=None)

Transfer matrix to GPU memory if CuPy or MLX is requested/available, returning a GPUMatrix wrapper. Otherwise returns the CPU matrix as-is.

Parameters:

Name	Type	Description	Default
mat	ndarray | csr_array	The matrix to (optionally) transfer.	required
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:get_backend).	None

Returns:

Type	Description
GPUMatrix | ndarray | csr_array	A :class:GPUMatrix wrapping mat if a GPU backend was resolved, otherwise mat unchanged.

Source code in parafac2/backend.py

def to_gpu(
    mat: np.ndarray | csr_array, backend: str | None = None
) -> GPUMatrix | np.ndarray | csr_array:
    """
    Transfer matrix to GPU memory if CuPy or MLX is requested/available,
    returning a GPUMatrix wrapper. Otherwise returns the CPU matrix as-is.

    Parameters
    ----------
    mat : np.ndarray | csr_array
        The matrix to (optionally) transfer.
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is auto-detected (see :func:`get_backend`).

    Returns
    -------
    GPUMatrix | np.ndarray | csr_array
        A :class:`GPUMatrix` wrapping ``mat`` if a GPU backend was resolved,
        otherwise ``mat`` unchanged.
    """
    chosen = get_backend(backend)
    if chosen == "cpu":
        return mat
    return GPUMatrix(mat, backend=chosen)

parafac2.parafac2

parafac2.parafac2

Core PARAFAC2 decomposition routines.

Implements PARAFAC2 initialization (randomized SVD), the alternating-least- squares fitting loop, CANDELINC-style compression, and standardization/storage of the fitted factors and per-condition projections. Operates directly on a single (optionally sparse) data matrix held in an AnnData object, avoiding per-condition copies.

parafac2_init(X, condition_unique_idxs, rank=3, means=None, random_state=None, n_oversamples=10, n_iter=2, norm_tensor=None)

Compute initial factors using randomized SVD.

Parameters:

Name	Type	Description	Default
X	Any	The (optionally sparse or GPU-backed) data matrix, stacked across all conditions, with shape (total_cells, n_genes).	required
condition_unique_idxs	ndarray	Integer array of length total_cells giving each row's condition index.	required
rank	int	The number of components to compute.	3
means	ndarray | None	Per-gene means to mean-center X by, or None to skip centering.	None
random_state	int | Generator | None	Seed or generator controlling the random projection used by the randomized SVD.	None
n_oversamples	int	Extra dimensions added to rank when forming the random projection, to improve the accuracy of the randomized SVD.	10
n_iter	int	Number of power iterations used to refine the random projection.	2
norm_tensor	float | None	Precomputed squared Frobenius norm of the mean-centered X. If None, it is computed via :func:~parafac2.utils.calc_norm_sq.	None

Returns:

Type	Description
tuple[list[ndarray], float]	The initial [A, B, C] factor matrices (with A all-ones, B the identity, and C the top right-singular vectors of the mean-centered X), and the squared Frobenius norm norm_tensor.

Source code in parafac2/parafac2.py

def parafac2_init(
    X: Any,
    condition_unique_idxs: np.ndarray,
    rank: int = 3,
    means: np.ndarray | None = None,
    random_state: int | np.random.Generator | None = None,
    n_oversamples: int = 10,
    n_iter: int = 2,
    norm_tensor: float | None = None,
) -> tuple[list[np.ndarray], float]:
    """Compute initial factors using randomized SVD.

    Parameters
    ----------
    X : Any
        The (optionally sparse or GPU-backed) data matrix, stacked across
        all conditions, with shape ``(total_cells, n_genes)``.
    condition_unique_idxs : np.ndarray
        Integer array of length ``total_cells`` giving each row's condition
        index.
    rank : int, default 3
        The number of components to compute.
    means : np.ndarray | None, default None
        Per-gene means to mean-center ``X`` by, or ``None`` to skip
        centering.
    random_state : int | np.random.Generator | None, default None
        Seed or generator controlling the random projection used by the
        randomized SVD.
    n_oversamples : int, default 10
        Extra dimensions added to ``rank`` when forming the random
        projection, to improve the accuracy of the randomized SVD.
    n_iter : int, default 2
        Number of power iterations used to refine the random projection.
    norm_tensor : float | None, default None
        Precomputed squared Frobenius norm of the mean-centered ``X``. If
        ``None``, it is computed via :func:`~parafac2.utils.calc_norm_sq`.

    Returns
    -------
    tuple[list[np.ndarray], float]
        The initial ``[A, B, C]`` factor matrices (with ``A`` all-ones,
        ``B`` the identity, and ``C`` the top right-singular vectors of the
        mean-centered ``X``), and the squared Frobenius norm ``norm_tensor``.
    """
    n_cond = int(np.amax(condition_unique_idxs)) + 1
    if norm_tensor is None:
        norm_tensor = calc_norm_sq(X, means)

    C = randomized_svd_right(
        X,
        means,
        n_components=rank,
        n_oversamples=n_oversamples,
        n_power_iter=n_iter,
        random_state=random_state,
    )

    factors = [
        np.ones((n_cond, rank), dtype=np.float64),
        np.eye(rank, dtype=np.float64),
        C,
    ]
    return factors, norm_tensor

parafac2_nd(X_in, rank, n_iter_max=100, tol=1e-06, random_state=None, callback=None, backend=None, normalize_slices=False, n_inner=1, compress=None)

The same interface as regular PARAFAC2 with optional CANDELINC compression.

If compress is specified (or if X_in is already a :class:~parafac2.compress.CompressedData), PARAFAC2 is fit in the compressed subspace (Bro's "compress-then-fit"), eliminating per-sweep raw data passes.

If normalize_slices is True, each condition's contribution to the factor updates is rescaled by the inverse of its (mean-centered) Frobenius norm. This prevents conditions with many more cells (or much higher variance) from dominating the shared factors, e.g. the A matrix. The weighting is computed from small per-condition summary statistics and applied to per-condition intermediates only, so X is never copied or modified. The reported error/R2X are unaffected, since they are still computed from the unweighted fit.

Parameters:

Name	Type	Description	Default
X_in	AnnData | CompressedData	Input dataset with the (optionally sparse) data matrix in X_in.X, condition labels in X_in.obs["condition_unique_idxs"], and optionally per-gene means in X_in.var["means"] (defaults to zero, i.e. no centering, if absent). Alternatively, a pre-compressed :class:~parafac2.compress.CompressedData object.	required
rank	int	The number of components to fit.	required
n_iter_max	int	Maximum number of ALS iterations.	100
tol	float	Convergence tolerance: iteration stops once the (non-negative) decrease in relative error between successive iterations drops below this value.	1e-6
random_state	int | None	Seed controlling the randomized SVD initialization.	None
callback	Callable[[int, float, list[ndarray]], None] | None	Optional callback invoked after each iteration with the iteration index, the relative error, and the current factor matrices.	None
backend	str | None	Compute backend to run matrix products on: one of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:~parafac2.backend.get_backend).	None
normalize_slices	bool	Whether to rescale each condition's contribution to the factor updates by the inverse of its Frobenius norm, as described above.	False
n_inner	int	Number of (projection, A, B) sub-iterations per sweep. These read the data only through the cached W, costing O(n_cells * rank^2) against the O(nnz * rank) of a raw-data product, so they are close to free on sparse inputs. Raising n_inner trades that cheap compute for fewer sweeps, and so for fewer raw-data passes: on structured test data, n_inner=2 cut the sweeps needed to reach a fixed error by ~20%, with little further gain beyond 3. Whether that is a net win depends on how strongly the raw-data products dominate, so it is worth benchmarking per dataset. The default of 1 reproduces the classic one-update-per-mode ALS sweep.	1
compress	int | tuple[int, int | None] | str | bool | None	Compression mode. If None or False (default), exact ALS is used. If "auto" or True, sets compression dimensions automatically based on rank. If an integer, sets both gene and cell compression dimensions to that value. If a tuple (L_g, L_c), sets dimensions separately (pass L_c=None for gene-only compression). Ignored if X_in is already a :class:~parafac2.compress.CompressedData.	None

Returns:

Type	Description
tuple[tuple[ndarray, list[ndarray], list[ndarray]], float]	A ((weights, factors, projections), R2X) tuple: the standardized weights and [A, B, C] factor matrices, the per-condition projection matrices P_k, and the final fraction of variance explained.

Source code in parafac2/parafac2.py

def parafac2_nd(
    X_in: anndata.AnnData | CompressedData,
    rank: int,
    n_iter_max: int = 100,
    tol: float = 1e-6,
    random_state: int | None = None,
    callback: Callable[[int, float, list[np.ndarray]], None] | None = None,
    backend: str | None = None,
    normalize_slices: bool = False,
    n_inner: int = 1,
    compress: int | tuple[int, int | None] | str | bool | None = None,
) -> tuple[tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], float]:
    r"""The same interface as regular PARAFAC2 with optional CANDELINC compression.

    If ``compress`` is specified (or if ``X_in`` is already a
    :class:`~parafac2.compress.CompressedData`), PARAFAC2 is fit in the
    compressed subspace (Bro's "compress-then-fit"), eliminating per-sweep raw
    data passes.

    If ``normalize_slices`` is True, each condition's contribution to the
    factor updates is rescaled by the inverse of its (mean-centered)
    Frobenius norm. This prevents conditions with many more cells (or much
    higher variance) from dominating the shared factors, e.g. the ``A``
    matrix. The weighting is computed from small per-condition summary
    statistics and applied to per-condition intermediates only, so ``X`` is
    never copied or modified. The reported error/R2X are unaffected, since
    they are still computed from the unweighted fit.

    Parameters
    ----------
    X_in : anndata.AnnData | CompressedData
        Input dataset with the (optionally sparse) data matrix in ``X_in.X``,
        condition labels in ``X_in.obs["condition_unique_idxs"]``, and
        optionally per-gene means in ``X_in.var["means"]`` (defaults to zero,
        i.e. no centering, if absent). Alternatively, a pre-compressed
        :class:`~parafac2.compress.CompressedData` object.
    rank : int
        The number of components to fit.
    n_iter_max : int, default 100
        Maximum number of ALS iterations.
    tol : float, default 1e-6
        Convergence tolerance: iteration stops once the (non-negative)
        decrease in relative error between successive iterations drops
        below this value.
    random_state : int | None, default None
        Seed controlling the randomized SVD initialization.
    callback : Callable[[int, float, list[np.ndarray]], None] | None, default None
        Optional callback invoked after each iteration with the iteration
        index, the relative error, and the current factor matrices.
    backend : str | None, default None
        Compute backend to run matrix products on: one of ``'mlx'``,
        ``'cupy'``, or ``'cpu'``. If ``None``, the first available
        accelerator is auto-detected (see
        :func:`~parafac2.backend.get_backend`).
    normalize_slices : bool, default False
        Whether to rescale each condition's contribution to the factor
        updates by the inverse of its Frobenius norm, as described above.
    n_inner : int, default 1
        Number of ``(projection, A, B)`` sub-iterations per sweep. These read
        the data only through the cached ``W``, costing ``O(n_cells *
        rank^2)`` against the ``O(nnz * rank)`` of a raw-data product, so
        they are close to free on sparse inputs. Raising ``n_inner`` trades
        that cheap compute for fewer sweeps, and so for fewer raw-data
        passes: on structured test data, ``n_inner=2`` cut the sweeps needed
        to reach a fixed error by ~20%, with little further gain beyond 3.
        Whether that is a net win depends on how strongly the raw-data
        products dominate, so it is worth benchmarking per dataset. The
        default of 1 reproduces the classic one-update-per-mode ALS sweep.
    compress : int | tuple[int, int | None] | str | bool | None, default None
        Compression mode. If ``None`` or ``False`` (default), exact ALS is
        used. If ``"auto"`` or ``True``, sets compression dimensions
        automatically based on ``rank``. If an integer, sets both gene and cell
        compression dimensions to that value. If a tuple ``(L_g, L_c)``, sets
        dimensions separately (pass ``L_c=None`` for gene-only compression).
        Ignored if ``X_in`` is already a
        :class:`~parafac2.compress.CompressedData`.

    Returns
    -------
    tuple[tuple[np.ndarray, list[np.ndarray], list[np.ndarray]], float]
        A ``((weights, factors, projections), R2X)`` tuple: the standardized
        weights and ``[A, B, C]`` factor matrices, the per-condition
        projection matrices ``P_k``, and the final fraction of variance
        explained.
    """
    # Verbose if this is not an automated build
    verbose = "CI" not in os.environ

    if isinstance(X_in, CompressedData):
        return _fit_parafac2_compressed(
            X_in,
            rank=rank,
            n_iter_max=n_iter_max,
            tol=tol,
            random_state=random_state,
            callback=callback,
            verbose=verbose,
        )

    if compress is not None and compress is not False:
        compressed = compress_dataset(
            X_in,
            L=compress,
            rank=rank,
            random_state=random_state,
            normalize_slices=normalize_slices,
            backend=backend,
        )
        return _fit_parafac2_compressed(
            compressed,
            rank=rank,
            n_iter_max=n_iter_max,
            tol=tol,
            random_state=random_state,
            callback=callback,
            verbose=verbose,
        )

    (
        X_mat,
        condition_unique_idxs,
        means,
        norm_tensor,
        slice_weights,
    ) = extract_dataset_info(X_in, normalize_slices=normalize_slices)

    X_raw = to_gpu(X_mat, backend=backend)

    factors, _ = parafac2_init(
        X_raw,
        condition_unique_idxs,
        rank=rank,
        means=means,
        random_state=random_state,
        norm_tensor=norm_tensor,
    )

    cond_slices = condition_slices(
        condition_unique_idxs, int(np.amax(condition_unique_idxs)) + 1
    )

    # W depends only on C, so it stays valid across the A and B updates and is
    # recomputed only once C changes. Each sweep therefore costs exactly two
    # raw-data products: this one and the X^T @ H inside the mode-2 update.
    W = calc_W(X_raw, means, factors[2])
    projections, S = project_data(W, factors, cond_slices)
    errs = [calc_err(S, factors, norm_tensor) / norm_tensor]

    tq = tqdm(range(n_iter_max), disable=(not verbose), delay=0.5)
    for iteration in tq:
        # The (P, A, B) block reads the data only through the cached W, so
        # extra inner passes buy convergence at no raw-data cost.
        for _ in range(n_inner):
            factors = parafac_update(factors, 0, S, slice_weights=slice_weights)
            projections, S = project_data(W, factors, cond_slices)
            factors = parafac_update(factors, 1, S, slice_weights=slice_weights)
            projections, S = project_data(W, factors, cond_slices)

        factors = parafac_update(
            factors,
            2,
            S,
            projections,
            X=X_raw,
            means=means,
            cond_slices=cond_slices,
            slice_weights=slice_weights,
        )

        # C changed, so refresh W; this also yields the projections and error
        # for the factors as they stand at the end of this sweep.
        W = calc_W(X_raw, means, factors[2])
        projections, S = project_data(W, factors, cond_slices)
        errs.append(calc_err(S, factors, norm_tensor) / norm_tensor)

        delta = errs[-2] - errs[-1]
        tq.set_postfix(error=errs[-1], R2X=1.0 - errs[-1], Δ=delta, refresh=False)
        if callback is not None:
            callback(iteration, errs[-1], factors)

        if 0 <= delta < tol:
            break

    R2X = 1 - errs[-1]

    # Standardize the results and return
    return standardize_pf2(factors, projections), R2X

store_pf2(X, parafac2_output)

Store the Pf2 results into the anndata object.

Parameters:

Name	Type	Description	Default
X	AnnData | CompressedData	The dataset the factorization was fit on. Must have X.obs["condition_unique_idxs"] set (as produced by :func:~parafac2.normalize.prepare_dataset or equivalent). If a :class:~parafac2.compress.CompressedData is provided, factors are written to its underlying .adata object.	required
parafac2_output	tuple[ndarray, list[ndarray], list[ndarray]]	The (weights, factors, projections) output of :func:parafac2_nd, where factors is the [A, B, C] factor matrices and projections is the per-condition projection matrices P_k.	required

Returns:

Type	Description
AnnData	The target AnnData object, mutated in place, with the weights in X.uns["Pf2_weights"], factors in X.uns["Pf2_A"]/X.uns["Pf2_B"]/X.varm["Pf2_C"], and per-cell projections in X.obsm["projections"] and X.obsm["weighted_projections"] (projections composed with B).

Source code in parafac2/parafac2.py

def store_pf2(
    X: anndata.AnnData | CompressedData,
    parafac2_output: tuple[np.ndarray, list[np.ndarray], list[np.ndarray]],
) -> anndata.AnnData:
    """Store the Pf2 results into the anndata object.

    Parameters
    ----------
    X : anndata.AnnData | CompressedData
        The dataset the factorization was fit on. Must have
        ``X.obs["condition_unique_idxs"]`` set (as produced by
        :func:`~parafac2.normalize.prepare_dataset` or equivalent). If a
        :class:`~parafac2.compress.CompressedData` is provided, factors are
        written to its underlying ``.adata`` object.
    parafac2_output : tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]
        The ``(weights, factors, projections)`` output of :func:`parafac2_nd`,
        where ``factors`` is the ``[A, B, C]`` factor matrices and
        ``projections`` is the per-condition projection matrices ``P_k``.

    Returns
    -------
    anndata.AnnData
        The target AnnData object, mutated in place, with the weights in
        ``X.uns["Pf2_weights"]``, factors in
        ``X.uns["Pf2_A"]``/``X.uns["Pf2_B"]``/``X.varm["Pf2_C"]``, and
        per-cell projections in ``X.obsm["projections"]`` and
        ``X.obsm["weighted_projections"]`` (projections composed with ``B``).
    """
    if isinstance(X, CompressedData):
        if X.adata is None:
            raise ValueError("CompressedData has no associated AnnData object.")
        target_adata = X.adata
        condition_unique_idxs = X.condition_unique_idxs
    else:
        target_adata = X
        condition_unique_idxs = target_adata.obs["condition_unique_idxs"]

    target_adata.uns["Pf2_weights"] = parafac2_output[0]
    target_adata.uns["Pf2_A"], target_adata.uns["Pf2_B"], target_adata.varm["Pf2_C"] = (
        parafac2_output[1]
    )

    target_adata.obsm["projections"] = np.zeros(
        (target_adata.shape[0], len(target_adata.uns["Pf2_weights"])), dtype=np.float32
    )
    for i, p in enumerate(parafac2_output[2]):
        target_adata.obsm["projections"][condition_unique_idxs == i, :] = p

    target_adata.obsm["weighted_projections"] = (
        target_adata.obsm["projections"] @ target_adata.uns["Pf2_B"]
    ).astype(np.float32, copy=False)

    return target_adata

parafac2.normalize

parafac2.normalize

Dataset preprocessing and normalization utilities for PARAFAC2 analysis.

This module provides functions to filter, normalize, and annotate single-cell gene expression datasets stored in AnnData objects prior to PARAFAC2 matrix factorization.

prepare_dataset(X, condition_name, geneThreshold)

Preprocess and normalize an AnnData dataset for PARAFAC2 factorization.

Performs quality control filtering of low-count cells and low-expression genes, normalizes total cell counts and gene sums, applies a log10 transformation, and computes metadata required by PARAFAC2 (condition indices and gene means).

Parameters:

Name	Type	Description	Default
X	AnnData	Input single-cell dataset with raw count matrix stored in X.X (must be a sparse matrix with non-negative values).	required
condition_name	str	Column name in X.obs identifying the sample or experimental condition grouping for each cell.	required
geneThreshold	float	Minimum threshold fraction for gene inclusion. Genes with total counts less than geneThreshold * total_cells are filtered out.	required

Returns:

Type	Description
AnnData	A filtered and normalized copy of the AnnData object. Contains the log-transformed normalized counts in X.X, integer condition codes in X.obs["condition_unique_idxs"], and per-gene mean expression values in X.var["means"].

Source code in parafac2/normalize.py

def prepare_dataset(
    X: anndata.AnnData, condition_name: str, geneThreshold: float
) -> anndata.AnnData:
    """Preprocess and normalize an AnnData dataset for PARAFAC2 factorization.

    Performs quality control filtering of low-count cells and low-expression
    genes, normalizes total cell counts and gene sums, applies a log10
    transformation, and computes metadata required by PARAFAC2 (condition
    indices and gene means).

    Parameters
    ----------
    X : anndata.AnnData
        Input single-cell dataset with raw count matrix stored in ``X.X``
        (must be a sparse matrix with non-negative values).
    condition_name : str
        Column name in ``X.obs`` identifying the sample or experimental
        condition grouping for each cell.
    geneThreshold : float
        Minimum threshold fraction for gene inclusion. Genes with total counts
        less than ``geneThreshold * total_cells`` are filtered out.

    Returns
    -------
    anndata.AnnData
        A filtered and normalized copy of the AnnData object. Contains the
        log-transformed normalized counts in ``X.X``, integer condition
        codes in ``X.obs["condition_unique_idxs"]``, and per-gene mean
        expression values in ``X.var["means"]``.
    """
    assert issparse(X.X)
    X_X_raw = cast("csr_array", X.X)
    assert np.amin(X_X_raw.data) >= 0.0

    # Filter out genes with too few reads, and cells with fewer than 10 counts
    cell_mask = np.ravel(X_X_raw.sum(axis=1)) > 10
    gene_mask = np.ravel(X_X_raw.sum(axis=0)) > (geneThreshold * X_X_raw.shape[0])

    # Subset and materialize actual AnnData object before modifying X.X
    if cell_mask.all() and gene_mask.all():
        X = X.copy()
    else:
        X = X[cell_mask, gene_mask].copy()

    # Convert subset to csr_array and float32 data
    X.X = csr_array(X.X)
    X_X = cast("csr_array", X.X)

    if X_X.dtype != np.float32:
        X_X.data = X_X.data.astype(np.float32)

    ## Normalize total counts per cell
    # Keep the counts on a reasonable scale to avoid accuracy issues
    counts_per_cell = np.ravel(X_X.sum(axis=1)).astype(np.float32, copy=False)
    counts_per_cell /= np.median(counts_per_cell)
    # In-place CSR row scaling
    X_X.data /= np.repeat(counts_per_cell, np.diff(X_X.indptr))

    # Scale genes by sum, in-place CSR column scaling
    gene_sums = np.ravel(X_X.sum(axis=0)).astype(np.float32, copy=False)
    X_X.data /= gene_sums[X_X.indices]

    # Transform values in-place to avoid nnz-sized temporaries
    X_X.data *= np.float32(1000.0)
    X_X.data += np.float32(1.0)
    np.log10(X_X.data, out=X_X.data)

    # Get the indices for subsetting the data
    X.obs["condition_unique_idxs"] = pd.Categorical(X.obs[condition_name]).codes

    # Pre-calculate gene means
    X.var["means"] = np.ravel(X_X.mean(axis=0))

    return X

parafac2.compress

parafac2.compress

CANDELINC compression for PARAFAC2.

Compresses the gene and cell modes of single-cell datasets prior to PARAFAC2 factorization, collapsing the problem to small dense per-condition cores.

CompressedData dataclass

Compressed representation of a multi-condition single-cell dataset.

Stores the small dense per-condition cores and the orthonormal bases for the gene mode (Q) and cell mode (Q_k), enabling fast rank sweeps without touching the raw data again.

Source code in parafac2/compress.py

@dataclass
class CompressedData:
    """Compressed representation of a multi-condition single-cell dataset.

    Stores the small dense per-condition cores and the orthonormal bases for
    the gene mode (``Q``) and cell mode (``Q_k``), enabling fast rank sweeps
    without touching the raw data again.
    """

    cores: list[np.ndarray]
    """List of length ``n_cond`` with dense cores ``Y_k`` of shape ``(L_c_k, L_g)``."""

    Q: np.ndarray
    """Gene projection basis of shape ``(n_genes, L_g)`` with orthonormal columns."""

    Q_k: list[np.ndarray | None] | None
    """Per-condition cell projection bases ``(n_cells_k, L_c_k)`` or ``None``."""

    condition_unique_idxs: np.ndarray
    """Integer condition indices for each cell."""

    norm_tensor: float
    """Total squared Frobenius norm of the original mean-centered dataset."""

    lost_var: float
    """Variance discarded by the compression projectors."""

    total_cells: int
    """Total number of cells across all conditions."""

    n_genes: int
    """Number of genes in the uncompressed dataset."""

    n_cond: int
    """Number of conditions."""

    slice_weights: np.ndarray | None = None
    """Optional per-condition slice weights for normalized ALS."""

    means: np.ndarray | None = None
    """Per-gene means subtracted during compression."""

    adata: anndata.AnnData | None = None
    """Reference to original AnnData object if available."""

    @property
    def L_g(self) -> int:
        """Gene compression dimension."""
        return self.Q.shape[1]

    @property
    def max_cell_dim(self) -> int:
        """Maximum cell dimension across cores."""
        return max(c.shape[0] for c in self.cores)

L_g property

Gene compression dimension.

Q instance-attribute

Gene projection basis of shape (n_genes, L_g) with orthonormal columns.

Q_k instance-attribute

Per-condition cell projection bases (n_cells_k, L_c_k) or None.

adata = None class-attribute instance-attribute

Reference to original AnnData object if available.

condition_unique_idxs instance-attribute

Integer condition indices for each cell.

cores instance-attribute

List of length n_cond with dense cores Y_k of shape (L_c_k, L_g).

lost_var instance-attribute

Variance discarded by the compression projectors.

max_cell_dim property

Maximum cell dimension across cores.

means = None class-attribute instance-attribute

Per-gene means subtracted during compression.

n_cond instance-attribute

Number of conditions.

n_genes instance-attribute

Number of genes in the uncompressed dataset.

norm_tensor instance-attribute

Total squared Frobenius norm of the original mean-centered dataset.

slice_weights = None class-attribute instance-attribute

Optional per-condition slice weights for normalized ALS.

total_cells instance-attribute

Total number of cells across all conditions.

compress_cells(X_c, condition_unique_idxs, L_c=None)

Compute per-condition cell-mode compression projectors Q_k and cores Y_k.

Parameters:

Name	Type	Description	Default
X_c	ndarray	Gene-compressed data matrix of shape (total_cells, L_g).	required
condition_unique_idxs	ndarray	Integer condition index for each cell.	required
L_c	int | None	Target cell subspace dimension per condition. If None, cell compression is skipped (cores are X_c slices).	None

Returns:

Type	Description
tuple[list[ndarray], list[ndarray | None] | None, float]	(cores, Q_k_list, norm_cores_sq) where each core Y_k has shape (L_c_k, L_g) and norm_cores_sq is the total squared Frobenius norm of all cores.

Source code in parafac2/compress.py

def compress_cells(
    X_c: np.ndarray,
    condition_unique_idxs: np.ndarray,
    L_c: int | None = None,
) -> tuple[list[np.ndarray], list[np.ndarray | None] | None, float]:
    """Compute per-condition cell-mode compression projectors Q_k and cores Y_k.

    Parameters
    ----------
    X_c : np.ndarray
        Gene-compressed data matrix of shape ``(total_cells, L_g)``.
    condition_unique_idxs : np.ndarray
        Integer condition index for each cell.
    L_c : int | None, default None
        Target cell subspace dimension per condition. If ``None``, cell
        compression is skipped (cores are ``X_c`` slices).

    Returns
    -------
    tuple[list[np.ndarray], list[np.ndarray | None] | None, float]
        ``(cores, Q_k_list, norm_cores_sq)`` where each core ``Y_k`` has
        shape ``(L_c_k, L_g)`` and ``norm_cores_sq`` is the total squared
        Frobenius norm of all cores.
    """
    n_cond = int(np.amax(condition_unique_idxs)) + 1
    cores: list[np.ndarray] = []
    Q_k_list: list[np.ndarray | None] | None = [] if L_c is not None else None
    norm_cores_sq = 0.0

    for i in range(n_cond):
        cond_i = condition_unique_idxs == i
        X_c_i = X_c[cond_i]
        n_k = X_c_i.shape[0]

        if L_c is None or n_k <= L_c:
            cores.append(X_c_i)
            if Q_k_list is not None:
                Q_k_list.append(None)
            norm_cores_sq += float(np.sum(X_c_i**2))
        else:
            # Thin SVD of (n_k, L_g) where L_g <= 100
            U_i, S_i, Vh_i = np.linalg.svd(X_c_i, full_matrices=False)
            L_k = min(n_k, L_c)
            Q_i = U_i[:, :L_k].astype(np.float64)
            Y_i = (S_i[:L_k, np.newaxis] * Vh_i[:L_k, :]).astype(np.float64)
            cores.append(Y_i)
            if Q_k_list is not None:
                Q_k_list.append(Q_i)
            norm_cores_sq += float(np.sum(S_i[:L_k] ** 2))

    return cores, Q_k_list, norm_cores_sq

compress_dataset(X_in, L='auto', rank=None, n_power_iter=2, random_state=None, normalize_slices=False, backend=None)

Compress an AnnData dataset in gene and cell modes.

Parameters:

Name	Type	Description	Default
X_in	AnnData	Input dataset with data in X_in.X, condition indices in X_in.obs["condition_unique_idxs"], and optional means in X_in.var["means"].	required
L	int | tuple[int, int | None] | str	Compression dimension(s). If "auto", picks dimensions based on rank (or default rank 30 if rank is None). If an int, sets both L_g = L and L_c = L. If a tuple (L_g, L_c), sets gene and cell dimensions individually (pass L_c=None for gene-only compression).	"auto"
rank	int | None	Expected maximum rank to fit on the compressed data. Used when L="auto".	None
n_power_iter	int	Number of power iterations for randomized SVD.	2
random_state	int | Generator | None	Random seed or generator.	None
normalize_slices	bool	Whether to precalculate slice weights for normalized ALS.	False
backend	str | None	Compute backend for raw matrix products.	None

Returns:

Type	Description
CompressedData	The compressed dataset ready for fast PARAFAC2 fitting.

Source code in parafac2/compress.py

def compress_dataset(
    X_in: anndata.AnnData,
    L: int | tuple[int, int | None] | str = "auto",
    rank: int | None = None,
    n_power_iter: int = 2,
    random_state: int | np.random.Generator | None = None,
    normalize_slices: bool = False,
    backend: str | None = None,
) -> CompressedData:
    """Compress an AnnData dataset in gene and cell modes.

    Parameters
    ----------
    X_in : anndata.AnnData
        Input dataset with data in ``X_in.X``, condition indices in
        ``X_in.obs["condition_unique_idxs"]``, and optional means in
        ``X_in.var["means"]``.
    L : int | tuple[int, int | None] | str, default "auto"
        Compression dimension(s). If ``"auto"``, picks dimensions based on
        ``rank`` (or default rank 30 if ``rank`` is None). If an int, sets
        both ``L_g = L`` and ``L_c = L``. If a tuple ``(L_g, L_c)``, sets
        gene and cell dimensions individually (pass ``L_c=None`` for
        gene-only compression).
    rank : int | None, default None
        Expected maximum rank to fit on the compressed data. Used when
        ``L="auto"``.
    n_power_iter : int, default 2
        Number of power iterations for randomized SVD.
    random_state : int | np.random.Generator | None, default None
        Random seed or generator.
    normalize_slices : bool, default False
        Whether to precalculate slice weights for normalized ALS.
    backend : str | None, default None
        Compute backend for raw matrix products.

    Returns
    -------
    CompressedData
        The compressed dataset ready for fast PARAFAC2 fitting.
    """
    (
        X_mat,
        condition_unique_idxs,
        means,
        norm_tensor,
        slice_weights,
    ) = extract_dataset_info(X_in, normalize_slices=normalize_slices)
    total_cells, n_genes = X_mat.shape
    n_cond = int(np.amax(condition_unique_idxs)) + 1

    # Determine L_g and L_c
    target_rank = rank if rank is not None else 30
    if isinstance(L, str) and L == "auto":
        L_g_val = min(n_genes, max(4 * target_rank, target_rank + 20))
        L_c_val: int | None = max(4 * target_rank, target_rank + 20)
    elif isinstance(L, tuple):
        L_g_val, L_c_val = L
        L_g_val = min(n_genes, L_g_val)
    elif isinstance(L, (int, np.integer)):
        L_g_val = min(n_genes, int(L))
        L_c_val = int(L)
    else:
        raise ValueError(f"Invalid compression parameter L: {L}")

    X_raw = to_gpu(X_mat, backend=backend)
    X_c, Q, _norm_Xc_sq = compress_genes(
        X_raw,
        means,
        L_g=L_g_val,
        n_power_iter=n_power_iter,
        random_state=random_state,
    )

    cores, Q_k, norm_cores_sq = compress_cells(
        X_c,
        condition_unique_idxs,
        L_c=L_c_val,
    )

    lost_var = float(np.maximum(0.0, norm_tensor - norm_cores_sq))

    return CompressedData(
        cores=cores,
        Q=Q,
        Q_k=Q_k,
        condition_unique_idxs=condition_unique_idxs,
        norm_tensor=norm_tensor,
        lost_var=lost_var,
        total_cells=total_cells,
        n_genes=n_genes,
        n_cond=n_cond,
        slice_weights=slice_weights,
        means=means,
        adata=X_in,
    )

compress_genes(X, means, L_g, n_power_iter=2, random_state=None)

Compute gene-mode compression projector Q and compressed matrix Xc.

Parameters:

Name	Type	Description	Default
X	Any	Stacked data matrix of shape (total_cells, n_genes).	required
means	ndarray | None	Per-gene means for centering.	required
L_g	int	Target gene subspace dimension.	required
n_power_iter	int	Number of power iterations for randomized SVD.	2
random_state	int | Generator | None	Random seed or generator.	None

Returns:

Type	Description
tuple[ndarray, ndarray, float]	(X_c, Q, norm_Xc_sq) where X_c = (X - 1 mu^T) @ Q of shape (total_cells, L_g), Q is (n_genes, L_g) orthonormal, and norm_Xc_sq is the squared Frobenius norm of X_c.

Source code in parafac2/compress.py

def compress_genes(
    X: Any,
    means: np.ndarray | None,
    L_g: int,
    n_power_iter: int = 2,
    random_state: int | np.random.Generator | None = None,
) -> tuple[np.ndarray, np.ndarray, float]:
    """Compute gene-mode compression projector Q and compressed matrix Xc.

    Parameters
    ----------
    X : Any
        Stacked data matrix of shape ``(total_cells, n_genes)``.
    means : np.ndarray | None
        Per-gene means for centering.
    L_g : int
        Target gene subspace dimension.
    n_power_iter : int, default 2
        Number of power iterations for randomized SVD.
    random_state : int | np.random.Generator | None, default None
        Random seed or generator.

    Returns
    -------
    tuple[np.ndarray, np.ndarray, float]
        ``(X_c, Q, norm_Xc_sq)`` where ``X_c = (X - 1 mu^T) @ Q`` of shape
        ``(total_cells, L_g)``, ``Q`` is ``(n_genes, L_g)`` orthonormal, and
        ``norm_Xc_sq`` is the squared Frobenius norm of ``X_c``.
    """
    _n_cells, n_genes = X.shape
    L_g = min(n_genes, L_g)

    Q = randomized_svd_right(
        X,
        means,
        n_components=L_g,
        n_oversamples=0,
        n_power_iter=n_power_iter,
        random_state=random_state,
    )
    X_c = calc_W(X, means, Q)
    norm_Xc_sq = float(np.sum(X_c**2))

    return X_c, Q, norm_Xc_sq

init_compressed_factors(cores, rank, random_state=None)

Initialize factor matrices [A, B, C_L] directly on compressed cores.

Source code in parafac2/compress.py

def init_compressed_factors(
    cores: list[np.ndarray],
    rank: int,
    random_state: int | np.random.Generator | None = None,
) -> list[np.ndarray]:
    """Initialize factor matrices [A, B, C_L] directly on compressed cores."""
    n_cond = len(cores)
    L_g = cores[0].shape[1]
    assert rank <= L_g, f"Rank {rank} exceeds compressed gene dimension {L_g}"

    # SVD of stacked cores to initialize C_L
    Y_stacked = np.concatenate(cores, axis=0)
    _, _, vh = np.linalg.svd(Y_stacked, full_matrices=False)
    C_L = vh[:rank, :].T.astype(np.float64)

    return [
        np.ones((n_cond, rank), dtype=np.float64),
        np.eye(rank, dtype=np.float64),
        C_L,
    ]

project_data_compressed(cores, factors, norm_tensor, mode, return_projections=False, slice_weights=None)

project_data_compressed(
    cores: list[np.ndarray],
    factors: list[np.ndarray],
    norm_tensor: float,
    mode: int,
    return_projections: Literal[False] = False,
    slice_weights: np.ndarray | None = None,
) -> tuple[np.ndarray, float]

project_data_compressed(
    cores: list[np.ndarray],
    factors: list[np.ndarray],
    norm_tensor: float,
    mode: int,
    return_projections: Literal[True],
    slice_weights: np.ndarray | None = None,
) -> list[np.ndarray]

Project compressed per-condition cores and accumulate MTTKRP and error.

Parameters:

Name	Type	Description	Default
cores	list[ndarray]	List of per-condition core matrices Y_k of shape (L_c_k, L_g).	required
factors	list[ndarray]	Current factor matrices [A, B, C_L] in the compressed space.	required
norm_tensor	float	Squared Frobenius norm of the original mean-centered tensor.	required
mode	int	Mode to update (0, 1, or 2).	required
return_projections	bool	Whether to return the list of projection matrices P_tilde_k.	False
slice_weights	ndarray | None	Optional per-condition slice weights.	None

Returns:

Type	Description
tuple[ndarray, float] | list[ndarray]	(mttkrp, norm_sq_err) or list of P_tilde_k.

Source code in parafac2/compress.py

def project_data_compressed(
    cores: list[np.ndarray],
    factors: list[np.ndarray],
    norm_tensor: float,
    mode: int,
    return_projections: bool = False,
    slice_weights: np.ndarray | None = None,
) -> tuple[np.ndarray, float] | list[np.ndarray]:
    """Project compressed per-condition cores and accumulate MTTKRP and error.

    Parameters
    ----------
    cores : list[np.ndarray]
        List of per-condition core matrices ``Y_k`` of shape ``(L_c_k, L_g)``.
    factors : list[np.ndarray]
        Current factor matrices ``[A, B, C_L]`` in the compressed space.
    norm_tensor : float
        Squared Frobenius norm of the original mean-centered tensor.
    mode : int
        Mode to update (0, 1, or 2).
    return_projections : bool, default False
        Whether to return the list of projection matrices ``P_tilde_k``.
    slice_weights : np.ndarray | None, default None
        Optional per-condition slice weights.

    Returns
    -------
    tuple[np.ndarray, float] | list[np.ndarray]
        ``(mttkrp, norm_sq_err)`` or list of ``P_tilde_k``.
    """
    A, B, C_L = factors
    rank = B.shape[0]
    n_cond = len(cores)

    norm_sq_err = norm_tensor + float(((A.T @ A) * (B.T @ B) * (C_L.T @ C_L)).sum())

    if mode == 0:
        mttkrp = np.zeros((n_cond, rank), dtype=np.float64)
    elif mode == 1:
        mttkrp = np.zeros((rank, rank), dtype=np.float64)
    else:
        mttkrp = np.zeros_like(C_L, dtype=np.float64)

    proj_list = []
    for i in range(n_cond):
        Y_i = cores[i]
        W_i = Y_i @ C_L  # (L_c_i, rank)
        M = W_i @ (B * A[i]).T  # (L_c_i, rank)
        proj = polar_factor(M)
        proj_list.append(proj)

        if return_projections:
            continue

        psc = proj.T @ W_i  # (rank, rank)
        m_i = np.sum(psc * B, axis=0)
        norm_sq_err -= 2.0 * float(np.dot(A[i], m_i))

        w_i = 1.0 if slice_weights is None else slice_weights[i]

        if mode == 0:
            mttkrp[i] = m_i * w_i
        elif mode == 1:
            mttkrp += psc * A[i] * w_i
        else:
            H_tilde_i = proj @ (B * A[i]) * w_i
            mttkrp += Y_i.T @ H_tilde_i

    if return_projections:
        return proj_list

    return mttkrp, float(norm_sq_err)

parafac2.utils

parafac2.utils

Low-level numerical routines supporting the PARAFAC2 fit.

Provides norm computation over (optionally mean-centered, optionally sparse) data, the per-condition projection step, the per-mode ALS factor update (which forms its own MTTKRP), and post-fit standardization of the factors and projections.

The fit touches the raw data through exactly two products, which together dominate runtime on single-cell-sized inputs:

• W = (X - 1 mu^T) @ C (:func:calc_W), which depends only on C.
• X^T @ H for the mode-2 MTTKRP (inside :func:parafac_update).

Everything else flows through the compressed per-condition slices S_k = P_k^T W_k, an (n_cond, rank, rank) array small enough to keep resident. In particular the mode-0 and mode-1 MTTKRPs and the reconstruction error are all functions of S alone, so the projections and both of those factor updates can be recomputed from a cached W without re-reading the data.

calc_W(X, means, C)

Compute W = (X - 1 mu^T) @ C, the first of the two raw-data products.

W depends only on C, so it stays valid across the A and B updates and only has to be recomputed once C changes.

The product is taken in X's own dtype. That matters: handing a float64 C to a float32 sparse X makes SciPy upcast the entire sparse matrix, doubling both the memory traffic that dominates this step and the peak memory. The result is widened to float64 afterwards, which is O(n_cells * rank) and so negligible beside the product itself.

Parameters:

Name	Type	Description	Default
X	Any	The (optionally sparse or GPU-backed) data matrix, stacked across all conditions, with shape (total_cells, n_genes).	required
means	ndarray | None	Per-gene means to mean-center X by, or None to skip centering.	required
C	ndarray	The current gene factor matrix, shape (n_genes, rank).	required

Returns:

Type	Description
ndarray	The float64 array W of shape (total_cells, rank).

Source code in parafac2/utils.py

def calc_W(X: Any, means: np.ndarray | None, C: np.ndarray) -> np.ndarray:
    """Compute ``W = (X - 1 mu^T) @ C``, the first of the two raw-data products.

    ``W`` depends only on ``C``, so it stays valid across the ``A`` and ``B``
    updates and only has to be recomputed once ``C`` changes.

    The product is taken in ``X``'s own dtype. That matters: handing a
    float64 ``C`` to a float32 sparse ``X`` makes SciPy upcast the entire
    sparse matrix, doubling both the memory traffic that dominates this step
    and the peak memory. The result is widened to float64 afterwards, which
    is ``O(n_cells * rank)`` and so negligible beside the product itself.

    Parameters
    ----------
    X : Any
        The (optionally sparse or GPU-backed) data matrix, stacked across all
        conditions, with shape ``(total_cells, n_genes)``.
    means : np.ndarray | None
        Per-gene means to mean-center ``X`` by, or ``None`` to skip centering.
    C : np.ndarray
        The current gene factor matrix, shape ``(n_genes, rank)``.

    Returns
    -------
    np.ndarray
        The float64 array ``W`` of shape ``(total_cells, rank)``.
    """
    C_op = np.ascontiguousarray(C, dtype=matrix_dtype(X))
    W = np.asarray(matmul(X, C_op), dtype=np.float64)
    if means is not None:
        W -= means @ C
    return W

calc_err(S, factors, norm_X_sq)

Return the squared reconstruction error from the compressed slices.

Uses the expansion ||X||^2 + Tr(A^T A * B^T B * C^T C) - 2 <A, diag(B^T S_k)>, so no raw-data pass is needed and the error is free to evaluate as often as desired (e.g. to monitor an inner iteration).

Parameters:

Name	Type	Description	Default
S	ndarray	The stacked compressed slices from :func:project_data.	required
factors	list[ndarray]	The current [A, B, C] factor matrices.	required
norm_X_sq	float	The squared Frobenius norm of the mean-centered X, as returned by :func:calc_norm_sq.	required

Returns:

Type	Description
float	The squared reconstruction error.

Source code in parafac2/utils.py

def calc_err(S: np.ndarray, factors: list[np.ndarray], norm_X_sq: float) -> float:
    """Return the squared reconstruction error from the compressed slices.

    Uses the expansion ``||X||^2 + Tr(A^T A * B^T B * C^T C) - 2 <A, diag(B^T
    S_k)>``, so no raw-data pass is needed and the error is free to evaluate
    as often as desired (e.g. to monitor an inner iteration).

    Parameters
    ----------
    S : np.ndarray
        The stacked compressed slices from :func:`project_data`.
    factors : list[np.ndarray]
        The current ``[A, B, C]`` factor matrices.
    norm_X_sq : float
        The squared Frobenius norm of the mean-centered ``X``, as returned by
        :func:`calc_norm_sq`.

    Returns
    -------
    float
        The squared reconstruction error.
    """
    A, B, C = factors
    norm_sq_err = norm_X_sq + float(((A.T @ A) * (B.T @ B) * (C.T @ C)).sum())
    norm_sq_err -= 2.0 * float(np.sum(A * np.einsum("kqr,qr->kr", S, B)))
    return norm_sq_err

calc_norm_sq(X, means=None)

Return the squared Frobenius norm of the mean-centered matrix.

Parameters:

Name	Type	Description	Default
X	ndarray | csr_array	The (dense or sparse) matrix to compute the norm of.	required
means	ndarray | None	Per-column means to subtract before computing the norm. If None or all-zero, X is used uncentered.	None

Returns:

Type	Description
float	sum((X - means) ** 2), computed without densifying a sparse X.

Source code in parafac2/utils.py

def calc_norm_sq(X: np.ndarray | csr_array, means: np.ndarray | None = None) -> float:
    """Return the squared Frobenius norm of the mean-centered matrix.

    Parameters
    ----------
    X : np.ndarray | csr_array
        The (dense or sparse) matrix to compute the norm of.
    means : np.ndarray | None, default None
        Per-column means to subtract before computing the norm. If ``None``
        or all-zero, ``X`` is used uncentered.

    Returns
    -------
    float
        ``sum((X - means) ** 2)``, computed without densifying a sparse
        ``X``.
    """
    if means is None or np.all(means == 0):
        if issparse(X):
            return float(np.sum(cast("csr_array", X).data ** 2))
        return float(np.sum(X**2))

    means_arr = np.asarray(means).ravel()
    if issparse(X):
        mat_csr = cast("csr_array", X)
        M = mat_csr.shape[0]
        term1 = np.sum(mat_csr.data**2)
        term2 = -2.0 * np.sum(mat_csr.data * means_arr[mat_csr.indices])
        term3 = M * np.sum(means_arr**2)
        return float(term1 + term2 + term3)
    return float(np.sum((X - means_arr) ** 2))

calc_slice_norms(X, means, condition_unique_idxs, n_cond)

Return the per-condition Frobenius norm of the mean-centered slices.

Parameters:

Name	Type	Description	Default
X	ndarray | csr_array	The (dense or sparse) matrix stacked across all conditions.	required
means	ndarray | None	Per-column means to subtract before computing each slice's norm, or None/all-zero to skip centering.	required
condition_unique_idxs	ndarray	Integer array assigning each row of X to a condition index in [0, n_cond).	required
n_cond	int	The total number of conditions.	required

Returns:

Type	Description
ndarray	Array of length n_cond with the Frobenius norm of each condition's (mean-centered) rows of X.

Source code in parafac2/utils.py

def calc_slice_norms(
    X: np.ndarray | csr_array,
    means: np.ndarray | None,
    condition_unique_idxs: np.ndarray,
    n_cond: int,
) -> np.ndarray:
    """Return the per-condition Frobenius norm of the mean-centered slices.

    Parameters
    ----------
    X : np.ndarray | csr_array
        The (dense or sparse) matrix stacked across all conditions.
    means : np.ndarray | None
        Per-column means to subtract before computing each slice's norm, or
        ``None``/all-zero to skip centering.
    condition_unique_idxs : np.ndarray
        Integer array assigning each row of ``X`` to a condition index in
        ``[0, n_cond)``.
    n_cond : int
        The total number of conditions.

    Returns
    -------
    np.ndarray
        Array of length ``n_cond`` with the Frobenius norm of each
        condition's (mean-centered) rows of ``X``.
    """
    idxs = np.asarray(condition_unique_idxs)
    counts = np.bincount(idxs, minlength=n_cond).astype(np.float64)

    if issparse(X):
        mat_csr = cast("csr_array", X)
        group_of_nnz = np.repeat(idxs, np.diff(mat_csr.indptr))
        sums_sq = np.bincount(
            group_of_nnz, weights=mat_csr.data.astype(np.float64) ** 2, minlength=n_cond
        )
        if means is None or np.all(means == 0):
            return np.sqrt(sums_sq)

        means_arr = np.asarray(means).ravel()
        cross = np.bincount(
            group_of_nnz,
            weights=mat_csr.data.astype(np.float64) * means_arr[mat_csr.indices],
            minlength=n_cond,
        )
        mean_sq_total = np.sum(means_arr**2)
        return np.sqrt(np.maximum(sums_sq - 2.0 * cross + counts * mean_sq_total, 0.0))

    means_arr = np.asarray(means).ravel() if means is not None else 0.0
    row_sums_sq = np.sum((np.asarray(X) - means_arr) ** 2, axis=1)
    return np.sqrt(np.bincount(idxs, weights=row_sums_sq, minlength=n_cond))

condition_slices(condition_unique_idxs, n_cond)

Return a per-condition row selector for each condition.

Computing condition_unique_idxs == i inside the per-condition loop costs O(n_cells) per condition, i.e. O(n_cells * n_cond) per pass over the data, plus a fancy-indexed copy each time. Precomputing the selectors once drops that to O(n_cells), and when the rows are already grouped by condition (the usual case, since conditions are concatenated) the selectors are plain slice objects, making W[sel] a zero-copy view.

Parameters:

Name	Type	Description	Default
condition_unique_idxs	ndarray	Integer array assigning each row to a condition in [0, n_cond).	required
n_cond	int	The total number of conditions.	required

Returns:

Type	Description
list[slice | ndarray]	One selector per condition: a slice when the condition's rows are contiguous, otherwise an integer index array.

Source code in parafac2/utils.py

def condition_slices(
    condition_unique_idxs: np.ndarray, n_cond: int
) -> list[slice | np.ndarray]:
    """Return a per-condition row selector for each condition.

    Computing ``condition_unique_idxs == i`` inside the per-condition loop
    costs ``O(n_cells)`` per condition, i.e. ``O(n_cells * n_cond)`` per pass
    over the data, plus a fancy-indexed copy each time. Precomputing the
    selectors once drops that to ``O(n_cells)``, and when the rows are
    already grouped by condition (the usual case, since conditions are
    concatenated) the selectors are plain ``slice`` objects, making
    ``W[sel]`` a zero-copy view.

    Parameters
    ----------
    condition_unique_idxs : np.ndarray
        Integer array assigning each row to a condition in ``[0, n_cond)``.
    n_cond : int
        The total number of conditions.

    Returns
    -------
    list[slice | np.ndarray]
        One selector per condition: a ``slice`` when the condition's rows are
        contiguous, otherwise an integer index array.
    """
    idxs = np.asarray(condition_unique_idxs)

    if idxs.size and np.all(np.diff(idxs) >= 0):
        starts = np.searchsorted(idxs, np.arange(n_cond), side="left")
        stops = np.searchsorted(idxs, np.arange(n_cond), side="right")
        return [slice(int(a), int(b)) for a, b in zip(starts, stops, strict=True)]

    order = np.argsort(idxs, kind="stable")
    bounds = np.searchsorted(idxs[order], np.arange(n_cond + 1))
    return [order[bounds[k] : bounds[k + 1]] for k in range(n_cond)]

extract_dataset_info(X_in, normalize_slices=False)

Extract matrix, condition indices, gene means, norm_sq, and optional slice weights.

Parameters:

Name	Type	Description	Default
X_in	AnnData	Input single-cell AnnData dataset.	required
normalize_slices	bool	Whether to calculate per-condition slice inverse-norm weights.	False

Returns:

Type	Description
tuple[ndarray | csr_array, ndarray, ndarray, float, ndarray | None]	The (X_mat, condition_unique_idxs, means, norm_tensor, slice_weights) tuple.

Source code in parafac2/utils.py

def extract_dataset_info(
    X_in: anndata.AnnData,
    normalize_slices: bool = False,
) -> tuple[np.ndarray | csr_array, np.ndarray, np.ndarray, float, np.ndarray | None]:
    """Extract matrix, condition indices, gene means, norm_sq, and optional slice weights.

    Parameters
    ----------
    X_in : anndata.AnnData
        Input single-cell AnnData dataset.
    normalize_slices : bool, default False
        Whether to calculate per-condition slice inverse-norm weights.

    Returns
    -------
    tuple[np.ndarray | csr_array, np.ndarray, np.ndarray, float, np.ndarray | None]
        The ``(X_mat, condition_unique_idxs, means, norm_tensor, slice_weights)`` tuple.
    """
    assert X_in.X is not None
    X_mat = cast("np.ndarray | csr_array", X_in.X)
    condition_unique_idxs = cast(
        "np.ndarray", X_in.obs["condition_unique_idxs"].to_numpy(dtype=int)
    )
    n_cond = int(np.amax(condition_unique_idxs)) + 1

    if "means" in X_in.var:
        means = X_in.var["means"].to_numpy()
    else:
        means = np.zeros(X_mat.shape[1])

    norm_tensor = calc_norm_sq(X_mat, means)

    slice_weights: np.ndarray | None = None
    if normalize_slices:
        slice_norms = calc_slice_norms(X_mat, means, condition_unique_idxs, n_cond)
        slice_weights = np.where(slice_norms > 1e-10, 1.0 / slice_norms, 1.0)

    return X_mat, condition_unique_idxs, means, norm_tensor, slice_weights

parafac_update(factors, mode, S, projections=None, *, X=None, means=None, cond_slices=None, slice_weights=None)

Form the MTTKRP for the requested mode and update that factor.

Modes 0 and 1 are built from the compressed slices S alone and cost O(n_cond * rank^2). Mode 2 is the only update that has to revisit the raw data, via X^T @ H with H_k = P_k B diag(a_k).

slice_weights, if given, is a per-condition scalar (e.g. an inverse Frobenius norm) applied only to the MTTKRP contributions. This rebalances how much each slice contributes to the factor updates without touching or copying X, and without affecting the reported error (which :func:calc_err computes from the unweighted S).

Parameters:

Name	Type	Description	Default
factors	list[ndarray]	The current [A, B, C] factor matrices; factors[mode] is replaced with the updated matrix.	required
mode	int	Which factor to update (index into factors).	required
S	ndarray	The stacked compressed slices from :func:project_data.	required
projections	list[ndarray] | None	The per-condition projections. Required for mode=2 only.	None
X	(Any, keyword - only)	The raw data matrix. Required for mode=2 only.	None
means	(ndarray | None, keyword - only)	Per-gene means to mean-center X by. Used for mode=2 only.	None
cond_slices	(list[slice | ndarray] | None, keyword - only)	Per-condition row selectors from :func:condition_slices. Required for mode=2 only.	None
slice_weights	(ndarray | None, keyword - only)	Optional per-condition scalar weights, as described above.	None

Returns:

Type	Description
list[ndarray]	factors, with factors[mode] updated by solving the normal equations factors[mode] @ v = mttkrp for the Gram-matrix product v of the other factors (falling back to a least-squares solve if v is singular).

Raises:

Type	Description
ValueError	If mode=2 is requested without projections, X, or cond_slices.

Source code in parafac2/utils.py

def parafac_update(
    factors: list[np.ndarray],
    mode: int,
    S: np.ndarray,
    projections: list[np.ndarray] | None = None,
    *,
    X: Any = None,
    means: np.ndarray | None = None,
    cond_slices: list[slice | np.ndarray] | None = None,
    slice_weights: np.ndarray | None = None,
) -> list[np.ndarray]:
    """
    Form the MTTKRP for the requested mode and update that factor.

    Modes 0 and 1 are built from the compressed slices ``S`` alone and cost
    ``O(n_cond * rank^2)``. Mode 2 is the only update that has to revisit the
    raw data, via ``X^T @ H`` with ``H_k = P_k B diag(a_k)``.

    ``slice_weights``, if given, is a per-condition scalar (e.g. an inverse
    Frobenius norm) applied only to the MTTKRP contributions. This rebalances
    how much each slice contributes to the factor updates without touching or
    copying ``X``, and without affecting the reported error (which
    :func:`calc_err` computes from the unweighted ``S``).

    Parameters
    ----------
    factors : list[np.ndarray]
        The current ``[A, B, C]`` factor matrices; ``factors[mode]`` is
        replaced with the updated matrix.
    mode : int
        Which factor to update (index into ``factors``).
    S : np.ndarray
        The stacked compressed slices from :func:`project_data`.
    projections : list[np.ndarray] | None, default None
        The per-condition projections. Required for ``mode=2`` only.
    X : Any, keyword-only, default None
        The raw data matrix. Required for ``mode=2`` only.
    means : np.ndarray | None, keyword-only, default None
        Per-gene means to mean-center ``X`` by. Used for ``mode=2`` only.
    cond_slices : list[slice | np.ndarray] | None, keyword-only, default None
        Per-condition row selectors from :func:`condition_slices`. Required
        for ``mode=2`` only.
    slice_weights : np.ndarray | None, keyword-only, default None
        Optional per-condition scalar weights, as described above.

    Returns
    -------
    list[np.ndarray]
        ``factors``, with ``factors[mode]`` updated by solving the normal
        equations ``factors[mode] @ v = mttkrp`` for the Gram-matrix product
        ``v`` of the other factors (falling back to a least-squares solve if
        ``v`` is singular).

    Raises
    ------
    ValueError
        If ``mode=2`` is requested without ``projections``, ``X``, or
        ``cond_slices``.
    """
    A, B, _C = factors
    rank = B.shape[0]

    if mode == 0:
        mttkrp = np.einsum("kqr,qr->kr", S, B)
        if slice_weights is not None:
            mttkrp = mttkrp * slice_weights[:, np.newaxis]
    elif mode == 1:
        A_w = A if slice_weights is None else A * slice_weights[:, np.newaxis]
        mttkrp = np.einsum("kqr,kr->qr", S, A_w)
    else:
        if projections is None or X is None or cond_slices is None:
            raise ValueError(
                "mode=2 needs `projections`, `X`, and `cond_slices` to form its MTTKRP."
            )
        # Build H^T directly so the dense operand of the X^T @ H product is
        # C-contiguous and shares X's dtype (see calc_W on why that matters).
        H_T = np.empty((rank, X.shape[0]), dtype=matrix_dtype(X))
        for k, sel in enumerate(cond_slices):
            w_k = 1.0 if slice_weights is None else slice_weights[k]
            H_T[:, sel] = (projections[k] @ (B * A[k]) * w_k).T

        mttkrp_T = np.asarray(rmatmul(H_T, X), dtype=np.float64)
        if means is not None:
            mttkrp_T -= np.outer(H_T.sum(axis=1), means)
        mttkrp = mttkrp_T.T

    return solve_factors(factors, mttkrp, mode)

polar_factor(M)

Compute the nearest orthonormal matrix to M via polar decomposition.

Source code in parafac2/utils.py

def polar_factor(M: np.ndarray) -> np.ndarray:
    """Compute the nearest orthonormal matrix to M via polar decomposition."""
    G = M.T @ M
    _, V = np.linalg.eigh(G)
    MV = M @ V
    col_norms = np.linalg.norm(MV, axis=0, keepdims=True)
    safe_norms = np.where(col_norms > 1e-10, col_norms, 1.0)
    return (MV / safe_norms) @ V.T

project_data(W, factors, cond_slices)

Compute each condition's projection matrix and compressed slice.

For condition k the projection P_k is the orthonormal polar factor of W_k diag(a_k) B^T, and the compressed slice is S_k = P_k^T W_k. Costs O(n_cells * rank^2) and touches no raw data, so it is roughly two orders of magnitude cheaper than :func:calc_W and can be repeated freely while W is cached.

Parameters:

Name	Type	Description	Default
W	ndarray	The cached (X - 1 mu^T) @ C from :func:calc_W.	required
factors	list[ndarray]	The current [A, B, C] factor matrices.	required
cond_slices	list[slice | ndarray]	Per-condition row selectors from :func:condition_slices.	required

Returns:

Type	Description
tuple[list[ndarray], ndarray]	The per-condition projections P_k (each (n_k, rank) with orthonormal columns), and the stacked compressed slices S with shape (n_cond, rank, rank).

Source code in parafac2/utils.py

def project_data(
    W: np.ndarray,
    factors: list[np.ndarray],
    cond_slices: list[slice | np.ndarray],
) -> tuple[list[np.ndarray], np.ndarray]:
    """Compute each condition's projection matrix and compressed slice.

    For condition ``k`` the projection ``P_k`` is the orthonormal polar
    factor of ``W_k diag(a_k) B^T``, and the compressed slice is
    ``S_k = P_k^T W_k``. Costs ``O(n_cells * rank^2)`` and touches no raw
    data, so it is roughly two orders of magnitude cheaper than
    :func:`calc_W` and can be repeated freely while ``W`` is cached.

    Parameters
    ----------
    W : np.ndarray
        The cached ``(X - 1 mu^T) @ C`` from :func:`calc_W`.
    factors : list[np.ndarray]
        The current ``[A, B, C]`` factor matrices.
    cond_slices : list[slice | np.ndarray]
        Per-condition row selectors from :func:`condition_slices`.

    Returns
    -------
    tuple[list[np.ndarray], np.ndarray]
        The per-condition projections ``P_k`` (each ``(n_k, rank)`` with
        orthonormal columns), and the stacked compressed slices ``S`` with
        shape ``(n_cond, rank, rank)``.
    """
    A, B = factors[0], factors[1]
    rank = B.shape[0]

    projections: list[np.ndarray] = []
    S = np.empty((len(cond_slices), rank, rank))

    for i, sel in enumerate(cond_slices):
        W_i = W[sel]
        M = W_i @ (B * A[i]).T  # (n_k, rank)
        proj = polar_factor(M)
        projections.append(proj)
        S[i] = proj.T @ W_i

    return projections, S

randomized_svd_right(X, means, n_components, n_oversamples=0, n_power_iter=2, random_state=None)

Compute the top right-singular vectors of the mean-centered matrix (X - 1 mu^T).

Parameters:

Name	Type	Description	Default
X	Any	The (optionally sparse or GPU-backed) data matrix of shape (total_cells, n_genes).	required
means	ndarray | None	Per-gene means for implicit centering, or None.	required
n_components	int	Number of right-singular vectors to return.	required
n_oversamples	int	Additional random test vectors for randomized SVD projection.	0
n_power_iter	int	Number of power iterations for subspace refinement.	2
random_state	int | Generator | None	Random seed or NumPy generator.	None

Returns:

Type	Description
ndarray	Array of shape (n_genes, n_components) with orthonormal columns.

Source code in parafac2/utils.py

def randomized_svd_right(
    X: Any,
    means: np.ndarray | None,
    n_components: int,
    n_oversamples: int = 0,
    n_power_iter: int = 2,
    random_state: int | np.random.Generator | None = None,
) -> np.ndarray:
    """Compute the top right-singular vectors of the mean-centered matrix ``(X - 1 mu^T)``.

    Parameters
    ----------
    X : Any
        The (optionally sparse or GPU-backed) data matrix of shape
        ``(total_cells, n_genes)``.
    means : np.ndarray | None
        Per-gene means for implicit centering, or ``None``.
    n_components : int
        Number of right-singular vectors to return.
    n_oversamples : int, default 0
        Additional random test vectors for randomized SVD projection.
    n_power_iter : int, default 2
        Number of power iterations for subspace refinement.
    random_state : int | np.random.Generator | None, default None
        Random seed or NumPy generator.

    Returns
    -------
    np.ndarray
        Array of shape ``(n_genes, n_components)`` with orthonormal columns.
    """
    rng = (
        random_state
        if isinstance(random_state, np.random.Generator)
        else np.random.default_rng(random_state)
    )
    n_genes = X.shape[1]
    l_dim = min(n_genes, n_components + n_oversamples)

    Omega = rng.normal(size=(n_genes, l_dim)).astype(np.float64)
    Y = np.asarray(matmul(X, Omega), dtype=np.float64)
    if means is not None:
        Y -= means @ Omega

    for _ in range(n_power_iter):
        Q, _ = np.linalg.qr(Y, mode="reduced")
        Z_T = np.asarray(rmatmul(Q.T, X), dtype=np.float64)
        if means is not None:
            Z_T -= np.outer(np.sum(Q.T, axis=1), means)
        Z = Z_T.T
        Q_z, _ = np.linalg.qr(Z, mode="reduced")
        Y = np.asarray(matmul(X, Q_z), dtype=np.float64)
        if means is not None:
            Y -= means @ Q_z

    Q, _ = np.linalg.qr(Y, mode="reduced")
    B = np.asarray(rmatmul(Q.T, X), dtype=np.float64)
    if means is not None:
        B -= np.outer(np.sum(Q.T, axis=1), means)

    _, _, vh = np.linalg.svd(B, full_matrices=False)
    return vh[:n_components, :].T.astype(np.float64)

solve_factors(factors, mttkrp, mode)

ALS factor update for a single mode using its precomputed MTTKRP.

Source code in parafac2/utils.py

def solve_factors(
    factors: list[np.ndarray],
    mttkrp: np.ndarray,
    mode: int,
) -> list[np.ndarray]:
    """ALS factor update for a single mode using its precomputed MTTKRP."""
    rank = factors[0].shape[1]
    v = np.ones((rank, rank))
    for i, factor in enumerate(factors):
        if i != mode:
            v *= factor.T @ factor

    try:
        factors[mode] = np.linalg.solve(v.T, mttkrp.T).T
    except np.linalg.LinAlgError:
        factors[mode] = np.linalg.lstsq(v.T, mttkrp.T, rcond=None)[0].T

    return factors

standardize_pf2(factors, projections)

Put a fitted PARAFAC2 model into a canonical, comparable form.

Reorders components by condition variance-to-mean ratio, normalizes and sign-flips the factors (via TensorLy's cp_normalize/cp_flip_sign), permutes components to maximize the diagonal of B (via linear-sum assignment), and flips signs so that B's diagonal is non-negative.

Parameters:

Name	Type	Description	Default
factors	list[ndarray]	The fitted [A, B, C] factor matrices.	required
projections	list[ndarray]	The fitted per-condition projection matrices P_k.	required

Returns:

Type	Description
tuple[ndarray, list[ndarray], list[ndarray]]	The (weights, factors, projections) triple after standardization, with components reordered/sign-flipped consistently across factors and projections.

Source code in parafac2/utils.py

def standardize_pf2(
    factors: list[np.ndarray], projections: list[np.ndarray]
) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]:
    """Put a fitted PARAFAC2 model into a canonical, comparable form.

    Reorders components by condition variance-to-mean ratio, normalizes and
    sign-flips the factors (via TensorLy's ``cp_normalize``/``cp_flip_sign``),
    permutes components to maximize the diagonal of ``B`` (via linear-sum
    assignment), and flips signs so that ``B``'s diagonal is non-negative.

    Parameters
    ----------
    factors : list[np.ndarray]
        The fitted ``[A, B, C]`` factor matrices.
    projections : list[np.ndarray]
        The fitted per-condition projection matrices ``P_k``.

    Returns
    -------
    tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]
        The ``(weights, factors, projections)`` triple after standardization,
        with components reordered/sign-flipped consistently across
        ``factors`` and ``projections``.
    """
    # Order components by condition variance-to-mean ratio
    mean_a = np.mean(factors[0], axis=0)
    gini = np.divide(
        np.var(factors[0], axis=0),
        mean_a,
        out=np.zeros_like(mean_a),
        where=np.abs(mean_a) > 1e-12,
    )
    gini_idx = np.argsort(gini)
    factors = [f[:, gini_idx] for f in factors]

    weights, factors = cp_flip_sign(cp_normalize((None, factors)), mode=1)

    # Order eigen-cells to maximize the diagonal of B
    _, col_ind = linear_sum_assignment(np.abs(factors[1].T), maximize=True)
    factors[1] = factors[1][col_ind, :]
    projections = [p[:, col_ind] for p in projections]

    # Flip the sign based on B
    signn = np.sign(np.diag(factors[1]))
    factors[1] *= signn[:, np.newaxis]
    projections = [p * signn for p in projections]

    return weights, factors, projections

parafac2.backend

parafac2.backend

GPU/CPU backend abstraction for matrix operations.

Provides a unified interface for performing dense and sparse (CSR) matrix multiplications on CPU (NumPy), Apple GPUs (MLX), or NVIDIA GPUs (CuPy). This lets the PARAFAC2 fit run its matrix products on whichever accelerator is available without copying data through an intermediate common format.

GPUMatrix

Wrapper for a single matrix (csr_array or np.ndarray) stored on GPU memory (CuPy or MLX) or CPU. Evaluates matrix products on the device and returns results as NumPy ndarrays.

Parameters:

Name	Type	Description	Default
mat	ndarray | csr_array	The matrix to wrap and transfer to the selected device.	required
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:get_backend).	None

Source code in parafac2/backend.py

class GPUMatrix:
    """
    Wrapper for a single matrix (csr_array or np.ndarray) stored on GPU memory
    (CuPy or MLX) or CPU. Evaluates matrix products on the device and returns
    results as NumPy ndarrays.

    Parameters
    ----------
    mat : np.ndarray | csr_array
        The matrix to wrap and transfer to the selected device.
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is auto-detected (see :func:`get_backend`).
    """

    __array_priority__ = 1000

    def __init__(self, mat: np.ndarray | csr_array, backend: str | None = None) -> None:
        """Transfer ``mat`` to the resolved backend's device memory."""
        self.backend = get_backend(backend)
        self.shape = mat.shape
        self.dtype = mat.dtype
        self.is_sparse = issparse(mat)

        if self.backend == "cupy":
            self.device_mat = _to_cupy_matrix(mat)
        elif self.backend == "mlx":
            self.device_mat = _to_mlx_matrix(mat)
        else:
            self.device_mat = mat

    def matmul(self, rhs: np.ndarray) -> np.ndarray:
        """Compute ``self @ rhs`` on the wrapped device.

        Parameters
        ----------
        rhs : np.ndarray
            The right-hand operand.

        Returns
        -------
        np.ndarray
            The product, as a NumPy array.
        """
        if self.backend == "cupy":
            return _matmul_cupy(self.device_mat, rhs)
        elif self.backend == "mlx":
            return _matmul_mlx(
                self.device_mat, rhs, is_sparse=self.is_sparse, shape=self.shape
            )
        return self.device_mat @ rhs

    def rmatmul(self, lhs: np.ndarray) -> np.ndarray:
        """Compute ``lhs @ self`` on the wrapped device.

        Parameters
        ----------
        lhs : np.ndarray
            The left-hand operand.

        Returns
        -------
        np.ndarray
            The product, as a NumPy array.
        """
        if self.backend == "cupy":
            return _rmatmul_cupy(lhs, self.device_mat)
        elif self.backend == "mlx":
            return _rmatmul_mlx(
                lhs, self.device_mat, is_sparse=self.is_sparse, shape=self.shape
            )
        return lhs @ self.device_mat

    def __matmul__(self, rhs: np.ndarray) -> np.ndarray:
        """Operator form of :meth:`matmul`, enabling ``gpu_matrix @ rhs``."""
        return self.matmul(rhs)

    def __rmatmul__(self, lhs: np.ndarray) -> np.ndarray:
        """Operator form of :meth:`rmatmul`, enabling ``lhs @ gpu_matrix``."""
        return self.rmatmul(lhs)

__init__(mat, backend=None)

Transfer mat to the resolved backend's device memory.

Source code in parafac2/backend.py

def __init__(self, mat: np.ndarray | csr_array, backend: str | None = None) -> None:
    """Transfer ``mat`` to the resolved backend's device memory."""
    self.backend = get_backend(backend)
    self.shape = mat.shape
    self.dtype = mat.dtype
    self.is_sparse = issparse(mat)

    if self.backend == "cupy":
        self.device_mat = _to_cupy_matrix(mat)
    elif self.backend == "mlx":
        self.device_mat = _to_mlx_matrix(mat)
    else:
        self.device_mat = mat

__matmul__(rhs)

Operator form of :meth:matmul, enabling gpu_matrix @ rhs.

Source code in parafac2/backend.py

def __matmul__(self, rhs: np.ndarray) -> np.ndarray:
    """Operator form of :meth:`matmul`, enabling ``gpu_matrix @ rhs``."""
    return self.matmul(rhs)

__rmatmul__(lhs)

Operator form of :meth:rmatmul, enabling lhs @ gpu_matrix.

Source code in parafac2/backend.py

def __rmatmul__(self, lhs: np.ndarray) -> np.ndarray:
    """Operator form of :meth:`rmatmul`, enabling ``lhs @ gpu_matrix``."""
    return self.rmatmul(lhs)

matmul(rhs)

Compute self @ rhs on the wrapped device.

Parameters:

Name	Type	Description	Default
rhs	ndarray	The right-hand operand.	required

Returns:

Type	Description
ndarray	The product, as a NumPy array.

Source code in parafac2/backend.py

def matmul(self, rhs: np.ndarray) -> np.ndarray:
    """Compute ``self @ rhs`` on the wrapped device.

    Parameters
    ----------
    rhs : np.ndarray
        The right-hand operand.

    Returns
    -------
    np.ndarray
        The product, as a NumPy array.
    """
    if self.backend == "cupy":
        return _matmul_cupy(self.device_mat, rhs)
    elif self.backend == "mlx":
        return _matmul_mlx(
            self.device_mat, rhs, is_sparse=self.is_sparse, shape=self.shape
        )
    return self.device_mat @ rhs

rmatmul(lhs)

Compute lhs @ self on the wrapped device.

Parameters:

Name	Type	Description	Default
lhs	ndarray	The left-hand operand.	required

Returns:

Type	Description
ndarray	The product, as a NumPy array.

Source code in parafac2/backend.py

def rmatmul(self, lhs: np.ndarray) -> np.ndarray:
    """Compute ``lhs @ self`` on the wrapped device.

    Parameters
    ----------
    lhs : np.ndarray
        The left-hand operand.

    Returns
    -------
    np.ndarray
        The product, as a NumPy array.
    """
    if self.backend == "cupy":
        return _rmatmul_cupy(lhs, self.device_mat)
    elif self.backend == "mlx":
        return _rmatmul_mlx(
            lhs, self.device_mat, is_sparse=self.is_sparse, shape=self.shape
        )
    return lhs @ self.device_mat

get_backend(backend=None)

Return the requested backend, or auto-detect the first available one.

Parameters:

Name	Type	Description	Default
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is chosen by attempting to import cupy then mlx.core, falling back to 'cpu' if neither is installed.	None

Returns:

Type	Description
str	The resolved backend name: 'mlx', 'cupy', or 'cpu'.

Raises:

Type	Description
ValueError	If backend is given but is not one of the supported names.

Source code in parafac2/backend.py

def get_backend(backend: str | None = None) -> str:
    """Return the requested backend, or auto-detect the first available one.

    Parameters
    ----------
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is chosen by attempting to import ``cupy``
        then ``mlx.core``, falling back to ``'cpu'`` if neither is
        installed.

    Returns
    -------
    str
        The resolved backend name: ``'mlx'``, ``'cupy'``, or ``'cpu'``.

    Raises
    ------
    ValueError
        If ``backend`` is given but is not one of the supported names.
    """
    if backend is not None:
        backend_lower = backend.lower()
        if backend_lower in ("mlx", "cupy", "cpu"):
            return backend_lower
        raise ValueError(
            f"Unknown backend '{backend}'. Supported backends: 'mlx', 'cupy', 'cpu'."
        )

    try:
        import cupy  # noqa: F401  # ty: ignore[unresolved-import]

        return "cupy"
    except ImportError:
        pass

    try:
        import mlx.core  # noqa: F401  # ty: ignore[unresolved-import]

        return "mlx"
    except ImportError:
        pass

    return "cpu"

matmul(mat, rhs)

Compute mat @ rhs, dispatching to the fastest available kernel.

Parameters:

Name	Type	Description	Default
mat	Any	A :class:GPUMatrix, SciPy sparse matrix, or dense NumPy array.	required
rhs	ndarray	The dense right-hand operand.	required

Returns:

Type	Description
ndarray	The product mat @ rhs.

Notes

rhs should already share mat's dtype. Handing a float64 rhs to a float32 sparse mat makes SciPy upcast the whole sparse matrix, which for single-cell-sized data is both slow and memory-hostile; see :func:~parafac2.utils.calc_W.

Source code in parafac2/backend.py

def matmul(mat: Any, rhs: np.ndarray) -> np.ndarray:
    """Compute ``mat @ rhs``, dispatching to the fastest available kernel.

    Parameters
    ----------
    mat : Any
        A :class:`GPUMatrix`, SciPy sparse matrix, or dense NumPy array.
    rhs : np.ndarray
        The dense right-hand operand.

    Returns
    -------
    np.ndarray
        The product ``mat @ rhs``.

    Notes
    -----
    ``rhs`` should already share ``mat``'s dtype. Handing a float64 ``rhs``
    to a float32 sparse ``mat`` makes SciPy upcast the *whole* sparse matrix,
    which for single-cell-sized data is both slow and memory-hostile; see
    :func:`~parafac2.utils.calc_W`.
    """
    if isinstance(mat, GPUMatrix):
        return mat.matmul(rhs)
    if _mkl_compatible(mat, rhs):
        return _get_mkl_dot()(mat, rhs)
    return mat @ rhs

matrix_dtype(mat)

Return the dtype of a :class:GPUMatrix, sparse matrix, or ndarray.

Source code in parafac2/backend.py

def matrix_dtype(mat: Any) -> np.dtype:
    """Return the dtype of a :class:`GPUMatrix`, sparse matrix, or ndarray."""
    return np.dtype(getattr(mat, "dtype", np.float64))

rmatmul(lhs, mat)

Compute lhs @ mat, dispatching to the fastest available kernel.

Parameters:

Name	Type	Description	Default
lhs	ndarray	The dense left-hand operand.	required
mat	Any	A :class:GPUMatrix, SciPy sparse matrix, or dense NumPy array.	required

Returns:

Type	Description
ndarray	The product lhs @ mat.

Source code in parafac2/backend.py

def rmatmul(lhs: np.ndarray, mat: Any) -> np.ndarray:
    """Compute ``lhs @ mat``, dispatching to the fastest available kernel.

    Parameters
    ----------
    lhs : np.ndarray
        The dense left-hand operand.
    mat : Any
        A :class:`GPUMatrix`, SciPy sparse matrix, or dense NumPy array.

    Returns
    -------
    np.ndarray
        The product ``lhs @ mat``.
    """
    if isinstance(mat, GPUMatrix):
        return mat.rmatmul(lhs)
    if _mkl_compatible(mat, lhs):
        return _get_mkl_dot()(lhs, mat)
    return lhs @ mat

to_gpu(mat, backend=None)

Transfer matrix to GPU memory if CuPy or MLX is requested/available, returning a GPUMatrix wrapper. Otherwise returns the CPU matrix as-is.

Parameters:

Name	Type	Description	Default
mat	ndarray | csr_array	The matrix to (optionally) transfer.	required
backend	str	One of 'mlx', 'cupy', or 'cpu'. If None, the first available accelerator is auto-detected (see :func:get_backend).	None

Returns:

Type	Description
GPUMatrix | ndarray | csr_array	A :class:GPUMatrix wrapping mat if a GPU backend was resolved, otherwise mat unchanged.

Source code in parafac2/backend.py

def to_gpu(
    mat: np.ndarray | csr_array, backend: str | None = None
) -> GPUMatrix | np.ndarray | csr_array:
    """
    Transfer matrix to GPU memory if CuPy or MLX is requested/available,
    returning a GPUMatrix wrapper. Otherwise returns the CPU matrix as-is.

    Parameters
    ----------
    mat : np.ndarray | csr_array
        The matrix to (optionally) transfer.
    backend : str, optional
        One of ``'mlx'``, ``'cupy'``, or ``'cpu'``. If ``None``, the first
        available accelerator is auto-detected (see :func:`get_backend`).

    Returns
    -------
    GPUMatrix | np.ndarray | csr_array
        A :class:`GPUMatrix` wrapping ``mat`` if a GPU backend was resolved,
        otherwise ``mat`` unchanged.
    """
    chosen = get_backend(backend)
    if chosen == "cpu":
        return mat
    return GPUMatrix(mat, backend=chosen)