From e40320f91d7f60c001ee5c4c5f63e5d9f237e7af Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Thu, 12 Mar 2026 10:50:30 +0100 Subject: [PATCH 01/19] minor preparation --- emu_mps/mps_backend_impl.py | 117 ++++++++++++++++---------------- test/emu_mps/test_end_to_end.py | 6 +- 2 files changed, 61 insertions(+), 62 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index a5a92c40..7bc4ff2a 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -101,12 +101,13 @@ class MPSBackendImpl: well_prepared_qubits_filter: Optional[torch.Tensor] hamiltonian: MPO state: MPS + left_baths: list[torch.Tensor] right_baths: list[torch.Tensor] - sweep_index: int - swipe_direction: SwipeDirection - timestep_index: int target_time: float results: Results + _swipe_direction = SwipeDirection.LEFT_TO_RIGHT + _sweep_index: int = 0 + _timestep_index: int = 0 def __init__(self, mps_config: MPSConfig, pulser_data: SequenceData): self.config = mps_config @@ -137,11 +138,8 @@ def __init__(self, mps_config: MPSConfig, pulser_data: SequenceData): ) self.hamiltonian_type = pulser_data.hamiltonian_type - self.left_baths: list[torch.Tensor] self.time = time.time() - self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT - self.sweep_index = 0 - self.timestep_index = 0 + self.results = Results( atom_order=optimat.permute_tuple( pulser_data.qubit_ids, self.qubit_permutation @@ -293,18 +291,18 @@ def init_noiseless_hamiltonian(self) -> None: def update_H(self) -> None: update_H( hamiltonian=self.hamiltonian, - omega=self.omega[self.timestep_index, :], - delta=self.delta[self.timestep_index, :], - phi=self.phi[self.timestep_index, :], + omega=self.omega[self._timestep_index, :], + delta=self.delta[self._timestep_index, :], + phi=self.phi[self._timestep_index, :], noise=self.lindblad_noise, ) def update_H_no_noise(self) -> None: update_H( hamiltonian=self.hamiltonian, - omega=self.omega[self.timestep_index, :], - delta=self.delta[self.timestep_index, :], - phi=self.phi[self.timestep_index, :], + omega=self.omega[self._timestep_index, :], + delta=self.delta[self._timestep_index, :], + phi=self.phi[self._timestep_index, :], noise=torch.zeros(self.dim, self.dim, dtype=dtype), # no noise ) @@ -330,7 +328,7 @@ def init(self) -> None: self.init_baths() def is_finished(self) -> bool: - return self.timestep_index >= self.timestep_count + return self._timestep_index >= self.timestep_count def _evolve( self, *indices: int, dt: float, orth_center_right: Optional[bool] = None @@ -383,8 +381,8 @@ def progress(self) -> None: """ Do one unit of simulation work given the current state. Update the state accordingly. - The state of the simulation is stored in self.sweep_index and - self.swipe_direction. + The state of the simulation is stored in self.__sweep_index and + self._swipe_direction. """ if self.is_finished(): return @@ -394,8 +392,8 @@ def progress(self) -> None: assert self.qubit_count >= 1 if 1 <= self.qubit_count <= 2: # Corner case: only 1 or 2 qubits - assert self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT - assert self.sweep_index == 0 + assert self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT + assert self._sweep_index == 0 if self.qubit_count == 1: self._evolve(0, dt=delta_time) @@ -405,70 +403,71 @@ def progress(self) -> None: self.sweep_complete() elif ( - self.sweep_index < self.qubit_count - 2 - and self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT + self._sweep_index < self.qubit_count - 2 + and self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT ): # Left-to-right swipe of TDVP self._evolve( - self.sweep_index, - self.sweep_index + 1, + self._sweep_index, + self._sweep_index + 1, dt=delta_time / 2, orth_center_right=True, ) self.left_baths.append( 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.state.factors[self._sweep_index], + self.hamiltonian.factors[self._sweep_index], + ).to(self.state.factors[self._sweep_index + 1].device) ) - self._evolve(self.sweep_index + 1, dt=-delta_time / 2) + self._evolve(self._sweep_index + 1, dt=-delta_time / 2) self.right_baths.pop() - self.sweep_index += 1 + self._sweep_index += 1 elif ( - self.sweep_index == self.qubit_count - 2 - and self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT + self._sweep_index == self.qubit_count - 2 + and self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT ): # Time-evolution of the rightmost 2 tensors self._evolve( - self.sweep_index, - self.sweep_index + 1, + self._sweep_index, + self._sweep_index + 1, dt=delta_time, orth_center_right=False, ) - self.swipe_direction = SwipeDirection.RIGHT_TO_LEFT + self._swipe_direction = SwipeDirection.RIGHT_TO_LEFT elif ( - 1 <= self.sweep_index and self.swipe_direction == SwipeDirection.RIGHT_TO_LEFT + 1 <= self._sweep_index + and self._swipe_direction == SwipeDirection.RIGHT_TO_LEFT ): # Right-to-left swipe of TDVP - assert self.sweep_index <= self.qubit_count - 2 + assert self._sweep_index <= self.qubit_count - 2 self.right_baths.append( 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.state.factors[self._sweep_index + 1], + self.hamiltonian.factors[self._sweep_index + 1], + ).to(self.state.factors[self._sweep_index].device) ) if not self.has_lindblad_noise: # Free memory because it won't be used anymore deallocate_tensor(self.right_baths[-2]) - self._evolve(self.sweep_index, dt=-delta_time / 2) + self._evolve(self._sweep_index, dt=-delta_time / 2) self.left_baths.pop() self._evolve( - self.sweep_index - 1, - self.sweep_index, + self._sweep_index - 1, + self._sweep_index, dt=delta_time / 2, orth_center_right=False, ) - self.sweep_index -= 1 + self._sweep_index -= 1 - if self.sweep_index == 0: + if self._sweep_index == 0: self.sweep_complete() - self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT + self._swipe_direction = SwipeDirection.LEFT_TO_RIGHT else: raise Exception("Didn't expect this") @@ -481,7 +480,7 @@ def sweep_complete(self) -> None: def timestep_complete(self) -> None: self.fill_results() - self.timestep_index += 1 + self._timestep_index += 1 interaction_matrix = self._get_interaction_matrix() # at new time is_the_same_matrix = torch.allclose( @@ -499,7 +498,7 @@ def timestep_complete(self) -> None: ) if not self.is_finished(): - self.target_time = self.target_times[self.timestep_index + 1] + self.target_time = self.target_times[self._timestep_index + 1] self.update_H() self.init_baths() @@ -712,7 +711,7 @@ def sweep_complete(self) -> None: if self.root_finder.is_converged(tolerance=1): self.do_random_quantum_jump() - self.target_time = self.target_times[self.timestep_index + 1] + self.target_time = self.target_times[self._timestep_index + 1] self.root_finder = None else: self.target_time = self.root_finder.get_next_abscissa() @@ -773,13 +772,13 @@ def progress(self) -> None: return # perform one two-site energy minimization and update - idx = self.sweep_index - assert self.swipe_direction in ( + idx = self._sweep_index + assert self._swipe_direction in ( SwipeDirection.LEFT_TO_RIGHT, SwipeDirection.RIGHT_TO_LEFT, ), "Unknown Swipe direction" - orth_center_right = self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT + orth_center_right = self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT new_L, new_R, energy = minimize_energy_pair( state_factors=self.state.factors[idx : idx + 2], ham_factors=self.hamiltonian.factors[idx : idx + 2], @@ -793,9 +792,9 @@ def progress(self) -> None: self.current_energy = energy # updating baths and orthogonality center - if self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT: + if self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT: self._left_to_right_update(idx) - elif self.swipe_direction == SwipeDirection.RIGHT_TO_LEFT: + elif self._swipe_direction == SwipeDirection.RIGHT_TO_LEFT: self._right_to_left_update(idx) else: raise Exception("Did not expect this") @@ -812,10 +811,10 @@ def _left_to_right_update(self, idx: int) -> None: ).to(self.state.factors[idx + 1].device) ) self.right_baths.pop() - self.sweep_index += 1 + self._sweep_index += 1 - if self.sweep_index == self.qubit_count - 2: - self.swipe_direction = SwipeDirection.RIGHT_TO_LEFT + if self._sweep_index == self.qubit_count - 2: + self._swipe_direction = SwipeDirection.RIGHT_TO_LEFT def _right_to_left_update(self, idx: int) -> None: if idx > 0: @@ -827,11 +826,11 @@ def _right_to_left_update(self, idx: int) -> None: ).to(self.state.factors[idx].device) ) self.left_baths.pop() - self.sweep_index -= 1 + self._sweep_index -= 1 - if self.sweep_index == 0: + if self._sweep_index == 0: self.state.orthogonalize(0) - self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT + self._swipe_direction = SwipeDirection.LEFT_TO_RIGHT self.sweep_count += 1 self.sweep_complete() @@ -847,9 +846,9 @@ def sweep_complete(self) -> None: # not converged for the current sweep. restart self.previous_energy = self.current_energy - assert self.sweep_index == 0 + assert self._sweep_index == 0 assert self.state.orthogonality_center == 0 - assert self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT + assert self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT self.current_energy = None diff --git a/test/emu_mps/test_end_to_end.py b/test/emu_mps/test_end_to_end.py index 4fec6076..3ac66190 100644 --- a/test/emu_mps/test_end_to_end.py +++ b/test/emu_mps/test_end_to_end.py @@ -677,8 +677,8 @@ def check_baths(impl: MPSBackendImpl): # the right baths administration happens properly when a quantum jump occurs. assert len(impl.right_baths) in [ - impl.state.num_sites - impl.sweep_index, - impl.state.num_sites - impl.sweep_index - 1, + impl.state.num_sites - impl._sweep_index, + impl.state.num_sites - impl._sweep_index - 1, ] expected_right_baths = right_baths( @@ -843,7 +843,7 @@ def save_simulation_mock_side_effect(self): self.last_save_time = time.time() + 999 return save_simulation_original(self) - assert self.timestep_index == 11 + assert self._timestep_index == 11 self.last_save_time = 0 # Trigger saving regardless of time save_simulation_original(self) From 725752e6540cf77738d47cee2eac6515f329e5da Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Thu, 12 Mar 2026 10:57:43 +0100 Subject: [PATCH 02/19] deep copy --- emu_mps/mps_backend_impl.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 7bc4ff2a..fd29fb25 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -103,6 +103,8 @@ class MPSBackendImpl: state: MPS left_baths: list[torch.Tensor] right_baths: list[torch.Tensor] + left_baths_compressed: list[torch.Tensor] + right_baths_compressed: list[torch.Tensor] target_time: float results: Results _swipe_direction = SwipeDirection.LEFT_TO_RIGHT @@ -311,6 +313,10 @@ def init_baths(self) -> None: torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device) ] self.right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2) + + self.left_baths_compressed = [t.clone() for t in self.left_baths] + self.right_baths_compressed = [t.clone() for t in self.right_baths] + assert len(self.right_baths) == self.qubit_count - 1 def get_current_right_bath(self) -> torch.Tensor: From adf7b310f592ce3ff965f56acb715ebdcf066fb0 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Fri, 13 Mar 2026 14:02:02 +0100 Subject: [PATCH 03/19] preliminary res --- emu_mps/mps_backend_impl.py | 103 +++++++++++++++++++++++++++++++++++- emu_mps/utils.py | 83 ++++++++++++++++++++++++++++- 2 files changed, 183 insertions(+), 3 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 71bd42e6..d43f8cf6 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -45,6 +45,25 @@ dtype = torch.complex128 +def tensor_memory_bytes(obj: list | tuple | torch.Tensor) -> float: + if isinstance(obj, torch.Tensor): + return obj.numel() * obj.element_size() + elif isinstance(obj, list): + return sum(tensor_memory_bytes(x) for x in obj) + elif isinstance(obj, tuple): + return sum(tensor_memory_bytes(x) for x in obj) + else: + return 0 + + +def human_bytes(n: float) -> str: + for unit in ["B", "KB", "MB", "GB", "TB"]: + if n < 1024: + return f"{n:.2f} {unit}" + n /= 1024 + return f"{n:.2f} PB" + + class Statistics(Observable): def __init__( self, @@ -103,6 +122,8 @@ class MPSBackendImpl: state: MPS left_baths: list[torch.Tensor] right_baths: list[torch.Tensor] + left_baths_compressed: list[list[torch.Tensor]] + right_baths_compressed: list[list[torch.Tensor]] target_time: float results: Results _swipe_direction: SwipeDirection = SwipeDirection.LEFT_TO_RIGHT @@ -312,8 +333,8 @@ def init_baths(self) -> None: ] self.right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2) - self.left_baths_compressed = [t.clone() for t in self.left_baths] - self.right_baths_compressed = [t.clone() for t in self.right_baths] + self.left_baths_compressed = [[t.clone()] for t in self.left_baths] + self.right_baths_compressed = [[t.clone()] for t in self.right_baths] assert len(self.right_baths) == self.qubit_count - 1 @@ -323,6 +344,12 @@ def get_current_right_bath(self) -> torch.Tensor: def get_current_left_bath(self) -> torch.Tensor: return self.left_baths[-1] + def get_current_right_bath_compressed(self) -> list[torch.Tensor]: + return self.right_baths_compressed[-1] + + def get_current_left_bath_compressed(self) -> list[torch.Tensor]: + return self.left_baths_compressed[-1] + def init(self) -> None: self.init_dark_qubits() self.init_initial_state(self.config.initial_state) @@ -413,6 +440,23 @@ def progress(self) -> None: else: self._right_to_left_update_tdvp(delta_time=delta_time) + # noncompressed = sum( + # tensor_memory_bytes(t) for t in [self.left_baths, self.right_baths] + # ) + # compressed = sum( + # tensor_memory_bytes(t) + # for t in [self.left_baths_compressed, self.right_baths_compressed] + # ) + + # ratio = compressed / noncompressed + + # hc = human_bytes(compressed) + # hnc = human_bytes(noncompressed) + # wf = human_bytes(tensor_memory_bytes(self.state.factors)) + # if (ratio > 1): + # print(f"MPS: {wf}, bath : {hnc}, compressed bath : {hc}") + # print(ratio) + self.save_simulation() def _left_to_right_update_tdvp(self, delta_time: float) -> None: @@ -434,6 +478,34 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None: self._evolve(self._sweep_index + 1, dt=-delta_time / 2) self.right_baths.pop() self._sweep_index += 1 + + # Operate with compressed bath + from emu_mps.utils import split_bath_node + + l, r = split_bath_node( + self.left_baths[-1], + max_error=self.config.precision, + max_rank=self.config.max_bond_dim, + orth_center_right=False, + preserve_norm=False, # only relevant for computing jump times + ) + self.left_baths_compressed.append(l) + self.right_baths_compressed.pop() + + assert len(self.left_baths_compressed) == len(self.left_baths) + assert len(self.right_baths_compressed) == len(self.right_baths) + + noncompressed = tensor_memory_bytes(self.left_baths[-1]) + compressed = tensor_memory_bytes(self.left_baths_compressed[-1]) + + hc = human_bytes(compressed) + hnc = human_bytes(noncompressed) + wf = human_bytes(tensor_memory_bytes(self.state.factors)) + # if (compressed / noncompressed > 1): + print( + f"step : {self._sweep_index} MPS: {wf}, bath : {hnc}, compressed bath : {hc}" + ) + else: # Time-evolution of the rightmost 2 tensors self._evolve( @@ -466,6 +538,33 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: ) self._sweep_index -= 1 + # Operate with compressed bath + from emu_mps.utils import split_bath_node + + l, r = split_bath_node( + self.right_baths[-1], + max_error=self.config.precision, + max_rank=self.config.max_bond_dim, + orth_center_right=False, + preserve_norm=False, # only relevant for computing jump times + ) + self.right_baths_compressed.append(l) + self.left_baths_compressed.pop() + + assert len(self.left_baths_compressed) == len(self.left_baths) + assert len(self.right_baths_compressed) == len(self.right_baths) + + noncompressed = tensor_memory_bytes(self.right_baths[-1]) + compressed = tensor_memory_bytes(self.right_baths_compressed[-1]) + + hc = human_bytes(compressed) + hnc = human_bytes(noncompressed) + wf = human_bytes(tensor_memory_bytes(self.state.factors)) + # if (compressed / noncompressed > 1): + print( + f"step : {self._sweep_index} MPS: {wf}, bath : {hnc}, compressed bath : {hc}" + ) + if self._sweep_index == 0: self.sweep_complete() self._swipe_direction = SwipeDirection.LEFT_TO_RIGHT diff --git a/emu_mps/utils.py b/emu_mps/utils.py index 6ead3b33..20a8625a 100644 --- a/emu_mps/utils.py +++ b/emu_mps/utils.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import List, Optional, cast import torch @@ -218,3 +218,84 @@ def tensor_trace(tensor: torch.Tensor, dim1: int, dim2: int) -> torch.Tensor: """ assert tensor.shape[dim1] == tensor.shape[dim2], "dimensions should match" return tensor.diagonal(offset=0, dim1=dim1, dim2=dim2).sum(-1) + + +def my_hermitean_svd( + m: torch.Tensor, + tol: float = 1e-5, +) -> torch.Tensor: + # Hermitian eigendecomposition + evals, Q = torch.linalg.eigh(m) + # keep only important modes + mask = evals.abs() >= tol + evals = evals[mask] + Q = Q[:, mask] + if evals.numel() == 0: + return cast(torch.Tensor, Q) + + evals = evals.to(Q) + sqrt_evals = torch.sqrt(evals) + + u_comp = Q * sqrt_evals.unsqueeze(0) + + # assert torch.allclose(m, u_comp @ u_comp.mH, atol=1e-3) + + return cast(torch.Tensor, u_comp) + + +def my_svd( + m: torch.Tensor, + tol: float = 1e-5, +) -> tuple[torch.Tensor, torch.Tensor]: + U, S, Vh = torch.linalg.svd(m, full_matrices=True) + + mask = S >= tol + + # S_sgn = torch.sgn(S) + # S = S @ S_sgn + # Vh = (S_sgn ).unsqueeze(1) * Vh + # pseudo_id = torch.abs(U @ torch.diag_embed(S) @ Vh) # up to sign + # id = torch.eye(pseudo_id.shape[0], dtype=pseudo_id.dtype) + # if(not torch.allclose(pseudo_id, id, atol=1e-5)): + # assert torch.allclose(pseudo_id, id, atol=1e-5) + + U_cut = U[:, mask] + S_cut = S[mask] + Vh_cut = Vh[mask, :] + + sqrt_S = torch.sqrt(S_cut) + U_cut = U_cut * sqrt_S.unsqueeze(0) + Vh_cut = (sqrt_S).unsqueeze(1) * Vh_cut + # assert torch.allclose(m, U @ Vh, atol=tol) + return U_cut, Vh_cut + + +def split_bath_node( + bath_node: torch.Tensor, + max_error: float = 1e-5, + max_rank: int = 1024, + orth_center_right: bool = True, + preserve_norm: bool = False, +) -> list[list[torch.Tensor]]: + assert bath_node.ndim == 3 + assert bath_node.shape[0] == bath_node.shape[2] + + slices = [bath_node[:, i, :] for i in range(bath_node.shape[1])] + + svd_slices_l: list[torch.Tensor] = [] + svd_slices_r: list[torch.Tensor] = [] + for s in slices: + assert torch.allclose(s, s.mH, atol=1e-5) + # u, v = my_svd(s, tol=max_error) + u_h = my_hermitean_svd(s, tol=max_error) + + # l_s, r_s = split_matrix(s, max_error, max_rank, orth_center_right, preserve_norm) + svd_slices_l.append(u_h) + svd_slices_r.append(u_h) + + # l: torch.Tensor = torch.stack(svd_slices_l, dim=0) + # r: torch.Tensor = torch.stack(svd_slices_r, dim=0) + assert len(svd_slices_l) == len(svd_slices_r) + assert len(svd_slices_l) == bath_node.shape[1] + + return [svd_slices_l, svd_slices_r] From 1b37a83b3432a6a173ac6c8ca2213dd659127c9a Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Mon, 23 Mar 2026 12:39:26 +0100 Subject: [PATCH 04/19] packed class --- emu_base/__init__.py | 26 ++--- emu_base/math/packed_tensor.py | 50 ++++++++++ test/emu_base/math/test_packed_tensor.py | 121 +++++++++++++++++++++++ 3 files changed, 185 insertions(+), 12 deletions(-) create mode 100644 emu_base/math/packed_tensor.py create mode 100644 test/emu_base/math/test_packed_tensor.py diff --git a/emu_base/__init__.py b/emu_base/__init__.py index b05a919b..a303df3b 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.7.2" diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py new file mode 100644 index 00000000..0ab60834 --- /dev/null +++ b/emu_base/math/packed_tensor.py @@ -0,0 +1,50 @@ +import torch + + +class PackedHermitianTensor: + """ + Pack a tensor of shape (a, b, a), Hermitian in axes 0 and 2, + into shape (b, a*(a+1)//2) by storing the lower triangle + of each (a, a) slice at fixed middle index. + + Lower-triangular packed order: + (0,0), + (1,0), (1,1), + (2,0), (2,1), (2,2), + ... + """ + + def __init__( + self, + h: torch.Tensor, + *, + check_hermitian: bool = True, + rtol: float = 1e-5, + atol: float = 1e-8, + ) -> None: + if h.ndim != 3 or h.shape[0] != h.shape[2]: + raise ValueError(f"Expected shape (n, b, n), got {tuple(h.shape)}") + + if check_hermitian and not torch.allclose( + h, h.transpose(0, 2).conj(), rtol=rtol, atol=atol + ): + raise ValueError("Tensor is not Hermitian in axes 0 and 2") + + self.n = h.shape[0] + self._ii, self._kk = torch.tril_indices(self.n, self.n, device=h.device) + self.packed = h[self._ii, :, self._kk].transpose(0, 1).contiguous() + + def unpack(self) -> torch.Tensor: + b = self.packed.shape[0] + vals = self.packed.transpose(0, 1) + + h = torch.zeros( + (self.n, b, self.n), + dtype=self.packed.dtype, + device=self.packed.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/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py new file mode 100644 index 00000000..abc5d4a8 --- /dev/null +++ b/test/emu_base/math/test_packed_tensor.py @@ -0,0 +1,121 @@ +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()) + + packed = PackedHermitianTensor(h) + h2 = packed.unpack() + + assert h2.shape == h.shape + assert torch.allclose(h2, h) + + +def test_roundtrip_minimal_shape(): + h = torch.tensor( + [ + [ + [ + 2.0 + 0.0j, + ] + ] + ] + ) # shape (1, 1, 1) + + packed = PackedHermitianTensor(h) + h2 = packed.unpack() + + assert packed.packed.shape == (1, 1) + 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()) + + packed = PackedHermitianTensor(h) + + assert packed.packed.shape == (b, n * (n + 1) // 2) + + +def test_roundtrip_real_symmetric(): + n, b = 4, 3 + x = torch.randn(n, b, n) + h = 0.5 * (x + x.transpose(0, 2)) + + packed = PackedHermitianTensor(h) + h2 = packed.unpack() + + assert h2.shape == h.shape + assert torch.allclose(h2, h) + + +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()) + + packed = PackedHermitianTensor(h) + h2 = packed.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) + + packed = PackedHermitianTensor(h, check_hermitian=False) + + assert packed.packed.shape == (2, 6) + + +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()) + + packed = PackedHermitianTensor(h) + + assert packed.packed.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"): + 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) From 785924a8533a5135b2ea3eb0de2fd62e0d68f8b8 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Tue, 31 Mar 2026 14:38:45 +0200 Subject: [PATCH 05/19] packed --- emu_base/math/packed_tensor.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 0ab60834..34407875 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -32,16 +32,16 @@ def __init__( self.n = h.shape[0] self._ii, self._kk = torch.tril_indices(self.n, self.n, device=h.device) - self.packed = h[self._ii, :, self._kk].transpose(0, 1).contiguous() + self._packed = h[self._ii, :, self._kk].transpose(0, 1).contiguous() def unpack(self) -> torch.Tensor: - b = self.packed.shape[0] - vals = self.packed.transpose(0, 1) + b = self._packed.shape[0] + vals = self._packed.transpose(0, 1) h = torch.zeros( (self.n, b, self.n), - dtype=self.packed.dtype, - device=self.packed.device, + dtype=self._packed.dtype, + device=self._packed.device, ) h[self._ii, :, self._kk] = vals From d781945227de13dc01afd5557bedae668aedf7eb Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Tue, 31 Mar 2026 15:41:20 +0200 Subject: [PATCH 06/19] rm unused functions --- emu_mps/mps_backend_impl.py | 174 +++++++++++------------------------- emu_mps/utils.py | 83 +---------------- 2 files changed, 51 insertions(+), 206 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index d43f8cf6..e6d5e9e6 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -12,12 +12,12 @@ from collections import Counter from enum import Enum, auto from types import MethodType -from typing import Any, Optional +from typing import Any, cast, Optional import torch from pulser.backend import EmulationConfig, Observable, Results, State -from emu_base import DEVICE_COUNT, SequenceData, get_max_rss +from emu_base import DEVICE_COUNT, SequenceData, get_max_rss, PackedHermitianTensor from emu_base.math.brents_root_finding import BrentsRootFinder from emu_base.utils import deallocate_tensor @@ -43,25 +43,7 @@ ) dtype = torch.complex128 - - -def tensor_memory_bytes(obj: list | tuple | torch.Tensor) -> float: - if isinstance(obj, torch.Tensor): - return obj.numel() * obj.element_size() - elif isinstance(obj, list): - return sum(tensor_memory_bytes(x) for x in obj) - elif isinstance(obj, tuple): - return sum(tensor_memory_bytes(x) for x in obj) - else: - return 0 - - -def human_bytes(n: float) -> str: - for unit in ["B", "KB", "MB", "GB", "TB"]: - if n < 1024: - return f"{n:.2f} {unit}" - n /= 1024 - return f"{n:.2f} PB" +Bath = torch.Tensor | PackedHermitianTensor class Statistics(Observable): @@ -120,10 +102,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_compressed: list[list[torch.Tensor]] - right_baths_compressed: list[list[torch.Tensor]] + left_baths: list[Bath] + right_baths: list[Bath] target_time: float results: Results _swipe_direction: SwipeDirection = SwipeDirection.LEFT_TO_RIGHT @@ -328,27 +308,33 @@ def update_H_no_noise(self) -> None: ) def init_baths(self) -> None: - self.left_baths = [ + _left_baths: list[torch.Tensor] = [ torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device) ] - self.right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2) + _right_baths: list[torch.Tensor] = right_baths( + self.state, self.hamiltonian, final_qubit=2 + ) - self.left_baths_compressed = [[t.clone()] for t in self.left_baths] - self.right_baths_compressed = [[t.clone()] for t in self.right_baths] + if isinstance(self, MPSBackendImpl): + self.left_baths = [PackedHermitianTensor(t) for t in _left_baths] + self.right_baths = [PackedHermitianTensor(t) for t in _right_baths] + else: + self.left_baths = list(_left_baths) + self.right_baths = list(_right_baths) 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] - - def get_current_right_bath_compressed(self) -> list[torch.Tensor]: - return self.right_baths_compressed[-1] - - def get_current_left_bath_compressed(self) -> list[torch.Tensor]: - return self.left_baths_compressed[-1] + lbath = self.left_baths[-1] + if isinstance(lbath, PackedHermitianTensor): + return lbath.unpack() + return lbath def init(self) -> None: self.init_dark_qubits() @@ -440,23 +426,6 @@ def progress(self) -> None: else: self._right_to_left_update_tdvp(delta_time=delta_time) - # noncompressed = sum( - # tensor_memory_bytes(t) for t in [self.left_baths, self.right_baths] - # ) - # compressed = sum( - # tensor_memory_bytes(t) - # for t in [self.left_baths_compressed, self.right_baths_compressed] - # ) - - # ratio = compressed / noncompressed - - # hc = human_bytes(compressed) - # hnc = human_bytes(noncompressed) - # wf = human_bytes(tensor_memory_bytes(self.state.factors)) - # if (ratio > 1): - # print(f"MPS: {wf}, bath : {hnc}, compressed bath : {hc}") - # print(ratio) - self.save_simulation() def _left_to_right_update_tdvp(self, delta_time: float) -> None: @@ -468,44 +437,20 @@ 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( - 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) - ) + 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) + item: Bath = lb + if isinstance(self, MPSBackendImpl): + item = PackedHermitianTensor(lb) + self.left_baths.append(item) + self._evolve(self._sweep_index + 1, dt=-delta_time / 2) self.right_baths.pop() self._sweep_index += 1 - # Operate with compressed bath - from emu_mps.utils import split_bath_node - - l, r = split_bath_node( - self.left_baths[-1], - max_error=self.config.precision, - max_rank=self.config.max_bond_dim, - orth_center_right=False, - preserve_norm=False, # only relevant for computing jump times - ) - self.left_baths_compressed.append(l) - self.right_baths_compressed.pop() - - assert len(self.left_baths_compressed) == len(self.left_baths) - assert len(self.right_baths_compressed) == len(self.right_baths) - - noncompressed = tensor_memory_bytes(self.left_baths[-1]) - compressed = tensor_memory_bytes(self.left_baths_compressed[-1]) - - hc = human_bytes(compressed) - hnc = human_bytes(noncompressed) - wf = human_bytes(tensor_memory_bytes(self.state.factors)) - # if (compressed / noncompressed > 1): - print( - f"step : {self._sweep_index} MPS: {wf}, bath : {hnc}, compressed bath : {hc}" - ) - else: # Time-evolution of the rightmost 2 tensors self._evolve( @@ -518,16 +463,22 @@ 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( - 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) - ) + 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) + item: Bath = rb + if isinstance(self, MPSBackendImpl): + item = PackedHermitianTensor(rb) + self.right_baths.append(item) + if not self.has_lindblad_noise: + # TODO this should be in Noise. Not in noiseless Base class # Free memory because it won't be used anymore - deallocate_tensor(self.right_baths[-2]) + to_deallocate: torch.Tensor = cast(torch.Tensor, self.right_baths[-2]) + deallocate_tensor(to_deallocate) + self._evolve(self._sweep_index, dt=-delta_time / 2) self.left_baths.pop() self._evolve( @@ -538,33 +489,6 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: ) self._sweep_index -= 1 - # Operate with compressed bath - from emu_mps.utils import split_bath_node - - l, r = split_bath_node( - self.right_baths[-1], - max_error=self.config.precision, - max_rank=self.config.max_bond_dim, - orth_center_right=False, - preserve_norm=False, # only relevant for computing jump times - ) - self.right_baths_compressed.append(l) - self.left_baths_compressed.pop() - - assert len(self.left_baths_compressed) == len(self.left_baths) - assert len(self.right_baths_compressed) == len(self.right_baths) - - noncompressed = tensor_memory_bytes(self.right_baths[-1]) - compressed = tensor_memory_bytes(self.right_baths_compressed[-1]) - - hc = human_bytes(compressed) - hnc = human_bytes(noncompressed) - wf = human_bytes(tensor_memory_bytes(self.state.factors)) - # if (compressed / noncompressed > 1): - print( - f"step : {self._sweep_index} MPS: {wf}, bath : {hnc}, compressed bath : {hc}" - ) - if self._sweep_index == 0: self.sweep_complete() self._swipe_direction = SwipeDirection.LEFT_TO_RIGHT @@ -874,10 +798,12 @@ def progress(self) -> None: ), "Unknown Swipe direction" orth_center_right = self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT + lbath: torch.Tensor = cast(torch.Tensor, self.left_baths[-1]) + rbath: torch.Tensor = cast(torch.Tensor, self.right_baths[-1]) 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/utils.py b/emu_mps/utils.py index 20a8625a..6ead3b33 100644 --- a/emu_mps/utils.py +++ b/emu_mps/utils.py @@ -1,4 +1,4 @@ -from typing import List, Optional, cast +from typing import List, Optional import torch @@ -218,84 +218,3 @@ def tensor_trace(tensor: torch.Tensor, dim1: int, dim2: int) -> torch.Tensor: """ assert tensor.shape[dim1] == tensor.shape[dim2], "dimensions should match" return tensor.diagonal(offset=0, dim1=dim1, dim2=dim2).sum(-1) - - -def my_hermitean_svd( - m: torch.Tensor, - tol: float = 1e-5, -) -> torch.Tensor: - # Hermitian eigendecomposition - evals, Q = torch.linalg.eigh(m) - # keep only important modes - mask = evals.abs() >= tol - evals = evals[mask] - Q = Q[:, mask] - if evals.numel() == 0: - return cast(torch.Tensor, Q) - - evals = evals.to(Q) - sqrt_evals = torch.sqrt(evals) - - u_comp = Q * sqrt_evals.unsqueeze(0) - - # assert torch.allclose(m, u_comp @ u_comp.mH, atol=1e-3) - - return cast(torch.Tensor, u_comp) - - -def my_svd( - m: torch.Tensor, - tol: float = 1e-5, -) -> tuple[torch.Tensor, torch.Tensor]: - U, S, Vh = torch.linalg.svd(m, full_matrices=True) - - mask = S >= tol - - # S_sgn = torch.sgn(S) - # S = S @ S_sgn - # Vh = (S_sgn ).unsqueeze(1) * Vh - # pseudo_id = torch.abs(U @ torch.diag_embed(S) @ Vh) # up to sign - # id = torch.eye(pseudo_id.shape[0], dtype=pseudo_id.dtype) - # if(not torch.allclose(pseudo_id, id, atol=1e-5)): - # assert torch.allclose(pseudo_id, id, atol=1e-5) - - U_cut = U[:, mask] - S_cut = S[mask] - Vh_cut = Vh[mask, :] - - sqrt_S = torch.sqrt(S_cut) - U_cut = U_cut * sqrt_S.unsqueeze(0) - Vh_cut = (sqrt_S).unsqueeze(1) * Vh_cut - # assert torch.allclose(m, U @ Vh, atol=tol) - return U_cut, Vh_cut - - -def split_bath_node( - bath_node: torch.Tensor, - max_error: float = 1e-5, - max_rank: int = 1024, - orth_center_right: bool = True, - preserve_norm: bool = False, -) -> list[list[torch.Tensor]]: - assert bath_node.ndim == 3 - assert bath_node.shape[0] == bath_node.shape[2] - - slices = [bath_node[:, i, :] for i in range(bath_node.shape[1])] - - svd_slices_l: list[torch.Tensor] = [] - svd_slices_r: list[torch.Tensor] = [] - for s in slices: - assert torch.allclose(s, s.mH, atol=1e-5) - # u, v = my_svd(s, tol=max_error) - u_h = my_hermitean_svd(s, tol=max_error) - - # l_s, r_s = split_matrix(s, max_error, max_rank, orth_center_right, preserve_norm) - svd_slices_l.append(u_h) - svd_slices_r.append(u_h) - - # l: torch.Tensor = torch.stack(svd_slices_l, dim=0) - # r: torch.Tensor = torch.stack(svd_slices_r, dim=0) - assert len(svd_slices_l) == len(svd_slices_r) - assert len(svd_slices_l) == bath_node.shape[1] - - return [svd_slices_l, svd_slices_r] From a17f92decf2fbb460ac8a8e20c23b430814a8684 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Tue, 31 Mar 2026 16:24:58 +0200 Subject: [PATCH 07/19] tests --- emu_base/math/packed_tensor.py | 10 ++++---- emu_mps/mps_backend_impl.py | 18 +++++++++----- test/emu_base/math/test_packed_tensor.py | 30 ++++++++++++------------ 3 files changed, 32 insertions(+), 26 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 34407875..39381957 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -32,16 +32,16 @@ def __init__( self.n = h.shape[0] self._ii, self._kk = torch.tril_indices(self.n, self.n, device=h.device) - self._packed = h[self._ii, :, self._kk].transpose(0, 1).contiguous() + self._packed_data = h[self._ii, :, self._kk].transpose(0, 1).contiguous() def unpack(self) -> torch.Tensor: - b = self._packed.shape[0] - vals = self._packed.transpose(0, 1) + b = self._packed_data.shape[0] + vals = self._packed_data.transpose(0, 1) h = torch.zeros( (self.n, b, self.n), - dtype=self._packed.dtype, - device=self._packed.device, + dtype=self._packed_data.dtype, + device=self._packed_data.device, ) h[self._ii, :, self._kk] = vals diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index e6d5e9e6..cc70f8e7 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -315,7 +315,7 @@ def init_baths(self) -> None: self.state, self.hamiltonian, final_qubit=2 ) - if isinstance(self, MPSBackendImpl): + if type(self) is MPSBackendImpl: self.left_baths = [PackedHermitianTensor(t) for t in _left_baths] self.right_baths = [PackedHermitianTensor(t) for t in _right_baths] else: @@ -443,7 +443,7 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None: self.hamiltonian.factors[self._sweep_index], ).to(self.state.factors[self._sweep_index + 1].device) item: Bath = lb - if isinstance(self, MPSBackendImpl): + if type(self) is MPSBackendImpl: item = PackedHermitianTensor(lb) self.left_baths.append(item) @@ -469,15 +469,21 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: self.hamiltonian.factors[self._sweep_index + 1], ).to(self.state.factors[self._sweep_index].device) item: Bath = rb - if isinstance(self, MPSBackendImpl): + if type(self) is MPSBackendImpl: item = PackedHermitianTensor(rb) self.right_baths.append(item) if not self.has_lindblad_noise: - # TODO this should be in Noise. Not in noiseless Base class + # TODO this should be in Noise? Not in noiseless Base class # Free memory because it won't be used anymore - to_deallocate: torch.Tensor = cast(torch.Tensor, self.right_baths[-2]) - deallocate_tensor(to_deallocate) + if type(self) is MPSBackendImpl: + todealloc_hermitean: PackedHermitianTensor = cast( + PackedHermitianTensor, self.right_baths[-2] + ) + deallocate_tensor(todealloc_hermitean._packed_data) + else: + todealloc: torch.Tensor = cast(torch.Tensor, self.right_baths[-2]) + deallocate_tensor(todealloc) self._evolve(self._sweep_index, dt=-delta_time / 2) self.left_baths.pop() diff --git a/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py index abc5d4a8..2caab818 100644 --- a/test/emu_base/math/test_packed_tensor.py +++ b/test/emu_base/math/test_packed_tensor.py @@ -8,8 +8,8 @@ def test_roundtrip(): x = torch.randn(n, b, n, dtype=torch.complex64) h = 0.5 * (x + x.transpose(0, 2).conj()) - packed = PackedHermitianTensor(h) - h2 = packed.unpack() + packedht = PackedHermitianTensor(h) + h2 = packedht.unpack() assert h2.shape == h.shape assert torch.allclose(h2, h) @@ -26,10 +26,10 @@ def test_roundtrip_minimal_shape(): ] ) # shape (1, 1, 1) - packed = PackedHermitianTensor(h) - h2 = packed.unpack() + tensorpacked = PackedHermitianTensor(h) + h2 = tensorpacked.unpack() - assert packed.packed.shape == (1, 1) + assert tensorpacked._packed_data.shape == (1, 1) assert torch.allclose(h2, h) @@ -52,9 +52,9 @@ def test_packed_shape(): x = torch.randn(n, b, n, dtype=torch.complex64) h = 0.5 * (x + x.transpose(0, 2).conj()) - packed = PackedHermitianTensor(h) + packedht = PackedHermitianTensor(h) - assert packed.packed.shape == (b, n * (n + 1) // 2) + assert packedht._packed_data.shape == (b, n * (n + 1) // 2) def test_roundtrip_real_symmetric(): @@ -62,8 +62,8 @@ def test_roundtrip_real_symmetric(): x = torch.randn(n, b, n) h = 0.5 * (x + x.transpose(0, 2)) - packed = PackedHermitianTensor(h) - h2 = packed.unpack() + packedht = PackedHermitianTensor(h) + h2 = packedht.unpack() assert h2.shape == h.shape assert torch.allclose(h2, h) @@ -74,8 +74,8 @@ def test_unpack_preserves_dtype(): x = torch.randn(n, b, n, dtype=torch.complex128) h = 0.5 * (x + x.transpose(0, 2).conj()) - packed = PackedHermitianTensor(h) - h2 = packed.unpack() + packedht = PackedHermitianTensor(h) + h2 = packedht.unpack() assert h2.dtype == h.dtype @@ -85,9 +85,9 @@ def test_skip_hermitian_check(): # class accepts structurally valid input without paying for symmetry checks. h = torch.randn(3, 2, 3, dtype=torch.complex64) - packed = PackedHermitianTensor(h, check_hermitian=False) + packedht = PackedHermitianTensor(h, check_hermitian=False) - assert packed.packed.shape == (2, 6) + assert packedht._packed_data.shape == (2, 6) def test_packed_is_contiguous(): @@ -95,9 +95,9 @@ def test_packed_is_contiguous(): x = torch.randn(n, b, n, dtype=torch.complex64) h = 0.5 * (x + x.transpose(0, 2).conj()) - packed = PackedHermitianTensor(h) + packedht = PackedHermitianTensor(h) - assert packed.packed.is_contiguous() + assert packedht._packed_data.is_contiguous() def test_custom_tolerance_controls_hermitian_check(): From 691d8707d16f664f157c037fdf3fb307fec902e4 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Tue, 31 Mar 2026 16:39:27 +0200 Subject: [PATCH 08/19] dealloc --- emu_mps/mps_backend_impl.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index cc70f8e7..9349cf3d 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -476,14 +476,11 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: if not self.has_lindblad_noise: # TODO this should be in Noise? Not in noiseless Base class # Free memory because it won't be used anymore - if type(self) is MPSBackendImpl: - todealloc_hermitean: PackedHermitianTensor = cast( - PackedHermitianTensor, self.right_baths[-2] - ) - deallocate_tensor(todealloc_hermitean._packed_data) - else: - todealloc: torch.Tensor = cast(torch.Tensor, self.right_baths[-2]) - deallocate_tensor(todealloc) + 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() From ccdb8d61808491dc1dab5198480eb2f2ec63b52e Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 1 Apr 2026 15:11:20 +0200 Subject: [PATCH 09/19] shrinked code --- emu_mps/mps_backend_impl.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 9349cf3d..2892744e 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -12,7 +12,7 @@ from collections import Counter from enum import Enum, auto from types import MethodType -from typing import Any, cast, Optional +from typing import Any, Optional import torch from pulser.backend import EmulationConfig, Observable, Results, State @@ -308,16 +308,14 @@ def update_H_no_noise(self) -> None: ) def init_baths(self) -> None: - _left_baths: list[torch.Tensor] = [ + _left_baths = [ torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device) ] - _right_baths: list[torch.Tensor] = right_baths( - self.state, self.hamiltonian, final_qubit=2 - ) + _right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2) - if type(self) is MPSBackendImpl: - self.left_baths = [PackedHermitianTensor(t) for t in _left_baths] - self.right_baths = [PackedHermitianTensor(t) for t in _right_baths] + if not self.has_lindblad_noise: + self.left_baths = list([PackedHermitianTensor(t) for t in _left_baths]) + self.right_baths = list([PackedHermitianTensor(t) for t in _right_baths]) else: self.left_baths = list(_left_baths) self.right_baths = list(_right_baths) @@ -443,7 +441,7 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None: self.hamiltonian.factors[self._sweep_index], ).to(self.state.factors[self._sweep_index + 1].device) item: Bath = lb - if type(self) is MPSBackendImpl: + if not self.has_lindblad_noise: item = PackedHermitianTensor(lb) self.left_baths.append(item) @@ -469,7 +467,7 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: self.hamiltonian.factors[self._sweep_index + 1], ).to(self.state.factors[self._sweep_index].device) item: Bath = rb - if type(self) is MPSBackendImpl: + if not self.has_lindblad_noise: item = PackedHermitianTensor(rb) self.right_baths.append(item) @@ -801,8 +799,8 @@ def progress(self) -> None: ), "Unknown Swipe direction" orth_center_right = self._swipe_direction == SwipeDirection.LEFT_TO_RIGHT - lbath: torch.Tensor = cast(torch.Tensor, self.left_baths[-1]) - rbath: torch.Tensor = cast(torch.Tensor, self.right_baths[-1]) + 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], From 726f5b6bb58e2ce58f34f4aa8db2a8c9506bbf1d Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Thu, 2 Apr 2026 10:49:11 +0200 Subject: [PATCH 10/19] dmrg test --- test/emu_mps/test_mps_backend_impl.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/test/emu_mps/test_mps_backend_impl.py b/test/emu_mps/test_mps_backend_impl.py index 61545e24..b37fd2ef 100644 --- a/test/emu_mps/test_mps_backend_impl.py +++ b/test/emu_mps/test_mps_backend_impl.py @@ -587,7 +587,7 @@ def test_left_to_right_update( ): mock_make_H.return_value = MagicMock(factors=[None] * 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) @@ -595,8 +595,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) @@ -619,7 +619,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 = ( From 6c87844d2968d2692195519ee8f946752eecea23 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 14:16:07 +0200 Subject: [PATCH 11/19] better naming for internals --- emu_base/math/packed_tensor.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 39381957..5f5cc1ad 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -3,9 +3,11 @@ class PackedHermitianTensor: """ - Pack a tensor of shape (a, b, a), Hermitian in axes 0 and 2, - into shape (b, a*(a+1)//2) by storing the lower triangle - of each (a, a) slice at fixed middle index. + Pack a tensor of shape (χ, m, χ), Hermitian in axes 0 and 2, + into shape (m, χ(χ+1)//2) by storing the lower triangle + of each (χ, χ) slice at fixed middle index. + The `PackedHermitianTensor` is used to represent left and right + baths in TDVP/DMRG algorithms. Lower-triangular packed order: (0,0), @@ -23,23 +25,23 @@ def __init__( atol: float = 1e-8, ) -> None: if h.ndim != 3 or h.shape[0] != h.shape[2]: - raise ValueError(f"Expected shape (n, b, n), got {tuple(h.shape)}") + raise ValueError(f"Expected shape (χ, m, χ), got {tuple(h.shape)}") if check_hermitian and not torch.allclose( h, h.transpose(0, 2).conj(), rtol=rtol, atol=atol ): raise ValueError("Tensor is not Hermitian in axes 0 and 2") - self.n = h.shape[0] - self._ii, self._kk = torch.tril_indices(self.n, self.n, device=h.device) + self.chi = h.shape[0] + self._ii, self._kk = torch.tril_indices(self.chi, self.chi, device=h.device) self._packed_data = h[self._ii, :, self._kk].transpose(0, 1).contiguous() def unpack(self) -> torch.Tensor: - b = self._packed_data.shape[0] + m = self._packed_data.shape[0] vals = self._packed_data.transpose(0, 1) h = torch.zeros( - (self.n, b, self.n), + (self.chi, m, self.chi), dtype=self._packed_data.dtype, device=self._packed_data.device, ) From f96b09558ee72bfa79527cbdab69e9fffaa9c7be Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 17:20:37 +0200 Subject: [PATCH 12/19] append shorter expression --- emu_mps/mps_backend_impl.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 2892744e..f9edda5e 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -440,10 +440,9 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None: self.state.factors[self._sweep_index], self.hamiltonian.factors[self._sweep_index], ).to(self.state.factors[self._sweep_index + 1].device) - item: Bath = lb - if not self.has_lindblad_noise: - item = PackedHermitianTensor(lb) - self.left_baths.append(item) + self.left_baths.append( + lb if self.has_lindblad_noise else PackedHermitianTensor(lb) + ) self._evolve(self._sweep_index + 1, dt=-delta_time / 2) self.right_baths.pop() @@ -466,10 +465,9 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: self.state.factors[self._sweep_index + 1], self.hamiltonian.factors[self._sweep_index + 1], ).to(self.state.factors[self._sweep_index].device) - item: Bath = rb - if not self.has_lindblad_noise: - item = PackedHermitianTensor(rb) - self.right_baths.append(item) + self.right_baths.append( + rb if self.has_lindblad_noise else PackedHermitianTensor(rb) + ) if not self.has_lindblad_noise: # TODO this should be in Noise? Not in noiseless Base class From ae6a3d7f78bd2098149e7a929c357b0f6e980924 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 17:57:33 +0200 Subject: [PATCH 13/19] tests and BathNode alias --- emu_mps/mps_backend_impl.py | 6 ++--- test/emu_base/math/test_packed_tensor.py | 33 +++--------------------- 2 files changed, 6 insertions(+), 33 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index f9edda5e..567922b9 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -43,7 +43,7 @@ ) dtype = torch.complex128 -Bath = torch.Tensor | PackedHermitianTensor +BathNode = torch.Tensor | PackedHermitianTensor class Statistics(Observable): @@ -102,8 +102,8 @@ class MPSBackendImpl: well_prepared_qubits_filter: Optional[torch.Tensor] hamiltonian: MPO state: MPS - left_baths: list[Bath] - right_baths: list[Bath] + left_baths: list[BathNode] + right_baths: list[BathNode] target_time: float results: Results _swipe_direction: SwipeDirection = SwipeDirection.LEFT_TO_RIGHT diff --git a/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py index 2caab818..fafca853 100644 --- a/test/emu_base/math/test_packed_tensor.py +++ b/test/emu_base/math/test_packed_tensor.py @@ -15,24 +15,6 @@ def test_roundtrip(): assert torch.allclose(h2, h) -def test_roundtrip_minimal_shape(): - h = torch.tensor( - [ - [ - [ - 2.0 + 0.0j, - ] - ] - ] - ) # shape (1, 1, 1) - - tensorpacked = PackedHermitianTensor(h) - h2 = tensorpacked.unpack() - - assert tensorpacked._packed_data.shape == (1, 1) - assert torch.allclose(h2, h) - - def test_rejects_non_hermitian(): h = torch.randn(3, 2, 3, dtype=torch.complex64) @@ -57,18 +39,6 @@ def test_packed_shape(): assert packedht._packed_data.shape == (b, n * (n + 1) // 2) -def test_roundtrip_real_symmetric(): - n, b = 4, 3 - x = torch.randn(n, b, n) - h = 0.5 * (x + x.transpose(0, 2)) - - packedht = PackedHermitianTensor(h) - h2 = packedht.unpack() - - assert h2.shape == h.shape - assert torch.allclose(h2, h) - - def test_unpack_preserves_dtype(): n, b = 4, 2 x = torch.randn(n, b, n, dtype=torch.complex128) @@ -85,8 +55,10 @@ def test_skip_hermitian_check(): # 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 == (2, 6) @@ -109,6 +81,7 @@ def test_custom_tolerance_controls_hermitian_check(): 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) From 652c3317120b1d3a547d68cb5c8815945b240c10 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 18:13:32 +0200 Subject: [PATCH 14/19] diff shape --- emu_base/math/packed_tensor.py | 10 +++++----- test/emu_base/math/test_packed_tensor.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 5f5cc1ad..3aba6a82 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -4,7 +4,7 @@ class PackedHermitianTensor: """ Pack a tensor of shape (χ, m, χ), Hermitian in axes 0 and 2, - into shape (m, χ(χ+1)//2) by storing the lower triangle + into shape (χ(χ+1)//2, m) by storing the lower triangle of each (χ, χ) slice at fixed middle index. The `PackedHermitianTensor` is used to represent left and right baths in TDVP/DMRG algorithms. @@ -33,15 +33,15 @@ def __init__( raise ValueError("Tensor is not Hermitian in axes 0 and 2") self.chi = h.shape[0] + self.m = h.shape[1] self._ii, self._kk = torch.tril_indices(self.chi, self.chi, device=h.device) - self._packed_data = h[self._ii, :, self._kk].transpose(0, 1).contiguous() + self._packed_data = h[self._ii, :, self._kk] def unpack(self) -> torch.Tensor: - m = self._packed_data.shape[0] - vals = self._packed_data.transpose(0, 1) + vals = self._packed_data h = torch.zeros( - (self.chi, m, self.chi), + (self.chi, self.m, self.chi), dtype=self._packed_data.dtype, device=self._packed_data.device, ) diff --git a/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py index fafca853..dc7225b2 100644 --- a/test/emu_base/math/test_packed_tensor.py +++ b/test/emu_base/math/test_packed_tensor.py @@ -36,7 +36,7 @@ def test_packed_shape(): packedht = PackedHermitianTensor(h) - assert packedht._packed_data.shape == (b, n * (n + 1) // 2) + assert packedht._packed_data.shape == (n * (n + 1) // 2, b) def test_unpack_preserves_dtype(): @@ -59,7 +59,7 @@ def test_skip_hermitian_check(): packedht = PackedHermitianTensor(h, check_hermitian=False) # Data loss happened - assert packedht._packed_data.shape == (2, 6) + assert packedht._packed_data.shape == (6, 2) def test_packed_is_contiguous(): From 8b0f2a31a807f976c9255743809e681c9be7f42d Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 18:20:33 +0200 Subject: [PATCH 15/19] docstring --- emu_base/math/packed_tensor.py | 16 +++++----------- test/emu_base/math/test_packed_tensor.py | 2 +- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 3aba6a82..7a0e1a5c 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -4,16 +4,10 @@ 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. + 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 in TDVP/DMRG algorithms. - - Lower-triangular packed order: - (0,0), - (1,0), (1,1), - (2,0), (2,1), (2,2), - ... + baths nodes in TDVP/DMRG algorithms. """ def __init__( @@ -33,15 +27,15 @@ def __init__( raise ValueError("Tensor is not Hermitian in axes 0 and 2") self.chi = h.shape[0] - self.m = h.shape[1] self._ii, self._kk = torch.tril_indices(self.chi, self.chi, device=h.device) self._packed_data = h[self._ii, :, self._kk] def unpack(self) -> torch.Tensor: vals = self._packed_data + m = vals.shape[1] h = torch.zeros( - (self.chi, self.m, self.chi), + (self.chi, m, self.chi), dtype=self._packed_data.dtype, device=self._packed_data.device, ) diff --git a/test/emu_base/math/test_packed_tensor.py b/test/emu_base/math/test_packed_tensor.py index dc7225b2..8bda7ef4 100644 --- a/test/emu_base/math/test_packed_tensor.py +++ b/test/emu_base/math/test_packed_tensor.py @@ -36,7 +36,7 @@ def test_packed_shape(): packedht = PackedHermitianTensor(h) - assert packedht._packed_data.shape == (n * (n + 1) // 2, b) + assert packedht._packed_data.shape == (n * (n + 1) / 2, b) def test_unpack_preserves_dtype(): From 63e695c2c147c2e5de59bf12a94740d45b4c9be7 Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 18:48:45 +0200 Subject: [PATCH 16/19] shorter init_baths --- emu_mps/mps_backend_impl.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 567922b9..e6422a31 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -308,18 +308,14 @@ def update_H_no_noise(self) -> None: ) def init_baths(self) -> None: - _left_baths = [ - torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device) - ] - _right_baths = right_baths(self.state, self.hamiltonian, final_qubit=2) - - if not self.has_lindblad_noise: - self.left_baths = list([PackedHermitianTensor(t) for t in _left_baths]) - self.right_baths = list([PackedHermitianTensor(t) for t in _right_baths]) - else: - self.left_baths = list(_left_baths) - self.right_baths = list(_right_baths) + pack = (lambda x: x) if self.has_lindblad_noise else PackedHermitianTensor + self.left_baths = [ + pack(torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device)) + ] + self.right_baths = [ + pack(t) for t in right_baths(self.state, self.hamiltonian, final_qubit=2) + ] assert len(self.right_baths) == self.qubit_count - 1 def get_current_right_bath(self) -> torch.Tensor: From 496cd4f00965707a49881214a4865976b531e64b Mon Sep 17 00:00:00 2001 From: Kemal Bidzhiev Date: Wed, 29 Apr 2026 18:55:30 +0200 Subject: [PATCH 17/19] rm comment --- emu_mps/mps_backend_impl.py | 1 - 1 file changed, 1 deletion(-) diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index e6422a31..668ee31e 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -466,7 +466,6 @@ def _right_to_left_update_tdvp(self, delta_time: float) -> None: ) if not self.has_lindblad_noise: - # TODO this should be in Noise? Not in noiseless Base class # Free memory because it won't be used anymore item = self.right_baths[-2] to_dealloc = ( From 48f3faaf3c4a90aae30abfc558b1d554d49c635e Mon Sep 17 00:00:00 2001 From: Anton Quelle Date: Mon, 27 Jul 2026 09:02:12 +0200 Subject: [PATCH 18/19] fix memory issue --- emu_mps/solver_utils.py | 9 ++++----- test/emu_mps/test_solver_utils.py | 2 +- 2 files changed, 5 insertions(+), 6 deletions(-) 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_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( From e83deaf353705177f4fe500917c071d51185bc7e Mon Sep 17 00:00:00 2001 From: Anton Quelle Date: Wed, 12 Aug 2026 16:34:33 +0200 Subject: [PATCH 19/19] change tensor init order --- emu_base/math/packed_tensor.py | 25 ++++++-- emu_mps/mps_backend_impl.py | 89 +++++++++++++++++++-------- test/emu_mps/test_mps_backend_impl.py | 6 +- 3 files changed, 88 insertions(+), 32 deletions(-) diff --git a/emu_base/math/packed_tensor.py b/emu_base/math/packed_tensor.py index 7a0e1a5c..63691ca1 100644 --- a/emu_base/math/packed_tensor.py +++ b/emu_base/math/packed_tensor.py @@ -12,23 +12,38 @@ class PackedHermitianTensor: def __init__( self, - h: torch.Tensor, + 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 check_hermitian and not torch.allclose( - h, h.transpose(0, 2).conj(), rtol=rtol, atol=atol + 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.chi = h.shape[0] 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[:, :] = h[self._ii, :, self._kk] + self._packed_data = self._packed_data.to(h.device) def unpack(self) -> torch.Tensor: vals = self._packed_data diff --git a/emu_mps/mps_backend_impl.py b/emu_mps/mps_backend_impl.py index 0762edbe..8b6b237d 100644 --- a/emu_mps/mps_backend_impl.py +++ b/emu_mps/mps_backend_impl.py @@ -342,17 +342,34 @@ def update_H_no_noise(self) -> None: ) def init_baths(self) -> None: - pack = (lambda x: x) if self.has_lindblad_noise else PackedHermitianTensor - - self.left_baths = [ - pack(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[PackedHermitianTensor] = [] - self.right_baths = [ - pack(t) for t in 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: @@ -468,14 +485,25 @@ def _left_to_right_update_tdvp(self, delta_time: float) -> None: dt=delta_time / 2, orth_center_right=True, ) - 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 if self.has_lindblad_noise else PackedHermitianTensor(lb) - ) + 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() @@ -493,14 +521,25 @@ 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: - 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 if self.has_lindblad_noise else PackedHermitianTensor(rb) - ) + 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 diff --git a/test/emu_mps/test_mps_backend_impl.py b/test/emu_mps/test_mps_backend_impl.py index 80b18081..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,7 +581,9 @@ 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, 1, 1)] * (QUBIT_COUNT - 1) mock_new_left.return_value = torch.zeros(1)