emu-mps 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

emu_mps/__init__.py

@@ -10,10 +10,10 @@ from emu_base import (
     StateResult,
     SecondMomentOfEnergy,
 )
+from .mps_config import MPSConfig
 from .mpo import MPO
 from .mps import MPS, inner
 from .mps_backend import MPSBackend
-from .mps_config import MPSConfig
 
 
 __all__ = [
@@ -35,4 +35,4 @@ __all__ = [
     "SecondMomentOfEnergy",
 ]
 
-__version__ = "1.2.2"
+__version__ = "1.2.4"

emu_mps/algebra.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import math
+
+from emu_mps import MPSConfig
 from emu_mps.utils import truncate_impl
 
 
 def add_factors(
-    left: list[torch.tensor], right: list[torch.tensor]
-) -> list[torch.tensor]:
+    left: list[torch.Tensor], right: list[torch.Tensor]
+) -> list[torch.Tensor]:
     """
     Direct sum algorithm implementation to sum two tensor trains (MPS/MPO).
     It assumes the left and right bond are along the dimension 0 and -1 of each tensor.
@@ -48,8 +52,8 @@ def add_factors(
 
 
 def scale_factors(
-    factors: list[torch.tensor], scalar: complex, *, which: int
-) -> list[torch.tensor]:
+    factors: list[torch.Tensor], scalar: complex, *, which: int
+) -> list[torch.Tensor]:
     """
     Returns a new list of factors where the tensor at the given index is scaled by `scalar`.
     """
@@ -57,10 +61,10 @@ def scale_factors(
 
 
 def zip_right_step(
-    slider: torch.tensor,
-    top: torch.tensor,
-    bottom: torch.tensor,
-) -> torch.tensor:
+    slider: torch.Tensor,
+    top: torch.Tensor,
+    bottom: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Returns a new `MPS/O` factor of the result of the multiplication MPO @ MPS/O,
     and the updated slider, performing a single step of the
@@ -113,11 +117,10 @@ def zip_right_step(
 
 
 def zip_right(
-    top_factors: list[torch.tensor],
-    bottom_factors: list[torch.tensor],
-    max_error: float = 1e-5,
-    max_rank: int = 1024,
-) -> list[torch.tensor]:
+    top_factors: list[torch.Tensor],
+    bottom_factors: list[torch.Tensor],
+    config: Optional[MPSConfig] = None,
+) -> list[torch.Tensor]:
     """
     Returns a new matrix product, resulting from applying `top` to `bottom`.
     The resulting factors are:
@@ -136,6 +139,8 @@ def zip_right(
         A final truncation sweep, from right to left,
         moves back the orthogonal center to the first element.
     """
+    config = config if config is not None else MPSConfig()
+
     if len(top_factors) != len(bottom_factors):
         raise ValueError("Cannot multiply two matrix products of different lengths.")
 
@@ -146,6 +151,6 @@ def zip_right(
         new_factors.append(res)
     new_factors[-1] @= slider[:, :, 0]
 
-    truncate_impl(new_factors, max_error=max_error, max_rank=max_rank)
+    truncate_impl(new_factors, config=config)
 
     return new_factors

emu_mps/constants.py

@@ -0,0 +1,4 @@
+import torch
+
+
+DEVICE_COUNT = torch.cuda.device_count()

emu_mps/hamiltonian.py

@@ -285,6 +285,7 @@ def _get_interactions_to_keep(interaction_matrix: torch.Tensor) -> list[torch.Te
     returns a list of bool valued tensors,
     indicating which interaction terms to keep for each bond in the MPO
     """
+    interaction_matrix = interaction_matrix.clone()
     nqubits = interaction_matrix.size(dim=1)
     middle = nqubits // 2
     interaction_matrix += torch.eye(
@@ -362,7 +363,7 @@ def make_H(
 
     interactions_to_keep = _get_interactions_to_keep(interaction_matrix)
 
-    cores = [_first_factor(interactions_to_keep[0].item())]
+    cores = [_first_factor(interactions_to_keep[0].item() != 0.0)]
 
     if nqubits > 2:
         for i in range(1, middle):
@@ -394,7 +395,7 @@ def make_H(
                 )
             )
     if nqubits == 2:
-        scale = interaction_matrix[0, 1]
+        scale = interaction_matrix[0, 1].item()
     elif interactions_to_keep[-1][0]:
         scale = 1.0
     else:

emu_mps/mpo.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import itertools
-from typing import Any, List, cast, Iterable, Optional
+from typing import Any, List, Iterable, Optional
 
 import torch
 
@@ -8,7 +8,8 @@ from emu_mps.algebra import add_factors, scale_factors, zip_right
 from emu_base.base_classes.operator import FullOp, QuditOp
 from emu_base import Operator, State
 from emu_mps.mps import MPS
-from emu_mps.utils import new_left_bath, assign_devices, DEVICE_COUNT
+from emu_mps.utils import new_left_bath, assign_devices
+from emu_mps.constants import DEVICE_COUNT
 
 
 def _validate_operator_targets(operations: FullOp, nqubits: int) -> None:
@@ -80,8 +81,7 @@ class MPO(Operator):
         factors = zip_right(
             self.factors,
             other.factors,
-            max_error=other.precision,
-            max_rank=other.max_bond_dim,
+            config=other.config,
         )
         return MPS(factors, orthogonality_center=0)
 
@@ -175,14 +175,15 @@ class MPO(Operator):
         Returns:
             the operator in MPO form.
         """
+        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
 
         _validate_operator_targets(operations, nqubits)
 
         basis = set(basis)
         if basis == {"r", "g"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "gg": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -197,9 +198,9 @@ class MPO(Operator):
                 ).reshape(1, 2, 2, 1),
             }
         elif basis == {"0", "1"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "00": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -218,26 +219,28 @@ class MPO(Operator):
 
         mpos = []
         for coeff, tensorop in operations:
-            # this function will recurse through the operators, and replace any definitions
+            # this function will recurse through the operators_with_tensors,
+            # and replace any definitions
             # in terms of strings by the computed tensor
             def replace_operator_string(op: QuditOp | torch.Tensor) -> torch.Tensor:
-                if isinstance(op, dict):
-                    for opstr, coeff in op.items():
-                        tensor = replace_operator_string(operators[opstr])
-                        operators[opstr] = tensor
-                        op[opstr] = tensor * coeff
-                    op = sum(cast(list[torch.Tensor], op.values()))
-                return op
+                if isinstance(op, torch.Tensor):
+                    return op
+
+                result = torch.zeros(1, 2, 2, 1, dtype=torch.complex128)
+                for opstr, coeff in op.items():
+                    tensor = replace_operator_string(operators_with_tensors[opstr])
+                    operators_with_tensors[opstr] = tensor
+                    result += tensor * coeff
+                return result
 
             factors = [
                 torch.eye(2, 2, dtype=torch.complex128).reshape(1, 2, 2, 1)
             ] * nqubits
 
-            for i, op in enumerate(tensorop):
-                tensorop[i] = (replace_operator_string(op[0]), op[1])
-
             for op in tensorop:
-                for i in op[1]:
-                    factors[i] = op[0]
+                factor = replace_operator_string(op[0])
+                for target_qubit in op[1]:
+                    factors[target_qubit] = factor
+
             mpos.append(coeff * MPO(factors, **kwargs))
         return sum(mpos[1:], start=mpos[0])

emu_mps/mps.py

@@ -7,15 +7,16 @@ from typing import Any, List, Optional, Iterable
 import torch
 
 from emu_base import State
+from emu_mps import MPSConfig
 from emu_mps.algebra import add_factors, scale_factors
 from emu_mps.utils import (
-    DEVICE_COUNT,
     apply_measurement_errors,
     assign_devices,
     truncate_impl,
     tensor_trace,
     n_operator,
 )
+from emu_mps.constants import DEVICE_COUNT
 
 
 class MPS(State):
@@ -27,17 +28,13 @@ class MPS(State):
     Only qubits are supported.
     """
 
-    DEFAULT_MAX_BOND_DIM: int = 1024
-    DEFAULT_PRECISION: float = 1e-5
-
     def __init__(
         self,
         factors: List[torch.Tensor],
         /,
         *,
         orthogonality_center: Optional[int] = None,
-        precision: float = DEFAULT_PRECISION,
-        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
+        config: Optional[MPSConfig] = None,
         num_gpus_to_use: Optional[int] = DEVICE_COUNT,
     ):
         """
@@ -53,14 +50,11 @@ class MPS(State):
                 of the data to this constructor, or some shared objects.
             orthogonality_center: the orthogonality center of the MPS, or None (in which case
                 it will be orthogonalized when needed)
-            precision: the precision with which to truncate here or in tdvp
-            max_bond_dim: the maximum bond dimension to allow
+            config: the emu-mps config object passed to the run method
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu, None=keep the existing device assignment.
         """
-        self.precision = precision
-        self.max_bond_dim = max_bond_dim
-
+        self.config = config if config is not None else MPSConfig()
         assert all(
             factors[i - 1].shape[2] == factors[i].shape[0] for i in range(1, len(factors))
         ), "The dimensions of consecutive tensors should match"
@@ -84,8 +78,7 @@ class MPS(State):
     def make(
         cls,
         num_sites: int,
-        precision: float = DEFAULT_PRECISION,
-        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
+        config: Optional[MPSConfig] = None,
         num_gpus_to_use: int = DEVICE_COUNT,
     ) -> MPS:
         """
@@ -93,11 +86,12 @@ class MPS(State):
 
         Args:
             num_sites: the number of qubits
-            precision: the precision with which to truncate here or in tdvp
-            max_bond_dim: the maximum bond dimension to allow
+            config: the MPSConfig
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu
         """
+        config = config if config is not None else MPSConfig()
+
         if num_sites <= 1:
             raise ValueError("For 1 qubit states, do state vector")
 
@@ -106,8 +100,7 @@ class MPS(State):
                 torch.tensor([[[1.0], [0.0]]], dtype=torch.complex128)
                 for _ in range(num_sites)
             ],
-            precision=precision,
-            max_bond_dim=max_bond_dim,
+            config=config,
             num_gpus_to_use=num_gpus_to_use,
             orthogonality_center=0,  # Arbitrary: every qubit is an orthogonality center.
         )
@@ -165,15 +158,11 @@ class MPS(State):
         """
         SVD based truncation of the state. Puts the orthogonality center at the first qubit.
         Calls orthogonalize on the last qubit, and then sweeps a series of SVDs right-left.
-        Uses self.precision and self.max_bond_dim for determining accuracy.
+        Uses self.config for determining accuracy.
         An in-place operation.
         """
         self.orthogonalize(self.num_sites - 1)
-        truncate_impl(
-            self.factors,
-            max_error=self.precision,
-            max_rank=self.max_bond_dim,
-        )
+        truncate_impl(self.factors, config=self.config)
         self.orthogonality_center = 0
 
     def get_max_bond_dim(self) -> int:
@@ -236,7 +225,7 @@ class MPS(State):
         num_qubits = len(self.factors)
 
         bitstring = ""
-        acc_mps_j: torch.tensor = self.factors[0]
+        acc_mps_j: torch.Tensor = self.factors[0]
 
         for qubit in range(num_qubits):
             # comp_basis is a projector: 0 is for ket |0> and 1 for ket |1>
@@ -303,7 +292,7 @@ class MPS(State):
         """
         Returns the sum of two MPSs, computed with a direct algorithm.
         The resulting MPS is orthogonalized on the first site and truncated
-        up to `self.precision`.
+        up to `self.config.precision`.
 
         Args:
             other: the other state
@@ -315,8 +304,7 @@ class MPS(State):
         new_tt = add_factors(self.factors, other.factors)
         result = MPS(
             new_tt,
-            precision=self.precision,
-            max_bond_dim=self.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=None,  # Orthogonality is lost.
         )
@@ -341,8 +329,7 @@ class MPS(State):
         factors = scale_factors(self.factors, scalar, which=which)
         return MPS(
             factors,
-            precision=self.precision,
-            max_bond_dim=self.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=self.orthogonality_center,
         )

emu_mps/mps_backend.py

@@ -1,8 +1,12 @@
-from pulser import Sequence
-
 from emu_base import Backend, BackendConfig, Results
 from emu_mps.mps_config import MPSConfig
-from emu_mps.mps_backend_impl import create_impl
+from emu_mps.mps_backend_impl import create_impl, MPSBackendImpl
+from pulser import Sequence
+import pickle
+import os
+import time
+import logging
+import pathlib
 
 
 class MPSBackend(Backend):
@@ -11,6 +15,32 @@ class MPSBackend(Backend):
     aka tensor trains.
     """
 
+    def resume(self, autosave_file: str | pathlib.Path) -> Results:
+        """
+        Resume simulation from autosave file.
+        Only resume simulations from data you trust!
+        Unpickling of untrusted data is not safe.
+        """
+        if isinstance(autosave_file, str):
+            autosave_file = pathlib.Path(autosave_file)
+
+        if not autosave_file.is_file():
+            raise ValueError(f"Not a file: {autosave_file}")
+
+        with open(autosave_file, "rb") as f:
+            impl: MPSBackendImpl = pickle.load(f)
+
+        impl.autosave_file = autosave_file
+        impl.last_save_time = time.time()
+        impl.config.init_logging()  # FIXME: might be best to take logger object out of config.
+
+        logging.getLogger("global_logger").warning(
+            f"Resuming simulation from file {autosave_file}\n"
+            f"Saving simulation state every {impl.config.autosave_dt} seconds"
+        )
+
+        return self._run(impl)
+
     def run(self, sequence: Sequence, mps_config: BackendConfig) -> Results:
         """
         Emulates the given sequence.
@@ -29,7 +59,14 @@ class MPSBackend(Backend):
         impl = create_impl(sequence, mps_config)
         impl.init()  # This is separate from the constructor for testing purposes.
 
+        return self._run(impl)
+
+    @staticmethod
+    def _run(impl: MPSBackendImpl) -> Results:
         while not impl.is_finished():
             impl.progress()
 
+        if impl.autosave_file.is_file():
+            os.remove(impl.autosave_file)
+
         return impl.results

emu_mps/mps_backend_impl.py

@@ -1,7 +1,11 @@
 import math
+import pathlib
 import random
+import uuid
 from resource import RUSAGE_SELF, getrusage
 from typing import Optional
+import pickle
+import os
 
 import torch
 import time
@@ -17,7 +21,6 @@ from emu_mps.noise import compute_noise_from_lindbladians, pick_well_prepared_qu
 from emu_mps.tdvp import (
     evolve_single,
     evolve_pair,
-    EvolveConfig,
     new_right_bath,
     right_baths,
 )
@@ -52,6 +55,7 @@ class MPSBackendImpl:
     def __init__(self, mps_config: MPSConfig, pulser_data: PulserData):
         self.config = mps_config
         self.target_time = float(self.config.dt)
+        self.pulser_data = pulser_data
         self.qubit_count = pulser_data.qubit_count
         assert self.qubit_count >= 2
         self.omega = pulser_data.omega
@@ -71,15 +75,17 @@ class MPSBackendImpl:
         self.tdvp_index = 0
         self.timestep_index = 0
         self.results = Results()
-
-        self.evolve_config = EvolveConfig(
-            exp_tolerance=self.config.precision * self.config.extra_krylov_tolerance,
-            norm_tolerance=self.config.precision * self.config.extra_krylov_tolerance,
-            max_krylov_dim=self.config.max_krylov_dim,
-            is_hermitian=not self.has_lindblad_noise,
-            max_error=self.config.precision,
-            max_rank=self.config.max_bond_dim,
+        self.autosave_file = self._get_autosave_filepath(self.config.autosave_prefix)
+        self.config.logger.warning(
+            f"""Will save simulation state to file "{self.autosave_file.name}"
+            every {self.config.autosave_dt} seconds.\n"""
+            f"""To resume: `MPSBackend().resume("{self.autosave_file}")`"""
         )
+        self.last_save_time = time.time()
+
+    @staticmethod
+    def _get_autosave_filepath(autosave_prefix: str) -> pathlib.Path:
+        return pathlib.Path(os.getcwd()) / (autosave_prefix + str(uuid.uuid1()) + ".dat")
 
     def init_dark_qubits(self) -> None:
         has_state_preparation_error: bool = (
@@ -112,8 +118,7 @@ class MPSBackendImpl:
         if initial_state is None:
             self.state = MPS.make(
                 self.qubit_count,
-                precision=self.config.precision,
-                max_bond_dim=self.config.max_bond_dim,
+                config=self.config,
                 num_gpus_to_use=self.config.num_gpus_to_use,
             )
             return
@@ -128,8 +133,7 @@ class MPSBackendImpl:
         initial_state = MPS(
             # Deep copy of every tensor of the initial state.
             [f.clone().detach() for f in initial_state.factors],
-            precision=self.config.precision,
-            max_bond_dim=self.config.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=self.config.num_gpus_to_use,
         )
         initial_state.truncate()
@@ -198,7 +202,8 @@ class MPSBackendImpl:
                 ham_factor=self.hamiltonian.factors[index],
                 baths=baths,
                 dt=dt,
-                config=self.evolve_config,
+                config=self.config,
+                is_hermitian=not self.has_lindblad_noise,
             )
         else:
             assert orth_center_right is not None
@@ -213,8 +218,9 @@ class MPSBackendImpl:
                 ham_factors=self.hamiltonian.factors[l : r + 1],
                 baths=baths,
                 dt=dt,
-                config=self.evolve_config,
+                config=self.config,
                 orth_center_right=orth_center_right,
+                is_hermitian=not self.has_lindblad_noise,
             )
 
             self.state.orthogonality_center = r if orth_center_right else l
@@ -292,7 +298,7 @@ class MPSBackendImpl:
             )
             if not self.has_lindblad_noise:
                 # Free memory because it won't be used anymore
-                self.right_baths[-2] = None
+                self.right_baths[-2] = torch.zeros(0)
 
             self._evolve(self.tdvp_index, dt=-delta_time / 2)
             self.left_baths.pop()
@@ -312,7 +318,7 @@ class MPSBackendImpl:
         else:
             raise Exception("Didn't expect this")
 
-        # TODO: checkpoint/autosave here
+        self.save_simulation()
 
     def tdvp_complete(self) -> None:
         self.current_time = self.target_time
@@ -343,14 +349,40 @@ class MPSBackendImpl:
         self.log_step_statistics(duration=time.time() - self.time)
         self.time = time.time()
 
+    def save_simulation(self) -> None:
+        if self.last_save_time > time.time() - self.config.autosave_dt:
+            return
+
+        basename = self.autosave_file
+        with open(basename.with_suffix(".new"), "wb") as file_handle:
+            pickle.dump(self, file_handle)
+
+        if basename.is_file():
+            os.rename(basename, basename.with_suffix(".bak"))
+
+        os.rename(basename.with_suffix(".new"), basename)
+        autosave_filesize = os.path.getsize(self.autosave_file) / 1e6
+
+        if basename.with_suffix(".bak").is_file():
+            os.remove(basename.with_suffix(".bak"))
+
+        self.last_save_time = time.time()
+
+        self.config.logger.debug(
+            f"Saved simulation state in file {self.autosave_file} ({autosave_filesize}MB)"
+        )
+
     def fill_results(self) -> None:
         normalized_state = 1 / self.state.norm() * self.state
 
+        current_time_int: int = round(self.current_time)
+        assert abs(self.current_time - current_time_int) < 1e-10
+
         if self.well_prepared_qubits_filter is None:
             for callback in self.config.callbacks:
                 callback(
                     self.config,
-                    self.current_time,
+                    current_time_int,
                     normalized_state,
                     self.hamiltonian,
                     self.results,
@@ -359,7 +391,7 @@ class MPSBackendImpl:
 
         full_mpo, full_state = None, None
         for callback in self.config.callbacks:
-            if self.current_time not in callback.evaluation_times:
+            if current_time_int not in callback.evaluation_times:
                 continue
 
             if full_mpo is None or full_state is None:
@@ -380,7 +412,7 @@ class MPSBackendImpl:
                     ),
                 )
 
-            callback(self.config, self.current_time, full_state, full_mpo, self.results)
+            callback(self.config, current_time_int, full_state, full_mpo, self.results)
 
     def log_step_statistics(self, *, duration: float) -> None:
         if self.state.factors[0].is_cuda:
@@ -424,13 +456,12 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
     """
 
     jump_threshold: float
-    aggregated_lindblad_ops: Optional[torch.Tensor]
+    aggregated_lindblad_ops: torch.Tensor
     norm_gap_before_jump: float
     root_finder: Optional[BrentsRootFinder]
 
     def __init__(self, config: MPSConfig, pulser_data: PulserData):
         super().__init__(config, pulser_data)
-        self.aggregated_lindblad_ops = None
         self.lindblad_ops = pulser_data.lindblad_ops
         self.root_finder = None
 
@@ -491,7 +522,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
                 for qubit in range(self.state.num_sites)
                 for op in self.lindblad_ops
             ],
-            weights=jump_operator_weights.reshape(-1),
+            weights=jump_operator_weights.reshape(-1).tolist(),
         )[0]
 
         self.state.apply(jumped_qubit_index, jump_operator)

emu_mps/mps_config.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 from emu_base import BackendConfig, State
-from emu_mps.utils import DEVICE_COUNT
+from emu_mps.constants import DEVICE_COUNT
 
 
 class MPSConfig(BackendConfig):
@@ -23,6 +23,10 @@ class MPSConfig(BackendConfig):
         num_gpus_to_use: during the simulation, distribute the state over this many GPUs
             0=all factors to cpu. As shown in the benchmarks, using multiple GPUs might
             alleviate memory pressure per GPU, but the runtime should be similar.
+        autosave_prefix: filename prefix for autosaving simulation state to file
+        autosave_dt: minimum time interval in seconds between two autosaves
+            Saving the simulation state is only possible at specific times,
+            therefore this interval is only a lower bound.
         kwargs: arguments that are passed to the base class
 
     Examples:
@@ -43,6 +47,8 @@ class MPSConfig(BackendConfig):
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int = DEVICE_COUNT,
+        autosave_prefix: str = "emu_mps_save_",
+        autosave_dt: int = 600,  # 10 minutes
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -62,3 +68,11 @@ class MPSConfig(BackendConfig):
                 and self.noise_model.amp_sigma != 0.0
             ):
                 raise NotImplementedError("Unsupported noise type: amp_sigma")
+
+        self.autosave_prefix = autosave_prefix
+        self.autosave_dt = autosave_dt
+
+        MIN_AUTOSAVE_DT = 10
+        assert (
+            self.autosave_dt > MIN_AUTOSAVE_DT
+        ), f"autosave_dt must be larger than {MIN_AUTOSAVE_DT} seconds"

emu_mps/tdvp.py

@@ -3,6 +3,7 @@ import torch
 from emu_base import krylov_exp
 from emu_mps import MPS, MPO
 from emu_mps.utils import split_tensor
+from emu_mps.mps_config import MPSConfig
 
 
 def new_right_bath(
@@ -10,7 +11,8 @@ def new_right_bath(
 ) -> torch.Tensor:
     bath = torch.tensordot(state, bath, ([2], [2]))
     bath = torch.tensordot(op.to(bath.device), bath, ([2, 3], [1, 3]))
-    return torch.tensordot(state.conj(), bath, ([1, 2], [1, 3]))
+    bath = torch.tensordot(state.conj(), bath, ([1, 2], [1, 3]))
+    return bath
 
 
 """
@@ -79,27 +81,6 @@ def apply_effective_Hamiltonian(
 _TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/ms, dt in ns
 
 
-class EvolveConfig:
-    def __init__(
-        self,
-        *,
-        exp_tolerance: float,
-        norm_tolerance: float,
-        max_krylov_dim: int,
-        is_hermitian: bool,
-        max_error: float,
-        max_rank: int
-    ) -> None:
-        self.exp_tolerance = exp_tolerance
-        self.norm_tolerance = norm_tolerance
-        self.max_krylov_dim = max_krylov_dim
-        self.is_hermitian = is_hermitian
-        self.max_error = (
-            max_error  # FIXME: max_error and max_rank are irrelevant for evolve_single
-        )
-        self.max_rank = max_rank
-
-
 def evolve_pair(
     *,
     state_factors: list[torch.Tensor],
@@ -107,7 +88,8 @@ def evolve_pair(
     ham_factors: list[torch.Tensor],
     dt: float,
     orth_center_right: bool,
-    config: EvolveConfig
+    is_hermitian: bool,
+    config: MPSConfig,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Time evolution of a pair of tensors of a tensor train using baths and truncated SVD.
@@ -154,16 +136,16 @@ def evolve_pair(
     evol = krylov_exp(
         op,
         combined_state_factors,
-        exp_tolerance=config.exp_tolerance,
-        norm_tolerance=config.norm_tolerance,
+        exp_tolerance=config.precision * config.extra_krylov_tolerance,
+        norm_tolerance=config.precision * config.extra_krylov_tolerance,
         max_krylov_dim=config.max_krylov_dim,
-        is_hermitian=config.is_hermitian,
+        is_hermitian=is_hermitian,
     ).reshape(left_bond_dim * 2, 2 * right_bond_dim)
 
     l, r = split_tensor(
         evol,
-        max_error=config.max_error,
-        max_rank=config.max_rank,
+        max_error=config.precision,
+        max_rank=config.max_bond_dim,
         orth_center_right=orth_center_right,
     )
 
@@ -178,7 +160,8 @@ def evolve_single(
     baths: tuple[torch.Tensor, torch.Tensor],
     ham_factor: torch.Tensor,
     dt: float,
-    config: EvolveConfig
+    is_hermitian: bool,
+    config: MPSConfig,
 ) -> torch.Tensor:
     """
     Time evolution of a single tensor of a tensor train using baths.
@@ -202,8 +185,8 @@ def evolve_single(
     return krylov_exp(
         op,
         state_factor,
-        exp_tolerance=config.exp_tolerance,
-        norm_tolerance=config.norm_tolerance,
+        exp_tolerance=config.precision * config.extra_krylov_tolerance,
+        norm_tolerance=config.precision * config.extra_krylov_tolerance,
         max_krylov_dim=config.max_krylov_dim,
-        is_hermitian=config.is_hermitian,
+        is_hermitian=is_hermitian,
     )

emu_mps/utils.py

@@ -3,7 +3,7 @@ import torch
 import random
 from collections import Counter
 
-DEVICE_COUNT = torch.cuda.device_count()
+from emu_mps import MPSConfig
 
 
 def new_left_bath(
@@ -12,15 +12,16 @@ def new_left_bath(
     # this order is more efficient than contracting the op first in general
     bath = torch.tensordot(bath, state.conj(), ([0], [0]))
     bath = torch.tensordot(bath, op.to(bath.device), ([0, 2], [0, 1]))
-    return torch.tensordot(bath, state, ([0, 2], [0, 1]))
+    bath = torch.tensordot(bath, state, ([0, 2], [0, 1]))
+    return bath
 
 
 def _determine_cutoff_index(d: torch.Tensor, max_error: float) -> int:
     assert max_error > 0
     squared_max_error = max_error * max_error
-    acc = 0
+    acc = 0.0
     for i in range(d.shape[0]):
-        acc += d[i]
+        acc += d[i].item()
         if acc > squared_max_error:
             return i
     return 0  # type: ignore[no-any-return]
@@ -60,9 +61,8 @@ def split_tensor(
 
 
 def truncate_impl(
-    factors: list[torch.tensor],
-    max_error: float = 1e-5,
-    max_rank: int = 1024,
+    factors: list[torch.Tensor],
+    config: MPSConfig,
 ) -> None:
     """
     Eigenvalues-based truncation of a matrix product.
@@ -78,8 +78,8 @@ def truncate_impl(
 
         l, r = split_tensor(
             factors[i].reshape(factor_shape[0], -1),
-            max_error=max_error,
-            max_rank=max_rank,
+            max_error=config.precision,
+            max_rank=config.max_bond_dim,
             orth_center_right=False,
         )
 
@@ -241,7 +241,7 @@ def apply_measurement_errors(
     return result
 
 
-n_operator = torch.tensor(
+n_operator: torch.Tensor = torch.tensor(
     [
         [0, 0],
         [0, 1],

emu_mps-1.2.4.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: emu-mps
+Version: 1.2.4
+Summary: Pasqal MPS based pulse emulator built on PyTorch
+Author-email: Anton Quelle <anton.quelle@pasqal.com>, Mauro Mendizabal <mauro.mendizabal-pico@pasqal.com>, Stefano Grava <stefano.grava@pasqal.com>, Pablo Le Henaff <pablo.le-henaff@pasqal.com>
+License: PASQAL OPEN-SOURCE SOFTWARE LICENSE AGREEMENT (MIT-derived)
+        
+        The author of the License is:
+          Pasqal, a Société par Actions Simplifiée (Simplified Joint Stock Company) registered under number 849 441 522 at the Registre du commerce et des sociétés (Trade and Companies Register) of Evry – France, headquartered at 7 rue Leonard de Vinci – 91300 – Massy – France, duly represented by its Président, M. Georges-Olivier REYMOND,
+        
+        Hereafter referred to as « the Licensor »
+        
+        - Permission is hereby granted, free of charge, to any person obtaining a copy of this software (the “Licensee”) and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. The Software is “as is”, without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose and non-infringement. In no event shall the authors or copyright holders be liable for any claim, damages or other liability, whether in an action of contract, tort or otherwise arising from, out of or in connection with the Software or the use or other dealings in the Software.
+        
+        - If use of the Software leads to the necessary use of any patent of the Licensor and/or any of its Affiliates (defined as a company owned or controlled by the Licensor), the Licensee is granted a royalty-free license, in any country where such patent is in force, to use the object of such patent; or use the process covered by such patent,
+        
+        - Such a patent license is granted for internal research or academic use of the Licensee's, which includes use by employees and students of the Licensee, acting on behalf of the Licensee, for research purposes only.
+        
+        - The License is governed by the laws of France. Any dispute relating to the License, notably its execution, performance and/or termination shall be brought to, heard and tried by the Tribunal Judiciaire de Paris, regardless of the rules of jurisdiction in the matter.
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=3.10
+Requires-Dist: emu-base==1.2.4
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <img src="docs/logos/LogoTaglineSoftGreen.svg">
+
+  # Emu-MPS
+</div>
+
+**Emu-mps** is a backend for the [Pulser low-level Quantum Programming toolkit](https://pulser.readthedocs.io) that lets you run Quantum Algorithms on a simulated device, using GPU acceleration if available. More in depth, emu-mps is designed to **emu**late the dynamics of programmable arrays of neutral atoms, with matrix product states (**mps**). While benchmarking is incomplete as of this writing, early results suggest that this design makes emu-mps faster and more memory-efficient than previous generations of quantum emulators at running simulations with large numbers of qubits.
+
+## Installation
+
+**Warning:** installing emu-mps will update pulser-core
+
+### Using `hatch`, `uv` or any pyproject-compatible Python manager
+
+To add `emu-mps` to your project, edit your `pyproject.toml` to add the line
+
+```toml
+  "emu-mps"
+```
+
+to the list of `dependencies`.
+
+
+### Using `pip` or `pipx`
+To install the `pipy` package using `pip` or `pipx`
+
+1. Create a `venv` if that's not done yet
+
+```sh
+$ python -m venv venv
+
+```
+
+2. Enter the venv
+
+If you're running Unix:
+
+```sh
+$ . venv/bin/activate
+```
+
+If you're running Windows:
+
+```sh
+C:\> /path/to/new/virtual/environment/Scripts/activate
+```
+
+3. Install the package
+
+```sh
+$ pip install emu-mps
+# or
+$ pipx install emu-mps
+```
+
+
+Join us on [Slack](https://pasqalworkspace.slack.com/archives/C07MUV5K7EU) or by [e-mail](mailto:emulation@pasqal.com) to give us feedback about how you plan to use Emu-MPS or if you require specific feature-upgrades.
+
+## Usage
+
+For the time being, the easiest way to learn how to use this package is to look
+at the [examples](examples/emu_mps_examples) and [notebooks](https://pasqal-io.github.io/emulators/latest/).
+
+See also the [full documentation](https://github.com/pasqal-io/emulators/blob/main/docs/index.md) for
+the API, information about contributing, benchmarks, etc.
+
+
+## Getting in touch
+
+- [Pasqal Community Portal](https://community.pasqal.com/) (forums, chat, tutorials, examples, code library).
+- [GitHub Repository](https://github.com/pasqal-io/quantum-evolution-kernel) (source code, issue tracker).
+- [Professional Support](https://www.pasqal.com/contact-us/) (if you need tech support, custom licenses, a variant of this library optimized for your workload, your own QPU, remote access to a QPU, ...)

emu_mps-1.2.4.dist-info/RECORD

@@ -0,0 +1,15 @@
+emu_mps/__init__.py,sha256=ou1n7ieIHlszqg9hBVYhvA8IBOpeduAWno3GC7X2dWE,646
+emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
+emu_mps/constants.py,sha256=41LYkKLUCz-oxPbd-j7nUDZuhIbUrnez6prT0uR0jcE,56
+emu_mps/hamiltonian.py,sha256=LALCdAQyYQgqFLVvtB9cHoUilT_7eU-6bAHRhu39lqM,15865
+emu_mps/mpo.py,sha256=7y6q0UIfyX9igQknqtgt6nymuVcgjHlH3-Qv7N7uOZE,8769
+emu_mps/mps.py,sha256=OjG_caqPOioCdOt-bFUkOf2xuNGnKzj0LaMc3EJCHi4,17855
+emu_mps/mps_backend.py,sha256=6fVaq-D4xyicYRjGjhqMEieC7---90LpfpbV7ZD7zkQ,2192
+emu_mps/mps_backend_impl.py,sha256=XCrmIjV-ZzI0nTOoaIaM6RGtfo4WHoJakVtoqoTZ7zI,20556
+emu_mps/mps_config.py,sha256=MxahrPDaOpfdB6SLG1610iDUOuLR04IaCjKQRk99ICY,3346
+emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
+emu_mps/tdvp.py,sha256=TH4CcBNczRURXYGPXndWKDs0jWXz_x9ozM961uGiSOw,5711
+emu_mps/utils.py,sha256=n9BcpuIz4Kl6EYlATaK8TKsyF-T7FTwbBo6KSAQYzl8,8066
+emu_mps-1.2.4.dist-info/METADATA,sha256=nAxhFZeRlvVOqTdfJ-szluFRpa0q8EwbmogczPni3kM,5465
+emu_mps-1.2.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-1.2.4.dist-info/RECORD,,

emu_mps-1.2.2.dist-info/METADATA

@@ -1,133 +0,0 @@
-Metadata-Version: 2.4
-Name: emu-mps
-Version: 1.2.2
-Summary: Pasqal MPS based pulse emulator built on PyTorch
-Author-email: Anton Quelle <anton.quelle@pasqal.com>, Mauro Mendizabal <mauro.mendizabal-pico@pasqal.com>, Stefano Grava <stefano.grava@pasqal.com>, Pablo Le Henaff <pablo.le-henaff@pasqal.com>
-License: Proprietary
-Classifier: License :: Other/Proprietary License
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: Implementation :: CPython
-Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: >=3.10
-Requires-Dist: emu-base==1.2.2
-Description-Content-Type: text/markdown
-
-<div align="center">
-  <img src="docs/logos/LogoTaglineSoftGreen.svg">
-
-  # Emu-MPS
-</div>
-
-**EMU-MPS** is a Pulser backend, designed to **EMU**late the dynamics of programmable arrays of neutral atoms, with matrix product states (**MPS**). It allows users to increase the number of qubits and reduce computation time.
-
-Join us on [Slack](https://pasqalworkspace.slack.com/archives/C07MUV5K7EU) or by [e-mail](mailto:emulation@pasqal.com) to give us feedback about how you plan to use Emu-MPS or if you require specific feature-upgrades.
-
-
-## Getting started
-
-You can install from source, or download the package from the private pypi registry that pasqal maintains in gitlab.
-For developers, we recommend installing from source, for users we recommend installing from the registry.
-
-**Warning:** installing emu-mps will update pulser-core
-
-We always recommend using a virtual environment.
-
-<details>
-  <summary>Click me to see how it is done</summary>
-
-  #### Create a virtual environment using python
-
-  ```
-  python -m venv .venv
-  ```
-
-  Or
-
-  ```
-  python -m venv /path/to/new/virtual/environment
-  ```
-
-  Replace `/path/to/new/virtual/environment` with your desired directory path.
-
-  Then activate the environment On linux or MacOS
-
-  ```
-  source /path/to/new/virtual/environment/bin/activate
-  ```
-
-  While on Windows it's
-
-  ```
-  C:\> /path/to/new/virtual/environment/Scripts/activate
-  ```
-
-  Remember to replace `/path/to/new/virtual/environment` with the actual path to your virtual environment. Once the environment is activated, you can clone emu_mps and install it using
-
-</details>
-
-### installing from the registry
-
-When pip is configured to know about the pasqal registry, Emu-MPS installs as
-
-```bash
-pip install emu-mps
-```
-When pip is not already configured, the easiest way to do so, is to add a file `~/.config/pip/pip.conf` containing:
-
-```
-[global]
-extra-index-url=https://gitlab.pasqal.com/api/v4/projects/597/packages/pypi/simple
-                possible.other.urls
-```
-
-As this shows, it is also possible to have multiple extra repositories configured. Note that the order is not important.
-
-It is also possible to add the `extra-index-url` to the `pip install` command directly, if you somehow don't want to create a `pip.conf` file.
-
-### installing from source
-git clone this [repository ](https://gitlab.pasqal.com/emulation/rydberg-atoms/emu-ct) or download
-
-
-Then, `cd` into the root folder of the repo and type
-
-```bash
-pip install -e .
-```
-
-<details>
-  <summary>Guidelines for developers </summary>
-  We recommend using an environment, git clone the repository, then inside the `emu_mps` folder
-
-```bash
-pip install -e .
-```
-
-  Also, the installation of pytest, nbmake, pre-commit.
-
-  Do not forget to run the unit test suite by simply running `pytest` command.
-
-  Another way can be using hatch.
-
-  #### virtual environment with `hatch`
-
-  ```bash
-  python -m pip install hatch
-  python -m hatch -v shell
-  ```
-
-  When inside the shell with development dependencies, install first the pre-commit hook:
-  ```
-  pre-commit install
-  ```
-</details>
-
-## Check the tutorial notebooks and example scripts
-
-For more information, you can check the tutorials and examples located in the [examples folder](https://gitlab.pasqal.com/emulation/rydberg-atoms/emu-ct/-/tree/main/examples?ref_type=heads)
-
-## Documentation
-
-Please check the [documentation](./docs/index.md) page for more info about contributing, the API, benchmarks, etc.
-
-
-![Code Coverage](https://img.shields.io/badge/Coverage-95%25-brightgreen.svg)

emu_mps-1.2.2.dist-info/RECORD

@@ -1,14 +0,0 @@
-emu_mps/__init__.py,sha256=O7sRnIvjoRsX-2lVy3NwydAIos69aCbYyFNpUGwAoTw,646
-emu_mps/algebra.py,sha256=fMkhie3iy-zEivE6G3EpJoo2o6Gde0rqrOwBUKOBkME,5384
-emu_mps/hamiltonian.py,sha256=Ojb-T2bpb-Ym99IbFsTC2WqQ4C06IbcYK4VusA2pt_0,15799
-emu_mps/mpo.py,sha256=RA9n50li1NJsIWCid-uOFgYV-usX4S7ZgmRdPLht8hs,8611
-emu_mps/mps.py,sha256=3Oi_gT5KdY46Q3l9uGifWyrb77DUQOHxnu_O9ky7MYo,18338
-emu_mps/mps_backend.py,sha256=ZCZMCsJz4-Oi0Ox2CqCOWqQWdE9ZJj6n_Jssd0g_1T0,972
-emu_mps/mps_backend_impl.py,sha256=dEhBVciS7bdZ0AymbgDVyNUfprxuv-GgUoRcmKBkLHY,19512
-emu_mps/mps_config.py,sha256=-ScHCNiiAwg_kMBSGOsXydjnbo_y3AqrvqpG3RzTYmU,2703
-emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
-emu_mps/tdvp.py,sha256=mDStFKgbG3OXsEG-Ja8NJKsTGyLxsjNS6sQXK7B33EY,6088
-emu_mps/utils.py,sha256=05s5TBE2zhsJwugbQFDCFQi7L0pVX0ukI-WIW2zN5TA,8052
-emu_mps-1.2.2.dist-info/METADATA,sha256=fYI3Qzl0Kdm0FxPR0zQtJOegTg05-zOYuiv6Ee74fQM,4203
-emu_mps-1.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-1.2.2.dist-info/RECORD,,