"""Helpers for exact Schmidt spectra and entanglement calculations."""
import numpy as np
from scipy.linalg import eigh as array_eigh
from scipy.sparse import issparse
__all__ = ["schmidt_probabilities", "dense_schmidt_probabilities"]
# Column/row-block width used when accumulating the Gram matrix from a sparse
# bipartition matrix. Blocking keeps the transient dense chunk small instead of
# materializing a full (near-dense) sparse product.
_GRAM_BLOCK = 1024
def _gram_from_sparse_blocked(psi_matrix, on_rows: bool) -> np.ndarray:
"""Accumulate the dense smaller-side Gram of a sparse bipartition matrix.
Parameters
----------
psi_matrix : scipy.sparse.spmatrix
Sparse subsystem/environment bipartition matrix ``A`` of shape
``(subsys_dim, env_dim)``.
on_rows : bool
If ``True`` build ``A A^H`` (size ``subsys_dim``); otherwise build
``A^H A`` (size ``env_dim``).
Returns
-------
numpy.ndarray
Dense Hermitian Gram matrix of side ``min(subsys_dim, env_dim)``.
Notes
-----
The Gram is formed as ``sum_b C_b C_b^H`` over column blocks (or
``sum_b C_b^H C_b`` over row blocks), so only one dense ``dim x block``
chunk and one dense ``dim x dim`` accumulator are alive at a time. This
avoids building the full sparse product (which would be ~20 bytes/entry of
near-dense storage on top of the dense copy).
"""
n_rows, n_cols = psi_matrix.shape
if on_rows:
dim, axis_len = n_rows, n_cols
mat = psi_matrix.tocsc()
else:
dim, axis_len = n_cols, n_rows
mat = psi_matrix.tocsr()
gram = np.zeros((dim, dim), dtype=np.complex128)
for start in range(0, axis_len, _GRAM_BLOCK):
stop = start + _GRAM_BLOCK
if on_rows:
chunk = mat[:, start:stop].toarray()
gram += chunk @ chunk.conj().T
else:
chunk = mat[start:stop, :].toarray()
gram += chunk.conj().T @ chunk
return gram
[docs]
def schmidt_probabilities(psi_matrix) -> np.ndarray:
"""Return the exact Schmidt probabilities of a bipartition matrix.
Parameters
----------
psi_matrix : numpy.ndarray or scipy.sparse.spmatrix
Subsystem/environment bipartition matrix of the state.
Returns
-------
numpy.ndarray
Non-negative Schmidt probabilities (squared singular values), i.e. the
eigenvalues of the reduced density matrix.
Notes
-----
The probabilities are the eigenvalues of the smaller of the two Gram
matrices ``A A^H`` / ``A^H A``, computed exactly with a Hermitian
eigensolver. The smaller side is chosen so the dense Gram has side
``min(subsys_dim, env_dim)``. For sparse ``A`` the Gram is accumulated in
blocks (see :func:`_gram_from_sparse_blocked`). The eigendecomposition runs
in place (``overwrite_a=True``) to avoid an extra working copy.
"""
n_rows, n_cols = psi_matrix.shape
on_rows = n_rows <= n_cols
if issparse(psi_matrix):
gram = _gram_from_sparse_blocked(psi_matrix, on_rows)
else:
mat = np.asarray(psi_matrix)
gram = mat @ mat.conj().T if on_rows else mat.conj().T @ mat
llambdas = array_eigh(gram, eigvals_only=True, overwrite_a=True)
return np.maximum(llambdas, 0.0)
# Backward-compatible alias: the routine now also accepts sparse input.
dense_schmidt_probabilities = schmidt_probabilities