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

emu_sv/__init__.py

@@ -36,4 +36,4 @@ __all__ = [
     "inner",
 ]
 
-__version__ = "2.0.4"
+__version__ = "2.1.1"

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,33 +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.4.dist-info → emu_sv-2.1.1.dist-info}/METADATA

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

{emu_sv-2.0.4.dist-info → emu_sv-2.1.1.dist-info}/RECORD

@@ -1,14 +1,14 @@
-emu_sv/__init__.py,sha256=Tzc6RlABZ1ZVKt1mPUes9djq0eiK6FTgPagEHuFhF6Q,702
+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=AkEtI6-SY20D0ORro3Kv8tHDRUc8gxejSiRa6d--vBE,4452
-emu_sv/sv_config.py,sha256=QRy0VbCugmY6TQZ48nD6RxPJbpu0wzN7-E1Sud7YxLQ,5106
-emu_sv/time_evolution.py,sha256=obV7DcHot0jtnEmjR1ilYiSyDcJ5rTvThRB8hFjP-2s,797
+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.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,,
+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,,