emu-mps 2.3.0__py3-none-any.whl → 2.4.1__py3-none-any.whl

emu_mps/__init__.py

@@ -9,12 +9,11 @@ from pulser.backend import (
     StateResult,
     EnergySecondMoment,
 )
-from .mps_config import MPSConfig
+from .mps_config import MPSConfig, Solver
 from .mpo import MPO
 from .mps import MPS, inner
 from .mps_backend import MPSBackend
 from .observables import EntanglementEntropy
-from emu_base import aggregate
 
 
 __all__ = [
@@ -23,6 +22,7 @@ __all__ = [
     "MPS",
     "inner",
     "MPSConfig",
+    "Solver",
     "MPSBackend",
     "StateResult",
     "BitStrings",
@@ -33,8 +33,7 @@ __all__ = [
     "Energy",
     "EnergyVariance",
     "EnergySecondMoment",
-    "aggregate",
     "EntanglementEntropy",
 ]
 
-__version__ = "2.3.0"
+__version__ = "2.4.1"

emu_mps/algebra.py

@@ -1,11 +1,8 @@
 from __future__ import annotations
 
-from typing import Optional
 
 import torch
 import math
-
-from emu_mps import MPSConfig
 from emu_mps.utils import truncate_impl
 
 
@@ -119,7 +116,8 @@ def zip_right_step(
 def zip_right(
     top_factors: list[torch.Tensor],
     bottom_factors: list[torch.Tensor],
-    config: Optional[MPSConfig] = None,
+    precision: float,
+    max_bond_dim: int,
 ) -> list[torch.Tensor]:
     """
     Returns a new matrix product, resulting from applying `top` to `bottom`.
@@ -139,8 +137,6 @@ 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.")
 
@@ -151,6 +147,6 @@ def zip_right(
         new_factors.append(res)
     new_factors[-1] @= slider[:, :, 0]
 
-    truncate_impl(new_factors, config=config)
+    truncate_impl(new_factors, precision=precision, max_bond_dim=max_bond_dim)
 
     return new_factors

emu_mps/custom_callback_implementations.py

@@ -25,9 +25,9 @@ def qubit_occupation_mps_impl(
     """
     Custom implementation of the occupation ❬ψ|nᵢ|ψ❭ for the EMU-MPS.
     """
-    op = torch.tensor(
-        [[[0.0, 0.0], [0.0, 1.0]]], dtype=torch.complex128, device=state.factors[0].device
-    )
+    dim = state.dim
+    op = torch.zeros(1, dim, dim, dtype=torch.complex128, device=state.factors[0].device)
+    op[0, 1, 1] = 1.0
     return state.expect_batch(op).real.view(-1).cpu()
 
 

emu_mps/mpo.py

@@ -7,7 +7,7 @@ from pulser.backend import State, Operator
 from emu_base import DEVICE_COUNT
 from emu_mps.algebra import add_factors, scale_factors, zip_right
 from pulser.backend.operator import FullOp, QuditOp
-from emu_mps.mps import MPS
+from emu_mps.mps import MPS, DEFAULT_MAX_BOND_DIM, DEFAULT_PRECISION
 from emu_mps.utils import new_left_bath, assign_devices
 
 
@@ -62,7 +62,8 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
         factors = zip_right(
             self.factors,
             other.factors,
-            config=other.config,
+            precision=other.precision,
+            max_bond_dim=other.max_bond_dim,
         )
         return MPS(factors, orthogonality_center=0, eigenstates=other.eigenstates)
 
@@ -107,7 +108,12 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
             the composed operator
         """
         assert isinstance(other, MPO), "MPO can only be applied to another MPO"
-        factors = zip_right(self.factors, other.factors)
+        factors = zip_right(
+            self.factors,
+            other.factors,
+            precision=DEFAULT_PRECISION,
+            max_bond_dim=DEFAULT_MAX_BOND_DIM,
+        )
         return MPO(factors)
 
     def expect(self, state: State) -> torch.Tensor:

emu_mps/mps.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 import math
 from collections import Counter
 from typing import List, Optional, Sequence, TypeVar, Mapping
@@ -8,16 +7,19 @@ import torch
 
 from pulser.backend.state import State, Eigenstate
 from emu_base import DEVICE_COUNT, apply_measurement_errors
-from emu_mps import MPSConfig
 from emu_mps.algebra import add_factors, scale_factors
 from emu_mps.utils import (
     assign_devices,
     truncate_impl,
     tensor_trace,
-    n_operator,
 )
 
+
 ArgScalarType = TypeVar("ArgScalarType")
+dtype = torch.complex128
+
+DEFAULT_PRECISION = 1e-5
+DEFAULT_MAX_BOND_DIM = 1024
 
 
 class MPS(State[complex, torch.Tensor]):
@@ -35,12 +37,14 @@ class MPS(State[complex, torch.Tensor]):
         /,
         *,
         orthogonality_center: Optional[int] = None,
-        config: Optional[MPSConfig] = None,
+        precision: float = DEFAULT_PRECISION,
+        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
         num_gpus_to_use: Optional[int] = DEVICE_COUNT,
         eigenstates: Sequence[Eigenstate] = ("r", "g"),
     ):
         """
-        This constructor creates a MPS directly from a list of tensors. It is for internal use only.
+        This constructor creates a MPS directly from a list of tensors. It is
+        for internal use only.
 
         Args:
             factors: the tensors for each site
@@ -52,12 +56,14 @@ class MPS(State[complex, torch.Tensor]):
                 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)
-            config: the emu-mps config object passed to the run method
+            precision: the precision with which to keep this MPS
+            max_bond_dim: the maximum bond dimension to allow for this MPS
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu, None=keep the existing device assignment.
         """
         super().__init__(eigenstates=eigenstates)
-        self.config = config if config is not None else MPSConfig()
+        self.precision = precision
+        self.max_bond_dim = max_bond_dim
         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"
@@ -69,6 +75,17 @@ class MPS(State[complex, torch.Tensor]):
         self.num_sites = len(factors)
         assert self.num_sites > 1  # otherwise, do state vector
 
+        self.dim = len(self.eigenstates)
+        assert all(factors[i].shape[1] == self.dim for i in range(self.num_sites)), (
+            "All tensors should have the same physical dimension as the number "
+            "of eigenstates"
+        )
+
+        self.n_operator = torch.zeros(
+            self.dim, self.dim, dtype=dtype, device=self.factors[0].device
+        )
+        self.n_operator[1, 1] = 1.0
+
         assert (orthogonality_center is None) or (
             0 <= orthogonality_center < self.num_sites
         ), "Invalid orthogonality center provided"
@@ -86,7 +103,8 @@ class MPS(State[complex, torch.Tensor]):
     def make(
         cls,
         num_sites: int,
-        config: Optional[MPSConfig] = None,
+        precision: float = DEFAULT_PRECISION,
+        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
         num_gpus_to_use: int = DEVICE_COUNT,
         eigenstates: Sequence[Eigenstate] = ["0", "1"],
     ) -> MPS:
@@ -95,21 +113,35 @@ class MPS(State[complex, torch.Tensor]):
 
         Args:
             num_sites: the number of qubits
-            config: the MPSConfig
+            precision: the precision with which to keep this MPS
+            max_bond_dim: the maximum bond dimension to allow for this MPS
             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")
 
-        return cls(
-            [
-                torch.tensor([[[1.0], [0.0]]], dtype=torch.complex128)
+        if len(eigenstates) == 2:
+            ground_state = [
+                torch.tensor([[[1.0], [0.0]]], dtype=dtype) for _ in range(num_sites)
+            ]
+
+        elif len(eigenstates) == 3:  # (g,r,x)
+            ground_state = [
+                torch.tensor([[[1.0], [0.0], [0.0]]], dtype=dtype)
                 for _ in range(num_sites)
-            ],
-            config=config,
+            ]
+
+        else:
+            raise ValueError(
+                "Unsupported basis provided. The supported "
+                "bases are:{('0','1'),('r','g'),('r','g','x')}"
+            )
+
+        return cls(
+            ground_state,
+            precision=precision,
+            max_bond_dim=max_bond_dim,
             num_gpus_to_use=num_gpus_to_use,
             orthogonality_center=0,  # Arbitrary: every qubit is an orthogonality center.
             eigenstates=eigenstates,
@@ -140,7 +172,8 @@ class MPS(State[complex, torch.Tensor]):
 
         for i in range(lr_swipe_start, desired_orthogonality_center):
             q, r = torch.linalg.qr(self.factors[i].view(-1, self.factors[i].shape[2]))
-            self.factors[i] = q.view(self.factors[i].shape[0], 2, -1)
+
+            self.factors[i] = q.view(self.factors[i].shape[0], self.dim, -1)
             self.factors[i + 1] = torch.tensordot(
                 r.to(self.factors[i + 1].device), self.factors[i + 1], dims=1
             )
@@ -153,9 +186,10 @@ class MPS(State[complex, torch.Tensor]):
 
         for i in range(rl_swipe_start, desired_orthogonality_center, -1):
             q, r = torch.linalg.qr(
-                self.factors[i].view(self.factors[i].shape[0], -1).mT,
+                self.factors[i].contiguous().view(self.factors[i].shape[0], -1).mT,
             )
-            self.factors[i] = q.mT.view(-1, 2, self.factors[i].shape[2])
+
+            self.factors[i] = q.mT.view(-1, self.dim, self.factors[i].shape[2])
             self.factors[i - 1] = torch.tensordot(
                 self.factors[i - 1], r.to(self.factors[i - 1].device), ([2], [1])
             )
@@ -172,7 +206,9 @@ class MPS(State[complex, torch.Tensor]):
         An in-place operation.
         """
         self.orthogonalize(self.num_sites - 1)
-        truncate_impl(self.factors, config=self.config)
+        truncate_impl(
+            self.factors, precision=self.precision, max_bond_dim=self.max_bond_dim
+        )
         self.orthogonality_center = 0
 
     def get_max_bond_dim(self) -> int:
@@ -182,7 +218,7 @@ class MPS(State[complex, torch.Tensor]):
         Returns:
             the largest bond dimension in the state
         """
-        return max((x.shape[2] for x in self.factors), default=0)
+        return max((factor.shape[2] for factor in self.factors), default=0)
 
     def sample(
         self,
@@ -206,8 +242,6 @@ class MPS(State[complex, torch.Tensor]):
         assert one_state in {None, "r", "1"}
         self.orthogonalize(0)
 
-        rnd_matrix = torch.rand(num_shots, self.num_sites).to(self.factors[0].device)
-
         bitstrings: Counter[str] = Counter()
 
         # Shots are performed in batches.
@@ -215,55 +249,45 @@ class MPS(State[complex, torch.Tensor]):
         max_batch_size = 32
 
         shots_done = 0
+
         while shots_done < num_shots:
             batch_size = min(max_batch_size, num_shots - shots_done)
             batched_accumulator = torch.ones(
-                batch_size, 1, dtype=torch.complex128, device=self.factors[0].device
+                batch_size, 1, dtype=dtype, device=self.factors[0].device
             )
 
-            batch_outcomes = torch.empty(batch_size, self.num_sites, dtype=torch.bool)
-
+            batch_outcomes = torch.empty(batch_size, self.num_sites, dtype=torch.int)
+            rangebatch = torch.arange(batch_size)
             for qubit, factor in enumerate(self.factors):
                 batched_accumulator = torch.tensordot(
                     batched_accumulator.to(factor.device), factor, dims=1
                 )
 
-                # Probability of measuring qubit == 0 for each shot in the batch
-                probas = (
-                    torch.linalg.vector_norm(batched_accumulator[:, 0, :], dim=1) ** 2
-                )
+                # Probabilities for each state in the basis
+                probn = torch.linalg.vector_norm(batched_accumulator, dim=2) ** 2
 
-                outcomes = (
-                    rnd_matrix[shots_done : shots_done + batch_size, qubit].to(
-                        factor.device
-                    )
-                    > probas
-                )
-                batch_outcomes[:, qubit] = outcomes
+                # list of: 0,1 for |g>,|r> or 0,1,2 for |g>,|r>,|x>
+                outcomes = torch.multinomial(probn, num_samples=1).reshape(-1)
 
-                # Batch collapse qubit
-                tmp = torch.stack((~outcomes, outcomes), dim=1).to(dtype=torch.complex128)
+                batch_outcomes[:, qubit] = outcomes
 
-                batched_accumulator = (
-                    torch.tensordot(batched_accumulator, tmp, dims=([1], [1]))
-                    .diagonal(dim1=0, dim2=2)
-                    .transpose(1, 0)
-                )
-                batched_accumulator /= torch.sqrt(
-                    (~outcomes) * probas + outcomes * (1 - probas)
-                ).unsqueeze(1)
+                # expected shape (batch_size, bond_dim)
+                batched_accumulator = batched_accumulator[rangebatch, outcomes, :]
 
             shots_done += batch_size
 
             for outcome in batch_outcomes:
-                bitstrings.update(["".join("0" if x == 0 else "1" for x in outcome)])
+                bitstrings.update(["".join("1" if x == 1 else "0" for x in outcome)])
 
-        if p_false_neg > 0 or p_false_pos > 0:
+        if p_false_neg > 0 or p_false_pos > 0 and self.dim == 2:
             bitstrings = apply_measurement_errors(
                 bitstrings,
                 p_false_pos=p_false_pos,
                 p_false_neg=p_false_neg,
             )
+        if p_false_pos > 0 and self.dim > 2:
+            raise NotImplementedError("Not implemented for qudits > 2 levels")
+
         return bitstrings
 
     def norm(self) -> torch.Tensor:
@@ -312,8 +336,11 @@ class MPS(State[complex, torch.Tensor]):
     def entanglement_entropy(self, mps_site: int) -> torch.Tensor:
         """
         Returns
-        the Von Neumann entanglement entropy of the state `mps` at the bond between sites b and b+1
+        the Von Neumann entanglement entropy of the state `mps` at the bond
+        between sites b and b+1
+
         S = -Σᵢsᵢ² log(sᵢ²)),
+
         where sᵢ are the singular values at the chosen bond.
         """
         self.orthogonalize(mps_site)
@@ -358,7 +385,8 @@ class MPS(State[complex, torch.Tensor]):
         new_tt = add_factors(self.factors, other.factors)
         result = MPS(
             new_tt,
-            config=self.config,
+            precision=self.precision,
+            max_bond_dim=self.max_bond_dim,
             num_gpus_to_use=None,
             orthogonality_center=None,  # Orthogonality is lost.
             eigenstates=self.eigenstates,
@@ -384,7 +412,8 @@ class MPS(State[complex, torch.Tensor]):
         factors = scale_factors(self.factors, scalar, which=which)
         return MPS(
             factors,
-            config=self.config,
+            precision=self.precision,
+            max_bond_dim=self.max_bond_dim,
             num_gpus_to_use=None,
             orthogonality_center=self.orthogonality_center,
             eigenstates=self.eigenstates,
@@ -412,25 +441,43 @@ class MPS(State[complex, torch.Tensor]):
         Returns:
             The resulting MPS representation of the state.s
         """
+
+        leak = ""
+        one = "r"
         basis = set(eigenstates)
+
         if basis == {"r", "g"}:
-            one = "r"
+            pass
         elif basis == {"0", "1"}:
             one = "1"
+        elif basis == {"g", "r", "x"}:
+            leak = "x"
         else:
             raise ValueError("Unsupported basis provided")
+        dim = len(eigenstates)
+        if dim == 2:
+            basis_0 = torch.tensor([[[1.0], [0.0]]], dtype=dtype)  # ground state
+            basis_1 = torch.tensor([[[0.0], [1.0]]], dtype=dtype)  # excited state
 
-        basis_0 = torch.tensor([[[1.0], [0.0]]], dtype=torch.complex128)  # ground state
-        basis_1 = torch.tensor([[[0.0], [1.0]]], dtype=torch.complex128)  # excited state
+        elif dim == 3:
+            basis_0 = torch.tensor([[[1.0], [0.0], [0.0]]], dtype=dtype)  # ground state
+            basis_1 = torch.tensor([[[0.0], [1.0], [0.0]]], dtype=dtype)  # excited state
+            basis_x = torch.tensor([[[0.0], [0.0], [1.0]]], dtype=dtype)  # leakage state
 
         accum_mps = MPS(
-            [torch.zeros((1, 2, 1), dtype=torch.complex128)] * n_qudits,
+            [torch.zeros((1, dim, 1), dtype=dtype)] * n_qudits,
             orthogonality_center=0,
             eigenstates=eigenstates,
         )
-
         for state, amplitude in amplitudes.items():
-            factors = [basis_1 if ch == one else basis_0 for ch in state]
+            factors = []
+            for ch in state:
+                if ch == one:
+                    factors.append(basis_1)
+                elif ch == leak:
+                    factors.append(basis_x)
+                else:
+                    factors.append(basis_0)
             accum_mps += amplitude * MPS(factors, eigenstates=eigenstates)
         norm = accum_mps.norm()
         if not math.isclose(1.0, norm, rel_tol=1e-5, abs_tol=0.0):
@@ -453,9 +500,7 @@ class MPS(State[complex, torch.Tensor]):
             else self.orthogonalize(0)
         )
 
-        result = torch.zeros(
-            self.num_sites, single_qubit_operators.shape[0], dtype=torch.complex128
-        )
+        result = torch.zeros(self.num_sites, single_qubit_operators.shape[0], dtype=dtype)
 
         center_factor = self.factors[orthogonality_center]
         for qubit_index in range(orthogonality_center, self.num_sites):
@@ -503,7 +548,7 @@ class MPS(State[complex, torch.Tensor]):
         )
 
     def get_correlation_matrix(
-        self, *, operator: torch.Tensor = n_operator
+        self, operator: torch.Tensor | None = None
     ) -> torch.Tensor:
         """
         Efficiently compute the symmetric correlation matrix
@@ -511,14 +556,18 @@ class MPS(State[complex, torch.Tensor]):
         in basis ("r", "g").
 
         Args:
-            operator: a 2x2 Torch tensor to use
+            operator: a 2x2 (or 3x3) Torch tensor to use
 
         Returns:
             the corresponding correlation matrix
         """
-        assert operator.shape == (2, 2)
 
-        result = torch.zeros(self.num_sites, self.num_sites, dtype=torch.complex128)
+        if operator is None:
+            operator = self.n_operator
+
+        assert operator.shape == (self.dim, self.dim), "Operator has wrong shape"
+
+        result = torch.zeros(self.num_sites, self.num_sites, dtype=dtype)
 
         for left in range(0, self.num_sites):
             self.orthogonalize(left)

emu_mps/mps_backend_impl.py

@@ -10,7 +10,6 @@ import uuid
 from copy import deepcopy
 from collections import Counter
 from enum import Enum, auto
-from resource import RUSAGE_SELF, getrusage
 from types import MethodType
 from typing import Any, Optional
 
@@ -18,7 +17,7 @@ import torch
 from pulser import Sequence
 from pulser.backend import EmulationConfig, Observable, Results, State
 
-from emu_base import DEVICE_COUNT, PulserData
+from emu_base import DEVICE_COUNT, PulserData, get_max_rss
 from emu_base.math.brents_root_finding import BrentsRootFinder
 from emu_base.utils import deallocate_tensor
 
@@ -26,9 +25,9 @@ from emu_mps.hamiltonian import make_H, update_H
 from emu_mps.mpo import MPO
 from emu_mps.mps import MPS
 from emu_mps.mps_config import MPSConfig
-from emu_base.noise import pick_dark_qubits
 from emu_base.jump_lindblad_operators import compute_noise_from_lindbladians
 import emu_mps.optimatrix as optimat
+from emu_mps.solver import Solver
 from emu_mps.solver_utils import (
     evolve_pair,
     evolve_single,
@@ -69,14 +68,7 @@ class Statistics(Observable):
         """Calculates the observable to store in the Results."""
         assert isinstance(state, MPS)
         duration = self.data[-1]
-        if state.factors[0].is_cuda:
-            max_mem_per_device = (
-                torch.cuda.max_memory_allocated(device) * 1e-6
-                for device in range(torch.cuda.device_count())
-            )
-            max_mem = max(max_mem_per_device)
-        else:
-            max_mem = getrusage(RUSAGE_SELF).ru_maxrss * 1e-3
+        max_mem = get_max_rss(state.factors[0].is_cuda)
 
         config.logger.info(
             f"step = {len(self.data)}/{self.timestep_count}, "
@@ -193,10 +185,9 @@ class MPSBackendImpl:
     def init_dark_qubits(self) -> None:
         # has_state_preparation_error
         if self.config.noise_model.state_prep_error > 0.0:
+            bad_atoms = self.pulser_data.hamiltonian.bad_atoms
             self.well_prepared_qubits_filter = torch.logical_not(
-                pick_dark_qubits(
-                    self.config.noise_model.state_prep_error, self.qubit_count
-                )
+                torch.tensor(list(bool(x) for x in bad_atoms.values()))
             )
         else:
             self.well_prepared_qubits_filter = None
@@ -218,7 +209,8 @@ class MPSBackendImpl:
         if initial_state is None:
             self.state = MPS.make(
                 self.qubit_count,
-                config=self.config,
+                precision=self.config.precision,
+                max_bond_dim=self.config.max_bond_dim,
                 num_gpus_to_use=self.config.num_gpus_to_use,
             )
             return
@@ -245,13 +237,15 @@ class MPSBackendImpl:
         initial_state = MPS(
             # Deep copy of every tensor of the initial state.
             [f.detach().clone() for f in initial_state.factors],
-            config=self.config,
+            precision=self.config.precision,
+            max_bond_dim=self.config.max_bond_dim,
             num_gpus_to_use=self.config.num_gpus_to_use,
             eigenstates=initial_state.eigenstates,
         )
         initial_state.truncate()
         initial_state *= 1 / initial_state.norm()
         self.state = initial_state
+        self.state.orthogonalize(0)
 
     def init_hamiltonian(self) -> None:
         """
@@ -699,18 +693,20 @@ class DMRGBackendImpl(MPSBackendImpl):
         mps_config: MPSConfig,
         pulser_data: PulserData,
         energy_tolerance: float = 1e-5,
-        max_sweeps: int = 999,
-        residual_tolerance: float = 1e-7,
+        max_sweeps: int = 2000,
     ):
+
+        if mps_config.noise_model.noise_types != ():
+            raise NotImplementedError(
+                "DMRG solver does not currently support noise types"
+                f"you are using: {mps_config.noise_model.noise_types}"
+            )
         super().__init__(mps_config, pulser_data)
-        self.init()
-        self.state.orthogonalize(0)
         self.previous_energy: Optional[float] = None
         self.current_energy: Optional[float] = None
         self.sweep_count: int = 0
         self.energy_tolerance: float = energy_tolerance
         self.max_sweeps: int = max_sweeps
-        self.residual_tolerance: float = residual_tolerance
 
     def convergence_check(self, energy_tolerance: float) -> bool:
         if self.previous_energy is None or self.current_energy is None:
@@ -728,63 +724,72 @@ class DMRGBackendImpl(MPSBackendImpl):
             SwipeDirection.RIGHT_TO_LEFT,
         ), "Unknown Swipe direction"
 
-        if self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT:
-            left_idx, right_idx = idx, idx + 1
-        elif self.swipe_direction == SwipeDirection.RIGHT_TO_LEFT:
-            left_idx, right_idx = idx - 1, idx
-
         orth_center_right = self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT
         new_L, new_R, energy = minimize_energy_pair(
-            state_factors=self.state.factors[left_idx : right_idx + 1],
-            ham_factors=self.hamiltonian.factors[left_idx : right_idx + 1],
+            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]),
             orth_center_right=orth_center_right,
             config=self.config,
-            residual_tolerance=self.residual_tolerance,
+            residual_tolerance=self.config.precision,
         )
-        self.state.factors[left_idx], self.state.factors[right_idx] = new_L, new_R
-        self.state.orthogonality_center = right_idx if orth_center_right else left_idx
+        self.state.factors[idx], self.state.factors[idx + 1] = new_L, new_R
+        self.state.orthogonality_center = idx + 1 if orth_center_right else idx
         self.current_energy = energy
 
         # updating baths and orthogonality center
         if self.swipe_direction == SwipeDirection.LEFT_TO_RIGHT:
+            self._left_to_right_update(idx)
+        elif self.swipe_direction == SwipeDirection.RIGHT_TO_LEFT:
+            self._right_to_left_update(idx)
+        else:
+            raise Exception("Did not expect this")
+
+        self.save_simulation()
+
+    def _left_to_right_update(self, idx: int) -> None:
+        if idx < self.qubit_count - 2:
             self.left_baths.append(
                 new_left_bath(
                     self.get_current_left_bath(),
-                    self.state.factors[left_idx],
-                    self.hamiltonian.factors[right_idx],
-                ).to(self.state.factors[right_idx].device)
+                    self.state.factors[idx],
+                    self.hamiltonian.factors[idx],
+                ).to(self.state.factors[idx + 1].device)
             )
             self.right_baths.pop()
             self.sweep_index += 1
 
-            if self.sweep_index == self.qubit_count - 1:
-                self.swipe_direction = SwipeDirection.RIGHT_TO_LEFT
+        if self.sweep_index == self.qubit_count - 2:
+            self.swipe_direction = SwipeDirection.RIGHT_TO_LEFT
 
-        elif self.swipe_direction == SwipeDirection.RIGHT_TO_LEFT:
+    def _right_to_left_update(self, idx: int) -> None:
+        if idx > 0:
             self.right_baths.append(
                 new_right_bath(
                     self.get_current_right_bath(),
-                    self.state.factors[right_idx],
-                    self.hamiltonian.factors[right_idx],
-                ).to(self.state.factors[left_idx].device)
+                    self.state.factors[idx + 1],
+                    self.hamiltonian.factors[idx + 1],
+                ).to(self.state.factors[idx].device)
             )
             self.left_baths.pop()
             self.sweep_index -= 1
 
-            if self.sweep_index == 0:
-                self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT
-                self.sweep_count += 1
-                self.sweep_complete()
+        if self.sweep_index == 0:
+            self.state.orthogonalize(0)
+            self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT
+            self.sweep_count += 1
+            self.sweep_complete()
 
     def sweep_complete(self) -> None:
         # This marks the end of one full sweep: checking convergence
         if self.convergence_check(self.energy_tolerance):
+            self.current_time = self.target_time
             self.timestep_complete()
         elif self.sweep_count + 1 > self.max_sweeps:
-            # not converged: restart a new sweep
+            # not converged
             raise RuntimeError(f"DMRG did not converge after {self.max_sweeps} sweeps")
         else:
+            # not converged for the current sweep. restart
             self.previous_energy = self.current_energy
 
         assert self.sweep_index == 0
@@ -798,5 +803,6 @@ def create_impl(sequence: Sequence, config: MPSConfig) -> MPSBackendImpl:
 
     if pulser_data.has_lindblad_noise:
         return NoisyMPSBackendImpl(config, pulser_data)
-
+    if config.solver == Solver.DMRG:
+        return DMRGBackendImpl(config, pulser_data)
     return MPSBackendImpl(config, pulser_data)

emu_mps/mps_config.py

@@ -4,6 +4,9 @@ from types import MethodType
 import copy
 
 from emu_base import DEVICE_COUNT
+from emu_mps.mps import MPS, DEFAULT_MAX_BOND_DIM, DEFAULT_PRECISION
+from emu_mps.mpo import MPO
+from emu_mps.solver import Solver
 from emu_mps.custom_callback_implementations import (
     energy_mps_impl,
     energy_second_moment_mps_impl,
@@ -32,15 +35,17 @@ class MPSConfig(EmulationConfig):
     See the API for that class for a list of available options.
 
     Args:
-        dt: the timestep size that the solver uses. Note that observables are
+        dt: The timestep size that the solver uses. Note that observables are
             only calculated if the evaluation_times are divisible by dt.
-        precision: up to what precision the state is truncated
-        max_bond_dim: the maximum bond dimension that the state is allowed to have.
+        precision: Up to what precision the state is truncated.
+            Defaults to `1e-5`.
+        max_bond_dim: The maximum bond dimension that the state is allowed to have.
+            Defaults to `1024`.
         max_krylov_dim:
-            the size of the krylov subspace that the Lanczos algorithm maximally builds
+            The size of the krylov subspace that the Lanczos algorithm maximally builds
         extra_krylov_tolerance:
-            the Lanczos algorithm uses this*precision as the convergence tolerance
-        num_gpus_to_use: during the simulation, distribute the state over this many GPUs
+            The Lanczos algorithm uses this*precision as the convergence tolerance
+        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.
         optimize_qubit_ordering: Optimize the register ordering. Improves performance and
@@ -50,10 +55,15 @@ class MPSConfig(EmulationConfig):
         log_level: How much to log. Set to `logging.WARN` to get rid of the timestep info.
         log_file: If specified, log to this file rather than stout.
         autosave_prefix: filename prefix for autosaving simulation state to file
-        autosave_dt: minimum time interval in seconds between two autosaves.
+        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
+        solver: Chooses the solver algorithm to run a sequence.
+            Two options are currently available:
+            ``TDVP``, which performs ordinary time evolution,
+            and ``DMRG``, which adiabatically follows the ground state
+            of a given adiabatic pulse.
+        kwargs: Arguments that are passed to the base class
 
     Examples:
         >>> num_gpus_to_use = 2 #use 2 gpus if available, otherwise 1 or cpu
@@ -65,13 +75,15 @@ class MPSConfig(EmulationConfig):
 
     # Whether to warn if unexpected kwargs are received
     _enforce_expected_kwargs: ClassVar[bool] = True
+    _state_type = MPS
+    _operator_type = MPO
 
     def __init__(
         self,
         *,
         dt: int = 10,
-        precision: float = 1e-5,
-        max_bond_dim: int = 1024,
+        precision: float = DEFAULT_PRECISION,
+        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int = DEVICE_COUNT,
@@ -81,6 +93,7 @@ class MPSConfig(EmulationConfig):
         log_file: pathlib.Path | None = None,
         autosave_prefix: str = "emu_mps_save_",
         autosave_dt: int = 600,  # 10 minutes
+        solver: Solver = Solver.TDVP,
         **kwargs: Any,
     ):
         kwargs.setdefault("observables", [BitStrings(evaluation_times=[1.0])])
@@ -97,6 +110,7 @@ class MPSConfig(EmulationConfig):
             log_file=log_file,
             autosave_prefix=autosave_prefix,
             autosave_dt=autosave_dt,
+            solver=solver,
             **kwargs,
         )
         if self.optimize_qubit_ordering:
@@ -144,6 +158,7 @@ class MPSConfig(EmulationConfig):
             "log_file",
             "autosave_prefix",
             "autosave_dt",
+            "solver",
         }
 
     def monkeypatch_observables(self) -> None:

emu_mps/observables.py

@@ -6,7 +6,18 @@ import torch
 
 
 class EntanglementEntropy(Observable):
-    """Entanglement Entropy subclass used only in emu_mps"""
+    """Entanglement Entropy of the state partition at qubit `mps_site`.
+
+    Args:
+        mps_site: the qubit index at which the bipartition is made.
+            All qubits with index $\\leq$ `mps_site` are put in the left partition.
+        evaluation_times: The relative times at which to store the state.
+            If left as `None`, uses the ``default_evaluation_times`` of the
+            backend's ``EmulationConfig``.
+        tag_suffix: An optional suffix to append to the tag. Needed if
+            multiple instances of the same observable are given to the
+            same EmulationConfig.
+    """
 
     def __init__(
         self,

emu_mps/solver.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class Solver(str, Enum):
+    TDVP = "tdvp"
+    DMRG = "dmrg"

emu_mps/solver_utils.py

@@ -51,9 +51,10 @@ def make_op(
         .view(left_ham_factor.shape[0], 4, 4, -1)
     )
 
-    op = lambda x: time_step * apply_effective_Hamiltonian(
-        x, combined_hamiltonian_factors, left_bath, right_bath
-    )
+    def op(x: torch.Tensor) -> torch.Tensor:
+        return time_step * apply_effective_Hamiltonian(
+            x, combined_hamiltonian_factors, left_bath, right_bath
+        )
 
     return combined_state_factors, right_device, op
 
@@ -205,17 +206,18 @@ def evolve_single(
 
     left_bath, right_bath = baths
 
-    op = (
-        lambda x: -_TIME_CONVERSION_COEFF
-        * 1j
-        * dt
-        * apply_effective_Hamiltonian(
-            x,
-            ham_factor,
-            left_bath,
-            right_bath,
+    def op(x: torch.Tensor) -> torch.Tensor:
+        return (
+            -_TIME_CONVERSION_COEFF
+            * 1j
+            * dt
+            * apply_effective_Hamiltonian(
+                x,
+                ham_factor,
+                left_bath,
+                right_bath,
+            )
         )
-    )
 
     return krylov_exp(
         op,

emu_mps/utils.py

@@ -1,8 +1,6 @@
 from typing import List, Optional
 import torch
 
-from emu_mps import MPSConfig
-
 
 def new_left_bath(
     bath: torch.Tensor, state: torch.Tensor, op: torch.Tensor
@@ -59,8 +57,7 @@ def split_tensor(
 
 
 def truncate_impl(
-    factors: list[torch.Tensor],
-    config: MPSConfig,
+    factors: list[torch.Tensor], precision: float, max_bond_dim: int
 ) -> None:
     """
     Eigenvalues-based truncation of a matrix product.
@@ -76,8 +73,8 @@ def truncate_impl(
 
         l, r = split_tensor(
             factors[i].view(factor_shape[0], -1),
-            max_error=config.precision,
-            max_rank=config.max_bond_dim,
+            max_error=precision,
+            max_rank=max_bond_dim,
             orth_center_right=False,
         )
 
@@ -208,15 +205,6 @@ def get_extended_site_index(
     raise ValueError(f"Index {desired_index} does not exist")
 
 
-n_operator: torch.Tensor = torch.tensor(
-    [
-        [0, 0],
-        [0, 1],
-    ],
-    dtype=torch.complex128,
-)
-
-
 def tensor_trace(tensor: torch.Tensor, dim1: int, dim2: int) -> torch.Tensor:
     """
     Contract two legs of a single tensor.

{emu_mps-2.3.0.dist-info → emu_mps-2.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: emu-mps
-Version: 2.3.0
+Version: 2.4.1
 Summary: Pasqal MPS based pulse emulator built on PyTorch
 Project-URL: Documentation, https://pasqal-io.github.io/emulators/
 Project-URL: Repository, https://github.com/pasqal-io/emulators
@@ -25,7 +25,7 @@ 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==2.3.0
+Requires-Dist: emu-base==2.4.1
 Description-Content-Type: text/markdown
 
 <div align="center">

emu_mps-2.4.1.dist-info/RECORD

@@ -0,0 +1,19 @@
+emu_mps/__init__.py,sha256=R4Fw8r-_GEeckaxyQ9leUKVfLoukcrwvaTQsTbzxhbs,708
+emu_mps/algebra.py,sha256=vi3d4xOBEsfQ7gsjYNns9AMLjJXXXppYVEThavNvAg0,5408
+emu_mps/custom_callback_implementations.py,sha256=WeczmO6qkvBIipvXLqX45i3D7M4ovOrepusIGs6d2Ts,2420
+emu_mps/hamiltonian.py,sha256=gOPxNOBmk6jRPPjevERuCP_scGv0EKYeAJ0uxooihes,15622
+emu_mps/mpo.py,sha256=2HNwN4Fz04QIVfPcPaMmt2q89ZBxN3K-vVeiFkOtqzs,8049
+emu_mps/mps.py,sha256=KEXrLdqhi5EvpdX-9J38ZfrRdt3oDmcGsVLghMPUfQw,21600
+emu_mps/mps_backend.py,sha256=bS83qFxvdoK-c12_1WaPw6O7xUc7vdWifZNHUzNP5sM,2091
+emu_mps/mps_backend_impl.py,sha256=Pcbn27lhg3n3Lzo3CGwhlWPqPi-8YV1ntkgl901AoUs,30400
+emu_mps/mps_config.py,sha256=QmwgU8INEnxrxZkhboYIsHGwml0UhgPkddT5zh8KVBU,8867
+emu_mps/observables.py,sha256=4C_ewkd3YkJP0xghTrGUTgXUGvJRCQcetb8cU0SjMl0,1900
+emu_mps/solver.py,sha256=M9xkHhlEouTBvoPw2UYVu6kij7CO4Z1FXw_SiGFtdgo,85
+emu_mps/solver_utils.py,sha256=EnNzEaUrtTMQbrWoqOy8vyDsQwlsfQCUc2HgOp4z8dk,8680
+emu_mps/utils.py,sha256=pW5N_EbbGiOviQpJCw1a0pVgEDObP_InceNaIqY5bHE,6982
+emu_mps/optimatrix/__init__.py,sha256=fBXQ7-rgDro4hcaBijCGhx3J69W96qcw5_3mWc7tND4,364
+emu_mps/optimatrix/optimiser.py,sha256=k9suYmKLKlaZ7ozFuIqvXHyCBoCtGgkX1mpen9GOdOo,6977
+emu_mps/optimatrix/permutations.py,sha256=9DDMZtrGGZ01b9F3GkzHR3paX4qNtZiPoI7Z_Kia3Lc,3727
+emu_mps-2.4.1.dist-info/METADATA,sha256=2SUt6GLqhz_6-2ewuVqDH9xvz3gx1ZROU975intvuwk,3587
+emu_mps-2.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-2.4.1.dist-info/RECORD,,

emu_mps-2.3.0.dist-info/RECORD

@@ -1,18 +0,0 @@
-emu_mps/__init__.py,sha256=vvVRxFPudsVALA-_YyqWf8rlIRhtI1WzBWakYk2vfB8,734
-emu_mps/algebra.py,sha256=78XP9HEbV3wGNUzIulcLU5HizW4XAYmcFdkCe1T1x-k,5489
-emu_mps/custom_callback_implementations.py,sha256=SZGKVyS8U5hy07L-3SqpWlCAqGGKFTlSlWexZwSmjrM,2408
-emu_mps/hamiltonian.py,sha256=gOPxNOBmk6jRPPjevERuCP_scGv0EKYeAJ0uxooihes,15622
-emu_mps/mpo.py,sha256=aWSVuEzZM-_7ZD5Rz3-tSJWX22ARP0tMIl3gUu-_4V4,7834
-emu_mps/mps.py,sha256=n5KPuWr2qCNW4Q1OKpUGrSuhu4IbVp1v1BPLvLILbGg,19960
-emu_mps/mps_backend.py,sha256=bS83qFxvdoK-c12_1WaPw6O7xUc7vdWifZNHUzNP5sM,2091
-emu_mps/mps_backend_impl.py,sha256=O3vbjw3jlgyj0hOYv3nMmAP0lINqxbt96mQPgwPZN5w,30198
-emu_mps/mps_config.py,sha256=PoSKZxJMhG6zfzgEjj4tIvyiyYRQywxkRgidh8MRBsA,8222
-emu_mps/observables.py,sha256=7GQDH5kyaVNrwckk2f8ZJRV9Ca4jKhWWDsOCqYWsoEk,1349
-emu_mps/solver_utils.py,sha256=VQ02_RxvPcjyXippuIY4Swpx4EdqtoJTt8Ie70GgdqU,8550
-emu_mps/utils.py,sha256=hgtaRUtBAzk76ab-S_wTVkvqfVOmaUks38zWame9GRQ,7132
-emu_mps/optimatrix/__init__.py,sha256=fBXQ7-rgDro4hcaBijCGhx3J69W96qcw5_3mWc7tND4,364
-emu_mps/optimatrix/optimiser.py,sha256=k9suYmKLKlaZ7ozFuIqvXHyCBoCtGgkX1mpen9GOdOo,6977
-emu_mps/optimatrix/permutations.py,sha256=9DDMZtrGGZ01b9F3GkzHR3paX4qNtZiPoI7Z_Kia3Lc,3727
-emu_mps-2.3.0.dist-info/METADATA,sha256=WLLpx31eshWCwR95aB_B4yoM8BRAKT81UR2mjVR3998,3587
-emu_mps-2.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-2.3.0.dist-info/RECORD,,