diff --git a/emu_base/__init__.py b/emu_base/__init__.py index 93b25677..2cc01ac4 100644 --- a/emu_base/__init__.py +++ b/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" diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py new file mode 100644 index 00000000..63691ca1 --- /dev/null +++ b/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 diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index e368967e..8b6b237d 100644 --- a/emu_mps/mps_backend_impl.py +++ b/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, diff --git a/emu_mps/solver_utils.py b/emu_mps/solver_utils.py index 53329da0..4e2ad756 100644 --- a/emu_mps/solver_utils.py +++ b/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 diff --git a/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py new file mode 100644 index 00000000..8bda7ef4 --- /dev/null +++ b/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) diff --git a/test/emu_mps/test_mps_backend_impl.py b/test/emu_mps/test_mps_backend_impl.py index 7ba336e4..ab34e28d 100644 --- a/test/emu_mps/test_mps_backend_impl.py +++ b/test/emu_mps/test_mps_backend_impl.py @@ -11,7 +11,7 @@ ) from emu_base import HamiltonianType -from emu_mps.mps_config import MPSConfig +from emu_mps import MPSConfig, MPO from pulser import NoiseModel import math import cmath @@ -581,9 +581,11 @@ def test_progress_at_right_mps_boundary( def test_left_to_right_update( mock_right_baths, mock_make_H, mock_update_H, mock_minimize, mock_new_left ): - mock_make_H.return_value = MagicMock(factors=[None] * QUBIT_COUNT) + mock_make_H.return_value = MPO( + [torch.eye(2, 2, dtype=dtype).reshape(1, 2, 2, 1)] * QUBIT_COUNT + ) mock_update_H.return_value = None - mock_right_baths.return_value = [torch.zeros(1)] * (QUBIT_COUNT - 1) + mock_right_baths.return_value = [torch.zeros(1, 1, 1)] * (QUBIT_COUNT - 1) mock_new_left.return_value = torch.zeros(1) mock_minimize.return_value = (torch.tensor([[1.0]]), torch.tensor([[2.0]]), 0.5) @@ -591,8 +593,8 @@ def test_left_to_right_update( dmrg.init() dmrg._sweep_index = 1 dmrg._swipe_direction = SwipeDirection.LEFT_TO_RIGHT - dmrg.left_baths = [torch.zeros(1)] - dmrg.right_baths = [torch.zeros(1)] * 3 + dmrg.left_baths = [torch.zeros(1, 1, 1)] + dmrg.right_baths = [torch.zeros(1, 1, 1)] * 3 dmrg._left_to_right_update(idx=1) @@ -615,7 +617,9 @@ def test_right_to_left_update( ) mock_update_H.return_value = None - mock_right_baths.return_value = [torch.zeros(1, dtype=dtype)] * (QUBIT_COUNT - 1) + mock_right_baths.return_value = [torch.zeros(1, 1, 1, dtype=dtype)] * ( + QUBIT_COUNT - 1 + ) mock_new_left.return_value = torch.zeros(1, 1, 1, dtype=dtype) mock_minimize.return_value = ( diff --git a/test/emu_mps/test_solver_utils.py b/test/emu_mps/test_solver_utils.py index f838d36b..cf4c1386 100644 --- a/test/emu_mps/test_solver_utils.py +++ b/test/emu_mps/test_solver_utils.py @@ -55,7 +55,7 @@ def test_right_baths_total_magnetization(): state = MPS([mps_factor] * 3, eigenstates=("0", "1")) obs = MPO([mpo_factor1, mpo_factor2, mpo_factor3]) - baths = right_baths(state, obs, 1) + baths = list(right_baths(state, obs, 1)) # The baths carry the information of the magnetization, so the baths have shape # (1,2,1), and L_i = [-i,1], which basically counts how magnetized the bath is. assert torch.allclose(