emu-sv 2.0.3__py3-none-any.whl → 2.1.1__py3-none-any.whl

emu_sv/__init__.py

@@ -16,6 +16,7 @@ from .dense_operator import DenseOperator
 from .sv_backend import SVBackend, SVConfig
 from .state_vector import StateVector, inner
 
+
 __all__ = [
     "__version__",
     "BitStrings",
@@ -35,4 +36,4 @@ __all__ = [
     "inner",
 ]
 
-__version__ = "2.0.3"
+__version__ = "2.1.1"

emu_sv/hamiltonian.py

@@ -3,34 +3,36 @@ from emu_sv.state_vector import StateVector
 
 
 class RydbergHamiltonian:
-    """
-    Representation of the Rydberg Hamiltonian with light-matter interaction:
+    """Represents the Rydberg Hamiltonian for a system of interacting qubits
+    driven by laser fields, including detuning, phase, and interaction terms.
+
+    The Hamiltonian is defined as:
 
-        H = ∑ⱼΩⱼ/2[cos(ϕⱼ)σˣⱼ + sin(ϕⱼ)σʸⱼ] - ∑ⱼΔⱼnⱼ + ∑ᵢ﹥ⱼUᵢⱼnᵢnⱼ
+        H = ∑ⱼ (Ωⱼ/2)[cos(ϕⱼ) σˣⱼ + sin(ϕⱼ) σʸⱼ] - ∑ⱼ Δⱼ nⱼ + ∑_{i>j} Uᵢⱼ nᵢ nⱼ
 
-    The Hamiltonian is parameterized by driving strengths or amplitudes Ωⱼ (`omegas`), detuning
-    values Δⱼ (`deltas`), phases ϕⱼ (`phis`) and interaction terms Uᵢⱼ (`interaction_matrix`).
-    Implements an efficient H*|ψ❭ as custom sparse matrix-vector multiplication.
+    where:
+        - Ωⱼ is the Rabi frequency on qubit j,
+        - Δⱼ is the detuning on qubit j,
+        - ϕⱼ is the laser phase on qubit j,
+        - Uᵢⱼ is the interaction strength between qubits i and j,
+        - nⱼ = |1⟩⟨1| is the number operator on qubit j.
 
     Attributes:
-        omegas (torch.Tensor): driving strength Ωⱼ for each qubit, scaled by a factor 1/2.
-        deltas (torch.Tensor): detuning values Δⱼ for each qubit.
-        phis (torch.Tensor): phase values ϕⱼ for each qubit.
-        interaction_matrix (torch.Tensor): matrix Uᵢⱼ representing pairwise Rydberg
-            interaction strengths between qubits.
+        omegas (torch.Tensor): vector of Rabi frequencies Ωⱼ / 2 for each qubit.
+        deltas (torch.Tensor): vector of detunings Δⱼ for each qubit.
+        phis (torch.Tensor): vector of phases ϕⱼ for each qubit.
+        interaction_matrix (torch.Tensor): matrix Uᵢⱼ for pairwise interactions.
+        device (torch.device): device on which all tensors are allocated.
+        diag (torch.Tensor): diagonal contribution to the Hamiltonian (detuning + interactions).
+        inds (torch.Tensor): index mapping for σˣ operations.
         nqubits (int): number of qubits in the system.
-        diag (torch.Tensor): diagonal elements of the Hamiltonian,
-            calculated based on `deltas` and `interaction_matrix`.
-        inds (torch.Tensor): index tensor used for vector manipulations
-            in matrix-vector multiplications.
 
     Methods:
-        __mul__(vec): implements matrix-vector multiplication with a state vector.
-        _create_diagonal(): constructs the diagonal elements of the Hamiltonian
-            based on `deltas` and `interaction_matrix`.
-        _apply_sigma_operators_complex(): apply all driving sigma operators,
-             with driving strenght `omegas` and phases `phis`.
-        _apply_sigma_operators_real(): only applies ∑ⱼ(Ωⱼ/2)σˣⱼ when all phases are zero (ϕⱼ=0).
+        __mul__(vec): Applies the Hamiltonian H to a state vector |ψ⟩.
+        _apply_sigma_operators_real(): Applies only σˣ terms (ϕⱼ = 0).
+        _apply_sigma_operators_complex(): Applies generalized σˣ/σʸ terms (ϕⱼ ≠ 0).
+        _create_diagonal(): Computes the diagonal part of H from Δⱼ and Uᵢⱼ.
+        expect(state): Computes ⟨ψ|H|ψ⟩ for a given StateVector.
     """
 
     def __init__(
@@ -107,7 +109,7 @@ class RydbergHamiltonian:
         Returns:
             the resulting state vector.
         """
-        c_omegas = self.omegas * torch.exp(1j * self.phis)
+        c_omegas = self.omegas * torch.exp(1.0j * self.phis)
 
         dim_to_act = 1
         for n, c_omega_n in enumerate(c_omegas):

emu_sv/lindblad_operator.py

@@ -0,0 +1,188 @@
+import torch
+from emu_base.jump_lindblad_operators import compute_noise_from_lindbladians
+
+
+dtype = torch.complex128
+sigmax = torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=dtype)
+sigmay = torch.tensor([[0.0, -1.0j], [1.0j, 0.0]], dtype=dtype)
+n_op = torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype)
+
+
+class RydbergLindbladian:
+    """
+    Apply the Lindblad superoperator ℒ to a density matrix 𝜌,  ℒ(𝜌).
+
+    This class implements
+    H @𝜌- H @ 𝜌 + i ∑ₖ − 1/2 Aₖ† Aₖ 𝜌 − 1/2 𝜌 Aₖ^† Aₖ + Aₖ 𝜌 Aₖ^†,
+    where A_k is a jump operator and H is the Rydberg Hamiltonian.
+    The complex -𝑖, will be multiplied in the evolution.
+
+    Attributes:
+        nqubits (int): number of qubits in the system.
+        omegas (torch.Tensor): amplited frequencies  Ωⱼ for each qubit, divided by 2.
+        deltas (torch.Tensor): detunings 𝛿ᵢ for each qubit.
+        phis (torch.Tensor): phases 𝜙ᵢ for each qubit.
+        interaction_matrix (torch.Tensor): interaction_matrix (torch.Tensor): matrix Uᵢⱼ
+            representing pairwise Rydberg interaction strengths between qubits.
+        pulser_linblads (list[torch.Tensor]): List of 2x2 local Lindblad (jump)
+            operators acting on each qubit.
+        device (torch.device): device on which tensors are allocated. cpu or gpu: cuda.
+        complex (bool): flag indicating whether any drive phase is nonzero
+            (i.e., complex Hamiltonian terms).
+        diag (torch.Tensor): precomputed diagonal interaction term for the density matrix evolution.
+
+    Methods:
+        apply_local_op_to_density_matrix(density_matrix, local_op, target_qubit):
+            Applies a local operator to the density matrix from the left: L @ ρ.
+
+        apply_density_matrix_to_local_op_T(density_matrix, local_op, target_qubit):
+            Applies a daggered local operator to the density matrix from the right: ρ @ L†.
+
+        __matmul__(density_matrix):
+            Applies the full Lindbladian superoperator to the input density matrix,
+            including coherent evolution and all dissipation channels.
+    """
+
+    def __init__(
+        self,
+        omegas: torch.Tensor,
+        deltas: torch.Tensor,
+        phis: torch.Tensor,
+        pulser_linblads: list[torch.Tensor],
+        interaction_matrix: torch.Tensor,
+        device: torch.device,
+    ):
+        self.nqubits: int = len(omegas)
+        self.omegas: torch.Tensor = omegas / 2.0
+        self.deltas: torch.Tensor = deltas
+        self.phis: torch.Tensor = phis
+        self.interaction_matrix: torch.Tensor = interaction_matrix
+        self.pulser_linblads: list[torch.Tensor] = pulser_linblads
+        self.device: torch.device = device
+        self.complex = self.phis.any()
+
+        self.diag: torch.Tensor = self._create_diagonal()
+
+    def _create_diagonal(self) -> torch.Tensor:
+        """
+        Return the diagonal elements of the Rydberg Hamiltonian matrix
+        concerning only the interaction
+
+            H.diag =  ∑ᵢ﹥ⱼUᵢⱼnᵢnⱼ
+        """
+        diag = torch.zeros(2**self.nqubits, dtype=dtype, device=self.device)
+
+        for i in range(self.nqubits):
+            diag = diag.view(2**i, 2, -1)
+            i_fixed = diag[:, 1, :]
+            for j in range(i + 1, self.nqubits):
+                i_fixed = i_fixed.view(2**i, 2 ** (j - i - 1), 2, -1)
+                # replacing i_j_fixed by i_fixed breaks the code :)
+                i_j_fixed = i_fixed[:, :, 1, :]
+                i_j_fixed += self.interaction_matrix[i, j]
+        return diag.view(-1)
+
+    def apply_local_op_to_density_matrix(
+        self,
+        density_matrix: torch.Tensor,
+        local_op: torch.Tensor,
+        target_qubit: int,
+    ) -> torch.Tensor:
+        """
+        Calculate a local operator (2x2) L being multiplied by a density matrix ρ
+        from the left
+        Return L @ ρ
+        """
+
+        orignal_shape = density_matrix.shape
+        density_matrix = density_matrix.view(2**target_qubit, 2, -1)
+        density_matrix = local_op @ density_matrix
+
+        return density_matrix.view(orignal_shape)
+
+    def apply_density_matrix_to_local_op_T(
+        self,
+        density_matrix: torch.Tensor,
+        local_op: torch.Tensor,
+        target_qubit: int,
+    ) -> torch.Tensor:
+        """
+        Calculates a density matrix ρ being multiplied by a daggered local (2x2)
+          operator L† from the right,
+
+        return:  ρ @L†
+        """
+
+        orignal_shape = density_matrix.shape
+
+        density_matrix = density_matrix.view(2 ** (target_qubit + self.nqubits), 2, -1)
+        density_matrix = local_op.conj() @ density_matrix
+
+        return density_matrix.view(orignal_shape)
+
+    def __matmul__(self, density_matrix: torch.Tensor) -> torch.Tensor:
+        """Apply the i*RydbergLindbladian operator to the density matrix ρ
+        in the following way:
+        Define and effective Hamiltonian
+        Heff = Hρ  -0.5i ∑ₖ Lₖ† Lₖ ρ
+        Then, the Lindblad operator applying to ρ is giving by
+         ℒ(𝜌) = Heff - Heff^†+i*∑ₖ Lₖ ρ Lₖ†
+        """
+
+        # compute -0.5i ∑ₖ Lₖ† Lₖ
+        sum_lindblad_local = compute_noise_from_lindbladians(self.pulser_linblads).to(
+            self.device
+        )
+
+        # apply all local terms:  Ωⱼ σₓ - δⱼ n - 0.5i (∑ₖ Lₖ† Lₖ) to each qubit
+        H_den_matrix = torch.zeros_like(density_matrix, dtype=dtype, device=self.device)
+
+        if not self.complex:
+            for qubit, (omega, delta) in enumerate(zip(self.omegas, self.deltas)):
+                H_q = (
+                    omega * sigmax.to(device=self.device)
+                    - delta * n_op.to(device=self.device)
+                    + sum_lindblad_local
+                )
+                H_den_matrix += self.apply_local_op_to_density_matrix(
+                    density_matrix, H_q, qubit
+                )
+        else:
+            for qubit, (omega, delta, phi) in enumerate(
+                zip(self.omegas, self.deltas, self.phis)
+            ):
+                H_q = (
+                    omega
+                    * (
+                        (
+                            torch.cos(phi) * sigmax.to(device=self.device)
+                            + torch.sin(phi) * sigmay.to(device=self.device)
+                        )
+                    )
+                    - delta * n_op.to(device=self.device)
+                    + sum_lindblad_local
+                )
+                H_den_matrix += self.apply_local_op_to_density_matrix(
+                    density_matrix, H_q, qubit
+                )
+
+        # apply the interaction terms  ∑ᵢⱼ Uᵢⱼ nᵢ nⱼ
+        H_den_matrix += self.diag.view(-1, 1) * density_matrix
+
+        # Heff - Heff^†=  [H, ρ] - 0.5i ∑ₖ Lₖ† Lₖρ - ρ 0.5i ∑ₖ Lₖ† Lₖρ
+        H_den_matrix = H_den_matrix - H_den_matrix.conj().T
+
+        # compute ∑ₖ Lₖ ρ Lₖ†, last part of the Lindblad operator
+        L_den_matrix_Ldag = sum(
+            self.apply_density_matrix_to_local_op_T(
+                self.apply_local_op_to_density_matrix(
+                    density_matrix, L.to(self.device), qubit
+                ),
+                L.to(self.device),
+                qubit,
+            )
+            for qubit in range(self.nqubits)
+            for L in self.pulser_linblads
+        )
+
+        return H_den_matrix + 1.0j * L_den_matrix_Ldag

emu_sv/sv_backend.py

@@ -9,7 +9,7 @@ from emu_base import PulserData
 
 from emu_sv.state_vector import StateVector
 from emu_sv.sv_config import SVConfig
-from emu_sv.time_evolution import do_time_step
+from emu_sv.time_evolution import EvolveStateVector
 
 
 _TIME_CONVERSION_COEFF = 0.001  # Omega and delta are given in rad/μs, dt in ns
@@ -54,10 +54,10 @@ class SVBackend(EmulatorBackend):
         else:
             state = StateVector.make(nqubits, gpu=self._config.gpu)
 
+        stepper = EvolveStateVector.apply
         for step in range(nsteps):
             dt = self.target_times[step + 1] - self.target_times[step]
-
-            state.vector, H = do_time_step(
+            state.vector, H = stepper(
                 dt * _TIME_CONVERSION_COEFF,
                 omega[step],
                 delta[step],

emu_sv/sv_config.py

@@ -37,6 +37,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:

emu_sv/time_evolution.py

@@ -1,32 +1,283 @@
 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
 
 
-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_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 _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, index: int, device: torch.device, nqubits: int):
+        self.index = index
+        self.shape = (2**index, 2, 2 ** (nqubits - index - 1))
+        diag = torch.zeros(
+            *self.shape,
+            dtype=torch.complex128,
+            device=device,
+        )
+        diag[:, 1, :] = -1.0
+        self.diag = diag.reshape(-1)
+
+    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+        return vec * self.diag
+
+
+class EvolveStateVector(torch.autograd.Function):
+    """Custom autograd implementation of a step in the time evolution."""
+
+    @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):
+        """
+        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,
+        )
+        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`
+
+        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, grad_state_in = None, None, None, None
+
+        ham = RydbergHamiltonian(
+            omegas=omegas,
+            deltas=deltas,
+            phis=phis,
+            interaction_matrix=interaction_matrix,
+            device=state.device,
+        )
+
+        if any(ctx.needs_input_grad[1:4]):
+            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):
+                # dh as per the docstring
+                dhd = DHDDeltaSparse(i, e_l.device, nqubits)
+                # compute the trace
+                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[5]:
+            op = lambda x: (1j * dt) * (ham * x)
+            grad_state_in = krylov_exp(op, grad_state_out, tolerance, tolerance)
+
+        return None, grad_omegas, grad_deltas, grad_phis, None, grad_state_in, None

{emu_sv-2.0.3.dist-info → emu_sv-2.1.1.dist-info}/METADATA

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

emu_sv-2.1.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+emu_sv/__init__.py,sha256=907RfJcGq7Nq6ayS0bsE5q7vrMTTEf41FBIXt651lvY,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=FrSBG8pacgucZ4MHKApfPh-VEFApsjnBzVb03GCcTpc,4493
+emu_sv/sv_config.py,sha256=q-cnyWwr_nNbpXI_m5vG51Wz_tyV5TL5M28uP2WctP4,5412
+emu_sv/time_evolution.py,sha256=pTmWzgI4AboRYklvCz4OLQNNN_RB1bOqJBXdsrFf6jk,10867
+emu_sv/utils.py,sha256=-axfQ2tqw0C7I9yw-28g7lytyk373DNBjDALh4kLBrM,302
+emu_sv-2.1.1.dist-info/METADATA,sha256=8QGaiNu0nRthfw2O717nMMoO60LCLnVaVrAPjp_t7dk,3513
+emu_sv-2.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+emu_sv-2.1.1.dist-info/RECORD,,

emu_sv-2.0.3.dist-info/RECORD

@@ -1,13 +0,0 @@
-emu_sv/__init__.py,sha256=6ijDylrpDM3KZkI1hvmsh_uUAOKowLmjvxaE7wrAmOY,701
-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=veJlJh_Q2_Fgc0IIfKPSWb6n_oem5WWGQUGDeepl924,6138
-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=48C0DL_SOu7Jdjk2QKBNPsevOpQlgsPYUHE7cScY-ZM,796
-emu_sv/utils.py,sha256=-axfQ2tqw0C7I9yw-28g7lytyk373DNBjDALh4kLBrM,302
-emu_sv-2.0.3.dist-info/METADATA,sha256=BYetdN_TW7rMls-VeLMSObauvcKONqh8an0FZqouJfw,3513
-emu_sv-2.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-emu_sv-2.0.3.dist-info/RECORD,,