Coverage for jaxquantum/core/qarray.py: 84%
874 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""New Qarray implementation with sparse support."""
3from __future__ import annotations
5from abc import ABC, abstractmethod
6from copy import deepcopy
7from enum import Enum
8from math import prod
9from numbers import Number
10from typing import TYPE_CHECKING, Generic, List, Literal, TypeVar, Union, overload
12import jax.numpy as jnp
13import jax.scipy as jsp
14from flax import struct
15from jax import Array, config, default_backend
16from jax.experimental import sparse
17from numpy import ndarray
19from jaxquantum.core.dims import (
20 Qdims,
21 Qtypes,
22 check_dims,
23 ket_from_op_dims,
24)
25from jaxquantum.core.settings import SETTINGS, _maybe_shard
26from jaxquantum.utils.utils import robust_isscalar
28if TYPE_CHECKING:
29 from jaxquantum.core.sparse_bcoo import SparseBCOOImpl
31config.update("jax_enable_x64", True)
33# Type variable for implementation types
34ImplT = TypeVar("ImplT", bound="QarrayImpl")
36# Module-level registry mapping impl_class -> QarrayImplType member
37_IMPL_REGISTRY: dict = {}
40class QarrayImplType(Enum):
41 """Enumeration of available Qarray storage backends.
43 Each member maps one-to-one with a concrete ``QarrayImpl`` subclass.
44 New backends should call ``QarrayImplType.register(MyImpl, QarrayImplType.MY_TYPE)``
45 immediately after defining their impl class.
47 Members:
48 DENSE: Standard JAX dense array (``jnp.ndarray``).
49 SPARSE_BCOO: JAX experimental BCOO sparse array.
50 SPARSE_DIA: Diagonal sparse array.
51 """
53 DENSE = "dense"
54 SPARSE_BCOO = "sparse_bcoo"
55 SPARSE_DIA = "sparse_dia"
57 @classmethod
58 def register(cls, impl_class, member):
59 """Register an implementation class with a QarrayImplType member.
61 Args:
62 impl_class: The concrete ``QarrayImpl`` subclass to register.
63 member: The ``QarrayImplType`` enum member to associate with it.
64 """
65 _IMPL_REGISTRY[impl_class] = member
67 @classmethod
68 def has(cls, x) -> bool:
69 """Return True if x corresponds to a member of QarrayImplType.
71 Accepts an existing ``QarrayImplType`` member, a string equal to the
72 member name or value (case-insensitive), or an implementation class
73 (e.g. ``DenseImpl``, ``SparseBCOOImpl``) that has been registered.
75 Args:
76 x: Value to test — a ``QarrayImplType``, ``str``, or impl class.
78 Returns:
79 True if ``x`` maps to a known ``QarrayImplType`` member.
80 """
81 if isinstance(x, cls):
82 return True
84 if isinstance(x, str):
85 xl = x.lower()
86 return any(xl == member.value or xl == member.name.lower() for member in cls)
88 # Try mapping from an implementation class to an enum member
89 try:
90 cls.from_impl_class(x)
91 return True
92 except Exception: # noqa: BLE001
93 return False
95 @classmethod
96 def from_impl_class(cls, impl_class) -> "QarrayImplType":
97 """Return the ``QarrayImplType`` member associated with *impl_class*.
99 Args:
100 impl_class: A concrete ``QarrayImpl`` subclass that has been
101 registered via :meth:`register`.
103 Returns:
104 The corresponding ``QarrayImplType`` member.
106 Raises:
107 ValueError: If *impl_class* is not in the registry.
108 """
109 if impl_class in _IMPL_REGISTRY:
110 return _IMPL_REGISTRY[impl_class]
111 raise ValueError(f"Unknown implementation class: {impl_class}")
113 def get_impl_class(self):
114 """Return the implementation class registered for this member.
116 Returns:
117 The concrete ``QarrayImpl`` subclass associated with this member.
119 Raises:
120 ValueError: If no class has been registered for this member.
121 """
122 for cls_key, member in _IMPL_REGISTRY.items():
123 if member is self:
124 return cls_key
125 raise ValueError(f"No impl class registered for {self}")
128def robust_asarray(data) -> Union[Array, sparse.BCOO]:
129 """Convert *data* to a JAX array, leaving sparse BCOO and SparseDiaData untouched.
131 Args:
132 data: Input data — any array-like, ``sparse.BCOO``, or ``SparseDiaData``.
134 Returns:
135 A ``jax.Array``, ``sparse.BCOO``, or ``SparseDiaData``.
136 """
137 if isinstance(data, sparse.BCOO):
138 return data
139 # SparseDiaData has a ``_is_sparse_dia`` marker; pass it through unchanged
140 if getattr(data, "_is_sparse_dia", False):
141 return data
142 return jnp.asarray(data)
145class QarrayImpl(ABC):
146 """Abstract base class defining the interface every storage backend must implement.
148 A ``QarrayImpl`` wraps a raw data array (dense ``jnp.ndarray`` or sparse
149 ``BCOO``) and provides the mathematical primitives used by ``Qarray``.
150 Concrete subclasses must implement every ``@abstractmethod``.
152 Attributes:
153 PROMOTION_ORDER: Integer priority used by ``_coerce`` to decide which
154 side to promote when operands have different types. Higher means
155 "more general" (``DenseImpl = 1``, ``SparseBCOOImpl = 0``).
156 """
158 PROMOTION_ORDER: int = 0 # override in subclasses; higher = more general
159 # Current hierarchy: SparseDiaImpl=0, SparseBCOOImpl=1, DenseImpl=2
161 @abstractmethod
162 def get_data(self) -> Array:
163 """Return the underlying raw data array."""
164 pass
166 @property
167 def data(self) -> Array:
168 """The underlying raw data array."""
169 return self.get_data()
171 @property
172 def impl_type(self) -> QarrayImplType:
173 """The ``QarrayImplType`` member corresponding to this instance."""
174 return QarrayImplType.from_impl_class(type(self))
176 @classmethod
177 @abstractmethod
178 def from_data(cls, data) -> "QarrayImpl":
179 """Wrap raw data in this impl type.
181 Args:
182 data: Raw array data (dense ``jnp.ndarray`` or ``sparse.BCOO``).
184 Returns:
185 A new instance of this implementation wrapping *data*.
186 """
187 pass
189 @abstractmethod
190 def matmul(self, other: "QarrayImpl") -> "QarrayImpl":
191 """Matrix multiplication with *other*.
193 Args:
194 other: Right-hand operand.
196 Returns:
197 Result of ``self @ other`` as a ``QarrayImpl``.
198 """
199 pass
201 @abstractmethod
202 def add(self, other: "QarrayImpl") -> "QarrayImpl":
203 """Element-wise addition with *other*.
205 Args:
206 other: Right-hand operand.
208 Returns:
209 Result of ``self + other`` as a ``QarrayImpl``.
210 """
211 pass
213 @abstractmethod
214 def sub(self, other: "QarrayImpl") -> "QarrayImpl":
215 """Element-wise subtraction of *other*.
217 Args:
218 other: Right-hand operand.
220 Returns:
221 Result of ``self - other`` as a ``QarrayImpl``.
222 """
223 pass
225 @abstractmethod
226 def mul(self, scalar) -> "QarrayImpl":
227 """Scalar multiplication.
229 Args:
230 scalar: Scalar value to multiply by.
232 Returns:
233 Result of ``scalar * self`` as a ``QarrayImpl``.
234 """
235 pass
237 @abstractmethod
238 def dag(self) -> "QarrayImpl":
239 """Conjugate transpose.
241 Returns:
242 The conjugate transpose of this array as a ``QarrayImpl``.
243 """
244 pass
246 @abstractmethod
247 def to_dense(self) -> "DenseImpl":
248 """Convert to a ``DenseImpl``.
250 Returns:
251 A ``DenseImpl`` wrapping the same data.
252 """
253 pass
255 @abstractmethod
256 def to_sparse_bcoo(self) -> "SparseBCOOImpl":
257 """Convert to a ``SparseBCOOImpl`` (BCOO).
259 Returns:
260 A ``SparseBCOOImpl`` wrapping the same data.
261 """
262 pass
264 def to_sparse_dia(self) -> "QarrayImpl":
265 """Convert to a ``SparseDiaImpl``.
267 Default implementation goes through dense and auto-detects diagonals.
268 Subclasses may override for a more direct path.
270 Returns:
271 A ``SparseDiaImpl`` wrapping the same data.
272 """
273 # Import here to avoid circular imports at module load time
274 from jaxquantum.core.sparse_dia import SparseDiaImpl
275 return SparseDiaImpl.from_data(self.to_dense()._data)
277 @abstractmethod
278 def shape(self) -> tuple:
279 """Shape of the underlying data array.
281 Returns:
282 Tuple of dimension sizes.
283 """
284 pass
286 @abstractmethod
287 def dtype(self):
288 """Data type of the underlying array.
290 Returns:
291 A numpy/JAX dtype object.
292 """
293 pass
295 @abstractmethod
296 def __deepcopy__(self, memo=None):
297 pass
299 @abstractmethod
300 def tidy_up(self, atol):
301 """Zero out values whose magnitude is below *atol*.
303 Args:
304 atol: Absolute tolerance threshold.
306 Returns:
307 A new ``QarrayImpl`` with small values zeroed.
308 """
309 pass
311 @abstractmethod
312 def kron(self, other: "QarrayImpl") -> "QarrayImpl":
313 """Kronecker (tensor) product with another implementation.
315 Args:
316 other: Right-hand operand. Mixed-type pairs are handled by
317 ``_coerce`` — the result has the higher ``PROMOTION_ORDER``
318 type (dense wins over sparse).
320 Returns:
321 A new ``QarrayImpl`` containing the Kronecker product.
322 """
323 pass
325 @classmethod
326 @abstractmethod
327 def _eye_data(cls, n: int, dtype=None):
328 """Create identity matrix data of size n.
330 Args:
331 n: Matrix size.
332 dtype: Optional data type for the identity entries.
334 Returns:
335 Raw identity matrix data in the format appropriate for this impl.
336 """
337 pass
339 @classmethod
340 def _scaled_identity(cls, n: int, scalar, dtype=None) -> QarrayImpl:
341 """Create an identity scaled over any leading scalar batch axes."""
342 scalar = jnp.asarray(scalar) + 0.0j
343 if scalar.ndim:
344 scalar = scalar.reshape(scalar.shape + (1, 1))
345 return cls.from_data(cls._eye_data(n, dtype=dtype) * scalar)
347 @classmethod
348 @abstractmethod
349 def can_handle_data(cls, arr) -> bool:
350 """Return True if *arr* is a raw data type natively handled by this impl.
352 Used by the module-level :func:`dag_data` dispatcher to route raw
353 arrays to the correct backend without any isinstance chain outside the
354 impl classes.
356 Args:
357 arr: Raw array — e.g. ``jnp.ndarray`` for ``DenseImpl`` or
358 ``sparse.BCOO`` for ``SparseBCOOImpl``.
360 Returns:
361 True if this impl can operate on *arr* without conversion.
362 """
363 pass
365 @classmethod
366 @abstractmethod
367 def dag_data(cls, arr):
368 """Conjugate transpose of raw data in this impl's native format.
370 Implementations must handle batched arrays (last two axes are
371 swapped) and must not densify sparse arrays.
373 Args:
374 arr: Raw array in this impl's native format.
376 Returns:
377 Conjugate transpose with the last two axes swapped.
378 """
379 pass
381 def _promote_to(self, target_cls: type) -> "QarrayImpl":
382 """Convert this impl to *target_cls* by passing through dense.
384 Args:
385 target_cls: The destination ``QarrayImpl`` subclass.
387 Returns:
388 An instance of *target_cls* holding equivalent data.
389 """
390 if isinstance(self, target_cls):
391 return self
392 return target_cls.from_data(self.to_dense()._data)
394 def _coerce(self, other: "QarrayImpl") -> "tuple[QarrayImpl, QarrayImpl]":
395 """Coerce *self* and *other* to the same implementation type.
397 The impl type with the higher ``PROMOTION_ORDER`` wins; the other side
398 is promoted via :meth:`_promote_to`.
400 Args:
401 other: The other operand.
403 Returns:
404 A pair ``(a, b)`` of the same ``QarrayImpl`` subclass, suitable
405 for a binary operation.
406 """
407 if type(self) is type(other):
408 return self, other
409 if self.PROMOTION_ORDER >= other.PROMOTION_ORDER:
410 return self, other._promote_to(type(self))
411 return self._promote_to(type(other)), other
414@struct.dataclass
415class DenseImpl(QarrayImpl):
416 """Dense implementation using JAX dense arrays.
418 Attributes:
419 _data: The underlying ``jnp.ndarray``.
420 """
422 _data: Array
424 PROMOTION_ORDER = 2 # noqa: RUF012 — not a struct field; no annotation intentional
426 @classmethod
427 def _make(cls, data) -> "DenseImpl":
428 """Construct a ``DenseImpl``, applying the configured default sharding.
430 All internal construction sites route through this so that every
431 Qarray (including intermediates produced by ``matmul``, ``kron``,
432 etc.) satisfies the user's sharding invariant.
433 """
434 return cls(_data=_maybe_shard(data))
436 @classmethod
437 def from_data(cls, data) -> "DenseImpl":
438 """Wrap *data* in a new ``DenseImpl``.
440 Args:
441 data: Array-like input data.
443 Returns:
444 A ``DenseImpl`` wrapping ``robust_asarray(data)``.
445 """
446 return cls._make(robust_asarray(data))
448 def get_data(self) -> Array:
449 """Return the underlying dense array."""
450 return self._data
452 def matmul(self, other: QarrayImpl) -> QarrayImpl:
453 """Matrix multiply ``self @ other``, coercing types as needed.
455 Args:
456 other: Right-hand operand.
458 Returns:
459 A ``DenseImpl`` containing the matrix product.
460 """
461 a, b = self._coerce(other)
462 if a is not self:
463 return a.matmul(b)
464 return DenseImpl._make(self._data @ b._data)
466 def add(self, other: QarrayImpl) -> QarrayImpl:
467 """Element-wise addition ``self + other``, coercing types as needed.
469 Args:
470 other: Right-hand operand.
472 Returns:
473 A ``DenseImpl`` containing the sum.
474 """
475 a, b = self._coerce(other)
476 if a is not self:
477 return a.add(b)
478 return DenseImpl._make(self._data + b._data)
480 def sub(self, other: QarrayImpl) -> QarrayImpl:
481 """Element-wise subtraction ``self - other``, coercing types as needed.
483 Args:
484 other: Right-hand operand.
486 Returns:
487 A ``DenseImpl`` containing the difference.
488 """
489 a, b = self._coerce(other)
490 if a is not self:
491 return a.sub(b)
492 return DenseImpl._make(self._data - b._data)
494 def mul(self, scalar) -> QarrayImpl:
495 """Scalar multiplication.
497 Args:
498 scalar: Scalar value.
500 Returns:
501 A ``DenseImpl`` with each element multiplied by *scalar*.
502 """
503 return DenseImpl._make(scalar * self._data)
505 def dag(self) -> QarrayImpl:
506 """Conjugate transpose.
508 Returns:
509 A ``DenseImpl`` containing the conjugate transpose.
510 """
511 return DenseImpl._make(jnp.moveaxis(jnp.conj(self._data), -1, -2))
513 def to_dense(self) -> "DenseImpl":
514 """Return self (already dense).
516 Returns:
517 This ``DenseImpl`` instance unchanged.
518 """
519 return self
521 def to_sparse_bcoo(self) -> "SparseBCOOImpl":
522 """Convert to a ``SparseBCOOImpl`` via ``BCOO.fromdense``.
524 Returns:
525 A ``SparseBCOOImpl`` wrapping a BCOO conversion of this array.
526 """
527 from jaxquantum.core.sparse_bcoo import SparseBCOOImpl
528 return SparseBCOOImpl(sparse.BCOO.fromdense(self._data))
530 def shape(self) -> tuple:
531 """Shape of the underlying dense array.
533 Returns:
534 Tuple of dimension sizes.
535 """
536 return self._data.shape
538 def dtype(self):
539 """Data type of the underlying dense array.
541 Returns:
542 The dtype of ``_data``.
543 """
544 return self._data.dtype
546 def frobenius_norm(self) -> float:
547 """Compute the Frobenius norm.
549 Returns:
550 The Frobenius norm as a scalar.
551 """
552 return jnp.sqrt(jnp.sum(jnp.abs(self._data) ** 2))
554 def real(self) -> QarrayImpl:
555 """Element-wise real part.
557 Returns:
558 A ``DenseImpl`` containing the real parts.
559 """
560 return DenseImpl._make(jnp.real(self._data))
562 def imag(self) -> QarrayImpl:
563 """Element-wise imaginary part.
565 Returns:
566 A ``DenseImpl`` containing the imaginary parts.
567 """
568 return DenseImpl._make(jnp.imag(self._data))
570 def conj(self) -> QarrayImpl:
571 """Element-wise complex conjugate.
573 Returns:
574 A ``DenseImpl`` containing the complex-conjugated values.
575 """
576 return DenseImpl._make(jnp.conj(self._data))
578 def __deepcopy__(self, memo=None):
579 return DenseImpl._make(deepcopy(self._data, memo))
581 def tidy_up(self, atol):
582 """Zero out real/imaginary parts whose magnitude is below *atol*.
584 Args:
585 atol: Absolute tolerance threshold.
587 Returns:
588 A new ``DenseImpl`` with small values zeroed.
589 """
590 data = self._data
591 data_re = jnp.real(data)
592 data_im = jnp.imag(data)
593 data_re_mask = jnp.abs(data_re) > atol
594 data_im_mask = jnp.abs(data_im) > atol
595 data_new = data_re * data_re_mask + 1j * data_im * data_im_mask
597 return DenseImpl._make(data_new)
599 def kron(self, other: "QarrayImpl") -> "QarrayImpl":
600 """Kronecker product using ``jnp.kron``.
602 Args:
603 other: Right-hand operand.
605 Returns:
606 A ``DenseImpl`` containing the Kronecker product.
607 """
608 a, b = self._coerce(other)
609 if a is not self:
610 return a.kron(b)
611 return DenseImpl._make(jnp.kron(self._data, b._data))
613 @classmethod
614 def _eye_data(cls, n: int, dtype=None):
615 """Create an ``n x n`` identity matrix as a dense JAX array.
617 Args:
618 n: Matrix size.
619 dtype: Optional data type.
621 Returns:
622 A ``jnp.ndarray`` identity matrix of shape ``(n, n)``.
623 """
624 return jnp.eye(n, dtype=dtype)
626 @classmethod
627 def can_handle_data(cls, arr) -> bool:
628 """Return True for any non-BCOO, non-SparseDIA array.
630 ``SparseDiaData`` objects carry a ``_is_sparse_dia`` marker so we can
631 exclude them without a direct type import (which would be circular).
633 Args:
634 arr: Raw array.
636 Returns:
637 True when *arr* is a plain dense array (not BCOO, not SparseDiaData).
638 """
639 return not isinstance(arr, sparse.BCOO) and not getattr(arr, "_is_sparse_dia", False)
641 @classmethod
642 def dag_data(cls, arr) -> Array:
643 """Conjugate transpose for dense arrays.
645 Swaps the last two axes via :func:`jnp.moveaxis` and conjugates all
646 elements. For 1-D inputs only conjugation is applied.
648 Args:
649 arr: Dense array.
651 Returns:
652 Conjugate transpose with the last two axes swapped.
653 """
654 if len(arr.shape) == 1:
655 return jnp.conj(arr)
656 return jnp.moveaxis(jnp.conj(arr), -1, -2)
659# Register implementation classes with the enum registry
660# SparseBCOOImpl is registered in sparse_bcoo.py after import
661QarrayImplType.register(DenseImpl, QarrayImplType.DENSE)
664@struct.dataclass
665class Qarray(Generic[ImplT]):
666 """Quantum array with a pluggable storage backend.
668 ``Qarray`` wraps a ``QarrayImpl`` together with quantum-mechanical
669 dimension metadata (``_qdims``) and optional batch dimensions
670 (``_bdims``). The default backend is dense (``DenseImpl``); pass
671 ``implementation="sparse_bcoo"`` (or ``QarrayImplType.SPARSE_BCOO``) to
672 store data as a JAX BCOO sparse array.
674 Attributes:
675 _impl: The storage backend holding the raw data.
676 _qdims: Quantum dimension metadata (bra/ket structure, Hilbert space
677 sizes).
678 _bdims: Tuple of batch dimension sizes (empty tuple = non-batched).
680 Example:
681 >>> import jaxquantum as jqt
682 >>> a = jqt.destroy(10, implementation="sparse_bcoo")
683 >>> a.is_sparse_bcoo
684 True
685 """
687 _impl: ImplT
688 _qdims: Qdims = struct.field(pytree_node=False)
689 _bdims: tuple[int] = struct.field(pytree_node=False)
691 # Initialization ----
692 @classmethod
693 def _from_impl(cls, impl, qdims, bdims=None):
694 """Wrap a canonical internal result without repeating public validation."""
695 if bdims is None:
696 shape = impl.shape()
697 bdims = shape[:-2] if qdims.qtype == Qtypes.oper else shape[:-1]
698 return cls(impl, qdims, tuple(bdims))
700 @classmethod
701 @overload
702 def create(cls, data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.DENSE] = QarrayImplType.DENSE) -> "Qarray[DenseImpl]":
703 ...
705 @classmethod
706 @overload
707 def create(cls, data, dims=None, bdims=None, qtype=None, implementation: Literal[QarrayImplType.SPARSE_BCOO] = ...) -> "Qarray[SparseBCOOImpl]":
708 ...
710 @classmethod
711 @overload
712 def create(cls, data, dims=None, bdims=None, qtype=None, implementation=...) -> "Qarray[DenseImpl]":
713 ...
715 @classmethod
716 def create(cls, data, dims=None, bdims=None, qtype=None, implementation=QarrayImplType.DENSE):
717 """Create a ``Qarray`` from raw data.
719 Handles shape normalisation, dimension inference, and tidying of small
720 values.
722 State vectors are stored with their Hilbert space on a **single trailing
723 axis** — a ket/bra of dimension ``N`` has data shape ``bdims + (N,)``
724 (never ``(N,1)`` / ``(1,N)``). Operators keep the last two axes:
725 ``bdims + (M, N)``. The ket/bra/oper distinction lives in ``_qdims``,
726 not in the data shape.
728 Legacy ``(N,1)`` / ``(1,N)`` inputs are still accepted and are squeezed
729 to ``(N,)`` on the way in. Because a 1‑D ``(N,)`` array (or a square
730 ``(N,N)`` batch of vectors) is shape-ambiguous, pass ``qtype`` to be
731 explicit.
733 Args:
734 data: Input data array (dense array-like or ``sparse.BCOO``).
735 dims: Quantum dimensions as ``((row_dims...), (col_dims...))`` or, as
736 a convenience, a flat list of Hilbert space sizes. Inferred from
737 *data* shape when ``None``.
738 bdims: Tuple of batch dimension sizes. Inferred from the leading
739 dimensions of *data* when ``None``.
740 qtype: Optional quantum type — ``"ket"`` / ``"bra"`` / ``"oper"`` (or
741 a :class:`Qtypes` member). When given it is authoritative and
742 resolves shape ambiguity (e.g. a bare ``(N,)`` array with
743 ``qtype="ket"`` behaves like a legacy ``(N,1)`` input).
744 implementation: Storage backend — ``QarrayImplType.DENSE``
745 (default) or ``QarrayImplType.SPARSE_BCOO``, or the equivalent
746 string ``"dense"`` / ``"sparse_bcoo"``.
748 Returns:
749 A new ``Qarray`` backed by the requested implementation.
750 """
751 data = robust_asarray(data)
753 # Whether the caller explicitly requested a qtype (vs. it being derived
754 # from a full ``dims`` tuple). An explicit qtype is authoritative and
755 # mismatches with the data raise a clear error.
756 qtype_explicit = qtype is not None
757 if qtype is not None and not isinstance(qtype, Qtypes):
758 # Accepts "ket" / "bra" / "oper" (case-sensitive enum value).
759 qtype = Qtypes.from_str(qtype)
761 # Empty array (e.g. an empty list of operators) — keep legacy behaviour.
762 if len(data.shape) == 1 and data.shape[0] == 0:
763 dims = ((), ()) if dims is None else (tuple(dims[0]), tuple(dims[1]))
764 bdims = (0,) if bdims is None else tuple(bdims)
765 check_dims(dims, bdims, data.shape)
766 qdims = Qdims(dims)
767 impl_class = QarrayImplType(implementation).get_impl_class()
768 impl = impl_class.from_data(data)
769 impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])
770 return cls(impl, qdims, bdims)
772 # Resolve full quantum dims + batch dims.
773 dims_is_full = dims is not None and isinstance(dims[0], (list, tuple))
774 if qtype is None and dims_is_full:
775 qtype = Qtypes.from_dims((tuple(dims[0]), tuple(dims[1])))
777 if qtype is not None:
778 # Explicit (or dims-derived) qtype: unambiguous layout.
779 data, dims, bdims = cls._qtype_layout(data, dims, bdims, qtype)
780 else:
781 # Auto: legacy shape heuristics (1‑D → ket, square → oper, etc.).
782 data, dims, bdims = cls._legacy_layout(data, dims, bdims)
784 qdims = Qdims(dims)
786 # Store vectors as bdims + (N,); operators as bdims + (M, N).
787 if qdims.qtype == Qtypes.oper:
788 space = (prod(dims[0]), prod(dims[1]))
789 else:
790 space = (prod(dims[0]) * prod(dims[1]),)
792 target_shape = tuple(bdims) + space
794 # An explicitly requested qtype is authoritative: error out if the data
795 # cannot be interpreted as that type rather than silently coercing it.
796 if qtype_explicit:
797 if qdims.qtype != qtype:
798 raise ValueError(
799 f"Requested qtype={qtype} is incompatible with the resolved "
800 f"dims {dims} (which is a {qdims.qtype})."
801 )
802 if prod(data.shape) != prod(target_shape):
803 raise ValueError(
804 f"Data of shape {tuple(data.shape)} is incompatible with "
805 f"qtype={qtype} and dims={dims}: cannot reshape to "
806 f"{target_shape}."
807 )
809 # Reshape to the canonical stored shape only when needed. Use varargs so
810 # this works for dense arrays, BCOO, and SparseDiaData alike.
811 if tuple(data.shape) != target_shape:
812 data = data.reshape(*target_shape)
814 check_dims(dims, bdims, data.shape)
816 # NOTE: Constantly tidying up on Qarray creation might be a bit overkill.
817 # It increases the compilation time, but only very slightly
818 # increased the runtime of the jit compiled function.
819 # We could instead use this tidy up where we think we need it.
821 impl_class = QarrayImplType(implementation).get_impl_class()
822 impl = impl_class.from_data(data)
823 impl = impl.tidy_up(SETTINGS["auto_tidyup_atol"])
825 return cls(impl, qdims, bdims)
827 @staticmethod
828 def _legacy_layout(data, dims, bdims):
829 """Auto-resolve dims/bdims from data shape (no explicit ``qtype``).
831 Mirrors the historical shape heuristics: a 1‑D array is a ket, a square
832 2‑D array is an operator, a non-square 2‑D array is a batch of kets, and
833 legacy ``(N,1)`` / ``(1,N)`` carry the orientation. Returns
834 ``(data, dims, bdims)`` where *data* may still carry a trailing singleton
835 for vectors — the caller collapses it to ``(N,)``.
836 """
837 if len(data.shape) == 1 and data.shape[0] > 0:
838 data = data.reshape(data.shape[0], 1)
840 if (
841 len(data.shape) >= 2
842 and data.shape[-2] != data.shape[-1]
843 and not (data.shape[-2] == 1 or data.shape[-1] == 1)
844 ):
845 data = data.reshape(*data.shape[:-1], data.shape[-1], 1)
847 if bdims is not None and len(data.shape) - len(bdims) == 1:
848 data = data.reshape(*data.shape[:-1], data.shape[-1], 1)
850 if bdims is None:
851 bdims = tuple(data.shape[:-2])
853 if dims is None:
854 dims = ((data.shape[-2],), (data.shape[-1],))
856 if not isinstance(dims[0], (list, tuple)):
857 # Only the Hilbert space dimensions were sent in.
858 if data.shape[-1] == 1:
859 dims = (tuple(dims), tuple([1 for _ in dims]))
860 elif data.shape[-2] == 1:
861 dims = (tuple([1 for _ in dims]), tuple(dims))
862 else:
863 dims = (tuple(dims), tuple(dims))
864 else:
865 dims = (tuple(dims[0]), tuple(dims[1]))
867 return data, dims, tuple(bdims)
869 @staticmethod
870 def _qtype_layout(data, dims, bdims, qtype):
871 """Resolve dims/bdims given an explicit ``qtype`` (no shape guessing).
873 Accepts modern ``(...,N)`` vectors as well as legacy ``(N,1)`` / ``(1,N)``
874 (and batched variants) and returns ``(data, dims_full, bdims)``. The
875 caller reshapes *data* to the canonical stored shape.
876 """
877 shape = tuple(data.shape)
878 dims_is_full = dims is not None and isinstance(dims[0], (list, tuple))
880 if dims_is_full:
881 dims_full = (tuple(dims[0]), tuple(dims[1]))
882 elif dims is not None:
883 sl = tuple(dims)
884 if qtype == Qtypes.ket:
885 dims_full = (sl, tuple(1 for _ in sl))
886 elif qtype == Qtypes.bra:
887 dims_full = (tuple(1 for _ in sl), sl)
888 else:
889 dims_full = (sl, sl)
890 else:
891 dims_full = None # inferred from shape below
893 if qtype == Qtypes.oper:
894 if dims_full is None:
895 if len(shape) < 2:
896 raise ValueError(
897 f"qtype='oper' needs 2-D data or explicit dims; "
898 f"got shape {shape}."
899 )
900 dims_full = ((shape[-2],), (shape[-1],))
901 if bdims is None:
902 bdims = shape[:-2]
903 return data, dims_full, tuple(bdims)
905 # Vector (ket / bra).
906 if bdims is not None:
907 bdims = tuple(bdims)
908 if dims_full is None:
909 n = 1
910 for d in shape[len(bdims):]:
911 n *= d
912 dims_full = ((n,), (1,)) if qtype == Qtypes.ket else ((1,), (n,))
913 return data, dims_full, bdims
915 # bdims unknown: peel a legacy orientation singleton, then the last axis
916 # is the space axis and the leading axes are batch dims.
917 s = list(shape)
918 if qtype == Qtypes.ket and len(s) >= 2 and s[-1] == 1:
919 s = s[:-1]
920 elif qtype == Qtypes.bra and len(s) >= 2 and s[-2] == 1:
921 s = s[:-2] + s[-1:]
922 inferred_n = s[-1] if s else 1
923 bdims = tuple(s[:-1])
924 if dims_full is None:
925 dims_full = ((inferred_n,), (1,)) if qtype == Qtypes.ket else ((1,), (inferred_n,))
926 return data, dims_full, bdims
928 @classmethod
929 @overload
930 def from_sparse_bcoo(cls, data, dims=None, bdims=None) -> "Qarray[SparseBCOOImpl]":
931 ...
933 @classmethod
934 def from_sparse_bcoo(cls, data, dims=None, bdims=None):
935 """Create a ``Qarray`` directly from a sparse BCOO array without densifying.
937 Args:
938 data: A ``sparse.BCOO`` or array-like to store as sparse BCOO.
939 dims: Quantum dimensions. Inferred when ``None``.
940 bdims: Batch dimensions. Inferred when ``None``.
942 Returns:
943 A ``Qarray[SparseBCOOImpl]``.
944 """
945 return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_BCOO)
947 @classmethod
948 def from_sparse_dia(cls, data, dims=None, bdims=None) -> "Qarray":
949 """Create a SparseDIA-backed ``Qarray``.
951 Accepts either a dense array-like (diagonals are auto-detected) or a
952 :class:`~jaxquantum.core.sparse_dia.SparseDiaData` container.
954 Args:
955 data: Dense array of shape (*batch, n, n) or a ``SparseDiaData``.
956 dims: Quantum dimensions ``((row_dims,), (col_dims,))``.
957 bdims: Batch dimension sizes.
959 Returns:
960 A ``Qarray`` backed by ``SparseDiaImpl``.
961 """
962 return cls.create(data, dims=dims, bdims=bdims, implementation=QarrayImplType.SPARSE_DIA)
964 @classmethod
965 @overload
966 def from_list(cls, qarr_list: List["Qarray[DenseImpl]"], qtype=None) -> "Qarray[DenseImpl]":
967 ...
969 @classmethod
970 @overload
971 def from_list(cls, qarr_list: List["Qarray[SparseBCOOImpl]"], qtype=None) -> "Qarray[SparseBCOOImpl]":
972 ...
974 @classmethod
975 def from_list(cls, qarr_list: List[Qarray], qtype=None) -> Qarray:
976 """Create a batched ``Qarray`` from a list of same-shaped ``Qarray`` objects.
978 The output implementation is determined by the element with the highest
979 ``PROMOTION_ORDER``: if all inputs are sparse the result is sparse; if
980 any input is dense (or types are mixed) all inputs are promoted to dense
981 and the result is dense.
983 Works for kets/bras (stacked into ``(len, *bdims, N)``) as well as
984 operators, regardless of whether the elements were originally created
985 from ``(N,)`` or legacy ``(N,1)`` / ``(1,N)`` arrays — they are all
986 stored as ``(N,)`` vectors by the time they reach here.
988 Args:
989 qarr_list: List of ``Qarray`` objects with identical ``dims`` and
990 ``bdims``. May be empty.
991 qtype: Optional quantum type ("ket"/"bra"/"oper" or a ``Qtypes``)
992 forwarded to :meth:`create`. Defaults to the qtype of the
993 first element, which is the correct choice in all normal cases;
994 pass it explicitly to override or to be defensive.
996 Returns:
997 A ``Qarray`` with an extra leading batch dimension of size
998 ``len(qarr_list)``.
1000 Raises:
1001 ValueError: If the elements have mismatched ``dims`` or ``bdims``.
1002 """
1003 if len(qarr_list) == 0:
1004 dims = ((), ())
1005 bdims = (0,)
1006 return cls.create(jnp.array([]), dims=dims, bdims=bdims)
1008 dims = qarr_list[0].dims
1009 bdims = qarr_list[0].bdims
1011 if not all(qarr.dims == dims and qarr.bdims == bdims for qarr in qarr_list):
1012 raise ValueError("All Qarrays in the list must have the same dimensions.")
1014 if qtype is None:
1015 qtype = qarr_list[0].qtype
1017 new_bdims = (len(qarr_list),) + bdims
1019 # Pick the target type: highest PROMOTION_ORDER wins (dense beats sparse).
1020 target_impl_type = max(
1021 (q.impl_type for q in qarr_list),
1022 key=lambda t: t.get_impl_class().PROMOTION_ORDER,
1023 )
1025 if target_impl_type == QarrayImplType.SPARSE_DIA:
1026 # All inputs are SparseDIA — batch without densifying.
1027 # Compute union of offsets across all operators, then remap each
1028 # operator's _diags rows into the union shape and stack.
1029 from jaxquantum.core.sparse_dia import SparseDiaData # lazy to avoid circular
1030 union_offsets = tuple(sorted(
1031 set().union(*[set(q._impl._offsets) for q in qarr_list])
1032 ))
1033 union_idx = {k: i for i, k in enumerate(union_offsets)}
1034 n = qarr_list[0]._impl._diags.shape[-1]
1035 dtype = jnp.result_type(*[q._impl._diags.dtype for q in qarr_list])
1036 remapped = []
1037 for q in qarr_list:
1038 row = jnp.zeros((len(union_offsets), n), dtype=dtype)
1039 for i_src, k in enumerate(q._impl._offsets):
1040 row = row.at[union_idx[k], :].set(q._impl._diags[i_src, :])
1041 remapped.append(row)
1042 stacked = jnp.stack(remapped, axis=0) # (n_ops, n_union_diags, N)
1043 raw = SparseDiaData(offsets=union_offsets, diags=stacked)
1044 return cls.create(raw, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_DIA)
1046 if target_impl_type == QarrayImplType.SPARSE_BCOO:
1047 # All inputs are sparse BCOO — stack via dense intermediates then re-sparsify.
1048 data = jnp.array([q.data.todense() for q in qarr_list])
1049 return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.SPARSE_BCOO)
1051 # Target is dense: promote any sparse inputs before stacking.
1052 data = jnp.array([q.to_dense().data for q in qarr_list])
1053 return cls.create(data, dims=dims, bdims=new_bdims, qtype=qtype, implementation=QarrayImplType.DENSE)
1055 @classmethod
1056 @overload
1057 def from_array(cls, qarr_arr: "Qarray[DenseImpl]") -> "Qarray[DenseImpl]":
1058 ...
1060 @classmethod
1061 @overload
1062 def from_array(cls, qarr_arr: "Qarray[SparseBCOOImpl]") -> "Qarray[SparseBCOOImpl]":
1063 ...
1065 @classmethod
1066 def from_array(cls, qarr_arr) -> Qarray:
1067 """Create a ``Qarray`` from a (possibly nested) list of ``Qarray`` objects.
1069 Args:
1070 qarr_arr: A ``Qarray`` (returned as-is) or a nested list of
1071 ``Qarray`` objects.
1073 Returns:
1074 A ``Qarray`` with batch dimensions matching the nesting structure
1075 of *qarr_arr*.
1076 """
1077 if isinstance(qarr_arr, Qarray):
1078 return qarr_arr
1080 bdims = ()
1081 lvl = qarr_arr
1082 while not isinstance(lvl, Qarray):
1083 bdims = bdims + (len(lvl),)
1084 if len(lvl) > 0:
1085 lvl = lvl[0]
1086 else:
1087 break
1089 def flat(lis):
1090 flatList = []
1091 for element in lis:
1092 if type(element) is list:
1093 flatList += flat(element)
1094 else:
1095 flatList.append(element)
1096 return flatList
1098 qarr_list = flat(qarr_arr)
1099 qarr = cls.from_list(qarr_list)
1100 qarr = qarr.reshape_bdims(*bdims)
1101 return qarr
1103 # Properties ----
1104 @property
1105 def qtype(self):
1106 """Quantum type of this array (ket, bra, or operator)."""
1107 return self._qdims.qtype
1109 @property
1110 def dtype(self):
1111 """Data type of the underlying storage array."""
1112 return self._impl.dtype()
1114 @property
1115 def dims(self):
1116 """Quantum dimensions as ``((row_dims...), (col_dims...))``."""
1117 return self._qdims.dims
1119 @property
1120 def bdims(self):
1121 """Tuple of batch dimension sizes (empty tuple = non-batched)."""
1122 return self._bdims
1124 @property
1125 def qdims(self):
1126 """The ``Qdims`` metadata object for this array."""
1127 return self._qdims
1129 @property
1130 def space_dims(self):
1131 """Hilbert space dimensions for the relevant side (ket row / bra col)."""
1132 if self.qtype in [Qtypes.oper, Qtypes.ket]:
1133 return self.dims[0]
1134 elif self.qtype == Qtypes.bra:
1135 return self.dims[1]
1136 else:
1137 # TODO: not reached for some reason
1138 raise ValueError("Unsupported qtype.")
1140 @property
1141 def data(self):
1142 """The raw underlying data (dense ``jnp.ndarray`` or ``sparse.BCOO``)."""
1143 return self._impl.data
1145 @property
1146 def shaped_data(self):
1147 """Data reshaped to ``bdims + dims[0] + dims[1]``."""
1148 return self.data.reshape(self.bdims + self.dims[0] + self.dims[1])
1150 @property
1151 def shape(self):
1152 """Shape of the underlying data array."""
1153 return self.data.shape
1155 @property
1156 def is_batched(self):
1157 """True if this array has one or more batch dimensions."""
1158 return len(self.bdims) > 0
1160 @property
1161 def is_sparse_bcoo(self):
1162 """True if the storage backend is ``SparseBCOOImpl`` (BCOO)."""
1163 return self._impl.impl_type == QarrayImplType.SPARSE_BCOO
1165 @property
1166 def is_dense(self):
1167 """True if the storage backend is ``DenseImpl``."""
1168 return self._impl.impl_type == QarrayImplType.DENSE
1170 @property
1171 def is_sparse_dia(self):
1172 """True if the storage backend is ``SparseDiaImpl``."""
1173 return self._impl.impl_type == QarrayImplType.SPARSE_DIA
1175 @property
1176 def impl_type(self):
1177 """The ``QarrayImplType`` member of the current storage backend."""
1178 return self._impl.impl_type
1180 def to_sparse_bcoo(self) -> "Qarray[SparseBCOOImpl]":
1181 """Return a BCOO-sparse-backed copy of this array.
1183 If the array is already sparse BCOO, returns self unchanged.
1185 Returns:
1186 A ``Qarray[SparseBCOOImpl]``.
1187 """
1188 if self.is_sparse_bcoo:
1189 return self
1190 new_impl = self._impl.to_sparse_bcoo()
1191 return Qarray(new_impl, self._qdims, self._bdims)
1193 def to_sparse_dia(self) -> "Qarray":
1194 """Return a SparseDIA-backed copy of this array.
1196 If the array is already SparseDIA, returns self unchanged.
1198 Returns:
1199 A ``Qarray[SparseDiaImpl]``.
1200 """
1201 if self.is_sparse_dia:
1202 return self
1203 new_impl = self._impl.to_sparse_dia()
1204 return Qarray(new_impl, self._qdims, self._bdims)
1206 def to_dense(self) -> "Qarray[DenseImpl]":
1207 """Return a dense-backed copy of this array.
1209 If the array is already dense, returns self unchanged.
1211 Returns:
1212 A ``Qarray[DenseImpl]``.
1213 """
1214 if self.is_dense:
1215 return self
1216 new_impl = self._impl.to_dense()
1217 return Qarray(new_impl, self._qdims, self._bdims)
1219 def __getitem__(self, index):
1220 if len(self.bdims) > 0:
1221 impl = type(self._impl).from_data(self.data[index])
1222 return Qarray._from_impl(impl, self._qdims)
1223 else:
1224 raise ValueError("Cannot index a non-batched Qarray.")
1226 def reshape_bdims(self, *args):
1227 """Reshape the batch dimensions of this ``Qarray``.
1229 Args:
1230 *args: New batch dimension sizes.
1232 Returns:
1233 A new ``Qarray`` with the requested batch shape.
1234 """
1235 new_bdims = tuple(args)
1237 if prod(new_bdims) == 0:
1238 new_shape = new_bdims
1239 elif self.qtype == Qtypes.oper:
1240 new_shape = new_bdims + (prod(self.dims[0]), -1)
1241 else:
1242 # Vectors keep a single trailing space axis (no (N,1)).
1243 new_shape = new_bdims + (prod(self.dims[0]) * prod(self.dims[1]),)
1245 impl = type(self._impl).from_data(self.data.reshape(new_shape))
1246 return Qarray._from_impl(
1247 impl,
1248 self._qdims,
1249 new_bdims,
1250 )
1252 def space_to_qdims(self, space_dims: List[int]):
1253 """Convert Hilbert space dimensions to full quantum dims tuple.
1255 Args:
1256 space_dims: Sequence of per-subsystem Hilbert space sizes, or a
1257 full ``((row_dims), (col_dims))`` tuple (returned unchanged).
1259 Returns:
1260 A ``((row_dims...), (col_dims...))`` tuple.
1262 Raises:
1263 ValueError: If ``self.qtype`` is not ket, bra, or oper.
1264 """
1265 if isinstance(space_dims[0], (list, tuple)):
1266 return space_dims
1268 if self.qtype == Qtypes.oper:
1269 return (tuple(space_dims), tuple(space_dims))
1270 elif self.qtype == Qtypes.ket:
1271 return (tuple(space_dims), tuple([1 for _ in range(len(space_dims))]))
1272 elif self.qtype == Qtypes.bra:
1273 return (tuple([1 for _ in range(len(space_dims))]), tuple(space_dims))
1274 else:
1275 raise ValueError("Unsupported qtype for space_to_qdims conversion.")
1277 def reshape_qdims(self, *args):
1278 """Reshape the quantum dimensions of the Qarray.
1280 Note that this does not take in qdims but rather the new Hilbert space
1281 dimensions.
1283 Args:
1284 *args: New Hilbert dimensions for the Qarray.
1286 Returns:
1287 Qarray: reshaped Qarray.
1288 """
1290 new_space_dims = tuple(args)
1291 current_space_dims = self.space_dims
1292 assert prod(new_space_dims) == prod(current_space_dims)
1294 new_qdims = self.space_to_qdims(new_space_dims)
1295 return Qarray._from_impl(self._impl, Qdims(new_qdims), self._bdims)
1297 def resize(self, new_shape):
1298 """Resize the Qarray to a new shape.
1300 TODO: review and maybe deprecate this method.
1302 Args:
1303 new_shape: Target shape tuple.
1305 Returns:
1306 A new ``Qarray`` with data resized via ``jnp.resize``.
1307 """
1308 dims = self.dims
1309 data = jnp.resize(self.data, new_shape)
1310 # Preserve implementation type
1311 implementation = self.impl_type
1312 return Qarray.create(
1313 data,
1314 dims=dims,
1315 implementation=implementation,
1316 )
1318 def __len__(self):
1319 """Length along the first batch dimension.
1321 Returns:
1322 Size of the leading batch dimension.
1324 Raises:
1325 ValueError: If the array is not batched.
1326 """
1327 if len(self.bdims) > 0:
1328 return self.data.shape[0]
1329 else:
1330 raise ValueError("Cannot get length of a non-batched Qarray.")
1332 def __eq__(self, other):
1333 if not isinstance(other, Qarray):
1334 raise ValueError( # noqa: TRY004
1335 "Cannot calculate equality of a Qarray with a non-Qarray."
1336 )
1338 if self.dims != other.dims:
1339 return False
1341 if self.bdims != other.bdims:
1342 return False
1344 if self.is_sparse_bcoo and other.is_sparse_bcoo:
1345 # Fast structural path: same sparsity pattern → compare values only (no todense)
1346 if (self.data.indices.shape == other.data.indices.shape
1347 and bool(jnp.all(self.data.indices == other.data.indices))):
1348 return bool(jnp.allclose(self.data.data, other.data.data))
1349 # Different patterns: fall back to dense comparison (unavoidable)
1350 return bool(jnp.all(self.data.todense() == other.data.todense()))
1352 # At least one dense: convert sparse side to dense for comparison
1353 self_data = self.data.todense() if hasattr(self.data, 'todense') else self.data
1354 other_data = other.data.todense() if hasattr(other.data, 'todense') else other.data
1355 return bool(jnp.all(self_data == other_data))
1357 def __ne__(self, other):
1358 return not self.__eq__(other)
1360 # Elementary Math ----
1361 def __matmul__(self, other):
1362 if not isinstance(other, Qarray):
1363 return NotImplemented
1365 _qdims_new = self._qdims @ other._qdims
1366 st, ot = self.qtype, other.qtype
1368 # Operator @ operator: keep the backend-native path (dense or sparse),
1369 # which preserves the storage implementation. Operators are 2-D-tailed,
1370 # so create() infers bdims/qtype from the result shape unambiguously.
1371 if st == Qtypes.oper and ot == Qtypes.oper:
1372 new_impl = self._impl.matmul(other._impl)
1373 return Qarray._from_impl(new_impl, _qdims_new)
1375 # bra @ oper == dag( oper^dag @ bra^dag ). Reusing the oper @ ket path
1376 # keeps a sparse operator sparse (no densification of the operator).
1377 if st == Qtypes.bra and ot == Qtypes.oper:
1378 return (other.dag() @ self.dag()).dag()
1380 # oper @ ket: matrix-vector product. Vectors are dense (see migration
1381 # scope); the operator may be sparse and acts natively on the dense
1382 # vector. The vector's single space axis is contracted in place — it is
1383 # never reshaped to (N,1) for storage.
1384 if st == Qtypes.oper and ot == Qtypes.ket:
1385 vec = other.to_dense().data # (..., N)
1386 if self.is_dense:
1387 out = jnp.einsum("...ij,...j->...i", self.data, vec)
1388 else:
1389 # Transient column only inside the sparse kernel; never stored.
1390 rhs = vec[..., None] # (..., N, 1)
1391 out = self._impl.matmul(DenseImpl._make(rhs)).data[..., 0]
1392 return Qarray._from_impl(DenseImpl._make(out), _qdims_new)
1394 # Vector ⊗/· vector — both operands are dense.
1395 a = self.to_dense().data
1396 b = other.to_dense().data
1397 if st == Qtypes.ket and ot == Qtypes.bra:
1398 # Outer product |a><b| → operator.
1399 out = jnp.einsum("...i,...j->...ij", a, b)
1400 return Qarray._from_impl(DenseImpl._make(out), _qdims_new)
1401 if st == Qtypes.bra and ot == Qtypes.ket:
1402 # Inner product <a|b> → 1x1 "ket" (qdims ((1,),(1,))).
1403 out = jnp.einsum("...i,...i->...", a, b)
1404 out = out.reshape(out.shape + (1,))
1405 return Qarray._from_impl(DenseImpl._make(out), _qdims_new)
1407 return NotImplemented
1409 def __mul__(self, other):
1410 if isinstance(other, Qarray):
1411 return self.__matmul__(other)
1413 other = other + 0.0j
1414 if not robust_isscalar(other) and len(other.shape) > 0: # not a scalar
1415 # Broadcast per-batch scalars against the stored data: vectors carry
1416 # a single trailing space axis, operators carry two.
1417 extra = (1,) if self.qtype in (Qtypes.ket, Qtypes.bra) else (1, 1)
1418 other = other.reshape(other.shape + extra)
1420 new_impl = self._impl.mul(other)
1421 return Qarray._from_impl(new_impl, self._qdims)
1423 def __rmul__(self, other):
1424 return self.__mul__(other)
1426 def __neg__(self):
1427 return self.__mul__(-1)
1429 def __truediv__(self, other):
1430 """Divide by a scalar.
1432 Args:
1433 other: Scalar divisor.
1435 Returns:
1436 A new ``Qarray`` with all elements divided by *other*.
1438 Raises:
1439 ValueError: If *other* is a ``Qarray``.
1440 """
1441 if isinstance(other, Qarray):
1442 raise ValueError("Cannot divide a Qarray by another Qarray.") # noqa: TRY004
1444 return self.__mul__(1 / other)
1446 def _scaled_identity(self, scalar):
1447 return type(self._impl)._scaled_identity(
1448 self.data.shape[-2],
1449 scalar,
1450 dtype=self.data.dtype,
1451 )
1453 def __add__(self, other):
1454 if isinstance(other, Qarray):
1455 if self.dims != other.dims:
1456 msg = (
1457 "Dimensions are incompatible: "
1458 + repr(self.dims)
1459 + " and "
1460 + repr(other.dims)
1461 )
1462 raise ValueError(msg)
1463 new_impl = self._impl.add(other._impl)
1464 return Qarray._from_impl(new_impl, self._qdims)
1466 if isinstance(other, Number) and other == 0:
1467 return self.copy()
1469 if self.qtype == Qtypes.oper:
1470 return Qarray._from_impl(
1471 self._impl.add(self._scaled_identity(other)),
1472 self._qdims,
1473 )
1475 return NotImplemented
1477 def __radd__(self, other):
1478 return self.__add__(other)
1480 def __sub__(self, other):
1481 if isinstance(other, Qarray):
1482 if self.dims != other.dims:
1483 msg = (
1484 "Dimensions are incompatible: "
1485 + repr(self.dims)
1486 + " and "
1487 + repr(other.dims)
1488 )
1489 raise ValueError(msg)
1490 new_impl = self._impl.sub(other._impl)
1491 return Qarray._from_impl(new_impl, self._qdims)
1493 if isinstance(other, Number) and other == 0:
1494 return self.copy()
1496 if self.qtype == Qtypes.oper:
1497 return Qarray._from_impl(
1498 self._impl.sub(self._scaled_identity(other)),
1499 self._qdims,
1500 )
1502 return NotImplemented
1504 def __rsub__(self, other):
1505 return self.__neg__().__add__(other)
1507 def __xor__(self, other):
1508 if not isinstance(other, Qarray):
1509 return NotImplemented
1510 return tensor(self, other)
1512 def __rxor__(self, other):
1513 if not isinstance(other, Qarray):
1514 return NotImplemented
1515 return tensor(other, self)
1517 def __pow__(self, other):
1518 if not isinstance(other, int):
1519 return NotImplemented
1521 return powm(self, other)
1523 # String Representation ----
1524 def _str_header(self):
1525 """Build the one-line header string for ``__str__`` and ``__repr__``."""
1526 impl_type = self.impl_type.value
1527 out = ", ".join(
1528 [
1529 "Quantum array: dims = " + str(self.dims),
1530 "bdims = " + str(self.bdims),
1531 "shape = " + str(self.data.shape),
1532 "type = " + str(self.qtype),
1533 "impl = " + impl_type,
1534 ]
1535 )
1536 return out
1538 def __str__(self):
1539 return self._str_header() + "\nQarray data =\n" + str(self.data)
1541 @property
1542 def header(self):
1543 """One-line header string describing dimensions, shape, and backend."""
1544 return self._str_header()
1546 def __repr__(self):
1547 return self.__str__()
1549 # Utilities ----
1550 def copy(self, memo=None):
1551 """Return a deep copy of this ``Qarray``.
1553 Args:
1554 memo: Optional memo dict forwarded to ``deepcopy``.
1556 Returns:
1557 A new ``Qarray`` with independent copies of all data.
1558 """
1559 return self.__deepcopy__(memo)
1561 def __deepcopy__(self, memo):
1562 """Need to override this when defining __getattr__."""
1564 return Qarray(
1565 _impl=deepcopy(self._impl, memo=memo),
1566 _qdims=deepcopy(self._qdims, memo=memo),
1567 _bdims=deepcopy(self._bdims, memo=memo),
1568 )
1570 def __getattr__(self, method_name):
1571 if "__" == method_name[:2]:
1572 # NOTE: we return NotImplemented for binary special methods logic in python, plus things like __jax_array__
1573 return lambda *args, **kwargs: NotImplemented
1575 modules = [jnp, jnp.linalg, jsp, jsp.linalg]
1577 method_f = None
1578 for mod in modules:
1579 method_f = getattr(mod, method_name, None)
1580 if method_f is not None:
1581 break
1583 if method_f is None:
1584 raise NotImplementedError(
1585 f"Method {method_name} does not exist. No backup method found in {modules}."
1586 )
1588 def func(*args, **kwargs):
1589 # For operations that might not be supported in sparse, convert to dense
1590 if self.is_sparse_bcoo:
1591 dense_self = self.to_dense()
1592 res = method_f(dense_self.data, *args, **kwargs)
1593 else:
1594 res = method_f(self.data, *args, **kwargs)
1596 if getattr(res, "shape", None) is None or res.shape != self.data.shape:
1597 return res
1598 else:
1599 # Preserve implementation type
1600 return Qarray.create(res, dims=self._qdims.dims, implementation=self.impl_type)
1602 return func
1604 # Conversions / Reshaping ----
1605 def dag(self):
1606 """Conjugate transpose of this array."""
1607 return dag(self)
1609 def to_dm(self):
1610 """Convert a ket to a density matrix via outer product."""
1611 return ket2dm(self)
1613 def is_dm(self):
1614 """Return True if this array is an operator (density-matrix type)."""
1615 return self.qtype == Qtypes.oper
1617 def is_vec(self):
1618 """Return True if this array is a ket or bra."""
1619 return self.qtype == Qtypes.ket or self.qtype == Qtypes.bra
1621 def to_ket(self):
1622 """Convert a bra to a ket (no-op for kets)."""
1623 return to_ket(self)
1625 def transpose(self, *args):
1626 """Transpose subsystem indices."""
1627 return transpose(self, *args)
1629 def keep_only_diag_elements(self):
1630 """Zero out all off-diagonal elements."""
1631 return keep_only_diag_elements(self)
1633 # Math Functions ----
1634 def unit(self):
1635 """Return the normalised (unit-norm) version of this array."""
1636 return unit(self)
1638 def norm(self):
1639 """Compute the norm of this array."""
1640 return norm(self)
1642 def frobenius_norm(self):
1643 """Compute the Frobenius norm directly from the implementation.
1645 Returns:
1646 The Frobenius norm as a scalar.
1647 """
1648 return self._impl.frobenius_norm()
1650 def real(self):
1651 """Element-wise real part.
1653 Returns:
1654 A new ``Qarray`` containing the real parts of each element.
1655 """
1656 new_impl = self._impl.real()
1657 return Qarray._from_impl(new_impl, self._qdims)
1659 def imag(self):
1660 """Element-wise imaginary part.
1662 Returns:
1663 A new ``Qarray`` containing the imaginary parts of each element.
1664 """
1665 new_impl = self._impl.imag()
1666 return Qarray._from_impl(new_impl, self._qdims)
1668 def conj(self):
1669 """Element-wise complex conjugate.
1671 Returns:
1672 A new ``Qarray`` containing the complex-conjugated elements.
1673 """
1674 new_impl = self._impl.conj()
1675 return Qarray._from_impl(new_impl, self._qdims)
1677 def expm(self):
1678 """Matrix exponential."""
1679 return expm(self)
1681 def powm(self, n):
1682 """Matrix power.
1684 Args:
1685 n: Exponent (integer or float).
1687 Returns:
1688 This array raised to the *n*-th matrix power.
1689 """
1690 return powm(self, n)
1692 def cosm(self):
1693 """Matrix cosine."""
1694 return cosm(self)
1696 def sinm(self):
1697 """Matrix sine."""
1698 return sinm(self)
1700 def tr(self, **kwargs):
1701 """Full trace."""
1702 return tr(self, **kwargs)
1704 def trace(self, **kwargs):
1705 """Full trace (alias for :meth:`tr`)."""
1706 return tr(self, **kwargs)
1708 def ptrace(self, indx):
1709 """Partial trace over subsystem *indx*.
1711 Args:
1712 indx: Index of the subsystem to trace out.
1714 Returns:
1715 Reduced density matrix.
1716 """
1717 return ptrace(self, indx)
1719 def eigenstates(self):
1720 """Eigenvalues and eigenstates of this operator."""
1721 return eigenstates(self)
1723 def eigenenergies(self):
1724 """Eigenvalues of this operator."""
1725 return eigenenergies(self)
1727 def eigenvalues(self):
1728 """Eigenvalues of this operator (alias for :meth:`eigenenergies`)."""
1729 return eigenenergies(self)
1731 def collapse(self, mode="sum"):
1732 """Collapse batch dimensions.
1734 Args:
1735 mode: Collapse strategy — currently only ``"sum"`` is supported.
1737 Returns:
1738 A non-batched ``Qarray``.
1739 """
1740 return collapse(self, mode=mode)
1743# Qarray operations ---------------------------------------------------------------------
1745def concatenate(qarr_list: List[Qarray], axis: int = 0) -> Qarray:
1746 """Concatenate a list of Qarrays along a specified axis.
1748 Args:
1749 qarr_list: List of Qarrays to concatenate.
1750 axis: Axis along which to concatenate. Default is 0.
1752 Returns:
1753 Concatenated Qarray.
1754 """
1756 non_empty_qarr_list = [qarr for qarr in qarr_list if len(qarr.data) != 0]
1758 if len(non_empty_qarr_list) == 0:
1759 return Qarray.from_list([])
1761 concatenated_data = jnp.concatenate(
1762 [qarr.data for qarr in non_empty_qarr_list], axis=axis
1763 )
1765 dims = non_empty_qarr_list[0].dims
1766 return Qarray.create(concatenated_data, dims=dims)
1769def collapse(qarr: Qarray, mode="sum") -> Qarray:
1770 """Collapse the batch dimensions of *qarr*.
1772 Args:
1773 qarr: Quantum array with optional batch dimensions.
1774 mode: Collapse strategy. Only ``"sum"`` is currently supported.
1776 Returns:
1777 A non-batched ``Qarray`` obtained by summing over all batch axes.
1778 """
1780 if mode == "sum":
1781 if len(qarr.bdims) == 0:
1782 return qarr
1784 batch_axes = list(range(len(qarr.bdims)))
1786 # Preserve implementation type
1787 implementation = qarr.impl_type
1788 return Qarray.create(jnp.sum(qarr.data, axis=batch_axes), dims=qarr.dims, implementation=implementation)
1791def transpose(qarr: Qarray, indices: List[int]) -> Qarray:
1792 """Transpose subsystem indices of the quantum array.
1794 Args:
1795 qarr: Input quantum array.
1796 indices: New ordering of subsystem indices.
1798 Returns:
1799 Transposed ``Qarray`` (converted to dense first).
1800 """
1802 qarr = qarr.to_dense()
1804 indices = list(indices)
1806 shaped_data = qarr.shaped_data
1807 dims = qarr.dims
1808 bdims_indxs = list(range(len(qarr.bdims)))
1810 reshape_indices = indices + [j + len(dims[0]) for j in indices]
1811 reshape_indices = bdims_indxs + [j + len(bdims_indxs) for j in reshape_indices]
1813 shaped_data = shaped_data.transpose(reshape_indices)
1814 new_dims = (
1815 tuple([dims[0][j] for j in indices]),
1816 tuple([dims[1][j] for j in indices]),
1817 )
1819 full_dims = prod(dims[0])
1820 full_data = shaped_data.reshape(*qarr.bdims, full_dims, -1)
1822 # Preserve implementation type
1823 implementation = qarr.impl_type
1824 return Qarray.create(full_data, dims=new_dims, implementation=implementation)
1827def unit(qarr: Qarray) -> Qarray:
1828 """Normalize *qarr* to unit norm.
1830 Args:
1831 qarr: Input quantum array.
1833 Returns:
1834 Normalized quantum array.
1835 """
1836 return qarr / qarr.norm()
1839def norm(qarr: Qarray) -> float:
1840 """Compute the norm of a quantum array.
1842 Sparse paths (no densification):
1844 * ket / bra — L2 norm via :meth:`SparseBCOOImpl.l2_norm_batched` (handles
1845 batch dimensions).
1846 * operator — trace norm assuming PSD (nuclear norm = tr(rho) for density
1847 matrices). This is exact for density matrices; for general non-PSD
1848 operators convert to dense first.
1850 Args:
1851 qarr: Input quantum array.
1853 Returns:
1854 The norm as a scalar (or batched array of scalars).
1855 """
1856 if qarr.qtype in [Qtypes.ket, Qtypes.bra] and qarr.is_sparse_bcoo:
1857 return qarr._impl.l2_norm_batched(qarr.bdims)
1859 if qarr.qtype == Qtypes.oper and qarr.is_sparse_bcoo:
1860 # Nuclear norm = trace for positive-semidefinite (density matrix) operators.
1861 # jnp.real strips any floating-point imaginary artefact.
1862 return jnp.real(qarr._impl.trace())
1864 if qarr.qtype == Qtypes.oper and qarr.is_sparse_dia:
1865 return jnp.real(qarr._impl.trace())
1867 qarr = qarr.to_dense()
1869 qdata = qarr.data
1871 if qarr.qtype == Qtypes.oper:
1872 if default_backend() == "cpu":
1873 return jnp.sum(jnp.linalg.svd(qdata, compute_uv=False), axis=-1)
1874 gram = qdata @ jnp.swapaxes(jnp.conj(qdata), -1, -2)
1875 values = jnp.linalg.eigvalsh(gram)
1876 return jnp.sum(jnp.sqrt(jnp.abs(values)), axis=-1)
1878 elif qarr.qtype in [Qtypes.ket, Qtypes.bra]:
1879 # Vectors store the Hilbert space on the single trailing axis.
1880 return jnp.sqrt(jnp.sum(jnp.abs(qdata) ** 2, axis=-1))
1883def tensor(*args, **kwargs) -> Qarray:
1884 """Tensor (Kronecker) product of two or more ``Qarray`` objects.
1886 Args:
1887 *args: ``Qarray`` objects to tensor together (left to right).
1888 **kwargs: Optional keyword arguments. Pass ``parallel=True`` to use
1889 an einsum-based batched outer product instead of ``jnp.kron``.
1891 Returns:
1892 The tensor product as a ``Qarray``. The output implementation is
1893 determined by the highest ``PROMOTION_ORDER`` among the inputs: all-sparse
1894 inputs → sparse output; any dense input → dense output. This holds for
1895 both ``parallel=True`` and ``parallel=False``.
1897 Note:
1898 ``parallel=True`` uses an einsum-based batched outer product. The
1899 einsum is always computed on dense data for efficiency, but the result
1900 is then wrapped in the appropriate backend (sparse when all inputs are
1901 sparse, dense otherwise). For the default (``parallel=False``) path
1902 each backend's ``kron`` method is used directly.
1903 """
1904 parallel = kwargs.pop("parallel", False)
1906 if parallel:
1907 # Determine target implementation: highest PROMOTION_ORDER wins.
1908 # All-sparse → sparse; any dense input → dense (same rule as non-parallel).
1909 target_impl_type = max(
1910 (arg.impl_type for arg in args),
1911 key=lambda t: t.get_impl_class().PROMOTION_ORDER,
1912 )
1913 # Einsum-based batched outer product (computed on dense data).
1914 dense_args = [arg.to_dense() for arg in args]
1915 # Vectors (kets/bras) carry one trailing space axis; operators carry two.
1916 # Tensoring requires a consistent qtype across args (mixed is invalid).
1917 is_vec = dense_args[0].qtype in (Qtypes.ket, Qtypes.bra)
1918 n_space = 1 if is_vec else 2
1919 data = dense_args[0].data
1920 dims_0 = dense_args[0].dims[0]
1921 dims_1 = dense_args[0].dims[1]
1922 for arg in dense_args[1:]:
1923 a, b = data, arg.data
1924 ba, bb = a.shape[:-n_space], b.shape[:-n_space]
1925 if len(ba) > len(bb):
1926 batch_dim = ba
1927 elif len(ba) == len(bb):
1928 batch_dim = ba if prod(ba) >= prod(bb) else bb
1929 else:
1930 batch_dim = bb
1932 if is_vec:
1933 # (..., N) ⊗ (..., M) -> (..., N*M)
1934 data = jnp.einsum("...i,...j->...ij", a, b).reshape(
1935 *batch_dim, a.shape[-1] * b.shape[-1]
1936 )
1937 else:
1938 # (..., M, N) ⊗ (..., K, L) -> (..., M*K, N*L)
1939 data = jnp.einsum("...ij,...kl->...ikjl", a, b).reshape(
1940 *batch_dim, a.shape[-2] * b.shape[-2], -1
1941 )
1942 dims_0 = dims_0 + arg.dims[0]
1943 dims_1 = dims_1 + arg.dims[1]
1944 impl = target_impl_type.get_impl_class().from_data(data)
1945 return Qarray._from_impl(impl, Qdims((dims_0, dims_1)))
1947 # Non-parallel: delegate to each impl's kron method.
1948 # All-sparse inputs stay sparse; mixed inputs promote to dense via _coerce.
1949 current_impl = args[0]._impl
1950 dims_0 = args[0].dims[0]
1951 dims_1 = args[0].dims[1]
1952 for arg in args[1:]:
1953 current_impl = current_impl.kron(arg._impl)
1954 dims_0 = dims_0 + arg.dims[0]
1955 dims_1 = dims_1 + arg.dims[1]
1956 return Qarray._from_impl(current_impl, Qdims((dims_0, dims_1)))
1959def tr(qarr: Qarray, **kwargs) -> Array:
1960 """Full trace of *qarr*.
1962 For sparse ``Qarray`` objects the trace is computed natively on the BCOO
1963 data using a masked scatter — no densification. Custom axis arguments
1964 are ignored for sparse (the last two dimensions are always the matrix
1965 dimensions in jaxquantum's convention).
1967 Args:
1968 qarr: Input quantum array.
1969 **kwargs: Forwarded to ``jnp.trace`` for dense arrays (e.g.
1970 ``axis1``, ``axis2``).
1972 Returns:
1973 The trace as a scalar (or batched array of scalars).
1974 """
1975 if qarr.is_sparse_bcoo:
1976 return qarr._impl.trace()
1977 if qarr.is_sparse_dia:
1978 return qarr._impl.trace()
1979 axis1 = kwargs.get("axis1", -2)
1980 axis2 = kwargs.get("axis2", -1)
1981 return jnp.trace(qarr.data, axis1=axis1, axis2=axis2, **kwargs)
1984def trace(qarr: Qarray, **kwargs) -> Array:
1985 """Full trace (alias for :func:`tr`).
1987 Args:
1988 qarr: Input quantum array.
1989 **kwargs: Forwarded to :func:`tr`.
1991 Returns:
1992 The trace as a scalar (or batched array of scalars).
1993 """
1994 return tr(qarr, **kwargs)
1997def expm_data(data: Array, **kwargs) -> Array:
1998 """Matrix exponential of a raw array.
2000 Args:
2001 data: Dense matrix array.
2002 **kwargs: Forwarded to ``jsp.linalg.expm``.
2004 Returns:
2005 The matrix exponential.
2006 """
2007 return jsp.linalg.expm(data, **kwargs)
2010def expm(qarr: Qarray, **kwargs) -> Qarray:
2011 """Matrix exponential of a ``Qarray``.
2013 Args:
2014 qarr: Input quantum array (converted to dense internally).
2015 **kwargs: Forwarded to ``jsp.linalg.expm``.
2017 Returns:
2018 A dense ``Qarray`` containing the matrix exponential.
2019 """
2020 dims = qarr.dims
2021 # Convert to dense for expm
2022 dense_data = qarr.to_dense().data
2023 data = expm_data(dense_data, **kwargs)
2024 return Qarray.create(data, dims=dims)
2027def powm(qarr: Qarray, n: Union[int, float], clip_eigvals=False) -> Qarray:
2028 """Matrix power of a ``Qarray``.
2030 Args:
2031 qarr: Input quantum array.
2032 n: Exponent. Integer powers use ``jnp.linalg.matrix_power``; float
2033 powers diagonalise the matrix.
2034 clip_eigvals: When ``True``, clip negative eigenvalues to zero before
2035 applying the float power (useful for nearly-PSD matrices).
2037 Returns:
2038 The *n*-th matrix power as a ``Qarray`` (stays SparseDIA for integer
2039 non-negative exponents when the input is SparseDIA).
2041 Raises:
2042 ValueError: If *n* is a float and the matrix has negative eigenvalues
2043 (and *clip_eigvals* is ``False``).
2044 """
2045 # SparseDIA fast path: binary exponentiation stays in SparseDIA format.
2046 if qarr.is_sparse_dia and isinstance(n, int) and n >= 0:
2047 new_impl = qarr._impl.powm(n)
2048 return Qarray.create(new_impl.data, dims=qarr.dims, implementation=new_impl.impl_type)
2050 # Convert to dense for powm
2051 dense_qarr = qarr.to_dense()
2053 if isinstance(n, int):
2054 data_res = jnp.linalg.matrix_power(dense_qarr.data, n)
2055 else:
2056 evalues, evectors = jnp.linalg.eig(dense_qarr.data)
2057 if clip_eigvals:
2058 evalues = jnp.maximum(evalues, 0)
2059 else:
2060 if not (evalues >= 0).all():
2061 raise ValueError(
2062 "Non-integer power of a matrix can only be "
2063 "computed if the matrix is positive semi-definite."
2064 "Got a matrix with a negative eigenvalue."
2065 )
2066 data_res = evectors * jnp.pow(evalues, n) @ jnp.linalg.inv(evectors)
2068 return Qarray.create(data_res, dims=qarr.dims)
2071def cosm_data(data: Array, **kwargs) -> Array:
2072 """Matrix cosine of a raw array.
2074 Args:
2075 data: Dense matrix array.
2076 **kwargs: Unused; kept for API consistency.
2078 Returns:
2079 The matrix cosine computed as ``(expm(i*A) + expm(-i*A)) / 2``.
2080 """
2081 return (expm_data(1j * data) + expm_data(-1j * data)) / 2
2084def cosm(qarr: Qarray) -> Qarray:
2085 """Matrix cosine of a ``Qarray``.
2087 Args:
2088 qarr: Input quantum array (converted to dense internally).
2090 Returns:
2091 A dense ``Qarray`` containing the matrix cosine.
2092 """
2093 dims = qarr.dims
2094 # Convert to dense for cosm
2095 dense_data = qarr.to_dense().data
2096 data = cosm_data(dense_data)
2097 return Qarray.create(data, dims=dims)
2100def sinm_data(data: Array, **kwargs) -> Array:
2101 """Matrix sine of a raw array.
2103 Args:
2104 data: Dense matrix array.
2105 **kwargs: Unused; kept for API consistency.
2107 Returns:
2108 The matrix sine computed as ``(expm(i*A) - expm(-i*A)) / (2i)``.
2109 """
2110 return (expm_data(1j * data) - expm_data(-1j * data)) / (2j)
2113def sinm(qarr: Qarray) -> Qarray:
2114 """Matrix sine of a ``Qarray``.
2116 Args:
2117 qarr: Input quantum array (converted to dense internally).
2119 Returns:
2120 A dense ``Qarray`` containing the matrix sine.
2121 """
2122 dims = qarr.dims
2123 # Convert to dense for sinm
2124 dense_data = qarr.to_dense().data
2125 data = sinm_data(dense_data)
2126 return Qarray.create(data, dims=dims)
2129def keep_only_diag_elements(qarr: Qarray) -> Qarray:
2130 """Zero out all off-diagonal elements of *qarr*.
2132 For sparse ``Qarray`` objects the off-diagonal stored values are zeroed
2133 in-place on the BCOO structure — no densification.
2135 Args:
2136 qarr: Non-batched input quantum array.
2138 Returns:
2139 A ``Qarray`` with only diagonal entries non-zero.
2141 Raises:
2142 ValueError: If *qarr* has batch dimensions.
2143 """
2144 if len(qarr.bdims) > 0:
2145 raise ValueError("Cannot keep only diagonal elements of a batched Qarray.")
2147 dims = qarr.dims
2148 if qarr.is_sparse_bcoo:
2149 new_impl = qarr._impl.keep_only_diag()
2150 return Qarray.create(new_impl.data, dims=dims, implementation=QarrayImplType.SPARSE_BCOO)
2151 if qarr.is_sparse_dia:
2152 from jaxquantum.core.sparse_dia import SparseDiaImpl
2153 impl = qarr._impl
2154 n = impl._diags.shape[-1]
2155 if 0 in impl._offsets:
2156 i = impl._offsets.index(0)
2157 main_diag = impl._diags[..., i:i + 1, :]
2158 else:
2159 main_diag = jnp.zeros((*impl._diags.shape[:-2], 1, n), dtype=impl._diags.dtype)
2160 new_impl = SparseDiaImpl(_offsets=(0,), _diags=main_diag)
2161 return Qarray.create(new_impl.get_data(), dims=dims, implementation=QarrayImplType.SPARSE_DIA)
2162 data = jnp.diag(jnp.diag(qarr.data))
2163 return Qarray.create(data, dims=dims)
2166def to_ket(qarr: Qarray) -> Qarray:
2167 """Convert *qarr* to a ket.
2169 Args:
2170 qarr: A ket (returned as-is) or bra (conjugate-transposed).
2172 Returns:
2173 The ket form of *qarr*.
2175 Raises:
2176 ValueError: If *qarr* is an operator.
2177 """
2178 if qarr.qtype == Qtypes.ket:
2179 return qarr
2180 elif qarr.qtype == Qtypes.bra:
2181 return qarr.dag()
2182 else:
2183 raise ValueError("Can only get ket from a ket or bra.")
2186def eigenstates(qarr: Qarray) -> Qarray:
2187 """Eigenstates of a quantum array.
2189 Args:
2190 qarr: Hermitian operator (converted to dense internally).
2192 Returns:
2193 A tuple ``(eigenvalues, eigenstates_qarray)`` where eigenvalues are
2194 sorted in ascending order.
2195 """
2196 # Convert to dense for eigenstates
2197 dense_qarr = qarr.to_dense()
2199 evals, evecs = jnp.linalg.eigh(dense_qarr.data)
2200 dims = ket_from_op_dims(qarr.dims)
2202 # numpy returns [batch, :, i] as the i-th eigenvector
2203 # we want [batch, i, :] as the i-th eigenvector
2204 evecs = jnp.swapaxes(evecs, -2, -1)
2205 evecs = Qarray._from_impl(
2206 DenseImpl._make(evecs),
2207 Qdims(dims),
2208 )
2210 return evals, evecs
2213def eigenenergies(qarr: Qarray) -> Array:
2214 """Eigenvalues of a quantum array.
2216 Args:
2217 qarr: Hermitian operator (converted to dense internally).
2219 Returns:
2220 Sorted eigenvalues as a JAX array.
2221 """
2222 # Convert to dense for eigenenergies
2223 dense_qarr = qarr.to_dense()
2224 evals = jnp.linalg.eigvalsh(dense_qarr.data)
2225 return evals
2228def ptrace(qarr: Qarray, indx) -> Qarray:
2229 """Partial trace over subsystem *indx*.
2231 Args:
2232 qarr: Input quantum array (converted to dense internally).
2233 indx: Index of the subsystem to trace out.
2235 Returns:
2236 Reduced density matrix as a ``Qarray``.
2237 """
2238 # Convert to dense for ptrace
2239 dense_qarr = qarr.to_dense()
2240 dense_qarr = ket2dm(dense_qarr)
2241 rho = dense_qarr.shaped_data
2242 dims = dense_qarr.dims
2244 Nq = len(dims[0])
2246 indxs = [indx, indx + Nq]
2247 for j in range(Nq):
2248 if j == indx:
2249 continue
2250 indxs.append(j)
2251 indxs.append(j + Nq)
2253 bdims = dense_qarr.bdims
2254 len_bdims = len(bdims)
2255 bdims_indxs = list(range(len_bdims))
2256 indxs = bdims_indxs + [j + len_bdims for j in indxs]
2257 rho = rho.transpose(indxs)
2259 for j in range(Nq - 1):
2260 rho = jnp.trace(rho, axis1=2 + len_bdims, axis2=3 + len_bdims)
2262 return Qarray.create(rho)
2265def dag(qarr: Qarray) -> Qarray:
2266 """Conjugate transpose of *qarr*.
2268 For ket/bra vectors (stored on a single axis) this is just a complex
2269 conjugate with the ``dims`` reversed — no axis swap, so the data stays
2270 1‑D. For operators it is the usual conjugate transpose.
2272 Args:
2273 qarr: Input quantum array.
2275 Returns:
2276 The conjugate transpose with swapped ``dims``.
2277 """
2278 dims = qarr.dims[::-1]
2279 if qarr.qtype in (Qtypes.ket, Qtypes.bra):
2280 new_impl = qarr._impl.conj()
2281 else:
2282 new_impl = qarr._impl.dag()
2283 # Infer batch dimensions from the transformed data because ``_bdims`` is
2284 # static PyTree metadata and can be stale under ``vmap``.
2285 return Qarray._from_impl(new_impl, Qdims(dims))
2288def dag_data(arr) -> Array:
2289 """Conjugate transpose of a raw array, dispatching to the right backend.
2291 Iterates through registered :class:`QarrayImpl` subclasses and delegates
2292 to the first one whose :meth:`~QarrayImpl.can_handle_data` returns True.
2293 Adding a new backend automatically extends this function — no changes
2294 required here.
2296 Args:
2297 arr: Input array (``jnp.ndarray``, ``sparse.BCOO``, or any type
2298 handled by a registered impl). For 1-D dense arrays only
2299 conjugation is applied (no transpose).
2301 Returns:
2302 Conjugate transpose with the last two axes swapped.
2304 Raises:
2305 TypeError: If no registered impl can handle *arr*.
2306 """
2307 for impl_class in _IMPL_REGISTRY:
2308 if impl_class.can_handle_data(arr):
2309 return impl_class.dag_data(arr)
2310 raise TypeError(f"dag_data: no registered impl can handle type {type(arr)}")
2313def ket2dm(qarr: Qarray) -> Qarray:
2314 """Convert a ket to a density matrix via outer product.
2316 Args:
2317 qarr: Ket, bra, or operator. Operators are returned unchanged.
2319 Returns:
2320 Density matrix ``|ψ⟩⟨ψ|``.
2321 """
2322 if qarr.qtype == Qtypes.oper:
2323 return qarr
2325 if qarr.qtype == Qtypes.bra:
2326 qarr = qarr.dag()
2328 return qarr @ qarr.dag()
2331# Data level operations
2332def is_dm_data(data: Array) -> bool:
2333 """Check whether *data* has the shape of a density matrix (square matrix).
2335 Args:
2336 data: Array to check.
2338 Returns:
2339 True if the last two dimensions are equal.
2340 """
2341 return data.shape[-2] == data.shape[-1]
2344def powm_data(data: Array, n: int) -> Array:
2345 """Integer matrix power of a raw array.
2347 Args:
2348 data: Dense square matrix array.
2349 n: Integer exponent.
2351 Returns:
2352 The *n*-th matrix power.
2353 """
2354 return jnp.linalg.matrix_power(data, n)
2357# Type aliases for readability
2358DenseQarray = Qarray[DenseImpl]
2359# SparseBCOOQarray and SparseDIAQarray are defined lazily (impls imported at runtime)
2360# Use Qarray[SparseBCOOImpl] / Qarray[SparseDiaImpl] once those modules are imported.
2362ARRAY_TYPES = (Array, ndarray, Qarray)