emu-mps 2.4.2__py3-none-any.whl → 2.4.4__py3-none-any.whl

emu_mps/__init__.py

@@ -36,4 +36,4 @@ __all__ = [
     "EntanglementEntropy",
 ]
 
-__version__ = "2.4.2"
+__version__ = "2.4.4"

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
                 ),
-                "rg": 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
                 ),
-                "gr": 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
                 ),
-                "10": 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
                 ),
-                "01": 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

@@ -2,7 +2,7 @@ from __future__ import annotations
 import math
 from collections import Counter
 from typing import List, Optional, Sequence, TypeVar, Mapping
-
+import logging
 import torch
 
 from pulser.backend.state import State, Eigenstate
@@ -481,7 +481,9 @@ class MPS(State[complex, torch.Tensor]):
             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):
-            print("\nThe state is not normalized, normalizing it for you.")
+            logging.getLogger("emulators").warning(
+                "\nThe state is not normalized, normalizing it for you."
+            )
             accum_mps *= 1 / norm
 
         return accum_mps, amplitudes

emu_mps/mps_backend.py

@@ -1,10 +1,10 @@
 from pulser.backend import EmulatorBackend, Results
 from emu_mps.mps_config import MPSConfig
+from emu_base import init_logging
 from emu_mps.mps_backend_impl import create_impl, MPSBackendImpl
 import pickle
 import os
 import time
-import logging
 import pathlib
 
 
@@ -34,9 +34,9 @@ class MPSBackend(EmulatorBackend):
 
         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.
+        init_logging(impl.config.log_level, impl.config.log_file)
 
-        logging.getLogger("global_logger").warning(
+        impl.config.logger.warning(
             f"Resuming simulation from file {autosave_file}\n"
             f"Saving simulation state every {impl.config.autosave_dt} seconds"
         )

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
         )
@@ -276,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
@@ -451,6 +453,7 @@ class MPSBackendImpl:
             self.hamiltonian = make_H(
                 interaction_matrix=self.full_interaction_matrix,
                 hamiltonian_type=self.hamiltonian_type,
+                dim=self.dim,
                 num_gpus_to_use=self.resolved_num_gpus,
             )
 
@@ -686,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,6 +3,7 @@ from types import MethodType
 
 import copy
 
+from emu_base import init_logging
 from emu_mps.mps import MPS, DEFAULT_MAX_BOND_DIM, DEFAULT_PRECISION
 from emu_mps.mpo import MPO
 from emu_mps.solver import Solver
@@ -24,7 +25,6 @@ from pulser.backend import (
 )
 import logging
 import pathlib
-import sys
 
 
 class MPSConfig(EmulationConfig):
@@ -90,7 +90,7 @@ class MPSConfig(EmulationConfig):
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int | None = None,
-        optimize_qubit_ordering: bool = False,
+        optimize_qubit_ordering: bool = True,
         interaction_cutoff: float = 0.0,
         log_level: int = logging.INFO,
         log_file: pathlib.Path | None = None,
@@ -116,8 +116,6 @@ class MPSConfig(EmulationConfig):
             solver=solver,
             **kwargs,
         )
-        if self.optimize_qubit_ordering:
-            self.check_permutable_observables()
 
         MIN_AUTOSAVE_DT = 10
         assert (
@@ -125,20 +123,8 @@ class MPSConfig(EmulationConfig):
         ), f"autosave_dt must be larger than {MIN_AUTOSAVE_DT} seconds"
 
         self.monkeypatch_observables()
+        self.logger = init_logging(log_level, log_file)
 
-        self.logger = logging.getLogger("global_logger")
-        if log_file is None:
-            logging.basicConfig(
-                level=log_level, format="%(message)s", stream=sys.stdout, force=True
-            )  # default to stream = sys.stderr
-        else:
-            logging.basicConfig(
-                level=log_level,
-                format="%(message)s",
-                filename=str(log_file),
-                filemode="w",
-                force=True,
-            )
         if (self.noise_model.runs != 1 and self.noise_model.runs is not None) or (
             self.noise_model.samples_per_run != 1
             and self.noise_model.samples_per_run is not None
@@ -146,6 +132,9 @@ class MPSConfig(EmulationConfig):
             self.logger.warning(
                 "Warning: The runs and samples_per_run values of the NoiseModel are ignored!"
             )
+        self._backend_options[
+            "optimize_qubit_ordering"
+        ] &= self.check_permutable_observables()
 
     def _expected_kwargs(self) -> set[str]:
         return super()._expected_kwargs() | {
@@ -191,21 +180,7 @@ class MPSConfig(EmulationConfig):
             obs_list.append(obs_copy)
         self.observables = tuple(obs_list)
 
-    def init_logging(self) -> None:
-        if self.log_file is None:
-            logging.basicConfig(
-                level=self.log_level, format="%(message)s", stream=sys.stdout, force=True
-            )  # default to stream = sys.stderr
-        else:
-            logging.basicConfig(
-                level=self.log_level,
-                format="%(message)s",
-                filename=str(self.log_file),
-                filemode="w",
-                force=True,
-            )
-
-    def check_permutable_observables(self) -> None:
+    def check_permutable_observables(self) -> bool:
         allowed_permutable_obs = set(
             [
                 "bitstrings",
@@ -221,10 +196,10 @@ class MPSConfig(EmulationConfig):
         actual_obs = set([obs._base_tag for obs in self.observables])
         not_allowed = actual_obs.difference(allowed_permutable_obs)
         if not_allowed:
-            raise ValueError(
+            self.logger.warning(
                 f"emu-mps allows only {allowed_permutable_obs} observables with"
                 " `optimize_qubit_ordering = True`."
                 f" you provided unsupported {not_allowed}"
-                " To use other observables, please set"
-                " `optimize_qubit_ordering = False` in `MPSConfig()`."
+                " using `optimize_qubit_ordering = False` instead."
             )
+        return not_allowed == set()

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.2.dist-info → emu_mps-2.4.4.dist-info}/METADATA

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

emu_mps-2.4.4.dist-info/RECORD

@@ -0,0 +1,19 @@
+emu_mps/__init__.py,sha256=xZGOemPMjLVSBvpfex6jcvAYl0SQRrlcVyP9n0j1nvo,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=9LcneNCGMLYQJfiTQrDzBKz-YQmcL76P9l-A-3_bXOI,21707
+emu_mps/mps_backend.py,sha256=dXcW5Fa2AAL_WlHUgr30WS2l3qiDJgBVebiVpyYWpPM,2064
+emu_mps/mps_backend_impl.py,sha256=LMvi7KpfnsNmbiwQHuCtoA2lsK3Pf1vETuipL7BlG5U,30688
+emu_mps/mps_config.py,sha256=btlygf6VDFna9TQus-zOmUqB6pfe_aoQmNoivfvxABk,8251
+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.4.dist-info/METADATA,sha256=6iwSnm9w0kSDuUXbyy81ZBEwktkc0mjQa8K2GnVP0E8,3587
+emu_mps-2.4.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-2.4.4.dist-info/RECORD,,

emu_mps-2.4.2.dist-info/RECORD

@@ -1,19 +0,0 @@
-emu_mps/__init__.py,sha256=fEovcc3ow-4zoEodViCEyn2uZWnb-uwVSG-uRBs-g7U,708
-emu_mps/algebra.py,sha256=VZ5uaX5PYWGqDCpRKKr819BLMMtT0pKDxc2HrlLCdgU,5423
-emu_mps/custom_callback_implementations.py,sha256=WeczmO6qkvBIipvXLqX45i3D7M4ovOrepusIGs6d2Ts,2420
-emu_mps/hamiltonian.py,sha256=gOPxNOBmk6jRPPjevERuCP_scGv0EKYeAJ0uxooihes,15622
-emu_mps/mpo.py,sha256=-kbElxfodWh3IQa8QDor4TgLJE-zsaMRI5PZ6dZpYck,8049
-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=B8hi21uHFUSVJ3k5J4A4fXqyQo4sNHSb2uddN3lUYPA,30558
-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=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.2.dist-info/METADATA,sha256=SMYIZXDve_o9aZ7K1EfB7V6s8FBqA_wltE5a8ecNUf4,3587
-emu_mps-2.4.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-2.4.2.dist-info/RECORD,,