emu-mps 1.2.7__py3-none-any.whl → 2.0.1__py3-none-any.whl

emu_mps/__init__.py

@@ -1,19 +1,20 @@
-from emu_base import (
-    Callback,
+from pulser.backend import (
     BitStrings,
     CorrelationMatrix,
     Energy,
     EnergyVariance,
     Expectation,
     Fidelity,
-    QubitDensity,
+    Occupation,
     StateResult,
-    SecondMomentOfEnergy,
+    EnergySecondMoment,
 )
 from .mps_config import MPSConfig
 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,16 +24,17 @@ __all__ = [
     "inner",
     "MPSConfig",
     "MPSBackend",
-    "Callback",
     "StateResult",
     "BitStrings",
-    "QubitDensity",
+    "Occupation",
     "CorrelationMatrix",
     "Expectation",
     "Fidelity",
     "Energy",
     "EnergyVariance",
-    "SecondMomentOfEnergy",
+    "EnergySecondMoment",
+    "aggregate",
+    "EntanglementEntropy",
 ]
 
-__version__ = "1.2.7"
+__version__ = "2.0.1"

emu_mps/custom_callback_implementations.py

@@ -0,0 +1,96 @@
+import torch
+
+from pulser.backend.default_observables import (
+    CorrelationMatrix,
+    EnergySecondMoment,
+    EnergyVariance,
+    Occupation,
+    Energy,
+)
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from emu_mps.mps_config import MPSConfig
+    from emu_mps.mps import MPS
+    from emu_mps.mpo import MPO
+
+
+def qubit_occupation_mps_impl(
+    self: Occupation,
+    *,
+    config: "MPSConfig",
+    state: "MPS",
+    hamiltonian: "MPO",
+) -> torch.Tensor:
+    """
+    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
+    )
+    return state.expect_batch(op).real.reshape(-1).cpu()
+
+
+def correlation_matrix_mps_impl(
+    self: CorrelationMatrix,
+    *,
+    config: "MPSConfig",
+    state: "MPS",
+    hamiltonian: "MPO",
+) -> torch.Tensor:
+    """
+    Custom implementation of the density-density correlation ❬ψ|nᵢnⱼ|ψ❭ for the EMU-MPS.
+
+    TODO: extend to arbitrary two-point correlation ❬ψ|AᵢBⱼ|ψ❭
+    """
+    return state.get_correlation_matrix().cpu()
+
+
+def energy_variance_mps_impl(
+    self: EnergyVariance,
+    *,
+    config: "MPSConfig",
+    state: "MPS",
+    hamiltonian: "MPO",
+) -> torch.Tensor:
+    """
+    Custom implementation of the energy variance ❬ψ|H²|ψ❭-❬ψ|H|ψ❭² for the EMU-MPS.
+    """
+    h_squared = hamiltonian @ hamiltonian
+    h_2 = h_squared.expect(state).cpu()
+    h = hamiltonian.expect(state).cpu()
+    en_var = h_2 - h**2
+    return en_var.real  # type: ignore[no-any-return]
+
+
+def energy_second_moment_mps_impl(
+    self: EnergySecondMoment,
+    *,
+    config: "MPSConfig",
+    state: "MPS",
+    hamiltonian: "MPO",
+) -> torch.Tensor:
+    """
+    Custom implementation of the second moment of energy ❬ψ|H²|ψ❭
+    for the EMU-MPS.
+    """
+    h_square = hamiltonian @ hamiltonian
+    h_2 = h_square.expect(state).cpu()
+    assert torch.allclose(h_2.imag, torch.zeros_like(h_2.imag), atol=1e-4)
+    return h_2.real
+
+
+def energy_mps_impl(
+    self: Energy,
+    *,
+    config: "MPSConfig",
+    state: "MPS",
+    hamiltonian: "MPO",
+) -> torch.Tensor:
+    """
+    Custom implementation of the second moment of energy ❬ψ|H²|ψ❭
+    for the EMU-MPS.
+    """
+    h = hamiltonian.expect(state)
+    assert torch.allclose(h.imag, torch.zeros_like(h.imag), atol=1e-4)
+    return h.real

emu_mps/mpo.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 import itertools
-from typing import Any, List, Iterable, Optional
+from typing import Any, List, Sequence, Optional
 
 import torch
 
+from pulser.backend import State, Operator
+from emu_base import DEVICE_COUNT
 from emu_mps.algebra import add_factors, scale_factors, zip_right
-from emu_base.base_classes.operator import FullOp, QuditOp
-from emu_base import Operator, State, DEVICE_COUNT
+from pulser.backend.operator import FullOp, QuditOp
 from emu_mps.mps import MPS
 from emu_mps.utils import new_left_bath, assign_devices
 
@@ -29,7 +30,7 @@ def _validate_operator_targets(operations: FullOp, nqubits: int) -> None:
             )
 
 
-class MPO(Operator):
+class MPO(Operator[complex, torch.Tensor, MPS]):
     """
     Matrix Product Operator.
 
@@ -61,7 +62,7 @@ class MPO(Operator):
     def __repr__(self) -> str:
         return "[" + ", ".join(map(repr, self.factors)) + "]"
 
-    def __mul__(self, other: State) -> MPS:
+    def apply_to(self, other: MPS) -> MPS:
         """
         Applies this MPO to the given MPS.
         The returned MPS is:
@@ -84,7 +85,7 @@ class MPO(Operator):
         )
         return MPS(factors, orthogonality_center=0)
 
-    def __add__(self, other: Operator) -> MPO:
+    def __add__(self, other: MPO) -> MPO:
         """
         Returns the sum of two MPOs, computed with a direct algorithm.
         The result is currently not truncated
@@ -113,7 +114,7 @@ class MPO(Operator):
         factors = scale_factors(self.factors, scalar, which=0)
         return MPO(factors)
 
-    def __matmul__(self, other: Operator) -> MPO:
+    def __matmul__(self, other: MPO) -> MPO:
         """
         Compose two operators. The ordering is that
         self is applied after other.
@@ -128,7 +129,7 @@ class MPO(Operator):
         factors = zip_right(self.factors, other.factors)
         return MPO(factors)
 
-    def expect(self, state: State) -> float | complex:
+    def expect(self, state: State) -> torch.Tensor:
         """
         Compute the expectation value of self on the given state.
 
@@ -151,17 +152,17 @@ class MPO(Operator):
                 state.factors[i + 1].device
             )
         acc = new_left_bath(acc, state.factors[n], self.factors[n])
-        return acc.item()  # type: ignore [no-any-return]
-
-    @staticmethod
-    def from_operator_string(
-        basis: Iterable[str],
-        nqubits: int,
-        operations: FullOp,
-        operators: dict[str, QuditOp] = {},
-        /,
+        return acc.reshape(1)[0].cpu()
+
+    @classmethod
+    def _from_operator_repr(
+        cls,
+        *,
+        eigenstates: Sequence[str],
+        n_qudits: int,
+        operations: FullOp[complex],
         **kwargs: Any,
-    ) -> MPO:
+    ) -> tuple[MPO, FullOp[complex]]:
         """
         See the base class
 
@@ -174,15 +175,16 @@ class MPO(Operator):
         Returns:
             the operator in MPO form.
         """
-        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
 
-        _validate_operator_targets(operations, nqubits)
+        _validate_operator_targets(operations, n_qudits)
 
-        basis = set(basis)
+        basis = set(eigenstates)
+
+        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 |= {
+            operators_with_tensors = {
                 "gg": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -199,7 +201,7 @@ class MPO(Operator):
         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 |= {
+            operators_with_tensors = {
                 "00": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -234,12 +236,12 @@ class MPO(Operator):
 
             factors = [
                 torch.eye(2, 2, dtype=torch.complex128).reshape(1, 2, 2, 1)
-            ] * nqubits
+            ] * n_qudits
 
             for op in tensorop:
                 factor = replace_operator_string(op[0])
                 for target_qubit in op[1]:
                     factors[target_qubit] = factor
 
-            mpos.append(coeff * MPO(factors, **kwargs))
-        return sum(mpos[1:], start=mpos[0])
+            mpos.append(coeff * cls(factors, **kwargs))
+        return sum(mpos[1:], start=mpos[0]), operations  # type: ignore[no-any-return]

emu_mps/mps.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import math
 from collections import Counter
-from typing import Any, List, Optional, Iterable
+from typing import List, Optional, Sequence, TypeVar, Mapping
 
 import torch
 
-from emu_base import State, DEVICE_COUNT
+from pulser.backend.state import State, Eigenstate
+from emu_base import DEVICE_COUNT
 from emu_mps import MPSConfig
 from emu_mps.algebra import add_factors, scale_factors
 from emu_mps.utils import (
@@ -17,8 +18,10 @@ from emu_mps.utils import (
     n_operator,
 )
 
+ArgScalarType = TypeVar("ArgScalarType")
 
-class MPS(State):
+
+class MPS(State[complex, torch.Tensor]):
     """
     Matrix Product State, aka tensor train.
 
@@ -53,6 +56,7 @@ class MPS(State):
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu, None=keep the existing device assignment.
         """
+        self._eigenstates = ["0", "1"]
         self.config = config if config is not None else MPSConfig()
         assert all(
             factors[i - 1].shape[2] == factors[i].shape[0] for i in range(1, len(factors))
@@ -73,6 +77,11 @@ class MPS(State):
         if num_gpus_to_use is not None:
             assign_devices(self.factors, min(DEVICE_COUNT, num_gpus_to_use))
 
+    @property
+    def n_qudits(self) -> int:
+        """The number of qudits in the state."""
+        return self.num_sites
+
     @classmethod
     def make(
         cls,
@@ -174,7 +183,12 @@ class MPS(State):
         return max((x.shape[2] for x in self.factors), default=0)
 
     def sample(
-        self, num_shots: int, p_false_pos: float = 0.0, p_false_neg: float = 0.0
+        self,
+        *,
+        num_shots: int,
+        one_state: Eigenstate | None = None,
+        p_false_pos: float = 0.0,
+        p_false_neg: float = 0.0,
     ) -> Counter[str]:
         """
         Samples bitstrings, taking into account the specified error rates.
@@ -182,11 +196,12 @@ class MPS(State):
         Args:
             num_shots: how many bitstrings to sample
             p_false_pos: the rate at which a 0 is read as a 1
-            p_false_neg: teh rate at which a 1 is read as a 0
+            p_false_neg: the rate at which a 1 is read as a 0
 
         Returns:
             the measured bitstrings, by count
         """
+        assert one_state in {None, "r", "1"}
         self.orthogonalize(0)
 
         rnd_matrix = torch.rand(num_shots, self.num_sites).to(self.factors[0].device)
@@ -249,19 +264,17 @@ class MPS(State):
             )
         return bitstrings
 
-    def norm(self) -> float:
+    def norm(self) -> torch.Tensor:
         """Computes the norm of the MPS."""
         orthogonality_center = (
             self.orthogonality_center
             if self.orthogonality_center is not None
             else self.orthogonalize(0)
         )
+        # the torch.norm function is not properly typed.
+        return self.factors[orthogonality_center].norm().cpu()  # type: ignore[no-any-return]
 
-        return float(
-            torch.linalg.norm(self.factors[orthogonality_center].to("cpu")).item()
-        )
-
-    def inner(self, other: State) -> float | complex:
+    def inner(self, other: State) -> torch.Tensor:
         """
         Compute the inner product between this state and other.
         Note that self is the left state in the inner product,
@@ -285,7 +298,33 @@ class MPS(State):
             acc = torch.tensordot(acc, other.factors[i].to(acc.device), dims=1)
             acc = torch.tensordot(self.factors[i].conj(), acc, dims=([0, 1], [0, 1]))
 
-        return acc.item()  # type: ignore[no-any-return]
+        return acc.reshape(1)[0].cpu()
+
+    def overlap(self, other: State, /) -> torch.Tensor:
+        """
+        Compute the overlap of this state and other. This is defined as
+        $|\\langle self | other \\rangle |^2$
+        """
+        return torch.abs(self.inner(other)) ** 2  # type: ignore[no-any-return]
+
+    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
+        S = -Σᵢsᵢ² log(sᵢ²)),
+        where sᵢ are the singular values at the chosen bond.
+        """
+        self.orthogonalize(mps_site)
+
+        # perform svd on reshaped matrix at site b
+        matrix = self.factors[mps_site].flatten(end_dim=1)
+        s = torch.linalg.svdvals(matrix)
+
+        s_e = torch.Tensor(torch.special.entr(s**2))
+        s_e = torch.sum(s_e)
+
+        self.orthogonalize(0)
+        return s_e.cpu()
 
     def get_memory_footprint(self) -> float:
         """
@@ -347,14 +386,13 @@ class MPS(State):
     def __imul__(self, scalar: complex) -> MPS:
         return self.__rmul__(scalar)
 
-    @staticmethod
-    def from_state_string(
+    @classmethod
+    def _from_state_amplitudes(
+        cls,
         *,
-        basis: Iterable[str],
-        nqubits: int,
-        strings: dict[str, complex],
-        **kwargs: Any,
-    ) -> MPS:
+        eigenstates: Sequence[str],
+        amplitudes: Mapping[str, complex],
+    ) -> tuple[MPS, Mapping[str, complex]]:
         """
         See the base class.
 
@@ -367,7 +405,8 @@ class MPS(State):
             The resulting MPS representation of the state.s
         """
 
-        basis = set(basis)
+        nqubits = len(next(iter(amplitudes.keys())))
+        basis = set(eigenstates)
         if basis == {"r", "g"}:
             one = "r"
         elif basis == {"0", "1"}:
@@ -381,18 +420,17 @@ class MPS(State):
         accum_mps = MPS(
             [torch.zeros((1, 2, 1), dtype=torch.complex128)] * nqubits,
             orthogonality_center=0,
-            **kwargs,
         )
 
-        for state, amplitude in strings.items():
+        for state, amplitude in amplitudes.items():
             factors = [basis_1 if ch == one else basis_0 for ch in state]
-            accum_mps += amplitude * MPS(factors, **kwargs)
+            accum_mps += amplitude * MPS(factors)
         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.")
             accum_mps *= 1 / norm
 
-        return accum_mps
+        return accum_mps, amplitudes
 
     def expect_batch(self, single_qubit_operators: torch.Tensor) -> torch.Tensor:
         """
@@ -460,7 +498,7 @@ class MPS(State):
 
     def get_correlation_matrix(
         self, *, operator: torch.Tensor = n_operator
-    ) -> list[list[float]]:
+    ) -> torch.Tensor:
         """
         Efficiently compute the symmetric correlation matrix
             C_ij = <self|operator_i operator_j|self>
@@ -474,7 +512,7 @@ class MPS(State):
         """
         assert operator.shape == (2, 2)
 
-        result = [[0.0 for _ in range(self.num_sites)] for _ in range(self.num_sites)]
+        result = torch.zeros(self.num_sites, self.num_sites, dtype=torch.complex128)
 
         for left in range(0, self.num_sites):
             self.orthogonalize(left)
@@ -486,7 +524,7 @@ class MPS(State):
             accumulator = torch.tensordot(
                 accumulator, self.factors[left].conj(), dims=([0, 2], [0, 1])
             )
-            result[left][left] = accumulator.trace().item().real
+            result[left, left] = accumulator.trace().item().real
             for right in range(left + 1, self.num_sites):
                 partial = torch.tensordot(
                     accumulator.to(self.factors[right].device),
@@ -497,7 +535,7 @@ class MPS(State):
                     partial, self.factors[right].conj(), dims=([0], [0])
                 )
 
-                result[left][right] = (
+                result[left, right] = (
                     torch.tensordot(
                         partial, operator.to(partial.device), dims=([0, 2], [0, 1])
                     )
@@ -505,13 +543,13 @@ class MPS(State):
                     .item()
                     .real
                 )
-                result[right][left] = result[left][right]
+                result[right, left] = result[left, right]
                 accumulator = tensor_trace(partial, 0, 2)
 
         return result
 
 
-def inner(left: MPS, right: MPS) -> float | complex:
+def inner(left: MPS, right: MPS) -> torch.Tensor:
     """
     Wrapper around MPS.inner.
 

emu_mps/mps_backend.py

@@ -1,7 +1,6 @@
-from emu_base import Backend, BackendConfig, Results
+from pulser.backend import EmulatorBackend, Results
 from emu_mps.mps_config import MPSConfig
 from emu_mps.mps_backend_impl import create_impl, MPSBackendImpl
-from pulser import Sequence
 import pickle
 import os
 import time
@@ -9,13 +8,16 @@ import logging
 import pathlib
 
 
-class MPSBackend(Backend):
+class MPSBackend(EmulatorBackend):
     """
     A backend for emulating Pulser sequences using Matrix Product States (MPS),
     aka tensor trains.
     """
 
-    def resume(self, autosave_file: str | pathlib.Path) -> Results:
+    default_config = MPSConfig()
+
+    @staticmethod
+    def resume(autosave_file: str | pathlib.Path) -> Results:
         """
         Resume simulation from autosave file.
         Only resume simulations from data you trust!
@@ -39,24 +41,18 @@ class MPSBackend(Backend):
             f"Saving simulation state every {impl.config.autosave_dt} seconds"
         )
 
-        return self._run(impl)
+        return MPSBackend._run(impl)
 
-    def run(self, sequence: Sequence, mps_config: BackendConfig) -> Results:
+    def run(self) -> Results:
         """
         Emulates the given sequence.
 
-        Args:
-            sequence: a Pulser sequence to simulate
-            mps_config: the backends config. Should be of type MPSConfig
-
         Returns:
             the simulation results
         """
-        assert isinstance(mps_config, MPSConfig)
-
-        self.validate_sequence(sequence)
+        assert isinstance(self._config, MPSConfig)
 
-        impl = create_impl(sequence, mps_config)
+        impl = create_impl(self._sequence, self._config)
         impl.init()  # This is separate from the constructor for testing purposes.
 
         return self._run(impl)

emu_mps/mps_backend_impl.py

@@ -2,16 +2,19 @@ import math
 import pathlib
 import random
 import uuid
+
 from resource import RUSAGE_SELF, getrusage
-from typing import Optional
+from typing import Optional, Any
+import typing
 import pickle
 import os
-
 import torch
 import time
 from pulser import Sequence
+from types import MethodType
 
-from emu_base import Results, State, PulserData, DEVICE_COUNT
+from pulser.backend import State, Observable, EmulationConfig, Results
+from emu_base import PulserData, DEVICE_COUNT
 from emu_base.math.brents_root_finding import BrentsRootFinder
 from emu_mps.hamiltonian import make_H, update_H
 from emu_mps.mpo import MPO
@@ -33,6 +36,56 @@ from emu_mps.utils import (
 from enum import Enum, auto
 
 
+class Statistics(Observable):
+    def __init__(
+        self,
+        evaluation_times: typing.Sequence[float] | None,
+        data: list[float],
+        timestep_count: int,
+    ):
+        super().__init__(evaluation_times=evaluation_times)
+        self.data = data
+        self.timestep_count = timestep_count
+
+    @property
+    def _base_tag(self) -> str:
+        return "statistics"
+
+    def apply(
+        self,
+        *,
+        config: EmulationConfig,
+        state: State,
+        **kwargs: Any,
+    ) -> dict:
+        """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
+
+        config.logger.info(
+            f"step = {len(self.data)}/{self.timestep_count}, "
+            + f"χ = {state.get_max_bond_dim()}, "
+            + f"|ψ| = {state.get_memory_footprint():.3f} MB, "
+            + f"RSS = {max_mem:.3f} MB, "
+            + f"Δt = {duration:.3f} s"
+        )
+
+        return {
+            "max_bond_dimension": state.get_max_bond_dim(),
+            "memory_footprint": state.get_memory_footprint(),
+            "RSS": max_mem,
+            "duration": duration,
+        }
+
+
 class SwipeDirection(Enum):
     LEFT_TO_RIGHT = auto()
     RIGHT_TO_LEFT = auto()
@@ -75,9 +128,14 @@ class MPSBackendImpl:
         self.swipe_direction = SwipeDirection.LEFT_TO_RIGHT
         self.tdvp_index = 0
         self.timestep_index = 0
-        self.results = Results()
+        self.results = Results(atom_order=(), total_duration=self.target_times[-1])
+        self.statistics = Statistics(
+            evaluation_times=[t / self.target_times[-1] for t in self.target_times],
+            data=[],
+            timestep_count=self.timestep_count,
+        )
         self.autosave_file = self._get_autosave_filepath(self.config.autosave_prefix)
-        self.config.logger.warning(
+        self.config.logger.debug(
             f"""Will save simulation state to file "{self.autosave_file.name}"
             every {self.config.autosave_dt} seconds.\n"""
             f"""To resume: `MPSBackend().resume("{self.autosave_file}")`"""
@@ -90,16 +148,26 @@ class MPSBackendImpl:
                 f"but only {DEVICE_COUNT if DEVICE_COUNT > 0 else 'cpu'} available"
             )
 
+    def __getstate__(self) -> dict:
+        for obs in self.config.observables:
+            obs.apply = MethodType(type(obs).apply, obs)  # type: ignore[method-assign]
+        d = self.__dict__.copy()
+        # mypy thinks the method below is an attribute, because of the __getattr__ override
+        d["results"] = self.results._to_abstract_repr()  # type: ignore[operator]
+        return d
+
+    def __setstate__(self, d: dict) -> None:
+        self.__dict__ = d
+        self.results = Results._from_abstract_repr(d["results"])  # type: ignore [attr-defined]
+        self.config.monkeypatch_observables()
+
     @staticmethod
     def _get_autosave_filepath(autosave_prefix: str) -> pathlib.Path:
         return pathlib.Path(os.getcwd()) / (autosave_prefix + str(uuid.uuid1()) + ".dat")
 
     def init_dark_qubits(self) -> None:
         # has_state_preparation_error
-        if (
-            self.config.noise_model is not None
-            and self.config.noise_model.state_prep_error > 0.0
-        ):
+        if self.config.noise_model.state_prep_error > 0.0:
             self.well_prepared_qubits_filter = pick_well_prepared_qubits(
                 self.config.noise_model.state_prep_error, self.qubit_count
             )
@@ -359,7 +427,14 @@ class MPSBackendImpl:
             )
             self.init_baths()
 
-        self.log_step_statistics(duration=time.time() - self.time)
+        self.statistics.data.append(time.time() - self.time)
+        self.statistics(
+            self.config,
+            self.current_time / self.target_times[-1],
+            self.state,
+            self.hamiltonian,
+            self.results,
+        )
         self.time = time.time()
 
     def save_simulation(self) -> None:
@@ -369,7 +444,6 @@ class MPSBackendImpl:
         basename = self.autosave_file
         with open(basename.with_suffix(".new"), "wb") as file_handle:
             pickle.dump(self, file_handle)
-
         if basename.is_file():
             os.rename(basename, basename.with_suffix(".bak"))
 
@@ -389,13 +463,14 @@ class MPSBackendImpl:
         normalized_state = 1 / self.state.norm() * self.state
 
         current_time_int: int = round(self.current_time)
+        fractional_time = self.current_time / self.target_times[-1]
         assert abs(self.current_time - current_time_int) < 1e-10
 
         if self.well_prepared_qubits_filter is None:
-            for callback in self.config.callbacks:
+            for callback in self.config.observables:
                 callback(
                     self.config,
-                    current_time_int,
+                    fractional_time,
                     normalized_state,
                     self.hamiltonian,
                     self.results,
@@ -403,63 +478,34 @@ class MPSBackendImpl:
             return
 
         full_mpo, full_state = None, None
-        for callback in self.config.callbacks:
-            if current_time_int not in callback.evaluation_times:
-                continue
-
-            if full_mpo is None or full_state is None:
-                # Only do this potentially expensive step once and when needed.
-                full_mpo = MPO(
-                    extended_mpo_factors(
-                        self.hamiltonian.factors, self.well_prepared_qubits_filter
-                    )
-                )
-                full_state = MPS(
-                    extended_mps_factors(
-                        normalized_state.factors, self.well_prepared_qubits_filter
-                    ),
-                    num_gpus_to_use=None,  # Keep the already assigned devices.
-                    orthogonality_center=get_extended_site_index(
-                        self.well_prepared_qubits_filter,
-                        normalized_state.orthogonality_center,
-                    ),
+        for callback in self.config.observables:
+            time_tol = 0.5 / self.target_times[-1] + 1e-10
+            if (
+                callback.evaluation_times is not None
+                and self.config.is_time_in_evaluation_times(
+                    fractional_time, callback.evaluation_times, tol=time_tol
                 )
+            ) or self.config.is_evaluation_time(fractional_time, tol=time_tol):
+
+                if full_mpo is None or full_state is None:
+                    # Only do this potentially expensive step once and when needed.
+                    full_mpo = MPO(
+                        extended_mpo_factors(
+                            self.hamiltonian.factors, self.well_prepared_qubits_filter
+                        )
+                    )
+                    full_state = MPS(
+                        extended_mps_factors(
+                            normalized_state.factors, self.well_prepared_qubits_filter
+                        ),
+                        num_gpus_to_use=None,  # Keep the already assigned devices.
+                        orthogonality_center=get_extended_site_index(
+                            self.well_prepared_qubits_filter,
+                            normalized_state.orthogonality_center,
+                        ),
+                    )
 
-            callback(self.config, current_time_int, full_state, full_mpo, self.results)
-
-    def log_step_statistics(self, *, duration: float) -> None:
-        if self.state.factors[0].is_cuda:
-            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
-
-        self.config.logger.info(
-            f"step = {self.timestep_index}/{self.timestep_count}, "
-            + f"χ = {self.state.get_max_bond_dim()}, "
-            + f"|ψ| = {self.state.get_memory_footprint():.3f} MB, "
-            + f"RSS = {max_mem:.3f} MB, "
-            + f"Δt = {duration:.3f} s"
-        )
-
-        if self.results.statistics is None:
-            assert self.timestep_index == 1
-            self.results.statistics = {"steps": []}
-
-        assert "steps" in self.results.statistics
-        assert len(self.results.statistics["steps"]) == self.timestep_index - 1
-
-        self.results.statistics["steps"].append(
-            {
-                "max_bond_dimension": self.state.get_max_bond_dim(),
-                "memory_footprint": self.state.get_memory_footprint(),
-                "RSS": max_mem,
-                "duration": duration,
-            }
-        )
+                callback(self.config, fractional_time, full_state, full_mpo, self.results)
 
 
 class NoisyMPSBackendImpl(MPSBackendImpl):

emu_mps/mps_config.py

@@ -1,16 +1,37 @@
-from typing import Any
+from typing import Any, ClassVar
+from types import MethodType
 
-from emu_base import BackendConfig, State, DEVICE_COUNT
+import copy
 
+from emu_base import DEVICE_COUNT
+from emu_mps.custom_callback_implementations import (
+    energy_mps_impl,
+    energy_second_moment_mps_impl,
+    energy_variance_mps_impl,
+    correlation_matrix_mps_impl,
+    qubit_occupation_mps_impl,
+)
+from pulser.backend import (
+    Occupation,
+    CorrelationMatrix,
+    Energy,
+    EnergySecondMoment,
+    EnergyVariance,
+    BitStrings,
+    EmulationConfig,
+)
+import logging
+import pathlib
+import sys
 
-class MPSConfig(BackendConfig):
+
+class MPSConfig(EmulationConfig):
     """
-    The configuration of the emu-ct MPSBackend. The kwargs passed to this class
+    The configuration of the emu-mps MPSBackend. The kwargs passed to this class
     are passed on to the base class.
     See the API for that class for a list of available options.
 
     Args:
-        initial_state: the initial state to use in the simulation
         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
@@ -36,42 +57,129 @@ class MPSConfig(BackendConfig):
         >>>     with_modulation=True) #the last arg is taken from the base class
     """
 
+    # Whether to warn if unexpected kwargs are received
+    _enforce_expected_kwargs: ClassVar[bool] = True
+
     def __init__(
         self,
         *,
-        initial_state: State | None = None,
         dt: int = 10,
         precision: float = 1e-5,
         max_bond_dim: int = 1024,
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int = DEVICE_COUNT,
+        interaction_cutoff: float = 0.0,
+        log_level: int = logging.INFO,
+        log_file: pathlib.Path | None = None,
         autosave_prefix: str = "emu_mps_save_",
         autosave_dt: int = 600,  # 10 minutes
         **kwargs: Any,
     ):
-        super().__init__(**kwargs)
-        self.initial_state = initial_state
-        self.dt = dt
-        self.precision = precision
-        self.max_bond_dim = max_bond_dim
-        self.max_krylov_dim = max_krylov_dim
-        self.num_gpus_to_use = num_gpus_to_use
-        self.extra_krylov_tolerance = extra_krylov_tolerance
+        kwargs.setdefault("observables", [BitStrings(evaluation_times=[1.0])])
+        super().__init__(
+            dt=dt,
+            precision=precision,
+            max_bond_dim=max_bond_dim,
+            max_krylov_dim=max_krylov_dim,
+            extra_krylov_tolerance=extra_krylov_tolerance,
+            num_gpus_to_use=num_gpus_to_use,
+            interaction_cutoff=interaction_cutoff,
+            log_level=log_level,
+            log_file=log_file,
+            autosave_prefix=autosave_prefix,
+            autosave_dt=autosave_dt,
+            **kwargs,
+        )
 
-        if self.noise_model is not None:
-            if "doppler" in self.noise_model.noise_types:
-                raise NotImplementedError("Unsupported noise type: doppler")
-            if (
-                "amplitude" in self.noise_model.noise_types
-                and self.noise_model.amp_sigma != 0.0
-            ):
-                raise NotImplementedError("Unsupported noise type: amp_sigma")
-
-        self.autosave_prefix = autosave_prefix
-        self.autosave_dt = autosave_dt
+        if "doppler" in self.noise_model.noise_types:
+            raise NotImplementedError("Unsupported noise type: doppler")
+        if (
+            "amplitude" in self.noise_model.noise_types
+            and self.noise_model.amp_sigma != 0.0
+        ):
+            raise NotImplementedError("Unsupported noise type: amp_sigma")
 
         MIN_AUTOSAVE_DT = 10
         assert (
             self.autosave_dt > MIN_AUTOSAVE_DT
         ), f"autosave_dt must be larger than {MIN_AUTOSAVE_DT} seconds"
+
+        self.monkeypatch_observables()
+
+        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
+        ):
+            self.logger.warning(
+                "Warning: The runs and samples_per_run values of the NoiseModel are ignored!"
+            )
+
+    def _expected_kwargs(self) -> set[str]:
+        return super()._expected_kwargs() | {
+            "dt",
+            "precision",
+            "max_bond_dim",
+            "max_krylov_dim",
+            "extra_krylov_tolerance",
+            "num_gpus_to_use",
+            "interaction_cutoff",
+            "log_level",
+            "log_file",
+            "autosave_prefix",
+            "autosave_dt",
+        }
+
+    def monkeypatch_observables(self) -> None:
+        obs_list = []
+        for _, obs in enumerate(self.observables):  # monkey patch
+            obs_copy = copy.deepcopy(obs)
+            if isinstance(obs, Occupation):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    qubit_occupation_mps_impl, obs_copy
+                )
+            elif isinstance(obs, EnergyVariance):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    energy_variance_mps_impl, obs_copy
+                )
+            elif isinstance(obs, EnergySecondMoment):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    energy_second_moment_mps_impl, obs_copy
+                )
+            elif isinstance(obs, CorrelationMatrix):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    correlation_matrix_mps_impl, obs_copy
+                )
+            elif isinstance(obs, Energy):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    energy_mps_impl, obs_copy
+                )
+            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,
+            )

emu_mps/observables.py

@@ -0,0 +1,40 @@
+from pulser.backend.state import State
+from pulser.backend.observable import Observable
+from emu_mps.mps import MPS
+from typing import Sequence, Any
+import torch
+
+
+class EntanglementEntropy(Observable):
+    """Entanglement Entropy subclass used only in emu_mps"""
+
+    def __init__(
+        self,
+        mps_site: int,
+        *,
+        evaluation_times: Sequence[float] | None = None,
+        tag_suffix: str | None = None,
+    ):
+        super().__init__(evaluation_times=evaluation_times, tag_suffix=tag_suffix)
+        self.mps_site = mps_site
+
+    @property
+    def _base_tag(self) -> str:
+        return "entanglement_entropy"
+
+    def _to_abstract_repr(self) -> dict[str, Any]:
+        repr = super()._to_abstract_repr()
+        repr["mps_site"] = self.mps_site
+        return repr
+
+    def apply(self, *, state: State, **kwargs: Any) -> torch.Tensor:
+        if not isinstance(state, MPS):
+            raise NotImplementedError(
+                "Entanglement entropy observable is only available for emu_mps emulator."
+            )
+        if not (0 <= self.mps_site <= len(state.factors) - 2):
+            raise ValueError(
+                f"Invalid bond index {self.mps_site}. "
+                f"Expected value in range 0 <= bond_index <= {len(state.factors)-2}."
+            )
+        return state.entanglement_entropy(self.mps_site)

emu_mps/optimatrix/optimiser.py

@@ -210,9 +210,9 @@ def minimize_bandwidth(input_matrix: np.ndarray, samples: int = 100) -> list[int
     # We are interested in strength of the interaction, not sign
 
     L = input_mat.shape[0]
-    rnd_permutations = itertools.chain(
+    rnd_permutations: itertools.chain[list[int]] = itertools.chain(
         [list(range(L))],  # First element is always the identity list
-        (np.random.permutation(L).tolist() for _ in range(samples)),
+        (np.random.permutation(L).tolist() for _ in range(samples)),  # type: ignore[misc]
     )
 
     opt_permutations_and_opt_bandwidth = (

{emu_mps-1.2.7.dist-info → emu_mps-2.0.1.dist-info}/METADATA

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

emu_mps-2.0.1.dist-info/RECORD

@@ -0,0 +1,19 @@
+emu_mps/__init__.py,sha256=GIR0FJOvmMHby7Vx_Da3w7YqqLpOxa0E8E3poxXsapY,734
+emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
+emu_mps/custom_callback_implementations.py,sha256=CUs0kW3HRaPE7UeFNQOFbeWJMsz4hS2q4rgS57BBp-A,2411
+emu_mps/hamiltonian.py,sha256=LcBs6CKBb643a1e9AAVtQoUfa4L_0dIhLOKecx5OOWs,15864
+emu_mps/mpo.py,sha256=wSonS6i3zEt3yRTgyZ7F6vT51pUKavFcLOxFFphBv8k,8793
+emu_mps/mps.py,sha256=KqAjo-nxgM-xQSg1NFNchwXKoPRcrKuuycFMsWr7iX8,19610
+emu_mps/mps_backend.py,sha256=_3rlg6XeI4fHaDiJRfPL6pDkX9k48hAHKXd8fkvkOFs,2004
+emu_mps/mps_backend_impl.py,sha256=M_do7QRBLAoHfwF_EpyMCb6g7w7BSs5hsPa5UE0z9Zs,22958
+emu_mps/mps_config.py,sha256=89nu5OhNUX31eAeeYvvKnAHegpPVD43jH5Nmp635HVU,6984
+emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
+emu_mps/observables.py,sha256=7GQDH5kyaVNrwckk2f8ZJRV9Ca4jKhWWDsOCqYWsoEk,1349
+emu_mps/tdvp.py,sha256=pIQ2NXA2Mrkp3elhqQbX3pdJVbtKkG3c5r9fFlJo7pI,5755
+emu_mps/utils.py,sha256=BqRJYAcXqprtZVJ0V_j954ON2bhTdtZiaTojsYyrWrg,8193
+emu_mps/optimatrix/__init__.py,sha256=lHWYNeThHp57ZrwTwXd0p8bNvcCv0w_AZ31iCWflBUo,226
+emu_mps/optimatrix/optimiser.py,sha256=7j9_jMQC-Uh2DzdIVB44InRzZO6AbbGhvmm7lC6N3tk,6737
+emu_mps/optimatrix/permutations.py,sha256=JRXGont8B4QgbkV9CzrA0w7uzLgBrmZ1J9aa0G52hPo,1979
+emu_mps-2.0.1.dist-info/METADATA,sha256=aGt2cFa9rplEaJAo8z4J5uoAkqq9WnQIlOhciUc6vkg,3505
+emu_mps-2.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_mps-2.0.1.dist-info/RECORD,,

emu_mps-1.2.7.dist-info/RECORD

@@ -1,17 +0,0 @@
-emu_mps/__init__.py,sha256=KSTZWIZSKHhjt0yt8-fS23rFFcFVQciNiXBgHS0pnHU,646
-emu_mps/algebra.py,sha256=ngPtTH-j2ZCBWoaJZXlkUyIlug7dY7Q92gzfnRlpPMA,5485
-emu_mps/hamiltonian.py,sha256=LcBs6CKBb643a1e9AAVtQoUfa4L_0dIhLOKecx5OOWs,15864
-emu_mps/mpo.py,sha256=H5vkJvz4AfXfnPbvgWznBWpMUO8LnGL3_NAP3IhxZzQ,8740
-emu_mps/mps.py,sha256=J0I4oQP_F1woEKmnOqnXPOWxx2Y1addxNjosL3yhYAY,18214
-emu_mps/mps_backend.py,sha256=6fVaq-D4xyicYRjGjhqMEieC7---90LpfpbV7ZD7zkQ,2192
-emu_mps/mps_backend_impl.py,sha256=XT2HccHWd6Y1gIAs070pBxjPUPIHBl-hFCuqXJaPS-E,21256
-emu_mps/mps_config.py,sha256=ydKN0OOaWCBcNd9V-4CU5ZZ4w1FRT-bbKyZQD2WCaME,3317
-emu_mps/noise.py,sha256=h4X2EFjoC_Ok0gZ8I9wN77RANXaVehTBbjkcbY_GAmY,784
-emu_mps/tdvp.py,sha256=pIQ2NXA2Mrkp3elhqQbX3pdJVbtKkG3c5r9fFlJo7pI,5755
-emu_mps/utils.py,sha256=BqRJYAcXqprtZVJ0V_j954ON2bhTdtZiaTojsYyrWrg,8193
-emu_mps/optimatrix/__init__.py,sha256=lHWYNeThHp57ZrwTwXd0p8bNvcCv0w_AZ31iCWflBUo,226
-emu_mps/optimatrix/optimiser.py,sha256=cVMdm2r_4OpbthcQuFMrJ9rNR9WEJRga9c_lHrJFkhw,6687
-emu_mps/optimatrix/permutations.py,sha256=JRXGont8B4QgbkV9CzrA0w7uzLgBrmZ1J9aa0G52hPo,1979
-emu_mps-1.2.7.dist-info/METADATA,sha256=-yHfBZrLmNsmc-tA-Yb0KfmxULGgVrcLVUbx4F37oA4,3505
-emu_mps-1.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_mps-1.2.7.dist-info/RECORD,,