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

emu_mps/__init__.py

@@ -36,4 +36,4 @@ __all__ = [
     "EntanglementEntropy",
 ]
 
-__version__ = "2.4.1"
+__version__ = "2.4.3"

emu_mps/algebra.py

@@ -49,7 +49,7 @@ def add_factors(
 
 
 def scale_factors(
-    factors: list[torch.Tensor], scalar: complex, *, which: int
+    factors: list[torch.Tensor], scalar: complex | torch.Tensor, *, which: int
 ) -> list[torch.Tensor]:
     """
     Returns a new list of factors where the tensor at the given index is scaled by `scalar`.

emu_mps/hamiltonian.py

@@ -4,35 +4,41 @@ to the Hamiltonian of a neutral atoms quantum processor.
 """
 
 from abc import abstractmethod, ABC
-from typing import Iterator
+from typing import Iterator, Literal, Union
 
+from pulser.channels.base_channel import States
 from emu_base import HamiltonianType
 import torch
 from emu_mps.mpo import MPO
 
+
 dtype = torch.complex128
 
 
+Eigenstate = Union[States, Literal["0", "1"]]
+
+
 class Operators:
     id = torch.eye(2, dtype=dtype)
+    id_3x3 = torch.eye(3, dtype=dtype)
     n = torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype)
     creation = torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=dtype)
     sx = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=dtype)
     sy = torch.tensor([[0.0, -0.5j], [0.5j, 0.0]], dtype=dtype)
-    pu = torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype)
 
 
 class HamiltonianMPOFactors(ABC):
-    def __init__(self, interaction_matrix: torch.Tensor):
+    def __init__(self, interaction_matrix: torch.Tensor, dim: int = 2):
         assert interaction_matrix.ndim == 2, "interaction matrix is not a matrix"
         assert (
             interaction_matrix.shape[0] == interaction_matrix.shape[1]
         ), "interaction matrix is not square"
-
+        self.dim = dim
         self.interaction_matrix = interaction_matrix.clone()
         self.interaction_matrix.fill_diagonal_(0.0)  # or assert
         self.qubit_count = self.interaction_matrix.shape[0]
         self.middle = self.qubit_count // 2
+        self.identity = Operators.id if self.dim == 2 else Operators.id_3x3
 
     def __iter__(self) -> Iterator[torch.Tensor]:
         yield self.first_factor()
@@ -72,10 +78,12 @@ class HamiltonianMPOFactors(ABC):
 class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
     def first_factor(self) -> torch.Tensor:
         has_right_interaction = self.interaction_matrix[0, 1:].any()
-        fac = torch.zeros(1, 2, 2, 3 if has_right_interaction else 2, dtype=dtype)
-        fac[0, :, :, 1] = Operators.id
+        fac = torch.zeros(
+            1, self.dim, self.dim, 3 if has_right_interaction else 2, dtype=dtype
+        )
+        fac[0, :, :, 1] = self.identity
         if has_right_interaction:
-            fac[0, :, :, 2] = Operators.n
+            fac[0, :2, :2, 2] = Operators.n
 
         return fac
 
@@ -86,18 +94,18 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
 
         fac = torch.zeros(
             int(current_left_interactions.sum().item() + 2),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(left_interactions_to_keep.sum().item() + int(has_right_interaction) + 2),
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
         if has_right_interaction:
-            fac[1, :, :, -1] = Operators.n
+            fac[1, :2, :2, -1] = Operators.n
 
-        fac[2:, :, :, 0] = (
+        fac[2:, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.n
         )
@@ -106,7 +114,7 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
         j = 2
         for current_left_interaction in current_left_interactions.nonzero().flatten():
             if left_interactions_to_keep[current_left_interaction]:
-                fac[i, :, :, j] = Operators.id
+                fac[i, :, :, j] = self.identity
                 j += 1
             i += 1
         return fac
@@ -118,21 +126,21 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
 
         fac = torch.zeros(
             int(current_left_interactions.sum().item() + 2),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(current_right_interactions.sum().item() + 2),
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
 
-        fac[2:, :, :, 0] = (
+        fac[2:, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.n
         )
 
-        fac[1, :, :, 2:] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 2:] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.n.unsqueeze(-1)
 
@@ -140,7 +148,7 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
             self.interaction_matrix[:n, n + 1 :][current_left_interactions, :][
                 :, None, None, current_right_interactions
             ]
-            * Operators.id[None, ..., None]
+            * self.identity[None, ..., None]
         )
 
         return fac
@@ -152,18 +160,18 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
 
         fac = torch.zeros(
             int(right_interactions_to_keep.sum().item() + int(has_left_interaction) + 2),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(current_right_interactions.sum().item() + 2),
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
         if has_left_interaction:
-            fac[2, :, :, 0] = Operators.n
+            fac[2, :2, :2, 0] = Operators.n
 
-        fac[1, :, :, 2:] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 2:] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.n.unsqueeze(-1)
 
@@ -171,20 +179,22 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
         j = 2
         for current_right_interaction in current_right_interactions.nonzero().flatten():
             if right_interactions_to_keep[current_right_interaction]:
-                fac[i, :, :, j] = Operators.id
+                fac[i, :, :, j] = self.identity
                 i += 1
             j += 1
         return fac
 
     def last_factor(self) -> torch.Tensor:
         has_left_interaction = self.interaction_matrix[-1, :-1].any()
-        fac = torch.zeros(3 if has_left_interaction else 2, 2, 2, 1, dtype=dtype)
-        fac[0, :, :, 0] = Operators.id
+        fac = torch.zeros(
+            3 if has_left_interaction else 2, self.dim, self.dim, 1, dtype=dtype
+        )
+        fac[0, :, :, 0] = self.identity
         if has_left_interaction:
             if self.qubit_count >= 3:
-                fac[2, :, :, 0] = Operators.n
+                fac[2, :2, :2, 0] = Operators.n
             else:
-                fac[2, :, :, 0] = self.interaction_matrix[0, 1] * Operators.n
+                fac[2, :2, :2, 0] = self.interaction_matrix[0, 1] * Operators.n
 
         return fac
 
@@ -192,11 +202,13 @@ class RydbergHamiltonianMPOFactors(HamiltonianMPOFactors):
 class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
     def first_factor(self) -> torch.Tensor:
         has_right_interaction = self.interaction_matrix[0, 1:].any()
-        fac = torch.zeros(1, 2, 2, 4 if has_right_interaction else 2, dtype=dtype)
-        fac[0, :, :, 1] = Operators.id
+        fac = torch.zeros(
+            1, self.dim, self.dim, 4 if has_right_interaction else 2, dtype=dtype
+        )
+        fac[0, :, :, 1] = self.identity
         if has_right_interaction:
-            fac[0, :, :, 2] = Operators.creation
-            fac[0, :, :, 3] = Operators.creation.T
+            fac[0, :2, :2, 2] = Operators.creation
+            fac[0, :2, :2, 3] = Operators.creation.T
 
         return fac
 
@@ -207,8 +219,8 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
 
         fac = torch.zeros(
             int(2 * current_left_interactions.sum().item() + 2),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(
                 2 * left_interactions_to_keep.sum().item()
                 + 2 * int(has_right_interaction)
@@ -217,17 +229,17 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
         if has_right_interaction:
-            fac[1, :, :, -2] = Operators.creation
-            fac[1, :, :, -1] = Operators.creation.T
+            fac[1, :2, :2, -2] = Operators.creation
+            fac[1, :2, :2, -1] = Operators.creation.T
 
-        fac[2::2, :, :, 0] = (
+        fac[2::2, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.creation.T
         )
-        fac[3::2, :, :, 0] = (
+        fac[3::2, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.creation
         )
@@ -236,8 +248,8 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
         j = 2
         for current_left_interaction in current_left_interactions.nonzero().flatten():
             if left_interactions_to_keep[current_left_interaction]:
-                fac[i, :, :, j] = Operators.id
-                fac[i + 1, :, :, j + 1] = Operators.id
+                fac[i, :, :, j] = self.identity
+                fac[i + 1, :, :, j + 1] = self.identity
                 j += 2
             i += 2
         return fac
@@ -249,28 +261,28 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
 
         fac = torch.zeros(
             int(2 * current_left_interactions.sum().item() + 2),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(2 * current_right_interactions.sum().item() + 2),
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
 
-        fac[2::2, :, :, 0] = (
+        fac[2::2, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.creation.T
         )
-        fac[3::2, :, :, 0] = (
+        fac[3::2, :2, :2, 0] = (
             self.interaction_matrix[:n][current_left_interactions, n, None, None]
             * Operators.creation
         )
 
-        fac[1, :, :, 2::2] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 2::2] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.creation.unsqueeze(-1)
-        fac[1, :, :, 3::2] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 3::2] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.creation.T.unsqueeze(-1)
 
@@ -278,13 +290,13 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
             self.interaction_matrix[:n, n + 1 :][current_left_interactions, :][
                 :, None, None, current_right_interactions
             ]
-            * Operators.id[None, ..., None]
+            * self.identity[None, ..., None]
         )
         fac[3::2, :, :, 3::2] = (
             self.interaction_matrix[:n, n + 1 :][current_left_interactions, :][
                 :, None, None, current_right_interactions
             ]
-            * Operators.id[None, ..., None]
+            * self.identity[None, ..., None]
         )
 
         return fac
@@ -300,22 +312,22 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
                 + 2 * int(has_left_interaction)
                 + 2
             ),
-            2,
-            2,
+            self.dim,
+            self.dim,
             int(2 * current_right_interactions.sum().item() + 2),
             dtype=dtype,
         )
 
-        fac[0, :, :, 0] = Operators.id
-        fac[1, :, :, 1] = Operators.id
+        fac[0, :, :, 0] = self.identity
+        fac[1, :, :, 1] = self.identity
         if has_left_interaction:
-            fac[2, :, :, 0] = Operators.creation.T
-            fac[3, :, :, 0] = Operators.creation
+            fac[2, :2, :2, 0] = Operators.creation.T
+            fac[3, :2, :2, 0] = Operators.creation
 
-        fac[1, :, :, 2::2] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 2::2] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.creation.unsqueeze(-1)
-        fac[1, :, :, 3::2] = self.interaction_matrix[n + 1 :][
+        fac[1, :2, :2, 3::2] = self.interaction_matrix[n + 1 :][
             None, None, current_right_interactions, n
         ] * Operators.creation.T.unsqueeze(-1)
 
@@ -323,23 +335,25 @@ class XYHamiltonianMPOFactors(HamiltonianMPOFactors):
         j = 2
         for current_right_interaction in current_right_interactions.nonzero().flatten():
             if right_interactions_to_keep[current_right_interaction]:
-                fac[i, :, :, j] = Operators.id
-                fac[i + 1, :, :, j + 1] = Operators.id
+                fac[i, :, :, j] = self.identity
+                fac[i + 1, :, :, j + 1] = self.identity
                 i += 2
             j += 2
         return fac
 
     def last_factor(self) -> torch.Tensor:
         has_left_interaction = self.interaction_matrix[-1, :-1].any()
-        fac = torch.zeros(4 if has_left_interaction else 2, 2, 2, 1, dtype=dtype)
-        fac[0, :, :, 0] = Operators.id
+        fac = torch.zeros(
+            4 if has_left_interaction else 2, self.dim, self.dim, 1, dtype=dtype
+        )
+        fac[0, :, :, 0] = self.identity
         if has_left_interaction:
             if self.qubit_count >= 3:
-                fac[2, :, :, 0] = Operators.creation.T
-                fac[3, :, :, 0] = Operators.creation
+                fac[2, :2, :2, 0] = Operators.creation.T
+                fac[3, :2, :2, 0] = Operators.creation
             else:
-                fac[2, :, :, 0] = self.interaction_matrix[0, 1] * Operators.creation.T
-                fac[3, :, :, 0] = self.interaction_matrix[0, 1] * Operators.creation
+                fac[2, :2, :2, 0] = self.interaction_matrix[0, 1] * Operators.creation.T
+                fac[3, :2, :2, 0] = self.interaction_matrix[0, 1] * Operators.creation
 
         return fac
 
@@ -348,6 +362,7 @@ def make_H(
     *,
     interaction_matrix: torch.Tensor,  # depends on Hamiltonian Type
     hamiltonian_type: HamiltonianType,
+    dim: int = 2,
     num_gpus_to_use: int | None,
 ) -> MPO:
     r"""
@@ -357,53 +372,58 @@ def make_H(
     The Hamiltonian H is given by:
     H = ∑ⱼΩⱼ[cos(ϕⱼ)σˣⱼ + sin(ϕⱼ)σʸⱼ] - ∑ⱼΔⱼnⱼ + ∑ᵢ﹥ⱼC⁶/rᵢⱼ⁶ nᵢnⱼ
 
-    If noise is considered, the Hamiltonian includes an additional term to support
-    the Monte Carlo WaveFunction algorithm:
+    If noise is considered, the Hamiltonian includes an additional term to
+    support the Monte Carlo WaveFunction algorithm:
     H = ∑ⱼΩⱼ[cos(ϕⱼ)σˣⱼ + sin(ϕⱼ)σʸⱼ] - ∑ⱼΔⱼnⱼ + ∑ᵢ﹥ⱼC⁶/rᵢⱼ⁶ nᵢnⱼ - 0.5i∑ₘ ∑ᵤ Lₘᵘ⁺ Lₘᵘ
-    where Lₘᵘ are the Lindblad operators representing the noise, m for noise channel
-    and u for the number of atoms
+    where Lₘᵘ are the Lindblad operators representing the noise,
+    m for noise channel and u for the number of atoms
 
-    make_H constructs an MPO of the appropriate size, but the single qubit terms are left at zero.
+    make_H constructs an MPO of the appropriate size, but the single qubit
+    terms are left at zero.
     To fill in the appropriate values, call update_H
 
     Args:
-        interaction_matrix (torch.Tensor): The interaction matrix describing the interactions
-        between qubits.
-        num_gpus_to_use (int): how many gpus to put the Hamiltonian on. See utils.assign_devices
+        interaction_matrix (torch.Tensor): The interaction matrix describing
+        the interactions between qubits.
+        hamiltonian_type: whether to use XY or Rydberg interation
+        dim: dimension of the basis (2 or 3)
+        num_gpus_to_use (int): how many gpus to put the Hamiltonian on.
+        See utils.assign_devices
     Returns:
-        MPO: A Matrix Product Operator (MPO) representing the specified Hamiltonian.
+        MPO: A Matrix Product Operator (MPO) representing the specified
+        Hamiltonian.
 
     Note:
     For more information about the Hamiltonian and its usage, refer to the
     [Pulser documentation](https://pulser.readthedocs.io/en/stable/conventions.html#hamiltonians).
 
     """
+
     if hamiltonian_type == HamiltonianType.Rydberg:
         return MPO(
-            list(RydbergHamiltonianMPOFactors(interaction_matrix)),
+            list(RydbergHamiltonianMPOFactors(interaction_matrix, dim=dim)),
             num_gpus_to_use=num_gpus_to_use,
         )
 
     if hamiltonian_type == HamiltonianType.XY:
         return MPO(
-            list(XYHamiltonianMPOFactors(interaction_matrix)),
+            list(XYHamiltonianMPOFactors(interaction_matrix, dim=dim)),
             num_gpus_to_use=num_gpus_to_use,
         )
 
-    raise ValueError(f"Unsupported hamiltonian type {hamiltonian_type}")
-
 
 def update_H(
     hamiltonian: MPO,
     omega: torch.Tensor,
     delta: torch.Tensor,
     phi: torch.Tensor,
-    noise: torch.Tensor = torch.zeros(2, 2),
+    noise: torch.Tensor,
 ) -> None:
     """
     The single qubit operators in the Hamiltonian,
-    corresponding to the omega, delta, phi parameters and the aggregated Lindblad operators
-    have a well-determined position in the factors of the Hamiltonian.
+    corresponding to the omega, delta, phi parameters and the aggregated
+    Lindblad operators have a well-determined position in the factors of
+    the Hamiltonian.
     This function updates this part of the factors to update the
     Hamiltonian with new parameters without rebuilding the entire thing.
     See make_H for details about the Hamiltonian.
@@ -416,20 +436,24 @@ def update_H(
         phi (torch.Tensor): The phase ϕⱼ corresponding to each qubit.
         noise (torch.Tensor, optional): The single-qubit noise
         term -0.5i∑ⱼLⱼ†Lⱼ applied to all qubits.
-        This can be computed using the `compute_noise_from_lindbladians` function.
+        This can be computed using the `compute_noise_from_lindbladians`
+        function.
         Defaults to a zero tensor.
     """
 
-    assert noise.shape == (2, 2)
+    assert noise.shape == (2, 2) or (3, 3)
     nqubits = omega.size(dim=0)
 
     a = torch.tensordot(omega * torch.cos(phi), Operators.sx, dims=0)
-    c = torch.tensordot(delta, Operators.pu, dims=0)
+    c = torch.tensordot(delta, Operators.n, dims=0)
     b = torch.tensordot(omega * torch.sin(phi), Operators.sy, dims=0)
 
-    single_qubit_terms = a + b - c + noise
     factors = hamiltonian.factors
 
+    single_qubit_terms = torch.stack(nqubits * [noise])
+
+    single_qubit_terms[:, :2, :2] += a + b - c
+
     factors[0][0, :, :, 0] = single_qubit_terms[0]
     for i in range(1, nqubits):
         factors[i][1, :, :, 0] = single_qubit_terms[i]

emu_mps/mpo.py

@@ -10,6 +10,8 @@ from pulser.backend.operator import FullOp, QuditOp
 from emu_mps.mps import MPS, DEFAULT_MAX_BOND_DIM, DEFAULT_PRECISION
 from emu_mps.utils import new_left_bath, assign_devices
 
+dtype = torch.complex128
+
 
 class MPO(Operator[complex, torch.Tensor, MPS]):
     """
@@ -30,7 +32,8 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
             raise ValueError("For 1 qubit states, do state vector")
         if factors[0].shape[0] != 1 or factors[-1].shape[-1] != 1:
             raise ValueError(
-                "The dimension of the left (right) link of the first (last) tensor should be 1"
+                "The dimension of the left (right) link of the first (last) "
+                "tensor should be 1"
             )
         assert all(
             factors[i - 1].shape[-1] == factors[i].shape[0]
@@ -164,42 +167,75 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
         """
 
         basis = set(eigenstates)
-
+        dim = len(basis)
         operators_with_tensors: dict[str, torch.Tensor | QuditOp]
         if basis == {"r", "g"}:
-            # operators_with_tensors will now contain the basis for single qubit ops,
-            # and potentially user defined strings in terms of these
+            # 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).view(
-                    1, 2, 2, 1
+                "gg": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=dtype).view(
+                    1, dim, dim, 1
                 ),
-                "gr": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128).view(
+                "rg": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=dtype).view(
                     1, 2, 2, 1
                 ),
-                "rg": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128).view(
+                "gr": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=dtype).view(
                     1, 2, 2, 1
                 ),
-                "rr": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.complex128).view(
-                    1, 2, 2, 1
+                "rr": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype).view(
+                    1, dim, dim, 1
                 ),
             }
         elif basis == {"0", "1"}:
-            # operators_with_tensors will now contain the basis for single qubit ops,
-            # and potentially user defined strings in terms of these
+            # 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).view(
-                    1, 2, 2, 1
+                "00": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=dtype).view(
+                    1, dim, dim, 1
                 ),
-                "01": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128).view(
+                "10": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=dtype).view(
                     1, 2, 2, 1
                 ),
-                "10": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128).view(
+                "01": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=dtype).view(
                     1, 2, 2, 1
                 ),
-                "11": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.complex128).view(
-                    1, 2, 2, 1
+                "11": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype).view(
+                    1, dim, dim, 1
                 ),
             }
+        elif basis == {"r", "g", "x"}:
+            # 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, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "gr": torch.tensor(
+                    [[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "rg": torch.tensor(
+                    [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "rr": torch.tensor(
+                    [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "xx": torch.tensor(
+                    [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "xg": torch.tensor(
+                    [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "xr": torch.tensor(
+                    [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "gx": torch.tensor(
+                    [[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+                "rx": torch.tensor(
+                    [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]], dtype=dtype
+                ).view(1, dim, dim, 1),
+            }
+
         else:
             raise ValueError("Unsupported basis provided")
 
@@ -212,16 +248,14 @@ class MPO(Operator[complex, torch.Tensor, MPS]):
                 if isinstance(op, torch.Tensor):
                     return op
 
-                result = torch.zeros(1, 2, 2, 1, dtype=torch.complex128)
+                result = torch.zeros(1, dim, dim, 1, dtype=dtype)
                 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).view(1, 2, 2, 1)
-            ] * n_qudits
+            factors = [torch.eye(dim, dim, dtype=dtype).view(1, dim, dim, 1)] * n_qudits
 
             for op in tensorop:
                 factor = replace_operator_string(op[0])

emu_mps/mps.py

@@ -394,7 +394,7 @@ class MPS(State[complex, torch.Tensor]):
         result.truncate()
         return result
 
-    def __rmul__(self, scalar: complex) -> MPS:
+    def __rmul__(self, scalar: complex | torch.Tensor) -> MPS:
         """
         Multiply an MPS by a scalar.
 
@@ -419,7 +419,7 @@ class MPS(State[complex, torch.Tensor]):
             eigenstates=self.eigenstates,
         )
 
-    def __imul__(self, scalar: complex) -> MPS:
+    def __imul__(self, scalar: complex | torch.Tensor) -> MPS:
         return self.__rmul__(scalar)
 
     @classmethod

emu_mps/mps_backend_impl.py

@@ -42,6 +42,8 @@ from emu_mps.utils import (
     new_left_bath,
 )
 
+dtype = torch.complex128
+
 
 class Statistics(Observable):
     def __init__(
@@ -117,7 +119,8 @@ class MPSBackendImpl:
         self.phi = pulser_data.phi
         self.timestep_count: int = self.omega.shape[0]
         self.has_lindblad_noise = pulser_data.has_lindblad_noise
-        self.lindblad_noise = torch.zeros(2, 2, dtype=torch.complex128)
+        self.dim = pulser_data.dim
+        self.lindblad_noise = torch.zeros(self.dim, self.dim, dtype=dtype)
         self.qubit_permutation = (
             optimat.minimize_bandwidth(pulser_data.full_interaction_matrix)
             if self.config.optimize_qubit_ordering
@@ -126,6 +129,7 @@ class MPSBackendImpl:
         self.full_interaction_matrix = optimat.permute_tensor(
             pulser_data.full_interaction_matrix, self.qubit_permutation
         )
+
         self.masked_interaction_matrix = optimat.permute_tensor(
             pulser_data.masked_interaction_matrix, self.qubit_permutation
         )
@@ -155,12 +159,16 @@ class MPSBackendImpl:
             f"""To resume: `MPSBackend().resume("{self.autosave_file}")`"""
         )
         self.last_save_time = time.time()
+        requested_num_gpus = self.config.num_gpus_to_use
 
-        if self.config.num_gpus_to_use > DEVICE_COUNT:
+        if requested_num_gpus is None:
+            requested_num_gpus = DEVICE_COUNT
+        elif requested_num_gpus > DEVICE_COUNT:
             self.config.logger.warning(
-                f"Requested to use {self.config.num_gpus_to_use} GPU(s) "
+                f"Requested to use {requested_num_gpus} GPU(s) "
                 f"but only {DEVICE_COUNT if DEVICE_COUNT > 0 else 'cpu'} available"
             )
+        self.resolved_num_gpus = requested_num_gpus
 
     def __getstate__(self) -> dict:
         d = self.__dict__.copy()
@@ -211,7 +219,7 @@ class MPSBackendImpl:
                 self.qubit_count,
                 precision=self.config.precision,
                 max_bond_dim=self.config.max_bond_dim,
-                num_gpus_to_use=self.config.num_gpus_to_use,
+                num_gpus_to_use=self.resolved_num_gpus,
             )
             return
 
@@ -239,7 +247,7 @@ class MPSBackendImpl:
             [f.detach().clone() for f in initial_state.factors],
             precision=self.config.precision,
             max_bond_dim=self.config.max_bond_dim,
-            num_gpus_to_use=self.config.num_gpus_to_use,
+            num_gpus_to_use=self.resolved_num_gpus,
             eigenstates=initial_state.eigenstates,
         )
         initial_state.truncate()
@@ -259,7 +267,7 @@ class MPSBackendImpl:
                 else self.full_interaction_matrix
             ),
             hamiltonian_type=self.hamiltonian_type,
-            num_gpus_to_use=self.config.num_gpus_to_use,
+            num_gpus_to_use=self.resolved_num_gpus,
         )
 
         update_H(
@@ -272,9 +280,7 @@ class MPSBackendImpl:
 
     def init_baths(self) -> None:
         self.left_baths = [
-            torch.ones(
-                1, 1, 1, dtype=torch.complex128, device=self.state.factors[0].device
-            )
+            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)
         assert len(self.right_baths) == self.qubit_count - 1
@@ -447,7 +453,8 @@ class MPSBackendImpl:
             self.hamiltonian = make_H(
                 interaction_matrix=self.full_interaction_matrix,
                 hamiltonian_type=self.hamiltonian_type,
-                num_gpus_to_use=self.config.num_gpus_to_use,
+                dim=self.dim,
+                num_gpus_to_use=self.resolved_num_gpus,
             )
 
         if not self.is_finished():
@@ -682,6 +689,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
             omega=self.omega[self.timestep_index - 1, :],  # Meh
             delta=self.delta[self.timestep_index - 1, :],
             phi=self.phi[self.timestep_index - 1, :],
+            noise=torch.zeros(self.dim, self.dim, dtype=dtype),  # no noise
         )
 
         super().fill_results()

emu_mps/mps_config.py

@@ -3,7 +3,6 @@ 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
@@ -45,8 +44,12 @@ class MPSConfig(EmulationConfig):
             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
-            0=all factors to cpu. As shown in the benchmarks, using multiple GPUs might
+        num_gpus_to_use: number of GPUs to be used in a given simulation.
+            - if it is set to a number `n > 0`, the state will be distributed across `n` GPUs.
+            - if it is set to `n = 0`, the entire simulation runs on the CPU.
+            - if it is `None` (the default value), the backend internally chooses the number of GPUs
+            based on the hardware availability during runtime.
+        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
             accuracy, but disables certain features.
@@ -86,7 +89,7 @@ class MPSConfig(EmulationConfig):
         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,
+        num_gpus_to_use: int | None = None,
         optimize_qubit_ordering: bool = False,
         interaction_cutoff: float = 0.0,
         log_level: int = logging.INFO,

emu_mps/solver_utils.py

@@ -5,7 +5,7 @@ from emu_base import krylov_exp
 from emu_base.math.krylov_energy_min import krylov_energy_minimization
 from emu_base.utils import deallocate_tensor
 from emu_mps import MPS, MPO
-from emu_mps.utils import split_tensor
+from emu_mps.utils import split_matrix
 from emu_mps.mps_config import MPSConfig
 
 
@@ -178,11 +178,12 @@ def evolve_pair(
         is_hermitian=is_hermitian,
     ).view(left_bond_dim * 2, 2 * right_bond_dim)
 
-    l, r = split_tensor(
+    l, r = split_matrix(
         evol,
         max_error=config.precision,
         max_rank=config.max_bond_dim,
         orth_center_right=orth_center_right,
+        preserve_norm=not is_hermitian,  # only relevant for computing jump times
     )
 
     return l.view(left_bond_dim, 2, -1), r.view(-1, 2, right_bond_dim).to(right_device)
@@ -262,7 +263,7 @@ def minimize_energy_pair(
     )
     updated_state = updated_state.view(left_bond_dim * 2, 2 * right_bond_dim)
 
-    l, r = split_tensor(
+    l, r = split_matrix(
         updated_state,
         max_error=config.precision,
         max_rank=config.max_bond_dim,

emu_mps/utils.py

@@ -23,11 +23,12 @@ def _determine_cutoff_index(d: torch.Tensor, max_error: float) -> int:
     return 0  # type: ignore[no-any-return]
 
 
-def split_tensor(
+def split_matrix(
     m: torch.Tensor,
     max_error: float = 1e-5,
     max_rank: int = 1024,
     orth_center_right: bool = True,
+    preserve_norm: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Computes a low-rank approximation split of m using the Eckart-Young-Mirsky theorem.
@@ -41,8 +42,11 @@ def split_tensor(
             d.shape[0] - max_rank,
         )
         left = q[:, max_bond:]
-        right = q.T.conj() @ m
-        right = right[max_bond:, :]
+        right = left.T.conj() @ m
+        if preserve_norm:
+            old_norm2 = torch.sum(d)
+            new_norm2 = torch.sum(d[max_bond:])
+            right *= torch.sqrt(old_norm2 / new_norm2)
     else:
         d, q = torch.linalg.eigh(m.T.conj() @ m)
         max_bond = max(
@@ -50,8 +54,11 @@ def split_tensor(
             d.shape[0] - max_rank,
         )
         right = q[:, max_bond:].T.conj_physical()
-        left = m @ q
-        left = left[:, max_bond:]
+        left = m @ q[:, max_bond:]
+        if preserve_norm:
+            old_norm2 = torch.sum(d)
+            new_norm2 = torch.sum(d[max_bond:])
+            left *= torch.sqrt(old_norm2 / new_norm2)
 
     return left, right
 
@@ -71,7 +78,7 @@ def truncate_impl(
     for i in range(len(factors) - 1, 0, -1):
         factor_shape = factors[i].shape
 
-        l, r = split_tensor(
+        l, r = split_matrix(
             factors[i].view(factor_shape[0], -1),
             max_error=precision,
             max_rank=max_bond_dim,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: emu-mps
-Version: 2.4.1
+Version: 2.4.3
 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.4.1
+Requires-Dist: emu-base==2.4.3
 Description-Content-Type: text/markdown
 
 <div align="center">

emu_mps-2.4.3.dist-info/RECORD

@@ -0,0 +1,19 @@
+emu_mps/__init__.py,sha256=avViacgzWmSfTtWIUvoxlaE_1p7LRXPb9Hh5sWv8ND4,708
+emu_mps/algebra.py,sha256=VZ5uaX5PYWGqDCpRKKr819BLMMtT0pKDxc2HrlLCdgU,5423
+emu_mps/custom_callback_implementations.py,sha256=WeczmO6qkvBIipvXLqX45i3D7M4ovOrepusIGs6d2Ts,2420
+emu_mps/hamiltonian.py,sha256=OM0bPNZV7J5Egk6aTUwt4GaWqiUOP68ujZBTuqlBY1k,16289
+emu_mps/mpo.py,sha256=WmGDGkCMhlODmydd0b09YcSRlsk6Bg5xYQ4rXSNJvnY,9703
+emu_mps/mps.py,sha256=1aMvY4NQSBeAwzTwHRegXncGQqm-1g05W00OimFNzt8,21630
+emu_mps/mps_backend.py,sha256=bS83qFxvdoK-c12_1WaPw6O7xUc7vdWifZNHUzNP5sM,2091
+emu_mps/mps_backend_impl.py,sha256=LMvi7KpfnsNmbiwQHuCtoA2lsK3Pf1vETuipL7BlG5U,30688
+emu_mps/mps_config.py,sha256=j7rho3edFzDxPO_VX7j5jc0Drw9wO2NWRymkqZ9hzmU,9128
+emu_mps/observables.py,sha256=4C_ewkd3YkJP0xghTrGUTgXUGvJRCQcetb8cU0SjMl0,1900
+emu_mps/solver.py,sha256=M9xkHhlEouTBvoPw2UYVu6kij7CO4Z1FXw_SiGFtdgo,85
+emu_mps/solver_utils.py,sha256=CklWjVvhngpyElpJyfUTncr4ErBP6VbUE3Z9nf8s3mM,8762
+emu_mps/utils.py,sha256=rL75H55hB5lDMjy8a_O2PpJq51iZKjSx91X4euxB3mY,7293
+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.3.dist-info/METADATA,sha256=hq605tOeTSfyZa9cNwSkVUH72t5ZTHfLo19Kll9-sxU,3587
+emu_mps-2.4.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-2.4.3.dist-info/RECORD,,

emu_mps-2.4.1.dist-info/RECORD

@@ -1,19 +0,0 @@
-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,,