emu-sv 2.0.4__py3-none-any.whl → 2.2.0__py3-none-any.whl

emu_sv/__init__.py

@@ -15,6 +15,7 @@ from pulser.backend import (
 from .dense_operator import DenseOperator
 from .sv_backend import SVBackend, SVConfig
 from .state_vector import StateVector, inner
+from .density_matrix_state import DensityMatrix
 
 
 __all__ = [
@@ -34,6 +35,7 @@ __all__ = [
     "StateResult",
     "StateVector",
     "inner",
+    "DensityMatrix",
 ]
 
-__version__ = "2.0.4"
+__version__ = "2.2.0"

emu_sv/custom_callback_implementations.py

@@ -8,11 +8,13 @@ from pulser.backend import (
     Occupation,
     Energy,
 )
-
+from emu_sv.density_matrix_state import DensityMatrix
 from emu_sv.state_vector import StateVector
 from emu_sv.dense_operator import DenseOperator
 from emu_sv.hamiltonian import RydbergHamiltonian
 
+dtype = torch.float64
+
 
 def qubit_occupation_sv_impl(
     self: Occupation,
@@ -25,7 +27,7 @@ def qubit_occupation_sv_impl(
     Custom implementation of the occupation ❬ψ|nᵢ|ψ❭ for the state vector solver.
     """
     nqubits = state.n_qudits
-    occupation = torch.zeros(nqubits, dtype=torch.float64, device=state.vector.device)
+    occupation = torch.zeros(nqubits, dtype=dtype, device=state.vector.device)
     for i in range(nqubits):
         state_tensor = state.vector.view(2**i, 2, -1)
         # nᵢ is a projector and therefore nᵢ == nᵢnᵢ
@@ -34,6 +36,33 @@ def qubit_occupation_sv_impl(
     return occupation.cpu()
 
 
+def qubit_occupation_sv_den_mat_impl(
+    self: Occupation,
+    *,
+    config: EmulationConfig,
+    state: DensityMatrix,
+    hamiltonian: DenseOperator,
+) -> torch.Tensor:
+    """
+    Custom implementation of the occupation nᵢ observable for density matrix.
+    The observable nᵢ is given by: I ⊗ ... ⊗  nᵢ ⊗ ...  ⊗I
+    where nᵢ is the occupation operator for qubit i.
+    The expectation value is given by: <nᵢ> = Tr(ρ nᵢ).
+
+    The output will be a tensor of size (nqubits,), where each element will be the
+    expectation value of the occupation operator for each qubit.
+    In case of 3 atoms, the output will be a tensor of size (3,), where each element
+    will be <nᵢ> = Tr(ρnᵢ), or [ <n₁>, <n₂>, <n₃> ].
+    """
+    nqubits = state.n_qudits
+    occupation = torch.zeros(nqubits, dtype=dtype, device=state.matrix.device)
+    diag_state_tensor = state.matrix.diagonal()
+    for i in range(nqubits):
+        state_tensor = diag_state_tensor.view(2**i, 2, 2 ** (nqubits - i - 1))[:, 1, :]
+        occupation[i] = state_tensor.sum().real
+    return occupation.cpu()
+
+
 def correlation_matrix_sv_impl(
     self: CorrelationMatrix,
     *,
@@ -47,24 +76,44 @@ def correlation_matrix_sv_impl(
     TODO: extend to arbitrary two-point correlation ❬ψ|AᵢBⱼ|ψ❭
     """
     nqubits = state.n_qudits
-    correlation = torch.zeros(
-        nqubits, nqubits, dtype=torch.float64, device=state.vector.device
-    )
+    correlation = torch.zeros(nqubits, nqubits, dtype=dtype, device=state.vector.device)
 
     for i in range(nqubits):
         select_i = state.vector.view(2**i, 2, -1)
         select_i = select_i[:, 1]
-        for j in range(i, nqubits):  # select the upper triangle
-            if i == j:
-                value = torch.linalg.vector_norm(select_i) ** 2
-                correlation[j, j] = value
-            else:
-                select_i = select_i.view(2**i, 2 ** (j - i - 1), 2, -1)
-                select_ij = select_i[:, :, 1, :]
-                value = torch.linalg.vector_norm(select_ij) ** 2
-                correlation[i, j] = value
-                correlation[j, i] = value
+        correlation[i, i] = torch.linalg.vector_norm(select_i) ** 2
+        for j in range(i + 1, nqubits):  # select the upper triangle
+            select_i = select_i.view(2**i, 2 ** (j - i - 1), 2, -1)
+            select_ij = select_i[:, :, 1, :]
+            correlation[i, j] = torch.linalg.vector_norm(select_ij) ** 2
+            correlation[j, i] = correlation[i, j]
+
+    return correlation.cpu()
 
+
+def correlation_matrix_sv_den_mat_impl(
+    self: CorrelationMatrix,
+    *,
+    config: EmulationConfig,
+    state: DensityMatrix,
+    hamiltonian: DenseOperator,
+) -> torch.Tensor:
+    """
+    Custom implementation of the density-density correlation <nᵢnⱼ> = Tr(ρ nᵢnⱼ)
+    in the case of Lindblad noise
+    """
+    nqubits = state.n_qudits
+    correlation = torch.zeros(nqubits, nqubits, dtype=dtype)
+    state_diag_matrix = state.matrix.diagonal()
+    for i in range(nqubits):  # applying ni
+        shapei = (2**i, 2, 2 ** (nqubits - i - 1))
+        state_diag_ni = state_diag_matrix.view(*shapei)[:, 1, :]
+        correlation[i, i] = state_diag_ni.sum().real  # diagonal
+        for j in range(i + 1, nqubits):
+            shapeij = (2**i, 2 ** (j - i - 1), 2, 2 ** (nqubits - 1 - j))
+            state_diag_ni_nj = state_diag_ni.view(*shapeij)[:, :, 1, :]
+            correlation[i, j] = state_diag_ni_nj.sum().real
+            correlation[j, i] = correlation[i, j]
     return correlation.cpu()
 
 

emu_sv/density_matrix_state.py

@@ -111,6 +111,7 @@ class DensityMatrix(State[complex, torch.Tensor]):
         cls: Type[DensityMatrixType],
         *,
         eigenstates: Sequence[Eigenstate],
+        n_qudits: int,
         amplitudes: Mapping[str, complex],
     ) -> tuple[DensityMatrix, Mapping[str, complex]]:
         """Transforms a state given by a string into a density matrix.
@@ -140,7 +141,7 @@ class DensityMatrix(State[complex, torch.Tensor]):
         """
 
         state_vector, amplitudes = StateVector._from_state_amplitudes(
-            eigenstates=eigenstates, amplitudes=amplitudes
+            eigenstates=eigenstates, n_qudits=n_qudits, amplitudes=amplitudes
         )
 
         return DensityMatrix.from_state_vector(state_vector), amplitudes

emu_sv/lindblad_operator.py

@@ -1,5 +1,5 @@
 import torch
-from emu_base.jump_lindblad_operators import compute_noise_from_lindbladians
+from emu_base import compute_noise_from_lindbladians, matmul_2x2_with_batched
 
 
 dtype = torch.complex128
@@ -17,6 +17,10 @@ class RydbergLindbladian:
     where A_k is a jump operator and H is the Rydberg Hamiltonian.
     The complex -𝑖, will be multiplied in the evolution.
 
+    Only works with effective noise channels, i.e., the jump or collapse
+    operators. For more information, see:
+    https://pulser.readthedocs.io/en/stable/tutorials/effective_noise.html
+
     Attributes:
         nqubits (int): number of qubits in the system.
         omegas (torch.Tensor): amplited frequencies  Ωⱼ for each qubit, divided by 2.
@@ -96,7 +100,10 @@ class RydbergLindbladian:
 
         orignal_shape = density_matrix.shape
         density_matrix = density_matrix.view(2**target_qubit, 2, -1)
-        density_matrix = local_op @ density_matrix
+        if density_matrix.is_cpu:
+            density_matrix = local_op @ density_matrix
+        else:
+            density_matrix = matmul_2x2_with_batched(local_op, density_matrix)
 
         return density_matrix.view(orignal_shape)
 
@@ -114,9 +121,11 @@ class RydbergLindbladian:
         """
 
         orignal_shape = density_matrix.shape
-
         density_matrix = density_matrix.view(2 ** (target_qubit + self.nqubits), 2, -1)
-        density_matrix = local_op.conj() @ density_matrix
+        if density_matrix.is_cpu:
+            density_matrix = local_op.conj() @ density_matrix
+        else:
+            density_matrix = matmul_2x2_with_batched(local_op.conj(), density_matrix)
 
         return density_matrix.view(orignal_shape)
 

emu_sv/state_vector.py

@@ -51,7 +51,7 @@ class StateVector(State[complex, torch.Tensor]):
     @property
     def n_qudits(self) -> int:
         """The number of qudits in the state."""
-        nqudits = math.log2(self.vector.reshape(-1).shape[0])
+        nqudits = math.log2(self.vector.view(-1).shape[0])
         return int(nqudits)
 
     def _normalize(self) -> None:
@@ -220,6 +220,7 @@ class StateVector(State[complex, torch.Tensor]):
         cls: Type[StateVectorType],
         *,
         eigenstates: Sequence[Eigenstate],
+        n_qudits: int,
         amplitudes: Mapping[str, complex],
     ) -> tuple[StateVector, Mapping[str, complex]]:
         """Transforms a state given by a string into a state vector.
@@ -257,8 +258,7 @@ class StateVector(State[complex, torch.Tensor]):
         else:
             raise ValueError("Unsupported basis provided")
 
-        nqubits = cls._validate_amplitudes(amplitudes=amplitudes, eigenstates=eigenstates)
-        accum_state = StateVector.zero(num_sites=nqubits, eigenstates=eigenstates)
+        accum_state = StateVector.zero(num_sites=n_qudits, eigenstates=eigenstates)
 
         for state, amplitude in amplitudes.items():
             bin_to_int = int(

emu_sv/sv_backend.py

@@ -1,23 +1,14 @@
-import torch
-from resource import RUSAGE_SELF, getrusage
-import time
-import typing
-
-from pulser.backend import EmulatorBackend, Results, Observable, State, EmulationConfig
-
-from emu_base import PulserData
-
-from emu_sv.state_vector import StateVector
+from pulser.backend import EmulatorBackend
+from pulser.backend import Results
 from emu_sv.sv_config import SVConfig
-from emu_sv.time_evolution import do_time_step
-
-
-_TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/μs, dt in ns
+from emu_sv.sv_backend_impl import create_impl
 
 
 class SVBackend(EmulatorBackend):
     """
     A backend for emulating Pulser sequences using state vectors and sparse matrices.
+    Noisy simulation is supported by solving the Lindblad equation and using effective
+    noise channel or jump operators
     """
 
     default_config = SVConfig()
@@ -31,112 +22,5 @@ class SVBackend(EmulatorBackend):
         """
         assert isinstance(self._config, SVConfig)
 
-        pulser_data = PulserData(
-            sequence=self._sequence, config=self._config, dt=self._config.dt
-        )
-        self.target_times = pulser_data.target_times
-        self.time = time.time()
-        omega, delta, phi = pulser_data.omega, pulser_data.delta, pulser_data.phi
-
-        nsteps = omega.shape[0]
-        nqubits = omega.shape[1]
-
-        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=nsteps,
-        )
-
-        if self._config.initial_state is not None:
-            state = self._config.initial_state
-            state = StateVector(state.vector.clone(), gpu=state.vector.is_cuda)
-        else:
-            state = StateVector.make(nqubits, gpu=self._config.gpu)
-
-        for step in range(nsteps):
-            dt = self.target_times[step + 1] - self.target_times[step]
-
-            state.vector, H = do_time_step(
-                dt * _TIME_CONVERSION_COEFF,
-                omega[step],
-                delta[step],
-                phi[step],
-                pulser_data.full_interaction_matrix,
-                state.vector,
-                self._config.krylov_tolerance,
-            )
-
-            # callbacks in observables and self.statistics in H
-            # have "# type: ignore[arg-type]" because H has it's own type
-            # meaning H is not inherited from Operator class.
-            # We decided that ignore[arg-type] is better compared to
-            # having many unused NotImplemented methods
-            for callback in self._config.observables:
-                callback(
-                    self._config,
-                    self.target_times[step + 1] / self.target_times[-1],
-                    state,
-                    H,  # type: ignore[arg-type]
-                    self.results,
-                )
-
-            self.statistics.data.append(time.time() - self.time)
-            self.statistics(
-                self._config,
-                self.target_times[step + 1] / self.target_times[-1],
-                state,
-                H,  # type: ignore[arg-type]
-                self.results,
-            )
-            self.time = time.time()
-            del H
-
-        return self.results
-
-
-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: typing.Any,
-    ) -> dict:
-        """Calculates the observable to store in the Results."""
-        assert isinstance(state, StateVector)
-        assert isinstance(config, SVConfig)
-        duration = self.data[-1]
-        if state.vector.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"RSS = {max_mem:.3f} MB, "
-            + f"Δt = {duration:.3f} s"
-        )
-
-        return {
-            "RSS": max_mem,
-            "duration": duration,
-        }
+        impl = create_impl(self._sequence, self._config)
+        return impl._run()

emu_sv/sv_backend_impl.py

@@ -0,0 +1,241 @@
+from abc import abstractmethod
+import time
+import typing
+
+from emu_sv.hamiltonian import RydbergHamiltonian
+from emu_sv.lindblad_operator import RydbergLindbladian
+from pulser import Sequence
+import torch
+from resource import RUSAGE_SELF, getrusage
+
+from pulser.backend import Results, Observable, State, EmulationConfig
+from emu_base import PulserData
+
+from emu_sv.state_vector import StateVector
+from emu_sv.density_matrix_state import DensityMatrix
+from emu_sv.sv_config import SVConfig
+from emu_sv.time_evolution import EvolveStateVector, EvolveDensityMatrix
+
+_TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/μs, dt in ns
+
+
+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: typing.Any,
+    ) -> dict:
+        """Calculates the observable to store in the Results."""
+        assert isinstance(state, StateVector | DensityMatrix)
+        assert isinstance(config, SVConfig)
+        duration = self.data[-1]
+        if isinstance(state, StateVector) and state.vector.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)
+        elif isinstance(state, DensityMatrix) and state.matrix.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"RSS = {max_mem:.3f} MB, "
+            + f"Δt = {duration:.3f} s"
+        )
+
+        return {
+            "RSS": max_mem,
+            "duration": duration,
+        }
+
+
+class BaseSVBackendImpl:
+    """
+    This class is used to handle the state vector and density matrix evolution.
+    """
+
+    def __init__(self, config: SVConfig, pulser_data: PulserData):
+        self._config = config
+        self._pulser_data = pulser_data
+        self.target_times = pulser_data.target_times
+        self.omega = pulser_data.omega
+        self.delta = pulser_data.delta
+        self.phi = pulser_data.phi
+        self.nsteps = pulser_data.omega.shape[0]
+        self.nqubits = pulser_data.omega.shape[1]
+        self.state: State
+        self.time = time.time()
+        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.nsteps,
+        )
+        self._current_H: None | RydbergLindbladian | RydbergHamiltonian = None
+
+        if self._config.initial_state is not None and (
+            self._config.initial_state.n_qudits != self.nqubits
+        ):
+            raise ValueError(
+                "Mismatch in number of atoms: initial state has "
+                f"{self._config.initial_state.n_qudits} and the sequence has {self.nqubits}"
+            )
+
+    def step(self, step_idx: int) -> None:
+        """One step of the evolution"""
+        dt = self._compute_dt(step_idx)
+        self._evolve_step(dt, step_idx)
+        self._apply_observables(step_idx)
+        self._save_statistics(step_idx)
+
+    def _compute_dt(self, step_idx: int) -> float:
+        return self.target_times[step_idx + 1] - self.target_times[step_idx]
+
+    @abstractmethod
+    def _evolve_step(self, dt: float, step_idx: int) -> None:
+        """One step evolution"""
+
+    def _apply_observables(self, step_idx: int) -> None:
+        norm_time = self.target_times[step_idx + 1] / self.target_times[-1]
+        for callback in self._config.observables:
+            callback(
+                self._config,
+                norm_time,
+                self.state,
+                self._current_H,  # type: ignore[arg-type]
+                self.results,
+            )
+
+    def _save_statistics(self, step_idx: int) -> None:
+        norm_time = self.target_times[step_idx + 1] / self.target_times[-1]
+        self.statistics.data.append(time.time() - self.time)
+        self.statistics(
+            self._config,
+            norm_time,
+            self.state,
+            self._current_H,  # type: ignore[arg-type]
+            self.results,
+        )
+        self.time = time.time()
+        self._current_H = None
+
+    def _run(self) -> Results:
+        for step in range(self.nsteps):
+            self.step(step)
+
+        return self.results
+
+
+class SVBackendImpl(BaseSVBackendImpl):
+
+    def __init__(self, config: SVConfig, pulser_data: PulserData):
+        """
+        For running sequences without noise. The state will evolve accoring
+        to e^(-iH t)
+
+        Args:
+            config: The configuration for the emulator.
+            pulser_data: The data for the sequence to be emulated.
+        """
+        super().__init__(config, pulser_data)
+
+        self.state: StateVector = (
+            StateVector.make(self.nqubits, gpu=self._config.gpu)
+            if self._config.initial_state is None
+            else StateVector(
+                self._config.initial_state.vector.clone(),
+                gpu=self._config.gpu,
+            )
+        )
+
+        self.stepper = EvolveStateVector.apply
+
+    def _evolve_step(self, dt: float, step_idx: int) -> None:
+        self.state.vector, self._current_H = self.stepper(
+            dt * _TIME_CONVERSION_COEFF,
+            self.omega[step_idx],
+            self.delta[step_idx],
+            self.phi[step_idx],
+            self._pulser_data.full_interaction_matrix,
+            self.state.vector,
+            self._config.krylov_tolerance,
+        )
+
+
+class NoisySVBackendImpl(BaseSVBackendImpl):
+
+    def __init__(self, config: SVConfig, pulser_data: PulserData):
+        """
+        Initializes the NoisySVBackendImpl, master equation version.
+        This class handles the Lindblad operators and
+        solves the Lindblad master equation
+
+        Args:
+            config: The configuration for the emulator.
+            pulser_data: The data for the sequence to be emulated.
+        """
+
+        super().__init__(config, pulser_data)
+
+        self.pulser_lindblads = pulser_data.lindblad_ops
+
+        self.state: DensityMatrix = (
+            DensityMatrix.make(self.nqubits, gpu=self._config.gpu)
+            if self._config.initial_state is None
+            else DensityMatrix(
+                self._config.initial_state.matrix.clone(), gpu=self._config.gpu
+            )
+        )
+
+    def _evolve_step(self, dt: float, step_idx: int) -> None:
+        self.state.matrix, self._current_H = EvolveDensityMatrix.evolve(
+            dt * _TIME_CONVERSION_COEFF,
+            self.omega[step_idx],
+            self.delta[step_idx],
+            self.phi[step_idx],
+            self._pulser_data.full_interaction_matrix,
+            self.state.matrix,
+            self._config.krylov_tolerance,
+            self.pulser_lindblads,
+        )
+
+
+def create_impl(sequence: Sequence, config: SVConfig) -> BaseSVBackendImpl:
+    """
+    Creates the backend implementation for the given sequence and config.
+
+    Args:
+        sequence: The sequence to be emulated.
+        config: configuration for the emulator.
+
+    Returns:
+        An instance of SVBackendImpl.
+    """
+    pulse_data = PulserData(sequence=sequence, config=config, dt=config.dt)
+    if pulse_data.has_lindblad_noise:
+        return NoisySVBackendImpl(config, pulse_data)
+    else:
+        return SVBackendImpl(config, pulse_data)

emu_sv/sv_config.py

@@ -7,9 +7,11 @@ from typing import Any, ClassVar
 
 from emu_sv.custom_callback_implementations import (
     correlation_matrix_sv_impl,
+    correlation_matrix_sv_den_mat_impl,
     energy_second_moment_sv_impl,
     energy_variance_sv_impl,
     qubit_occupation_sv_impl,
+    qubit_occupation_sv_den_mat_impl,
 )
 
 from pulser.backend import (
@@ -37,6 +39,10 @@ class SVConfig(EmulationConfig):
             the Lanczos algorithm uses this as the convergence tolerance
         gpu: Use 1 gpu if True, and a GPU is available, otherwise, cpu.
             Will cause errors if True when a gpu is not available
+        interaction_cutoff: Set interaction coefficients below this value to `0`.
+            Potentially improves runtime and memory consumption.
+        log_level: How much to log. Set to `logging.WARN` to get rid of the timestep info.
+        log_file: If specified, log to this file rather than stout.
         kwargs: arguments that are passed to the base class
 
     Examples:
@@ -97,6 +103,10 @@ class SVConfig(EmulationConfig):
                 "Warning: The runs and samples_per_run "
                 "values of the NoiseModel are ignored!"
             )
+        if "SPAM" in self.noise_model.noise_types:
+            raise NotImplementedError(
+                "SPAM errors are currently not supported in emu-sv."
+            )
 
     def _expected_kwargs(self) -> set[str]:
         return super()._expected_kwargs() | {
@@ -111,24 +121,41 @@ class SVConfig(EmulationConfig):
 
     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_sv_impl, obs_copy
+                    (
+                        qubit_occupation_sv_impl
+                        if self.noise_model.noise_types == ()
+                        else qubit_occupation_sv_den_mat_impl
+                    ),
+                    obs_copy,
+                )
+            if isinstance(obs, CorrelationMatrix):
+                obs_copy.apply = MethodType(  # type: ignore[method-assign]
+                    (
+                        correlation_matrix_sv_impl
+                        if self.noise_model.noise_types == ()
+                        else correlation_matrix_sv_den_mat_impl
+                    ),
+                    obs_copy,
                 )
-            elif isinstance(obs, EnergyVariance):
+            if isinstance(obs, EnergyVariance):
+                if self.noise_model.noise_types != ():
+                    raise Exception("Not implemented for density matrix")
                 obs_copy.apply = MethodType(  # type: ignore[method-assign]
                     energy_variance_sv_impl, obs_copy
                 )
             elif isinstance(obs, EnergySecondMoment):
+                if self.noise_model.noise_types != ():
+                    raise Exception("Not implemented for density matrix")
+
                 obs_copy.apply = MethodType(  # type: ignore[method-assign]
                     energy_second_moment_sv_impl, obs_copy
                 )
-            elif isinstance(obs, CorrelationMatrix):
-                obs_copy.apply = MethodType(  # type: ignore[method-assign]
-                    correlation_matrix_sv_impl, obs_copy
-                )
             obs_list.append(obs_copy)
         self.observables = tuple(obs_list)
 

emu_sv/time_evolution.py

@@ -1,33 +1,366 @@
 import torch
-
+from typing import Any, no_type_check
 from emu_base.math.krylov_exp import krylov_exp
+from emu_base.math.double_krylov import double_krylov
 from emu_sv.hamiltonian import RydbergHamiltonian
+from emu_sv.lindblad_operator import RydbergLindbladian
+
 
+def _apply_omega_real(
+    result: torch.Tensor,
+    i: int,
+    inds: torch.Tensor,
+    source: torch.Tensor,
+    alpha: complex,
+) -> None:
+    """Accumulate to `result` the application of ασˣᵢ on `source`"""
+    result.index_add_(i, inds, source, alpha=alpha)
 
-def do_time_step(
-    dt: float,
-    omegas: torch.Tensor,
-    deltas: torch.Tensor,
-    phis: torch.Tensor,
-    full_interaction_matrix: torch.Tensor,
-    state_vector: torch.Tensor,
-    krylov_tolerance: float,
-) -> tuple[torch.Tensor, RydbergHamiltonian]:
-    ham = RydbergHamiltonian(
-        omegas=omegas,
-        deltas=deltas,
-        phis=phis,
-        interaction_matrix=full_interaction_matrix,
-        device=state_vector.device,
+
+def _apply_omega_complex(
+    result: torch.Tensor,
+    i: int,
+    inds: torch.Tensor,
+    source: torch.Tensor,
+    alpha: complex,
+) -> None:
+    """Accumulate to `result` the application of ασ⁺ᵢ + α*σ⁻ᵢ on `source`"""
+    result.index_add_(i, inds[0], source.select(i, 0).unsqueeze(i), alpha=alpha)
+    result.index_add_(
+        i,
+        inds[1],
+        source.select(i, 1).unsqueeze(2),
+        alpha=alpha.conjugate(),
     )
-    op = lambda x: -1j * dt * (ham * x)
 
-    return (
-        krylov_exp(
+
+class DHDOmegaSparse:
+    """
+    Derivative of the RydbergHamiltonian respect to Omega.
+        ∂H/∂Ωₖ = 0.5[cos(ϕₖ)σˣₖ + sin(ϕₖ)σʸₖ]
+
+    If ϕₖ=0, simplifies to ∂H/∂Ωₖ = 0.5σˣₖ
+    """
+
+    def __init__(self, index: int, device: torch.device, nqubits: int, phi: torch.Tensor):
+        self.index = index
+        self.shape = (2**index, 2, 2 ** (nqubits - index - 1))
+        self.inds = torch.tensor([1, 0], device=device)  # flips the state, for 𝜎ₓ
+        self.alpha = 0.5 * torch.exp(1j * phi).item()
+        if phi.is_nonzero():
+            self._apply_sigmas = _apply_omega_complex
+        else:  # ∂H/∂Ωₖ = 0.5σˣₖ
+            self._apply_sigmas = _apply_omega_real
+
+    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+        vec = vec.view(vec.shape[0], *self.shape)  # add batch dimension
+        result = torch.zeros_like(vec)
+        self._apply_sigmas(result, 2, self.inds, vec, alpha=self.alpha)
+        return result.view(vec.shape[0], -1)
+
+
+class DHDPhiSparse:
+    """
+    Derivative of the RydbergHamiltonian respect to Phi.
+        ∂H/∂ϕₖ = 0.5Ωₖ[cos(ϕₖ+π/2)σˣₖ + sin(ϕₖ+π/2)σʸₖ]
+    """
+
+    def __init__(
+        self,
+        index: int,
+        device: torch.device,
+        nqubits: int,
+        omega: torch.Tensor,
+        phi: torch.Tensor,
+    ):
+        self.index = index
+        self.shape = (2**index, 2, 2 ** (nqubits - index - 1))
+        self.alpha = 0.5 * (omega * torch.exp(1j * (phi + torch.pi / 2))).item()
+        self.inds = torch.tensor([1, 0], device=device)  # flips the state, for 𝜎ₓ
+
+    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+        vec = vec.view(vec.shape[0], *self.shape)  # add batch dimension
+        result = torch.zeros_like(vec)
+        _apply_omega_complex(result, 2, self.inds, vec, alpha=self.alpha)
+        return result.view(vec.shape[0], -1)
+
+
+class DHDDeltaSparse:
+    """
+    Derivative of the Rydberg Hamiltonian respect to Delta:
+        ∂H/∂Δᵢ = -nᵢ
+    """
+
+    def __init__(self, i: int, nqubits: int):
+        self.nqubits = nqubits
+        self.shape = (2**i, 2, 2 ** (nqubits - i - 1))
+
+    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+        result = vec.clone()
+        result = result.view(vec.shape[0], *self.shape)
+        result[:, :, 0] = 0.0
+        return -result.view(vec.shape[0], 2**self.nqubits)
+
+
+class DHDUSparse:
+    """
+    Derivative of the Rydberg Hamiltonian respect to the interaction matrix:
+        ∂H/∂Uᵢⱼ = nᵢnⱼ
+    """
+
+    def __init__(self, i: int, j: int, nqubits: int):
+        self.shape = (2**i, 2, 2 ** (j - i - 1), 2, 2 ** (nqubits - j - 1))
+        self.nqubits = nqubits
+
+    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+        result = vec.clone()
+        result = result.view(vec.shape[0], *self.shape)
+        result[:, :, 0] = 0.0
+        result[:, :, 1, :, 0] = 0.0
+        return result.view(vec.shape[0], 2**self.nqubits)
+
+
+class EvolveStateVector(torch.autograd.Function):
+    """Custom autograd implementation of a step in the time evolution."""
+
+    @staticmethod
+    def evolve(
+        dt: float,
+        omegas: torch.Tensor,
+        deltas: torch.Tensor,
+        phis: torch.Tensor,
+        interaction_matrix: torch.Tensor,
+        state: torch.Tensor,
+        krylov_tolerance: float,
+    ) -> tuple[torch.Tensor, RydbergHamiltonian]:
+        ham = RydbergHamiltonian(
+            omegas=omegas,
+            deltas=deltas,
+            phis=phis,
+            interaction_matrix=interaction_matrix,
+            device=state.device,
+        )
+        op = lambda x: -1j * dt * (ham * x)
+        res = krylov_exp(
             op,
-            state_vector,
+            state,
             norm_tolerance=krylov_tolerance,
             exp_tolerance=krylov_tolerance,
-        ),
-        ham,
-    )
+            is_hermitian=True,
+        )
+        return res, ham
+
+    @staticmethod
+    def forward(
+        ctx: Any,
+        dt: float,
+        omegas: torch.Tensor,
+        deltas: torch.Tensor,
+        phis: torch.Tensor,
+        interaction_matrix: torch.Tensor,
+        state: torch.Tensor,
+        krylov_tolerance: float,
+    ) -> tuple[torch.Tensor, RydbergHamiltonian]:
+        """
+        Returns the time evolved state
+            |ψ(t+dt)〉= exp(-i dt H)|ψ(t)〉
+        under the Hamiltonian H built from the input Tensor parameters, omegas, deltas, phis and
+        the interaction matrix.
+
+        Args:
+            ctx (Any): context object to stash information for backward computation.
+            dt (float): timestep
+            omegas (torch.Tensor): 1D tensor of driving strengths for each qubit.
+            deltas (torch.Tensor): 1D tensor of detuning values for each qubit.
+            phis (torch.Tensor): 1D tensor of phase values for each qubit.
+            interaction_matrix (torch.Tensor): matrix representing the interaction
+                strengths between each pair of qubits.
+            state (Tensor): input state to be evolved
+            krylov_tolerance (float):
+        """
+        res, ham = EvolveStateVector.evolve(
+            dt, omegas, deltas, phis, interaction_matrix, state, krylov_tolerance
+        )
+        ctx.save_for_backward(omegas, deltas, phis, interaction_matrix, state)
+        ctx.dt = dt
+        ctx.tolerance = krylov_tolerance
+        return res, ham
+
+    # mypy complains and I don't know why
+    # backward expects same number of gradients as output of forward, gham is unused
+    @no_type_check
+    @staticmethod
+    def backward(ctx: Any, grad_state_out: torch.Tensor, gham: None) -> tuple[
+        None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        None,
+    ]:
+        """
+        In the backward pass we receive a Tensor containing the gradient of the loss L
+        with respect to the output
+            |gψ(t+dt)〉= ∂L/∂|ψ(t+dt)〉,
+        and return the gradients of the loss with respect to the input tensors parameters
+            - gΩⱼ = ∂L/∂Ωⱼ =〈gψ(t+dt)|dU(H,∂H/∂Ωⱼ)|ψ(t)〉
+            - gΔⱼ = ∂L/∂Δⱼ =  ...
+            - |gψ(t)〉= ∂L/∂|ψ(t)〉= exp(i dt H)|gψ(t+dt)〉
+
+        Args:
+            ctx (Any): context object to stash information for backward computation.
+            grad_state_out (torch.Tensor): |gψ(t+dt)〉
+
+        Return:
+            grad_omegas (torch.Tensor): 1D tensor of gradients with respect to Ωⱼ for each qubit.
+            grad_deltas (torch.Tensor): 1D tensor of gradients with respect to Δⱼ for each qubit.
+            grad_phis (torch.Tensor): 1D tensor of gradients with respect to φⱼ for each qubit.
+            grad_state_in (torch.Tensor): 1D tensor gradient with respect to the input state.
+
+        Notes:
+        Gradients are obtained by matching the total variations
+            〈gψ(t+dt)|d|ψ(t+dt)〉= ∑ⱼgΔⱼ*dΔⱼ + ∑ⱼgΩⱼ*dΩⱼ + ∑ⱼgφ*dφⱼ +〈gψ(t)|d|ψ(t)〉  (1)
+
+        For the exponential map U = exp(-i dt H), differentiating reads:
+            d|ψ(t+dt)〉= dU|ψ(t)〉+ Ud|ψ(t)〉
+            dU = ∑ⱼdU(H,∂H/∂Δⱼ) + ∑ⱼdU(H,∂H/∂Ωⱼ) + ∑ⱼdU(H,∂H/∂φⱼ)                      (2)
+
+        where dU(H,E) is the Fréchet derivative of the exponential map
+        along the direction E:
+        - https://eprints.maths.manchester.ac.uk/1218/1/covered/MIMS_ep2008_26.pdf
+        - https://en.wikipedia.org/wiki/Derivative_of_the_exponential_map
+
+        Substituting (2) into (1) leads to the expressions of the gradients
+        with respect to the input tensors above.
+
+        Variations with respect to the Hamiltonian parameters are computed as
+            gΩ = 〈gψ(t+dt)|dU(H,∂H/∂Ω)|ψ(t)〉
+               = Tr( -i dt ∂H/∂Ω @ dU(H,|ψ(t)〉〈gψ(t+dt)|) ),
+        where under the trace sign, ∂H/∂Ω and |ψ(t)〉〈gψ(t+dt)| can be switched.
+
+        - The Fréchet derivative is computed in a Arnoldi-Gram-Schmidt
+        decomposition in the `double_krylov` method:
+            dU(H,|a〉〈b|) = Va @ dS @ Vb*
+        where Va,Vb are orthogonal Krylov basis associated
+        with |a〉and |b〉respectively.
+
+        - The action of the derivatives of the Hamiltonian with
+        respect to the input parameters are implemented separately in
+            - ∂H/∂Ω:  `DHDOmegaSparse`
+            - ∂H/∂Δ:  `DHDDeltaSparse`
+            - ∂H/∂φ:  `DHDPhiSparse`
+            - ∂H/∂Uᵢⱼ `DHDUSparse`
+
+        Then, the resulting gradient respect to a generic parameter reads:
+            gΩ = Tr( -i dt ∂H/∂Ω @ Vs @ dS @ Vg* )
+        """
+        omegas, deltas, phis, interaction_matrix, state = ctx.saved_tensors
+        dt = ctx.dt
+        tolerance = ctx.tolerance
+        nqubits = len(omegas)
+
+        grad_omegas, grad_deltas, grad_phis = None, None, None
+        grad_int_mat = None
+        grad_state_in = None
+
+        ham = RydbergHamiltonian(
+            omegas=omegas,
+            deltas=deltas,
+            phis=phis,
+            interaction_matrix=interaction_matrix,
+            device=state.device,
+        )
+
+        if any(ctx.needs_input_grad[1:5]):
+            op = lambda x: -1j * dt * (ham * x)
+            lanczos_vectors_state, dS, lanczos_vectors_grad = double_krylov(
+                op, state, grad_state_out, tolerance
+            )
+            # TODO: explore returning directly the basis in matrix form
+            Vs = torch.stack(lanczos_vectors_state)
+            del lanczos_vectors_state
+            Vg = torch.stack(lanczos_vectors_grad)
+            del lanczos_vectors_grad
+            e_l = dS.mT @ Vs
+
+        if ctx.needs_input_grad[1]:
+            grad_omegas = torch.zeros_like(omegas)
+            for i in range(nqubits):
+                # dh as per the docstring
+                dho = DHDOmegaSparse(i, e_l.device, nqubits, phis[i])
+                # compute the trace
+                v = dho @ e_l
+                grad_omegas[i] = (-1j * dt * torch.tensordot(Vg.conj(), v)).real
+
+        if ctx.needs_input_grad[2]:
+            grad_deltas = torch.zeros_like(deltas)
+            for i in range(nqubits):
+                dhd = DHDDeltaSparse(i, nqubits)
+                v = dhd @ e_l
+                grad_deltas[i] = (-1j * dt * torch.tensordot(Vg.conj(), v)).real
+
+        if ctx.needs_input_grad[3]:
+            grad_phis = torch.zeros_like(phis)
+            for i in range(nqubits):
+                dhp = DHDPhiSparse(i, e_l.device, nqubits, omegas[i], phis[i])
+                v = dhp @ e_l
+                grad_phis[i] = (-1j * dt * torch.tensordot(Vg.conj(), v)).real
+
+        if ctx.needs_input_grad[4]:
+            grad_int_mat = torch.zeros_like(interaction_matrix)
+            for i in range(nqubits):
+                for j in range(i + 1, nqubits):
+                    dhu = DHDUSparse(i, j, nqubits)
+                    v = dhu @ e_l
+                    grad_int_mat[i, j] = (-1j * dt * torch.tensordot(Vg.conj(), v)).real
+
+        if ctx.needs_input_grad[5]:
+            op = lambda x: (1j * dt) * (ham * x)
+            grad_state_in = krylov_exp(op, grad_state_out.detach(), tolerance, tolerance)
+
+        return (
+            None,
+            grad_omegas,
+            grad_deltas,
+            grad_phis,
+            grad_int_mat,
+            grad_state_in,
+            None,
+        )
+
+
+class EvolveDensityMatrix:
+    """Evolution of a density matrix under a Lindbladian operator."""
+
+    @staticmethod
+    def evolve(
+        dt: float,
+        omegas: torch.Tensor,
+        deltas: torch.Tensor,
+        phis: torch.Tensor,
+        full_interaction_matrix: torch.Tensor,
+        density_matrix: torch.Tensor,
+        krylov_tolerance: float,
+        pulser_lindblads: list[torch.Tensor],
+    ) -> tuple[torch.Tensor, RydbergLindbladian]:
+        ham = RydbergLindbladian(
+            omegas=omegas,
+            deltas=deltas,
+            phis=phis,
+            pulser_linblads=pulser_lindblads,
+            interaction_matrix=full_interaction_matrix,
+            device=density_matrix.device,
+        )
+
+        op = lambda x: -1j * dt * (ham @ x)
+        return (
+            krylov_exp(
+                op,
+                density_matrix,
+                norm_tolerance=krylov_tolerance,
+                exp_tolerance=krylov_tolerance,
+                is_hermitian=False,
+            ),
+            ham,
+        )

{emu_sv-2.0.4.dist-info → emu_sv-2.2.0.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: emu-sv
-Version: 2.0.4
+Version: 2.2.0
 Summary: Pasqal State Vector based pulse emulator built on PyTorch
 Project-URL: Documentation, https://pasqal-io.github.io/emulators/
 Project-URL: Repository, https://github.com/pasqal-io/emulators
 Project-URL: Issues, https://github.com/pasqal-io/emulators/issues
-Author-email: Anton Quelle <anton.quelle@pasqal.com>, Mauro Mendizabal <mauro.mendizabal-pico@pasqal.com>, Stefano Grava <stefano.grava@pasqal.com>, Pablo Le Henaff <pablo.le-henaff@pasqal.com>
+Author-email: Kemal Bidzhiev <kemal.bidzhiev@pasqal.com>, Stefano Grava <stefano.grava@pasqal.com>, Pablo Le Henaff <pablo.le-henaff@pasqal.com>, Mauro Mendizabal <mauro.mendizabal-pico@pasqal.com>, Elie Merhej <elie.merhej@pasqal.com>, Anton Quelle <anton.quelle@pasqal.com>
 License: PASQAL OPEN-SOURCE SOFTWARE LICENSE AGREEMENT (MIT-derived)
         
         The author of the License is:
@@ -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.0.4
+Requires-Dist: emu-base==2.2.0
 Description-Content-Type: text/markdown
 
 <div align="center">

emu_sv-2.2.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+emu_sv/__init__.py,sha256=vVU7zIUsHPgFXUBWMn1lUtkWUg-avdFRqLUrtbmPvDI,771
+emu_sv/custom_callback_implementations.py,sha256=j_G5x-xSnFSyKE81QN_DKoiTaF9JvNWaM1N6k2_t0Oo,5501
+emu_sv/dense_operator.py,sha256=NfgzVpnNitc5ZSM4RlfpAc5Ls2wFPNsTxdeFdhJSg1o,6909
+emu_sv/density_matrix_state.py,sha256=5W_UKIAYHb0k3ryRLQ2dbFUgrb5ju5jceDGAekM2gNE,7035
+emu_sv/hamiltonian.py,sha256=CqNGuWJlO2ZljK47wt130s-5uKiOldQUsC3tjwk1mKA,6106
+emu_sv/lindblad_operator.py,sha256=Rlxh24dXVUAutSrSacNO2ilNVlxKix8pfFt7h2CfVOg,7893
+emu_sv/state_vector.py,sha256=zKHCdgl_eRIOPE4qVKO53ig9UyYTQ7a_guNFXgynU7g,9753
+emu_sv/sv_backend.py,sha256=-soOkSEzEBK1dCKnYnbtvYjmNZtZra1_4jP3H1ROOtM,737
+emu_sv/sv_backend_impl.py,sha256=mdPWBLDwH0q7EEwQTmLNLLx5tycMmsCQbUifIHvciMk,8059
+emu_sv/sv_config.py,sha256=ixMTgDXKll4bXsYtAe4a_9Gng0bhwCgS42KKMwZCFHI,6308
+emu_sv/time_evolution.py,sha256=_VH4f2RF6lGKzO08WxTYJ5Aw8_pTTMRKcyMnIuxH03I,13382
+emu_sv/utils.py,sha256=-axfQ2tqw0C7I9yw-28g7lytyk373DNBjDALh4kLBrM,302
+emu_sv-2.2.0.dist-info/METADATA,sha256=aw71GDNhy2P9-zjXjXszkLyaOY95qCpPkOmG2lAmYZg,3595
+emu_sv-2.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_sv-2.2.0.dist-info/RECORD,,

emu_sv-2.0.4.dist-info/RECORD

@@ -1,14 +0,0 @@
-emu_sv/__init__.py,sha256=Tzc6RlABZ1ZVKt1mPUes9djq0eiK6FTgPagEHuFhF6Q,702
-emu_sv/custom_callback_implementations.py,sha256=zvsSiDIc56gwybKq87VFZyKsniTDye6-oFd2-R0shpg,3447
-emu_sv/dense_operator.py,sha256=NfgzVpnNitc5ZSM4RlfpAc5Ls2wFPNsTxdeFdhJSg1o,6909
-emu_sv/density_matrix_state.py,sha256=6UBLUXaJaUdzOhflrKolcnH8737JszX7sry1WmbyakI,6993
-emu_sv/hamiltonian.py,sha256=CqNGuWJlO2ZljK47wt130s-5uKiOldQUsC3tjwk1mKA,6106
-emu_sv/lindblad_operator.py,sha256=KmaNCahpAi8SIXh-TrFD-ggmGpa1zklp8DMWVK9Y_J4,7433
-emu_sv/state_vector.py,sha256=lqSbv4BMtDtgY0YUPuhIUNJxrlVa7vUWuN_XqwpG5sQ,9823
-emu_sv/sv_backend.py,sha256=AkEtI6-SY20D0ORro3Kv8tHDRUc8gxejSiRa6d--vBE,4452
-emu_sv/sv_config.py,sha256=QRy0VbCugmY6TQZ48nD6RxPJbpu0wzN7-E1Sud7YxLQ,5106
-emu_sv/time_evolution.py,sha256=obV7DcHot0jtnEmjR1ilYiSyDcJ5rTvThRB8hFjP-2s,797
-emu_sv/utils.py,sha256=-axfQ2tqw0C7I9yw-28g7lytyk373DNBjDALh4kLBrM,302
-emu_sv-2.0.4.dist-info/METADATA,sha256=paHoVgW22OONoxvlwaypB-UF01zi0giZqUtvAz7fhmw,3513
-emu_sv-2.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_sv-2.0.4.dist-info/RECORD,,