Bath compression

emu_base/__init__.py

@@ -1,26 +1,28 @@
 from .constants import DEVICE_COUNT
-from .pulser_adapter import PulserData, HamiltonianType, SequenceData
-from .math.brents_root_finding import find_root_brents
-from .math.krylov_exp import krylov_exp, DEFAULT_MAX_KRYLOV_DIM
 from .jump_lindblad_operators import compute_noise_from_lindbladians
+from .math.brents_root_finding import find_root_brents
+from .math.krylov_exp import DEFAULT_MAX_KRYLOV_DIM, krylov_exp
 from .math.matmul import matmul_2x2_with_batched
-from .utils import get_max_rss, apply_measurement_errors, unix_like, init_logging
+from .math.packed_tensor import PackedHermitianTensor
+from .pulser_adapter import HamiltonianType, PulserData, SequenceData
+from .utils import apply_measurement_errors, get_max_rss, init_logging, unix_like
 
 __all__ = [
     "__version__",
-    "get_max_rss",
-    "compute_noise_from_lindbladians",
-    "matmul_2x2_with_batched",
+    "DEFAULT_MAX_KRYLOV_DIM",
+    "DEVICE_COUNT",
+    "HamiltonianType",
+    "PackedHermitianTensor",
     "PulserData",
     "SequenceData",
+    "apply_measurement_errors",
+    "compute_noise_from_lindbladians",
     "find_root_brents",
+    "get_max_rss",
+    "init_logging",
     "krylov_exp",
-    "HamiltonianType",
-    "DEFAULT_MAX_KRYLOV_DIM",
-    "DEVICE_COUNT",
-    "apply_measurement_errors",
+    "matmul_2x2_with_batched",
     "unix_like",
-    "init_logging",
 ]
 
 __version__ = "2.9.1"

emu_base/math/packed_tensor.py

@@ -0,0 +1,61 @@
+import torch
+
+
+class PackedHermitianTensor:
+    """
+    Pack a tensor of shape (χ, m, χ), Hermitian in axes 0 and 2,
+    into shape (χ(χ+1)/2, m) by storing the lower triangle
+    of each (χ, χ) slice at fixed middle index m.
+    The `PackedHermitianTensor` is used to represent left and right
+    baths nodes in TDVP/DMRG algorithms.
+    """
+
+    def __init__(
+        self,
+        chi: int,
+        m: int,
+        *,
+        check_hermitian: bool = True,
+        rtol: float = 1e-5,
+        atol: float = 1e-8,
+    ) -> None:
+        self._packed_data = torch.zeros(
+            int(chi * (chi + 1) / 2), m, dtype=torch.complex128
+        )
+        self.check_hermitian = check_hermitian
+        self.rtol = rtol
+        self.atol = atol
+        self.chi = chi
+        self.m = m
+
+    def pack(self, h: torch.Tensor) -> None:
+        if h.ndim != 3 or h.shape[0] != h.shape[2]:
+            raise ValueError(f"Expected shape (χ, m, χ), got {tuple(h.shape)}")
+        if self.chi != h.shape[0] or self.m != h.shape[1]:
+            raise ValueError(
+                f"Initialized for ({self.chi},{self.m},{self.chi}), got {h.shape}"
+            )
+
+        if self.check_hermitian and not torch.allclose(
+            h, h.transpose(0, 2).conj(), rtol=self.rtol, atol=self.atol
+        ):
+            raise ValueError("Tensor is not Hermitian in axes 0 and 2")
+
+        self._ii, self._kk = torch.tril_indices(self.chi, self.chi, device=h.device)
+        self._packed_data[:, :] = h[self._ii, :, self._kk]
+        self._packed_data = self._packed_data.to(h.device)
+
+    def unpack(self) -> torch.Tensor:
+        vals = self._packed_data
+        m = vals.shape[1]
+
+        h = torch.zeros(
+            (self.chi, m, self.chi),
+            dtype=self._packed_data.dtype,
+            device=self._packed_data.device,
+        )
+        h[self._ii, :, self._kk] = vals
+
+        offdiag = self._ii != self._kk
+        h[self._kk[offdiag], :, self._ii[offdiag]] = vals[offdiag].conj()
+        return h

emu_mps/mps_backend_impl.py

@@ -17,7 +17,13 @@
 import torch
 from pulser.backend import EmulationConfig, Observable, Results, State, AggregationMethod
 
-from emu_base import DEVICE_COUNT, SequenceData, get_max_rss, HamiltonianType
+from emu_base import (
+    DEVICE_COUNT,
+    SequenceData,
+    get_max_rss,
+    HamiltonianType,
+    PackedHermitianTensor,
+)
 from emu_base.math.brents_root_finding import BrentsRootFinder
 from emu_base.utils import deallocate_tensor
 
@@ -43,6 +49,7 @@
 )
 
 dtype = torch.complex128
+BathNode = torch.Tensor | PackedHermitianTensor
 
 
 class Statistics(Observable):
@@ -104,8 +111,8 @@ class MPSBackendImpl:
     well_prepared_qubits_filter: Optional[torch.Tensor]
     hamiltonian: MPO
     state: MPS
-    left_baths: list[torch.Tensor]
-    right_baths: list[torch.Tensor]
+    left_baths: list[BathNode]
+    right_baths: list[BathNode]
     target_time: float
     results: Results
     _swipe_direction: SwipeDirection = SwipeDirection.LEFT_TO_RIGHT
@@ -335,20 +342,47 @@ def update_H_no_noise(self) -> None:
         )
 
     def init_baths(self) -> None:
-        self.left_baths = [
-            torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device)
-        ]
         # clean up the memory in right baths, since otherwise we temporarily have
         # both the old and the new in memory, which can cause OOM errors
-        self.right_baths: list[torch.Tensor] = []
-        self.right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2)
+        self.right_baths = []
+
+        if self.has_lindblad_noise:
+            self.left_baths = [
+                torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device)
+            ]
+            self.right_baths = [
+                t for t in right_baths(self.state, self.hamiltonian, final_qubit=2)
+            ]
+        else:
+            self.left_baths = [PackedHermitianTensor(1, 1)]
+            self.left_baths[0].pack(  # type: ignore[union-attr]
+                torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device)
+            )
+            self.right_baths = [PackedHermitianTensor(1, 1)]
+            self.right_baths += [
+                PackedHermitianTensor(
+                    self.state.factors[i].shape[0], self.hamiltonian.factors[i].shape[0]
+                )
+                for i in range(len(self.state.factors) - 1, 1, -1)
+            ]
+            for i, b in enumerate(
+                right_baths(self.state, self.hamiltonian, final_qubit=2)
+            ):
+                self.right_baths[i].pack(b)  # type: ignore[union-attr]
+
         assert len(self.right_baths) == self.qubit_count - 1
 
     def get_current_right_bath(self) -> torch.Tensor:
-        return self.right_baths[-1]
+        rbath = self.right_baths[-1]
+        if isinstance(rbath, PackedHermitianTensor):
+            return rbath.unpack()
+        return rbath
 
     def get_current_left_bath(self) -> torch.Tensor:
-        return self.left_baths[-1]
+        lbath = self.left_baths[-1]
+        if isinstance(lbath, PackedHermitianTensor):
+            return lbath.unpack()
+        return lbath
 
     def init(self) -> None:
         self.init_dark_qubits()
@@ -451,16 +485,30 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None:
                 dt=delta_time / 2,
                 orth_center_right=True,
             )
-            self.left_baths.append(
-                new_left_bath(
+            if self.has_lindblad_noise:
+                lb = new_left_bath(
                     self.get_current_left_bath(),
                     self.state.factors[self._sweep_index],
                     self.hamiltonian.factors[self._sweep_index],
                 ).to(self.state.factors[self._sweep_index + 1].device)
-            )
+                self.left_baths.append(lb)
+            else:
+                packed = PackedHermitianTensor(
+                    self.state.factors[self._sweep_index].shape[-1],
+                    self.hamiltonian.factors[self._sweep_index].shape[-1],
+                )
+                lb = new_left_bath(
+                    self.get_current_left_bath(),
+                    self.state.factors[self._sweep_index],
+                    self.hamiltonian.factors[self._sweep_index],
+                ).to(self.state.factors[self._sweep_index + 1].device)
+                packed.pack(lb)
+                self.left_baths.append(packed)
+
             self._evolve(self._sweep_index + 1, dt=-delta_time / 2)
             self.right_baths.pop()
             self._sweep_index += 1
+
         else:
             # Time-evolution of the rightmost 2 tensors
             self._evolve(
@@ -473,16 +521,34 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None:
 
     def _right_to_left_update_tdvp(self, delta_time: float) -> None:
         if self._sweep_index > 0:
-            self.right_baths.append(
-                new_right_bath(
+            if self.has_lindblad_noise:
+                rb = new_right_bath(
                     self.get_current_right_bath(),
                     self.state.factors[self._sweep_index + 1],
                     self.hamiltonian.factors[self._sweep_index + 1],
                 ).to(self.state.factors[self._sweep_index].device)
-            )
+                self.right_baths.append(rb)
+            else:
+                packed = PackedHermitianTensor(
+                    self.state.factors[self._sweep_index + 1].shape[0],
+                    self.hamiltonian.factors[self._sweep_index + 1].shape[0],
+                )
+                rb = new_right_bath(
+                    self.get_current_right_bath(),
+                    self.state.factors[self._sweep_index + 1],
+                    self.hamiltonian.factors[self._sweep_index + 1],
+                ).to(self.state.factors[self._sweep_index].device)
+                packed.pack(rb)
+                self.right_baths.append(packed)
+
             if not self.has_lindblad_noise:
                 # Free memory because it won't be used anymore
-                deallocate_tensor(self.right_baths[-2])
+                item = self.right_baths[-2]
+                to_dealloc = (
+                    item._packed_data if isinstance(item, PackedHermitianTensor) else item
+                )
+                deallocate_tensor(to_dealloc)
+
             self._evolve(self._sweep_index, dt=-delta_time / 2)
             self.left_baths.pop()
             self._evolve(
@@ -802,10 +868,12 @@ def progress(self) -> None:
         ), "Unknown Swipe direction"
 
         orth_center_right = self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT
+        lbath = self.get_current_left_bath()
+        rbath = self.get_current_right_bath()
         new_L, new_R, energy = minimize_energy_pair(
             state_factors=self.state.factors[idx : idx + 2],
             ham_factors=self.hamiltonian.factors[idx : idx + 2],
-            baths=(self.left_baths[-1], self.right_baths[-1]),
+            baths=(lbath, rbath),
             orth_center_right=orth_center_right,
             config=self.config,
             residual_tolerance=self.config.precision,

emu_mps/solver_utils.py

@@ -1,5 +1,5 @@
 import torch
-from typing import Callable, Sequence
+from typing import Callable, Sequence, Iterator
 
 from emu_base import krylov_exp
 from emu_base.math.krylov_energy_min import krylov_energy_minimization
@@ -144,15 +144,14 @@ def new_right_bath(
 """
 
 
-def right_baths(state: MPS, op: MPO, final_qubit: int) -> list[torch.Tensor]:
+def right_baths(state: MPS, op: MPO, final_qubit: int) -> Iterator[torch.Tensor]:
     state_factor = state.factors[-1]
     bath = torch.ones(1, 1, 1, device=state_factor.device, dtype=state_factor.dtype)
-    baths = [bath]
+    yield bath
     for i in range(len(state.factors) - 1, final_qubit - 1, -1):
         bath = new_right_bath(bath, state.factors[i], op.factors[i])
         bath = bath.to(state.factors[i - 1].device)
-        baths.append(bath)
-    return baths
+        yield bath
 
 
 _TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/μs, dt in ns

test/emu_base/math/test_packed_tensor.py

@@ -0,0 +1,94 @@
+import pytest
+import torch
+from emu_base import PackedHermitianTensor
+
+
+def test_roundtrip():
+    n, b = 4, 3
+    x = torch.randn(n, b, n, dtype=torch.complex64)
+    h = 0.5 * (x + x.transpose(0, 2).conj())
+
+    packedht = PackedHermitianTensor(h)
+    h2 = packedht.unpack()
+
+    assert h2.shape == h.shape
+    assert torch.allclose(h2, h)
+
+
+def test_rejects_non_hermitian():
+    h = torch.randn(3, 2, 3, dtype=torch.complex64)
+
+    with pytest.raises(ValueError, match="not Hermitian"):
+        PackedHermitianTensor(h)
+
+
+def test_rejects_wrong_shape():
+    h = torch.randn(3, 2, 4)
+
+    with pytest.raises(ValueError, match="Expected shape"):
+        PackedHermitianTensor(h)
+
+
+def test_packed_shape():
+    n, b = 4, 3
+    x = torch.randn(n, b, n, dtype=torch.complex64)
+    h = 0.5 * (x + x.transpose(0, 2).conj())
+
+    packedht = PackedHermitianTensor(h)
+
+    assert packedht._packed_data.shape == (n * (n + 1) / 2, b)
+
+
+def test_unpack_preserves_dtype():
+    n, b = 4, 2
+    x = torch.randn(n, b, n, dtype=torch.complex128)
+    h = 0.5 * (x + x.transpose(0, 2).conj())
+
+    packedht = PackedHermitianTensor(h)
+    h2 = packedht.unpack()
+
+    assert h2.dtype == h.dtype
+
+
+def test_skip_hermitian_check():
+    # check_hermitian=False is needed for hot paths: this test ensures the
+    # class accepts structurally valid input without paying for symmetry checks.
+    h = torch.randn(3, 2, 3, dtype=torch.complex64)
+
+    # the tensor above should not pass for `check_hermitian=True`
+    packedht = PackedHermitianTensor(h, check_hermitian=False)
+
+    # Data loss happened
+    assert packedht._packed_data.shape == (6, 2)
+
+
+def test_packed_is_contiguous():
+    n, b = 4, 3
+    x = torch.randn(n, b, n, dtype=torch.complex64)
+    h = 0.5 * (x + x.transpose(0, 2).conj())
+
+    packedht = PackedHermitianTensor(h)
+
+    assert packedht._packed_data.is_contiguous()
+
+
+def test_custom_tolerance_controls_hermitian_check():
+    n, b = 4, 2
+    x = torch.randn(n, b, n, dtype=torch.complex64)
+    h = 0.5 * (x + x.transpose(0, 2).conj())
+
+    h_perturbed = h.clone()
+    h_perturbed[0, 0, 1] += 1e-4
+
+    with pytest.raises(ValueError, match="not Hermitian"):
+        # Not convertable within given tolerance
+        PackedHermitianTensor(h_perturbed, atol=1e-8)
+
+    PackedHermitianTensor(h_perturbed, atol=1e-3)
+
+
+def test_rejects_non_3d_input():
+    h = torch.randn(3, 3)
+
+    with pytest.raises(ValueError, match="Expected shape"):
+        PackedHermitianTensor(h)