emu-sv 1.0.1__py3-none-any.whl → 2.0.1__py3-none-any.whl

emu_sv/__init__.py

@@ -1,42 +1,38 @@
-from emu_sv.state_vector import StateVector, inner
-from emu_sv.dense_operator import DenseOperator
-from emu_sv.sv_backend import SVBackend, SVConfig
-from emu_base.base_classes import Results
-from emu_base.base_classes.callback import AggregationType
+from pulser.backend.results import Results
 
-from emu_base.base_classes import (
-    Callback,
+from pulser.backend import (
     BitStrings,
     CorrelationMatrix,
     Energy,
     EnergyVariance,
+    EnergySecondMoment,
     Expectation,
-    QubitDensity,
-    StateResult,
-    SecondMomentOfEnergy,
     Fidelity,
+    Occupation,
+    StateResult,
 )
 
+from .dense_operator import DenseOperator
+from .sv_backend import SVBackend, SVConfig
+from .state_vector import StateVector, inner
+
 __all__ = [
     "__version__",
-    "StateVector",
-    "DenseOperator",
-    "inner",
-    "SVBackend",
-    "SVConfig",
-    "Callback",
     "BitStrings",
     "CorrelationMatrix",
+    "DenseOperator",
     "Energy",
+    "EnergySecondMoment",
     "EnergyVariance",
     "Expectation",
     "Fidelity",
-    "QubitDensity",
-    "StateResult",
-    "SecondMomentOfEnergy",
-    "AggregationType",
+    "Occupation",
     "Results",
+    "SVBackend",
+    "SVConfig",
+    "StateResult",
+    "StateVector",
+    "inner",
 ]
 
-
-__version__ = "1.0.1"
+__version__ = "2.0.1"

emu_sv/custom_callback_implementations.py

@@ -1,93 +1,115 @@
-import math
 import torch
 
-from emu_base.base_classes.config import BackendConfig
-from emu_base.base_classes.default_callbacks import (
-    QubitDensity,
-    EnergyVariance,
-    SecondMomentOfEnergy,
+from pulser.backend import (
     CorrelationMatrix,
+    EmulationConfig,
+    EnergySecondMoment,
+    EnergyVariance,
+    Occupation,
+    Energy,
 )
-from emu_base.base_classes.operator import Operator
 
-from emu_sv import StateVector
+from emu_sv.state_vector import StateVector
+from emu_sv.dense_operator import DenseOperator
 from emu_sv.hamiltonian import RydbergHamiltonian
 
 
-def qubit_density_sv_impl(
-    self: QubitDensity, config: BackendConfig, t: int, state: StateVector, H: Operator
+def qubit_occupation_sv_impl(
+    self: Occupation,
+    *,
+    config: EmulationConfig,
+    state: StateVector,
+    hamiltonian: DenseOperator,
 ) -> torch.Tensor:
     """
-    Custom implementation of the qubit density ❬ψ|nᵢ|ψ❭ for the state vector solver.
+    Custom implementation of the occupation ❬ψ|nᵢ|ψ❭ for the state vector solver.
     """
-    nqubits = int(math.log2(len(state.vector)))
-    state_tensor = state.vector.reshape((2,) * nqubits)
-
-    qubit_density = torch.zeros(nqubits, dtype=torch.float64, device=state_tensor.device)
+    nqubits = state.n_qudits
+    occupation = torch.zeros(nqubits, dtype=torch.float64, device=state.vector.device)
     for i in range(nqubits):
-        qubit_density[i] = state_tensor.select(i, 1).norm() ** 2
-    return qubit_density
+        state_tensor = state.vector.view(2**i, 2, -1)
+        # nᵢ is a projector and therefore nᵢ == nᵢnᵢ
+        # ❬ψ|nᵢ|ψ❭ == ❬ψ|nᵢnᵢ|ψ❭ == ❬ψ|nᵢ * nᵢ|ψ❭ == ❬ϕ|ϕ❭ == |ϕ|**2
+        occupation[i] = torch.linalg.vector_norm(state_tensor[:, 1]) ** 2
+    return occupation.cpu()
 
 
 def correlation_matrix_sv_impl(
     self: CorrelationMatrix,
-    config: BackendConfig,
-    t: int,
+    *,
+    config: EmulationConfig,
     state: StateVector,
-    H: Operator,
+    hamiltonian: DenseOperator,
 ) -> torch.Tensor:
     """
-    Custom implementation of the density-density correlation ❬ψ|nᵢnⱼ|ψ❭ for the state vector solver.
-
+    Custom implementation of the density-density correlation ❬ψ|nᵢnⱼ|ψ❭
+      for the state vector solver.
     TODO: extend to arbitrary two-point correlation ❬ψ|AᵢBⱼ|ψ❭
     """
-    nqubits = int(math.log2(len(state.vector)))
-    state_tensor = state.vector.reshape((2,) * nqubits)
-
+    nqubits = state.n_qudits
     correlation = torch.zeros(
-        nqubits, nqubits, dtype=torch.float64, device=state_tensor.device
+        nqubits, nqubits, dtype=torch.float64, device=state.vector.device
     )
 
     for i in range(nqubits):
-        select_i = state_tensor.select(i, 1)
+        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 = select_i.norm() ** 2
+                value = torch.linalg.vector_norm(select_i) ** 2
+                correlation[j, j] = value
             else:
-                value = select_i.select(j - 1, 1).norm() ** 2
+                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, j] = value
-            correlation[j, i] = value
-    return correlation
+    return correlation.cpu()
 
 
 def energy_variance_sv_impl(
     self: EnergyVariance,
-    config: BackendConfig,
-    t: int,
+    *,
+    config: EmulationConfig,
     state: StateVector,
-    H: RydbergHamiltonian,
+    hamiltonian: RydbergHamiltonian,
 ) -> torch.Tensor:
     """
     Custom implementation of the energy variance ❬ψ|H²|ψ❭-❬ψ|H|ψ❭² for the state vector solver.
     """
-    hstate = H * state.vector
+    hstate = hamiltonian * state.vector
     h_squared = torch.vdot(hstate, hstate).real
     energy = torch.vdot(state.vector, hstate).real
-    energy_variance: torch.Tensor = h_squared - energy**2
-    return energy_variance
+    en_var: torch.Tensor = h_squared - energy**2
+    return en_var.cpu()
 
 
-def second_moment_sv_impl(
-    self: SecondMomentOfEnergy,
-    config: BackendConfig,
-    t: int,
+def energy_second_moment_sv_impl(
+    self: EnergySecondMoment,
+    *,
+    config: EmulationConfig,
     state: StateVector,
-    H: RydbergHamiltonian,
+    hamiltonian: RydbergHamiltonian,
 ) -> torch.Tensor:
     """
     Custom implementation of the second moment of energy ❬ψ|H²|ψ❭
     for the state vector solver.
     """
-    hstate = H * state.vector
-    return torch.vdot(hstate, hstate).real
+    hstate = hamiltonian * state.vector
+    en_2_mom: torch.Tensor = torch.vdot(hstate, hstate).real
+    return en_2_mom.cpu()
+
+
+def energy_sv_impl(
+    self: Energy,
+    *,
+    config: EmulationConfig,
+    state: StateVector,
+    hamiltonian: RydbergHamiltonian,
+) -> torch.Tensor:
+    """
+    Custom implementation of the energy ❬ψ|H|ψ❭ for the state vector solver.
+    """
+    en: torch.Tensor = hamiltonian.expect(state)
+    return en

emu_sv/dense_operator.py

@@ -1,12 +1,22 @@
 from __future__ import annotations
+
 import itertools
-from typing import Any, Iterable
+from functools import reduce
 
 import torch
-from emu_base.base_classes.operator import FullOp, QuditOp
-from emu_base import Operator, State, DEVICE_COUNT
+
+from typing import Sequence, Type
+
+from emu_base import DEVICE_COUNT
 from emu_sv.state_vector import StateVector
 
+from pulser.backend import (
+    Operator,
+    State,
+)
+from pulser.backend.operator import FullOp, QuditOp
+from pulser.backend.state import Eigenstate
+
 dtype = torch.complex128
 
 
@@ -29,8 +39,8 @@ def _validate_operator_targets(operations: FullOp, nqubits: int) -> None:
             )
 
 
-class DenseOperator(Operator):
-    """Operators in EMU-SV are dense matrices"""
+class DenseOperator(Operator[complex, torch.Tensor, StateVector]):
+    """DenseOperator in EMU-SV use dense matrices"""
 
     def __init__(
         self,
@@ -46,154 +56,145 @@ class DenseOperator(Operator):
 
     def __matmul__(self, other: Operator) -> DenseOperator:
         """
-        Apply this operator to a other. The ordering is that
-        self is applied after other.
+        Compose two DenseOperators via matrix multiplication.
 
         Args:
-            other: the operator to compose with self
+            other: a DenseOperator instance.
 
         Returns:
-            the composed operator
+            A new DenseOperator representing the product `self @ other`.
         """
         assert isinstance(
             other, DenseOperator
-        ), "DenseOperator can only be multiplied with Operator"
-
+        ), "DenseOperator can only be multiplied with a DenseOperator."
         return DenseOperator(self.matrix @ other.matrix)
 
     def __add__(self, other: Operator) -> DenseOperator:
         """
-        Returns the sum of two matrices
+        Element-wise addition of two DenseOperators.
 
         Args:
-            other: the other operator
+            other: a DenseOperator instance.
 
         Returns:
-            the summed operator
+            A new DenseOperator representing the sum.
         """
-        assert isinstance(other, DenseOperator), "MPO can only be added to another MPO"
-
+        assert isinstance(
+            other, DenseOperator
+        ), "DenseOperator can only be added to another DenseOperator."
         return DenseOperator(self.matrix + other.matrix)
 
     def __rmul__(self, scalar: complex) -> DenseOperator:
         """
-        Multiply a DenseOperator by scalar.
+        Scalar multiplication of the DenseOperator.
 
         Args:
-            scalar: the scale factor to multiply with
+            scalar: a number to scale the operator.
 
         Returns:
-            the scaled MPO
+            A new DenseOperator scaled by the given scalar.
         """
 
-        return DenseOperator(self.matrix * scalar)
+        return DenseOperator(scalar * self.matrix)
 
-    def __mul__(self, other: State) -> StateVector:
+    def apply_to(self, other: State) -> StateVector:
         """
-        Applies this DenseOperator to the given StateVector.
+        Apply the DenseOperator to a given StateVector.
 
         Args:
-            other: the state to apply this operator to
+            other: a StateVector instance.
 
         Returns:
-            the resulting state
+            A new StateVector after applying the operator.
         """
         assert isinstance(
             other, StateVector
-        ), "DenseOperator can only be applied to another DenseOperator"
+        ), "DenseOperator can only be applied to a StateVector."
 
         return StateVector(self.matrix @ other.vector)
 
-    def expect(self, state: State) -> float | complex:
+    def expect(self, state: State) -> torch.Tensor:
         """
-        Compute the expectation value of self on the given state.
+        Compute the expectation value of the operator with respect to a state.
 
         Args:
-            state: the state with which to compute
+            state: a StateVector instance.
 
         Returns:
-            the expectation
+            The expectation value as a float or complex number.
         """
         assert isinstance(
             state, StateVector
-        ), "currently, only expectation values of StateVectors are \
-        supported"
-
-        return torch.vdot(state.vector, self.matrix @ state.vector).item()
-
-    @staticmethod
-    def from_operator_string(
-        basis: Iterable[str],
-        nqubits: int,
-        operations: FullOp,
-        operators: dict[str, QuditOp] = {},
-        /,
-        **kwargs: Any,
-    ) -> DenseOperator:
+        ), "Only expectation values of StateVectors are supported."
+
+        return torch.vdot(state.vector, self.apply_to(state).vector).cpu()
+
+    @classmethod
+    def _from_operator_repr(
+        cls: Type[DenseOperator],
+        *,
+        eigenstates: Sequence[Eigenstate],
+        n_qudits: int,
+        operations: FullOp[complex],
+    ) -> tuple[DenseOperator, FullOp[complex]]:
         """
-        See the base class
+        Construct a DenseOperator from an operator representation.
 
         Args:
-            basis: the eigenstates in the basis to use e.g. ('r', 'g')
-            nqubits: how many qubits there are in the state
-            operations: which bitstrings make up the state with what weight
-            operators: additional symbols to be used in operations
+            eigenstates: the eigenstates of the basis to use, e.g. ("r", "g") or ("0", "1").
+            n_qudits: number of qudits in the system.
+            operations: which bitstrings make up the state with what weight.
 
         Returns:
-            the operator in MPO form.
+            A DenseOperator instance corresponding to the given representation.
         """
 
-        _validate_operator_targets(operations, nqubits)
+        _validate_operator_targets(operations, n_qudits)
+        assert len(set(eigenstates)) == 2, "Only qubits are supported in EMU-SV."
 
-        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
+        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict()
 
-        basis = set(basis)
-        if basis == {"r", "g"}:
+        if set(eigenstates) == {"r", "g"}:
             # operators_with_tensors will now contain the basis for single qubit ops,
-            # and potentially user defined strings in terms of these
+            # and potentially user defined strings in terms of {r, g} or {0, 1}
             operators_with_tensors |= {
-                "gg": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128),
-                "gr": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128),
-                "rg": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128),
-                "rr": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.complex128),
-            }
-        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 |= {
-                "00": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128),
-                "01": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.complex128),
-                "10": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=torch.complex128),
-                "11": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=torch.complex128),
+                "gg": torch.tensor([[1.0, 0.0], [0.0, 0.0]], dtype=dtype),
+                "gr": torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=dtype),
+                "rg": torch.tensor([[0.0, 1.0], [0.0, 0.0]], dtype=dtype),
+                "rr": torch.tensor([[0.0, 0.0], [0.0, 1.0]], dtype=dtype),
             }
+        elif set(eigenstates) == {"0", "1"}:
+            raise NotImplementedError(
+                "{'0','1'} basis is related to XY Hamiltonian, which is not implemented"
+            )
         else:
-            raise ValueError("Unsupported basis provided")
-
-        accum_res = torch.zeros(2**nqubits, 2**nqubits, dtype=torch.complex128)
-        for coeff, tensorop in operations:
-            # this function will recurse through the operators_with_tensors,
-            # and replace any definitions in terms of strings by the computed matrix
-            def replace_operator_string(op: QuditOp | torch.Tensor) -> torch.Tensor:
-                if isinstance(op, torch.Tensor):
-                    return op
-
-                result = torch.zeros(2, 2, dtype=torch.complex128)
-                for opstr, coeff in op.items():
-                    tensor = replace_operator_string(operators_with_tensors[opstr])
+            raise ValueError("An unsupported basis of eigenstates has been provided.")
+
+        accum_res = torch.zeros(2**n_qudits, 2**n_qudits, dtype=dtype)
+        for coeff, oper_torch_with_target_qubits in operations:
+
+            def build_torch_operator_from_string(
+                oper: QuditOp | torch.Tensor,
+            ) -> torch.Tensor:
+                if isinstance(oper, torch.Tensor):
+                    return oper
+
+                result = torch.zeros((2, 2), dtype=dtype)
+                for opstr, coeff in oper.items():
+                    tensor = build_torch_operator_from_string(
+                        operators_with_tensors[opstr]
+                    )
                     operators_with_tensors[opstr] = tensor
                     result += tensor * coeff
                 return result
 
-            total_op_per_qubit = [torch.eye(2, 2, dtype=torch.complex128)] * nqubits
+            single_qubit_gates = [torch.eye(2, dtype=dtype) for _ in range(n_qudits)]
 
-            for op in tensorop:
-                factor = replace_operator_string(op[0])
-                for target_qubit in op[1]:
-                    total_op_per_qubit[target_qubit] = factor
+            for operator_torch, target_qubits in oper_torch_with_target_qubits:
+                factor = build_torch_operator_from_string(operator_torch)
+                for target_qubit in target_qubits:
+                    single_qubit_gates[target_qubit] = factor
 
-            dense_op = total_op_per_qubit[0]
-            for single_qubit_operator in total_op_per_qubit[1:]:
-                dense_op = torch.kron(dense_op, single_qubit_operator)
+            accum_res += coeff * reduce(torch.kron, single_qubit_gates)
 
-            accum_res += coeff * dense_op
-        return DenseOperator(accum_res)
+        return DenseOperator(accum_res), operations