sparse_dia

Sparse diagonal (SparseDIA) backend for Qarray.

Stores only the diagonal values of a matrix, making quantum operators with small numbers of non-zero diagonals (annihilation, creation, number, Kerr…) far cheaper than Dense or BCOO:

• Memory: O(d * n) where d = number of stored diagonals, n = matrix size
• No index arrays (unlike BCOO which stores (row, col) per non-zero)
• _offsets is static Python metadata (pytree_node=False), so JAX unrolls all loops over diagonals at compile time — only static slices, no dynamic indexing or scatter/gather.

Padding convention (Convention A):

For diagonal at offset k (k ≥ 0): diags[..., i, j] = A[j-k, j] for j ∈ [k, n-1]; zeros at [0:k] For diagonal at offset k (k < 0): diags[..., i, j] = A[j-k, j] for j ∈ [0, n+k-1]; zeros at [n+k:]

Unified access formula (holds for any k, out-of-range slots are zero): A[i, i+k] = diags[..., diag_idx, i+k]

This makes every matrix operation a set of aligned slice multiplications.

Some improvements (_dia_slice helper, integer matrix power, diagonal-range pruning, offset detection) were identified by studying the dynamiqs library (https://github.com/dynamiqs/dynamiqs).

SparseDiaData

Lightweight pytree-compatible container for sparse-diagonal raw data.

Returned by SparseDiaImpl.get_data() and consumed by SparseDiaImpl.from_data(). Registered as a JAX pytree via Flax's @struct.dataclass; offsets is not a pytree leaf (it is static compile-time metadata).

Attributes:

Name	Type	Description
offsets	tuple	Static tuple of diagonal offsets (pytree_node=False).
diags	Array	JAX array of shape (*batch, n_diags, n) containing the padded diagonal values.

Source code in jaxquantum/core/sparse_dia.py

@struct.dataclass
class SparseDiaData:
    """Lightweight pytree-compatible container for sparse-diagonal raw data.

    Returned by ``SparseDiaImpl.get_data()`` and consumed by
    ``SparseDiaImpl.from_data()``.  Registered as a JAX pytree via Flax's
    ``@struct.dataclass``; ``offsets`` is *not* a pytree leaf (it is static
    compile-time metadata).

    Attributes:
        offsets: Static tuple of diagonal offsets (pytree_node=False).
        diags:   JAX array of shape (*batch, n_diags, n) containing the
                 padded diagonal values.
    """

    offsets: tuple = struct.field(pytree_node=False)
    diags: Array

    # Class-level marker (not a dataclass field — no type annotation).
    # Used by DenseImpl.can_handle_data to exclude SparseDiaData without
    # a direct import (which would be circular).
    _is_sparse_dia = True

    @property
    def shape(self) -> tuple:
        """Shape of the represented square matrix (*batch, n, n)."""
        n = self.diags.shape[-1]
        return (*self.diags.shape[:-2], n, n)

    @property
    def dtype(self):
        """Dtype of the stored diagonal values."""
        return self.diags.dtype

    def __mul__(self, scalar):
        return SparseDiaData(offsets=self.offsets, diags=self.diags * scalar)

    def __rmul__(self, scalar):
        return SparseDiaData(offsets=self.offsets, diags=scalar * self.diags)

    def __getitem__(self, index):
        """Index into the batch dimension(s), preserving offsets."""
        return SparseDiaData(offsets=self.offsets, diags=self.diags[index])

    def __len__(self):
        """Number of elements along the leading batch dimension."""
        return self.shape[0]

    def reshape(self, *new_shape):
        """Reshape batch dimensions while preserving diagonal structure.

        ``new_shape`` must end with ``(N, N)`` (the matrix dims are unchanged).
        Only the leading batch dims are reshaped.
        """
        new_batch = new_shape[:-2]
        n = self.diags.shape[-1]
        new_diags = self.diags.reshape(*new_batch, len(self.offsets), n)
        return SparseDiaData(offsets=self.offsets, diags=new_diags)

    def __matmul__(self, other):
        """SparseDIA @ dense → dense (used by mesolve ODE RHS)."""
        # _sparsedia_matmul_dense is defined later in this module; Python
        # resolves the name at call time so forward reference is fine.
        return _sparsedia_matmul_dense(self.offsets, self.diags, other)

    def __rmatmul__(self, other):
        """dense @ SparseDIA → dense (used by mesolve ODE RHS)."""
        return _sparsedia_rmatmul_dense(other, self.offsets, self.diags)

dtype property

Dtype of the stored diagonal values.

shape property

Shape of the represented square matrix (*batch, n, n).

__getitem__(index)

Index into the batch dimension(s), preserving offsets.

Source code in jaxquantum/core/sparse_dia.py

def __getitem__(self, index):
    """Index into the batch dimension(s), preserving offsets."""
    return SparseDiaData(offsets=self.offsets, diags=self.diags[index])

__len__()

Number of elements along the leading batch dimension.

Source code in jaxquantum/core/sparse_dia.py

def __len__(self):
    """Number of elements along the leading batch dimension."""
    return self.shape[0]

__matmul__(other)

SparseDIA @ dense → dense (used by mesolve ODE RHS).

Source code in jaxquantum/core/sparse_dia.py

def __matmul__(self, other):
    """SparseDIA @ dense → dense (used by mesolve ODE RHS)."""
    # _sparsedia_matmul_dense is defined later in this module; Python
    # resolves the name at call time so forward reference is fine.
    return _sparsedia_matmul_dense(self.offsets, self.diags, other)

__rmatmul__(other)

dense @ SparseDIA → dense (used by mesolve ODE RHS).

Source code in jaxquantum/core/sparse_dia.py

def __rmatmul__(self, other):
    """dense @ SparseDIA → dense (used by mesolve ODE RHS)."""
    return _sparsedia_rmatmul_dense(other, self.offsets, self.diags)

reshape(*new_shape)

Reshape batch dimensions while preserving diagonal structure.

new_shape must end with (N, N) (the matrix dims are unchanged). Only the leading batch dims are reshaped.

Source code in jaxquantum/core/sparse_dia.py

def reshape(self, *new_shape):
    """Reshape batch dimensions while preserving diagonal structure.

    ``new_shape`` must end with ``(N, N)`` (the matrix dims are unchanged).
    Only the leading batch dims are reshaped.
    """
    new_batch = new_shape[:-2]
    n = self.diags.shape[-1]
    new_diags = self.diags.reshape(*new_batch, len(self.offsets), n)
    return SparseDiaData(offsets=self.offsets, diags=new_diags)

SparseDiaImpl

Bases: QarrayImpl

Sparse-diagonal backend storing only diagonal values.

Data layout::

_offsets  : tuple[int, ...]          — static (pytree_node=False)
_diags    : Array[*batch, n_diags, n] — JAX-traced values

For offset k, the convention is: * k ≥ 0 : valid data at _diags[..., i, k:], zeros at [0:k] * k < 0 : valid data at _diags[..., i, :n+k], zeros at [n+k:]

In both cases: A[row, row+k] = _diags[..., i, row+k]

Source code in jaxquantum/core/sparse_dia.py

@struct.dataclass
class SparseDiaImpl(QarrayImpl):
    """Sparse-diagonal backend storing only diagonal values.

    Data layout::

        _offsets  : tuple[int, ...]          — static (pytree_node=False)
        _diags    : Array[*batch, n_diags, n] — JAX-traced values

    For offset k, the convention is:
        * k ≥ 0 : valid data at ``_diags[..., i, k:]``, zeros at ``[0:k]``
        * k < 0 : valid data at ``_diags[..., i, :n+k]``, zeros at ``[n+k:]``

    In both cases: ``A[row, row+k] = _diags[..., i, row+k]``
    """

    _offsets: tuple = struct.field(pytree_node=False)
    _diags: Array

    PROMOTION_ORDER = 0  # noqa: RUF012 — not a struct field

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    @classmethod
    def from_data(cls, data) -> "SparseDiaImpl":
        """Wrap *data* in a new ``SparseDiaImpl``.

        Accepts either a :class:`SparseDiaData` container (direct wrap) or
        a dense array-like (auto-detect non-zero diagonals via numpy, safe
        to call before JIT).

        Args:
            data: A :class:`SparseDiaData` or dense array of shape
                (*batch, n, n).

        Returns:
            A new ``SparseDiaImpl`` instance.
        """
        if isinstance(data, SparseDiaData):
            return cls(_offsets=data.offsets, _diags=data.diags)
        offsets, diags_np = _dense_to_sparsedia(np.asarray(data))
        return cls(_offsets=offsets, _diags=jnp.array(diags_np))

    @classmethod
    def from_diags(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
        """Directly construct from sorted offsets and padded diagonal array.

        This is the preferred factory when diagonal structure is known in
        advance (e.g., when building ``destroy`` or ``create`` operators).

        Args:
            offsets: Tuple of integer diagonal offsets (need not be sorted;
                will be sorted internally).
            diags:   JAX array of shape (*batch, n_diags, n) with padded
                     diagonal values matching *offsets*.

        Returns:
            A new ``SparseDiaImpl`` instance.
        """
        return cls(_offsets=tuple(sorted(offsets)), _diags=diags)

    # ------------------------------------------------------------------
    # QarrayImpl abstract methods
    # ------------------------------------------------------------------

    def get_data(self) -> SparseDiaData:
        """Return a :class:`SparseDiaData` container with the raw diagonal data."""
        return SparseDiaData(offsets=self._offsets, diags=self._diags)

    def shape(self) -> tuple:
        """Shape of the represented square matrix (including batch dims)."""
        n = self._diags.shape[-1]
        return (*self._diags.shape[:-2], n, n)

    def dtype(self):
        """Dtype of the stored diagonal values."""
        return self._diags.dtype

    def __deepcopy__(self, memo=None):
        return SparseDiaImpl(
            _offsets=deepcopy(self._offsets),
            _diags=self._diags,
        )

    # ------------------------------------------------------------------
    # Arithmetic
    # ------------------------------------------------------------------

    def mul(self, scalar) -> "SparseDiaImpl":
        """Scalar multiplication — scales all diagonal values."""
        return SparseDiaImpl(_offsets=self._offsets, _diags=scalar * self._diags)

    def neg(self) -> "SparseDiaImpl":
        """Negation."""
        return SparseDiaImpl(_offsets=self._offsets, _diags=-self._diags)

    def add(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise addition.

        SparseDIA + SparseDIA stays SparseDIA (union of offsets, static).
        Otherwise coerces to the higher-order type.
        """
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_add(self, other)
        a, b = self._coerce(other)
        if a is not self:
            return a.add(b)
        return a.add(b)

    def sub(self, other: QarrayImpl) -> QarrayImpl:
        """Element-wise subtraction."""
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_add(self, other, subtract=True)
        a, b = self._coerce(other)
        if a is not self:
            return a.sub(b)
        return a.sub(b)

    def matmul(self, other: QarrayImpl) -> QarrayImpl:
        """Matrix multiplication.

        * SparseDIA @ SparseDIA → SparseDIA  (O(d₁·d₂·n))
        * SparseDIA @ Dense    → Dense       (O(d·n²), no densification of self)
        * Others               → coerce then delegate
        """
        if isinstance(other, DenseImpl):
            return DenseImpl(_sparsedia_matmul_dense(
                self._offsets, self._diags, other._data
            ))
        if isinstance(other, SparseDiaImpl):
            offsets, diags = _sparsedia_matmul_sparsedia(
                self._offsets, self._diags,
                other._offsets, other._diags,
            )
            return SparseDiaImpl(_offsets=offsets, _diags=diags)
        a, b = self._coerce(other)
        if a is not self:
            return a.matmul(b)
        return a.matmul(b)

    def dag(self) -> "SparseDiaImpl":
        """Conjugate transpose without densification.

        Negates every offset and rearranges the stored values so that the
        padding convention remains consistent.
        """
        new_offsets = tuple(-k for k in self._offsets)
        new_diags = jnp.zeros_like(self._diags)
        for i, k in enumerate(self._offsets):
            s = _dia_slice(k)    # valid data slice for offset k
            sm = _dia_slice(-k)  # valid data slice for offset -k (the new position)
            new_diags = new_diags.at[..., i, sm].set(jnp.conj(self._diags[..., i, s]))
        return SparseDiaImpl(_offsets=new_offsets, _diags=new_diags)

    def kron(self, other: QarrayImpl) -> QarrayImpl:
        """Kronecker product.

        SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair
        (kA, kB) is ``kA * m + kB`` where m = dim(B).  Fully vectorised —
        no loops at JAX level.
        """
        if isinstance(other, SparseDiaImpl):
            return _sparsedia_kron(self, other)
        a, b = self._coerce(other)
        if a is not self:
            return a.kron(b)
        return a.kron(b)

    def tidy_up(self, atol) -> "SparseDiaImpl":
        """Zero diagonal values whose magnitude is below *atol*."""
        diags = self._diags
        real_part = jnp.where(jnp.abs(jnp.real(diags)) < atol, 0.0, jnp.real(diags))
        if jnp.issubdtype(diags.dtype, jnp.complexfloating):
            imag_part = jnp.where(jnp.abs(jnp.imag(diags)) < atol, 0.0, jnp.imag(diags))
            new_diags = (real_part + 1j * imag_part).astype(diags.dtype)
        else:
            new_diags = real_part.astype(diags.dtype)
        return SparseDiaImpl(_offsets=self._offsets, _diags=new_diags)

    # ------------------------------------------------------------------
    # Conversions
    # ------------------------------------------------------------------

    def to_dense(self) -> "DenseImpl":
        """Convert to a ``DenseImpl`` by summing diagonal contributions."""
        n = self._diags.shape[-1]
        batch_shape = self._diags.shape[:-2]
        result = jnp.zeros((*batch_shape, n, n), dtype=self._diags.dtype)
        for i, k in enumerate(self._offsets):
            s = _dia_slice(k)
            length = n - abs(k)
            if length <= 0:
                continue
            vals = self._diags[..., i, s]
            row_idx = jnp.arange(length) + max(-k, 0)
            col_idx = row_idx + k
            result = result.at[..., row_idx, col_idx].set(vals)
        return DenseImpl(result)

    def to_sparse_bcoo(self) -> "SparseBCOOImpl":
        """Convert to a ``SparseBCOOImpl`` (BCOO) via dense."""
        return self.to_dense().to_sparse_bcoo()

    def to_sparse_dia(self) -> "SparseDiaImpl":
        """Return self (already SparseDIA)."""
        return self

    # ------------------------------------------------------------------
    # Class-method interface
    # ------------------------------------------------------------------

    @classmethod
    def _eye_data(cls, n: int, dtype=None):
        """Return an n×n identity as a dense JAX array.

        ``from_data`` will automatically convert it to SparseDIA format
        when the implementation type is ``SPARSE_DIA``.
        """
        return jnp.eye(n, dtype=dtype)

    @classmethod
    def can_handle_data(cls, arr) -> bool:
        """Return True only for :class:`SparseDiaData` objects."""
        return isinstance(arr, SparseDiaData)

    @classmethod
    def dag_data(cls, arr: SparseDiaData) -> SparseDiaData:
        """Conjugate transpose of raw :class:`SparseDiaData` without densification."""
        impl = SparseDiaImpl(_offsets=arr.offsets, _diags=arr.diags)
        result = impl.dag()
        return result.get_data()

    # ------------------------------------------------------------------
    # Extra sparse-native methods (no densification)
    # ------------------------------------------------------------------

    def trace(self):
        """Compute trace directly from the main diagonal (offset 0).

        Returns:
            Scalar trace (sum of main diagonal values).
        """
        if 0 in self._offsets:
            i = self._offsets.index(0)
            return jnp.sum(self._diags[..., i, :], axis=-1)
        return jnp.zeros(self._diags.shape[:-2], dtype=self._diags.dtype)

    def frobenius_norm(self):
        """Frobenius norm computed directly from stored diagonal values."""
        return jnp.sqrt(jnp.sum(jnp.abs(self._diags) ** 2))

    def real(self) -> "SparseDiaImpl":
        """Element-wise real part of stored values."""
        return SparseDiaImpl(
            _offsets=self._offsets,
            _diags=jnp.real(self._diags).astype(self._diags.dtype),
        )

    def imag(self) -> "SparseDiaImpl":
        """Element-wise imaginary part of stored values."""
        return SparseDiaImpl(
            _offsets=self._offsets,
            _diags=jnp.imag(self._diags).astype(self._diags.dtype),
        )

    def conj(self) -> "SparseDiaImpl":
        """Element-wise complex conjugate of stored values."""
        return SparseDiaImpl(_offsets=self._offsets, _diags=jnp.conj(self._diags))

    def powm(self, n: int) -> "SparseDiaImpl":
        """Integer matrix power staying SparseDIA via binary exponentiation.

        Uses O(log n) SparseDIA @ SparseDIA multiplications rather than
        densifying.  A^0 returns the identity operator.

        Args:
            n: Non-negative integer exponent.

        Returns:
            A ``SparseDiaImpl`` equal to this matrix raised to the *n*-th power.

        Raises:
            ValueError: If *n* is negative.
        """
        if n < 0:
            raise ValueError("powm requires n >= 0")
        if n == 0:
            size = self._diags.shape[-1]
            eye_diags = jnp.ones((*self._diags.shape[:-2], 1, size), dtype=self._diags.dtype)
            return SparseDiaImpl(_offsets=(0,), _diags=eye_diags)
        if n == 1:
            return self
        half = self.powm(n // 2)
        squared = half.matmul(half)  # SparseDIA @ SparseDIA → SparseDIA
        return squared if n % 2 == 0 else self.matmul(squared)

add(other)

Element-wise addition.

SparseDIA + SparseDIA stays SparseDIA (union of offsets, static). Otherwise coerces to the higher-order type.

Source code in jaxquantum/core/sparse_dia.py

def add(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise addition.

    SparseDIA + SparseDIA stays SparseDIA (union of offsets, static).
    Otherwise coerces to the higher-order type.
    """
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_add(self, other)
    a, b = self._coerce(other)
    if a is not self:
        return a.add(b)
    return a.add(b)

can_handle_data(arr) classmethod

Return True only for :class:SparseDiaData objects.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def can_handle_data(cls, arr) -> bool:
    """Return True only for :class:`SparseDiaData` objects."""
    return isinstance(arr, SparseDiaData)

conj()

Element-wise complex conjugate of stored values.

Source code in jaxquantum/core/sparse_dia.py

def conj(self) -> "SparseDiaImpl":
    """Element-wise complex conjugate of stored values."""
    return SparseDiaImpl(_offsets=self._offsets, _diags=jnp.conj(self._diags))

dag()

Conjugate transpose without densification.

Negates every offset and rearranges the stored values so that the padding convention remains consistent.

Source code in jaxquantum/core/sparse_dia.py

def dag(self) -> "SparseDiaImpl":
    """Conjugate transpose without densification.

    Negates every offset and rearranges the stored values so that the
    padding convention remains consistent.
    """
    new_offsets = tuple(-k for k in self._offsets)
    new_diags = jnp.zeros_like(self._diags)
    for i, k in enumerate(self._offsets):
        s = _dia_slice(k)    # valid data slice for offset k
        sm = _dia_slice(-k)  # valid data slice for offset -k (the new position)
        new_diags = new_diags.at[..., i, sm].set(jnp.conj(self._diags[..., i, s]))
    return SparseDiaImpl(_offsets=new_offsets, _diags=new_diags)

dag_data(arr) classmethod

Conjugate transpose of raw :class:SparseDiaData without densification.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def dag_data(cls, arr: SparseDiaData) -> SparseDiaData:
    """Conjugate transpose of raw :class:`SparseDiaData` without densification."""
    impl = SparseDiaImpl(_offsets=arr.offsets, _diags=arr.diags)
    result = impl.dag()
    return result.get_data()

dtype()

Dtype of the stored diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def dtype(self):
    """Dtype of the stored diagonal values."""
    return self._diags.dtype

frobenius_norm()

Frobenius norm computed directly from stored diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def frobenius_norm(self):
    """Frobenius norm computed directly from stored diagonal values."""
    return jnp.sqrt(jnp.sum(jnp.abs(self._diags) ** 2))

from_data(data) classmethod

Wrap data in a new SparseDiaImpl.

Accepts either a :class:SparseDiaData container (direct wrap) or a dense array-like (auto-detect non-zero diagonals via numpy, safe to call before JIT).

Parameters:

Name	Type	Description	Default
data		A :class:SparseDiaData or dense array of shape (*batch, n, n).	required

Returns:

Type	Description
'SparseDiaImpl'	A new SparseDiaImpl instance.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def from_data(cls, data) -> "SparseDiaImpl":
    """Wrap *data* in a new ``SparseDiaImpl``.

    Accepts either a :class:`SparseDiaData` container (direct wrap) or
    a dense array-like (auto-detect non-zero diagonals via numpy, safe
    to call before JIT).

    Args:
        data: A :class:`SparseDiaData` or dense array of shape
            (*batch, n, n).

    Returns:
        A new ``SparseDiaImpl`` instance.
    """
    if isinstance(data, SparseDiaData):
        return cls(_offsets=data.offsets, _diags=data.diags)
    offsets, diags_np = _dense_to_sparsedia(np.asarray(data))
    return cls(_offsets=offsets, _diags=jnp.array(diags_np))

from_diags(offsets, diags) classmethod

Directly construct from sorted offsets and padded diagonal array.

This is the preferred factory when diagonal structure is known in advance (e.g., when building destroy or create operators).

Parameters:

Name	Type	Description	Default
offsets	tuple	Tuple of integer diagonal offsets (need not be sorted; will be sorted internally).	required
diags	Array	JAX array of shape (batch, n_diags, n) with padded diagonal values matching *offsets.	required

Returns:

Type	Description
'SparseDiaImpl'	A new SparseDiaImpl instance.

Source code in jaxquantum/core/sparse_dia.py

@classmethod
def from_diags(cls, offsets: tuple, diags: Array) -> "SparseDiaImpl":
    """Directly construct from sorted offsets and padded diagonal array.

    This is the preferred factory when diagonal structure is known in
    advance (e.g., when building ``destroy`` or ``create`` operators).

    Args:
        offsets: Tuple of integer diagonal offsets (need not be sorted;
            will be sorted internally).
        diags:   JAX array of shape (*batch, n_diags, n) with padded
                 diagonal values matching *offsets*.

    Returns:
        A new ``SparseDiaImpl`` instance.
    """
    return cls(_offsets=tuple(sorted(offsets)), _diags=diags)

get_data()

Return a :class:SparseDiaData container with the raw diagonal data.

Source code in jaxquantum/core/sparse_dia.py

def get_data(self) -> SparseDiaData:
    """Return a :class:`SparseDiaData` container with the raw diagonal data."""
    return SparseDiaData(offsets=self._offsets, diags=self._diags)

imag()

Element-wise imaginary part of stored values.

Source code in jaxquantum/core/sparse_dia.py

def imag(self) -> "SparseDiaImpl":
    """Element-wise imaginary part of stored values."""
    return SparseDiaImpl(
        _offsets=self._offsets,
        _diags=jnp.imag(self._diags).astype(self._diags.dtype),
    )

kron(other)

Kronecker product.

SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair (kA, kB) is kA * m + kB where m = dim(B). Fully vectorised — no loops at JAX level.

Source code in jaxquantum/core/sparse_dia.py

def kron(self, other: QarrayImpl) -> QarrayImpl:
    """Kronecker product.

    SparseDIA ⊗ SparseDIA stays SparseDIA: output offset for pair
    (kA, kB) is ``kA * m + kB`` where m = dim(B).  Fully vectorised —
    no loops at JAX level.
    """
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_kron(self, other)
    a, b = self._coerce(other)
    if a is not self:
        return a.kron(b)
    return a.kron(b)

matmul(other)

Matrix multiplication.

• SparseDIA @ SparseDIA → SparseDIA (O(d₁·d₂·n))
• SparseDIA @ Dense → Dense (O(d·n²), no densification of self)
• Others → coerce then delegate

Source code in jaxquantum/core/sparse_dia.py

def matmul(self, other: QarrayImpl) -> QarrayImpl:
    """Matrix multiplication.

    * SparseDIA @ SparseDIA → SparseDIA  (O(d₁·d₂·n))
    * SparseDIA @ Dense    → Dense       (O(d·n²), no densification of self)
    * Others               → coerce then delegate
    """
    if isinstance(other, DenseImpl):
        return DenseImpl(_sparsedia_matmul_dense(
            self._offsets, self._diags, other._data
        ))
    if isinstance(other, SparseDiaImpl):
        offsets, diags = _sparsedia_matmul_sparsedia(
            self._offsets, self._diags,
            other._offsets, other._diags,
        )
        return SparseDiaImpl(_offsets=offsets, _diags=diags)
    a, b = self._coerce(other)
    if a is not self:
        return a.matmul(b)
    return a.matmul(b)

mul(scalar)

Scalar multiplication — scales all diagonal values.

Source code in jaxquantum/core/sparse_dia.py

def mul(self, scalar) -> "SparseDiaImpl":
    """Scalar multiplication — scales all diagonal values."""
    return SparseDiaImpl(_offsets=self._offsets, _diags=scalar * self._diags)

neg()

Negation.

Source code in jaxquantum/core/sparse_dia.py

def neg(self) -> "SparseDiaImpl":
    """Negation."""
    return SparseDiaImpl(_offsets=self._offsets, _diags=-self._diags)

powm(n)

Integer matrix power staying SparseDIA via binary exponentiation.

Uses O(log n) SparseDIA @ SparseDIA multiplications rather than densifying. A^0 returns the identity operator.

Parameters:

Name	Type	Description	Default
n	int	Non-negative integer exponent.	required

Returns:

Type	Description
'SparseDiaImpl'	A SparseDiaImpl equal to this matrix raised to the n-th power.

Raises:

Type	Description
ValueError	If n is negative.

Source code in jaxquantum/core/sparse_dia.py

def powm(self, n: int) -> "SparseDiaImpl":
    """Integer matrix power staying SparseDIA via binary exponentiation.

    Uses O(log n) SparseDIA @ SparseDIA multiplications rather than
    densifying.  A^0 returns the identity operator.

    Args:
        n: Non-negative integer exponent.

    Returns:
        A ``SparseDiaImpl`` equal to this matrix raised to the *n*-th power.

    Raises:
        ValueError: If *n* is negative.
    """
    if n < 0:
        raise ValueError("powm requires n >= 0")
    if n == 0:
        size = self._diags.shape[-1]
        eye_diags = jnp.ones((*self._diags.shape[:-2], 1, size), dtype=self._diags.dtype)
        return SparseDiaImpl(_offsets=(0,), _diags=eye_diags)
    if n == 1:
        return self
    half = self.powm(n // 2)
    squared = half.matmul(half)  # SparseDIA @ SparseDIA → SparseDIA
    return squared if n % 2 == 0 else self.matmul(squared)

real()

Element-wise real part of stored values.

Source code in jaxquantum/core/sparse_dia.py

def real(self) -> "SparseDiaImpl":
    """Element-wise real part of stored values."""
    return SparseDiaImpl(
        _offsets=self._offsets,
        _diags=jnp.real(self._diags).astype(self._diags.dtype),
    )

shape()

Shape of the represented square matrix (including batch dims).

Source code in jaxquantum/core/sparse_dia.py

def shape(self) -> tuple:
    """Shape of the represented square matrix (including batch dims)."""
    n = self._diags.shape[-1]
    return (*self._diags.shape[:-2], n, n)

sub(other)

Element-wise subtraction.

Source code in jaxquantum/core/sparse_dia.py

def sub(self, other: QarrayImpl) -> QarrayImpl:
    """Element-wise subtraction."""
    if isinstance(other, SparseDiaImpl):
        return _sparsedia_add(self, other, subtract=True)
    a, b = self._coerce(other)
    if a is not self:
        return a.sub(b)
    return a.sub(b)

tidy_up(atol)

Zero diagonal values whose magnitude is below atol.

Source code in jaxquantum/core/sparse_dia.py

def tidy_up(self, atol) -> "SparseDiaImpl":
    """Zero diagonal values whose magnitude is below *atol*."""
    diags = self._diags
    real_part = jnp.where(jnp.abs(jnp.real(diags)) < atol, 0.0, jnp.real(diags))
    if jnp.issubdtype(diags.dtype, jnp.complexfloating):
        imag_part = jnp.where(jnp.abs(jnp.imag(diags)) < atol, 0.0, jnp.imag(diags))
        new_diags = (real_part + 1j * imag_part).astype(diags.dtype)
    else:
        new_diags = real_part.astype(diags.dtype)
    return SparseDiaImpl(_offsets=self._offsets, _diags=new_diags)

to_dense()

Convert to a DenseImpl by summing diagonal contributions.

Source code in jaxquantum/core/sparse_dia.py

def to_dense(self) -> "DenseImpl":
    """Convert to a ``DenseImpl`` by summing diagonal contributions."""
    n = self._diags.shape[-1]
    batch_shape = self._diags.shape[:-2]
    result = jnp.zeros((*batch_shape, n, n), dtype=self._diags.dtype)
    for i, k in enumerate(self._offsets):
        s = _dia_slice(k)
        length = n - abs(k)
        if length <= 0:
            continue
        vals = self._diags[..., i, s]
        row_idx = jnp.arange(length) + max(-k, 0)
        col_idx = row_idx + k
        result = result.at[..., row_idx, col_idx].set(vals)
    return DenseImpl(result)

to_sparse_bcoo()

Convert to a SparseBCOOImpl (BCOO) via dense.

Source code in jaxquantum/core/sparse_dia.py

def to_sparse_bcoo(self) -> "SparseBCOOImpl":
    """Convert to a ``SparseBCOOImpl`` (BCOO) via dense."""
    return self.to_dense().to_sparse_bcoo()

to_sparse_dia()

Return self (already SparseDIA).

Source code in jaxquantum/core/sparse_dia.py

def to_sparse_dia(self) -> "SparseDiaImpl":
    """Return self (already SparseDIA)."""
    return self

trace()

Compute trace directly from the main diagonal (offset 0).

Returns:

Type	Description
	Scalar trace (sum of main diagonal values).

Source code in jaxquantum/core/sparse_dia.py

def trace(self):
    """Compute trace directly from the main diagonal (offset 0).

    Returns:
        Scalar trace (sum of main diagonal values).
    """
    if 0 in self._offsets:
        i = self._offsets.index(0)
        return jnp.sum(self._diags[..., i, :], axis=-1)
    return jnp.zeros(self._diags.shape[:-2], dtype=self._diags.dtype)