emu-sv 1.0.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

emu_sv/density_matrix_state.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+from collections import Counter
+import math
+from typing import Mapping, TypeVar, Type, Sequence
+import torch
+from pulser.backend import State
+from emu_base import DEVICE_COUNT
+from emu_sv.state_vector import StateVector
+from emu_sv.utils import index_to_bitstring
+from pulser.backend.state import Eigenstate
+
+DensityMatrixType = TypeVar("DensityMatrixType", bound="DensityMatrix")
+
+dtype = torch.complex128
+
+
+class DensityMatrix(State[complex, torch.Tensor]):
+    """Represents a density matrix in a computational basis."""
+
+    # for the moment no need to check positivity and trace 1
+    def __init__(
+        self,
+        matrix: torch.Tensor,
+        *,
+        gpu: bool = True,
+    ):
+        # NOTE: this accepts also zero matrices.
+
+        device = "cuda" if gpu and DEVICE_COUNT > 0 else "cpu"
+        self.matrix = matrix.to(dtype=dtype, device=device)
+
+    @property
+    def n_qudits(self) -> int:
+        """The number of qudits in the state."""
+        nqudits = math.log2(self.matrix.shape[0])
+        return int(nqudits)
+
+    @classmethod
+    def make(cls, n_atoms: int, gpu: bool = True) -> DensityMatrix:
+        """Creates the density matrix of the ground state |000...0>"""
+        result = torch.zeros(2**n_atoms, 2**n_atoms, dtype=dtype)
+        result[0, 0] = 1.0
+        return cls(result, gpu=gpu)
+
+    def __add__(self, other: State) -> DensityMatrix:
+        raise NotImplementedError("Not implemented")
+
+    def __rmul__(self, scalar: complex) -> DensityMatrix:
+        raise NotImplementedError("Not implemented")
+
+    def _normalize(self) -> None:
+        # NOTE: use this in the callbacks
+        """Normalize the density matrix state"""
+        matrix_trace = torch.trace(self.matrix)
+        if not torch.allclose(matrix_trace, torch.tensor(1.0, dtype=torch.float64)):
+            self.matrix = self.matrix / matrix_trace
+
+    def overlap(self, other: State) -> torch.Tensor:
+        """
+        Compute Tr(self^† @ other). The type of other must be DensityMatrix.
+
+        Args:
+            other: the other state
+
+        Returns:
+            the inner product
+
+        Example:
+        >>> density_bell_state = (1/2* torch.tensor([[1, 0, 0, 1], [0, 0, 0, 0],
+        ... [0, 0, 0, 0], [1, 0, 0, 1]],dtype=torch.complex128))
+        >>> density_c = DensityMatrix(density_bell_state, gpu=False)
+        >>> density_c.overlap(density_c)
+        tensor(1.+0.j, dtype=torch.complex128)
+        """
+
+        assert isinstance(
+            other, DensityMatrix
+        ), "Other state also needs to be a DensityMatrix"
+        assert (
+            self.matrix.shape == other.matrix.shape
+        ), "States do not have the same number of sites"
+
+        return torch.vdot(self.matrix.flatten(), other.matrix.flatten())
+
+    @classmethod
+    def from_state_vector(cls, state: StateVector) -> DensityMatrix:
+        """Convert a state vector to a density matrix.
+        This function takes a state vector |ψ❭ and returns the corresponding
+        density matrix ρ = |ψ❭❬ψ| representing the pure state |ψ❭.
+        Example:
+           >>> from emu_sv import StateVector
+           >>> import math
+           >>> bell_state_vec = 1 / math.sqrt(2) * torch.tensor(
+           ... [1.0, 0.0, 0.0, 1.0j],dtype=torch.complex128)
+           >>> bell_state = StateVector(bell_state_vec, gpu=False)
+           >>> density = DensityMatrix.from_state_vector(bell_state)
+           >>> print(density.matrix)
+           tensor([[0.5000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.0000-0.5000j],
+                   [0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j],
+                   [0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j],
+                   [0.0000+0.5000j, 0.0000+0.0000j, 0.0000+0.0000j, 0.5000+0.0000j]],
+                  dtype=torch.complex128)
+        """
+
+        return cls(
+            torch.outer(state.vector, state.vector.conj()), gpu=state.vector.is_cuda
+        )
+
+    @classmethod
+    def _from_state_amplitudes(
+        cls: Type[DensityMatrixType],
+        *,
+        eigenstates: Sequence[Eigenstate],
+        amplitudes: Mapping[str, complex],
+    ) -> tuple[DensityMatrix, Mapping[str, complex]]:
+        """Transforms a state given by a string into a density matrix.
+
+        Construct a state from the pulser abstract representation
+        https://pulser.readthedocs.io/en/stable/conventions.html
+
+        Args:
+            basis: A tuple containing the basis states (e.g., ('r', 'g')).
+            nqubits: the number of qubits.
+            strings: A dictionary mapping state strings to complex or floats amplitudes.
+
+        Returns:
+            The resulting state.
+
+        Examples:
+            >>> eigenstates = ("r","g")
+            >>> n = 2
+            >>> dense_mat=DensityMatrix.from_state_amplitudes(eigenstates=eigenstates,
+            ... amplitudes={"rr":1.0,"gg":1.0})
+            >>> print(dense_mat.matrix)
+            tensor([[0.5000+0.j, 0.0000+0.j, 0.0000+0.j, 0.5000+0.j],
+                    [0.0000+0.j, 0.0000+0.j, 0.0000+0.j, 0.0000+0.j],
+                    [0.0000+0.j, 0.0000+0.j, 0.0000+0.j, 0.0000+0.j],
+                    [0.5000+0.j, 0.0000+0.j, 0.0000+0.j, 0.5000+0.j]],
+                   dtype=torch.complex128)
+        """
+
+        state_vector, amplitudes = StateVector._from_state_amplitudes(
+            eigenstates=eigenstates, amplitudes=amplitudes
+        )
+
+        return DensityMatrix.from_state_vector(state_vector), amplitudes
+
+    def sample(
+        self,
+        num_shots: int = 1000,
+        one_state: Eigenstate | None = None,
+        p_false_pos: float = 0.0,
+        p_false_neg: float = 0.0,
+    ) -> Counter[str]:
+        """
+        Samples bitstrings, taking into account the specified error rates.
+
+        Args:
+            num_shots: how many bitstrings to sample
+            p_false_pos: the rate at which a 0 is read as a 1
+            p_false_neg: teh rate at which a 1 is read as a 0
+
+        Returns:
+            the measured bitstrings, by count
+
+        Example:
+        >>> import math
+        >>> torch.manual_seed(1234)
+        >>> from emu_sv import StateVector
+        >>> bell_vec = 1 / math.sqrt(2) * torch.tensor(
+        ... [1.0, 0.0, 0.0, 1.0j],dtype=torch.complex128)
+        >>> bell_state_vec = StateVector(bell_vec)
+        >>> bell_density = DensityMatrix.from_state_vector(bell_state_vec)
+        >>> bell_density.sample(1000)
+         Counter({'00': 517, '11': 483})
+        """
+
+        assert p_false_neg == p_false_pos == 0.0, "Error rates must be 0.0"
+
+        probabilities = torch.abs(self.matrix.diagonal())
+
+        outcomes = torch.multinomial(probabilities, num_shots, replacement=True)
+
+        # Convert outcomes to bitstrings and count occurrences
+        counts = Counter(
+            [index_to_bitstring(self.n_qudits, outcome) for outcome in outcomes]
+        )
+
+        return counts
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

emu_sv/hamiltonian.py

@@ -71,15 +71,14 @@ class RydbergHamiltonian:
             the resulting state vector.
         """
         # (-∑ⱼΔⱼnⱼ + ∑ᵢ﹥ⱼUᵢⱼnᵢnⱼ)|ψ❭
-        diag_result = self.diag * vec
+        result = self.diag * vec
         # ∑ⱼΩⱼ/2[cos(ϕⱼ)σˣⱼ + sin(ϕⱼ)σʸⱼ]|ψ❭
-        sigma_result = self._apply_sigma_operators(vec)
-        result: torch.Tensor
-        result = diag_result + sigma_result
-
+        self._apply_sigma_operators(result, vec)
         return result
 
-    def _apply_sigma_operators_real(self, vec: torch.Tensor) -> torch.Tensor:
+    def _apply_sigma_operators_real(
+        self, result: torch.Tensor, vec: torch.Tensor
+    ) -> None:
         """
         Apply the ∑ⱼ(Ωⱼ/2)σˣⱼ operator to the input vector |ψ❭.
 
@@ -89,18 +88,16 @@ class RydbergHamiltonian:
         Returns:
             the resulting state vector.
         """
-        result = torch.zeros_like(vec)
-
         dim_to_act = 1
         for n, omega_n in enumerate(self.omegas):
             shape_n = (2**n, 2, 2 ** (self.nqubits - n - 1))
-            vec = vec.reshape(shape_n)
-            result = result.reshape(shape_n)
+            vec = vec.view(shape_n)
+            result = result.view(shape_n)
             result.index_add_(dim_to_act, self.inds, vec, alpha=omega_n)
 
-        return result.reshape(-1)
-
-    def _apply_sigma_operators_complex(self, vec: torch.Tensor) -> torch.Tensor:
+    def _apply_sigma_operators_complex(
+        self, result: torch.Tensor, vec: torch.Tensor
+    ) -> None:
         """
         Apply the ∑ⱼΩⱼ/2[cos(ϕⱼ)σˣⱼ + sin(ϕⱼ)σʸⱼ] operator to the input vector |ψ❭.
 
@@ -111,13 +108,12 @@ class RydbergHamiltonian:
             the resulting state vector.
         """
         c_omegas = self.omegas * torch.exp(1j * self.phis)
-        result = torch.zeros_like(vec)
 
         dim_to_act = 1
         for n, c_omega_n in enumerate(c_omegas):
             shape_n = (2**n, 2, 2 ** (self.nqubits - n - 1))
-            vec = vec.reshape(shape_n)
-            result = result.reshape(shape_n)
+            vec = vec.view(shape_n)
+            result = result.view(shape_n)
             result.index_add_(
                 dim_to_act, self.inds[0], vec[:, 0, :].unsqueeze(1), alpha=c_omega_n
             )
@@ -128,8 +124,6 @@ class RydbergHamiltonian:
                 alpha=c_omega_n.conj(),
             )
 
-        return result.reshape(-1)
-
     def _create_diagonal(self) -> torch.Tensor:
         """
         Return the diagonal elements of the Rydberg Hamiltonian matrix
@@ -139,19 +133,21 @@ class RydbergHamiltonian:
         diag = torch.zeros(2**self.nqubits, dtype=torch.complex128, device=self.device)
 
         for i in range(self.nqubits):
-            diag = diag.reshape(2**i, 2, -1)
+            diag = diag.view(2**i, 2, -1)
             i_fixed = diag[:, 1, :]
             i_fixed -= self.deltas[i]
             for j in range(i + 1, self.nqubits):
-                i_fixed = i_fixed.reshape(2**i, 2 ** (j - i - 1), 2, -1)
+                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.reshape(-1)
+        return diag.view(-1)
 
     def expect(self, state: StateVector) -> torch.Tensor:
         """Return the energy expectation value E=❬ψ|H|ψ❭"""
         assert isinstance(
             state, StateVector
         ), "Currently, only expectation values of StateVectors are supported"
-        return torch.vdot(state.vector, self * state.vector)
+        en = torch.vdot(state.vector, self * state.vector)
+        assert torch.allclose(en.imag, torch.zeros_like(en.imag), atol=1e-8)
+        return en.real

emu_sv/state_vector.py

@@ -1,18 +1,23 @@
 from __future__ import annotations
 
-from collections import Counter
-from typing import Any, Iterable
 import math
+from collections import Counter
+from typing import Sequence, Type, TypeVar, Mapping
 
+import torch
 
-from emu_base import State, DEVICE_COUNT
+from emu_sv.utils import index_to_bitstring
 
-import torch
+from emu_base import DEVICE_COUNT
+from pulser.backend import State
+from pulser.backend.state import Eigenstate
 
+StateVectorType = TypeVar("StateVectorType", bound="StateVector")
+# Default tensor data type
 dtype = torch.complex128
 
 
-class StateVector(State):
+class StateVector(State[complex, torch.Tensor]):
     """
     Represents a quantum state vector in a computational basis.
 
@@ -41,13 +46,26 @@ class StateVector(State):
         device = "cuda" if gpu and DEVICE_COUNT > 0 else "cpu"
         self.vector = vector.to(dtype=dtype, device=device)
 
+    @property
+    def n_qudits(self) -> int:
+        """The number of qudits in the state."""
+        nqudits = math.log2(self.vector.reshape(-1).shape[0])
+        return int(nqudits)
+
     def _normalize(self) -> None:
-        # NOTE: use this in the callbacks
-        """Checks if the input is normalized or not"""
-        norm_state = torch.linalg.vector_norm(self.vector)
+        """Normalizes the state vector to ensure it has unit norm.
 
-        if not torch.allclose(norm_state, torch.tensor(1.0, dtype=torch.float64)):
-            self.vector = self.vector / norm_state
+        If the vector norm is not 1, the method scales the vector
+        to enforce normalization.
+
+        Note:
+            This method is intended to be used in callbacks.
+        """
+
+        norm = torch.linalg.vector_norm(self.vector)
+
+        if not torch.allclose(norm, torch.ones_like(norm)):
+            self.vector = self.vector / norm
 
     @classmethod
     def zero(cls, num_sites: int, gpu: bool = True) -> StateVector:
@@ -62,7 +80,7 @@ class StateVector(State):
             The zero state
 
         Examples:
-            >>> StateVector.zero(2)
+            >>> StateVector.zero(2,gpu=False)
             tensor([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=torch.complex128)
         """
 
@@ -73,8 +91,7 @@ class StateVector(State):
     @classmethod
     def make(cls, num_sites: int, gpu: bool = True) -> StateVector:
         """
-        Returns a State vector in ground state |000..0>.
-        The vector in the output of StateVector has the shape (2,)*number of qubits
+        Returns a State vector in the ground state |00..0>.
 
         Args:
             num_sites: the number of qubits
@@ -84,17 +101,19 @@ class StateVector(State):
             The described state
 
         Examples:
-            >>> StateVector.make(2)
+            >>> StateVector.make(2,gpu=False)
             tensor([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=torch.complex128)
+
+
         """
 
         result = cls.zero(num_sites=num_sites, gpu=gpu)
         result.vector[0] = 1.0
         return result
 
-    def inner(self, other: State) -> float | complex:
+    def inner(self, other: State) -> torch.Tensor:
         """
-        Compute <self, other>. The type of other must be StateVector.
+        Compute <self|other>. The type of other must be StateVector.
 
         Args:
             other: the other state
@@ -102,17 +121,21 @@ class StateVector(State):
         Returns:
             the inner product
         """
-        assert isinstance(
-            other, StateVector
-        ), "Other state also needs to be a StateVector"
+        assert isinstance(other, StateVector), "Other state must be a StateVector"
         assert (
             self.vector.shape == other.vector.shape
-        ), "States do not have the same number of sites"
+        ), "States do not have the same shape"
 
-        return torch.vdot(self.vector, other.vector).item()
+        # by our internal convention inner and norm return to cpu
+        return torch.vdot(self.vector, other.vector).cpu()
 
     def sample(
-        self, num_shots: int = 1000, p_false_pos: float = 0.0, p_false_neg: float = 0.0
+        self,
+        *,
+        num_shots: int = 1000,
+        one_state: Eigenstate | None = None,
+        p_false_pos: float = 0.0,
+        p_false_neg: float = 0.0,
     ) -> Counter[str]:
         """
         Samples bitstrings, taking into account the specified error rates.
@@ -125,25 +148,20 @@ class StateVector(State):
         Returns:
             the measured bitstrings, by count
         """
+        assert p_false_neg == p_false_pos == 0.0, "Error rates must be 0.0"
 
         probabilities = torch.abs(self.vector) ** 2
-        probabilities /= probabilities.sum()  # multinomial does not normalize the input
 
         outcomes = torch.multinomial(probabilities, num_shots, replacement=True)
 
         # Convert outcomes to bitstrings and count occurrences
-        counts = Counter([self._index_to_bitstring(outcome) for outcome in outcomes])
+        counts = Counter(
+            [index_to_bitstring(self.n_qudits, outcome) for outcome in outcomes]
+        )
 
         # NOTE: false positives and negatives
         return counts
 
-    def _index_to_bitstring(self, index: int) -> str:
-        """
-        Convert an integer index into its corresponding bitstring representation.
-        """
-        nqubits = int(math.log2(self.vector.reshape(-1).shape[0]))
-        return format(index, f"0{nqubits}b")
-
     def __add__(self, other: State) -> StateVector:
         """Sum of two state vectors
 
@@ -153,11 +171,8 @@ class StateVector(State):
         Returns:
             The summed state
         """
-        assert isinstance(
-            other, StateVector
-        ), "Other state also needs to be a StateVector"
-        result = self.vector + other.vector
-        return StateVector(result)
+        assert isinstance(other, StateVector), "`Other` state can only be a StateVector"
+        return StateVector(self.vector + other.vector, gpu=self.vector.is_cuda)
 
     def __rmul__(self, scalar: complex) -> StateVector:
         """Scalar multiplication
@@ -168,62 +183,68 @@ class StateVector(State):
         Returns:
             The scaled state
         """
-        result = scalar * self.vector
+        return StateVector(scalar * self.vector, gpu=self.vector.is_cuda)
 
-        return StateVector(result)
-
-    def norm(self) -> float | complex:
+    def norm(self) -> torch.Tensor:
         """Returns the norm of the state
 
         Returns:
             the norm of the state
         """
-        norm: float | complex = torch.linalg.vector_norm(self.vector).item()
-        return norm
+        nrm: torch.Tensor = torch.linalg.vector_norm(self.vector).cpu()
+        return nrm
 
     def __repr__(self) -> str:
         return repr(self.vector)
 
-    @staticmethod
-    def from_state_string(
+    @classmethod
+    def _from_state_amplitudes(
+        cls: Type[StateVectorType],
         *,
-        basis: Iterable[str],
-        nqubits: int,
-        strings: dict[str, complex],
-        **kwargs: Any,
-    ) -> StateVector:
+        eigenstates: Sequence[Eigenstate],
+        amplitudes: Mapping[str, complex],
+    ) -> tuple[StateVector, Mapping[str, complex]]:
         """Transforms a state given by a string into a state vector.
 
         Construct a state from the pulser abstract representation
         https://pulser.readthedocs.io/en/stable/conventions.html
 
         Args:
-            basis: A tuple containing the basis states (e.g., ('r', 'g')).
-            nqubits: the number of qubits.
-            strings: A dictionary mapping state strings to complex or floats amplitudes.
+            eigenstates: A tuple containing the basis states (e.g., ('r', 'g')).
+            amplitudes: A dictionary mapping state strings to complex or floats amplitudes.
 
         Returns:
-            The resulting state.
+            The normalised resulting state.
 
         Examples:
             >>> basis = ("r","g")
             >>> n = 2
-            >>> st=StateVector.from_state_string(basis=basis,nqubits=n,strings={"rr":1.0,"gg":1.0})
+            >>> st=StateVector.from_state_string(basis=basis,
+                ... nqubits=n,strings={"rr":1.0,"gg":1.0},gpu=False)
+            >>> st = StateVector.from_state_amplitudes(
+            ...     eigenstates=basis,
+            ...     amplitudes={"rr": 1.0, "gg": 1.0}
+            ... )
             >>> print(st)
-            tensor([0.7071+0.j, 0.0000+0.j, 0.0000+0.j, 0.7071+0.j], dtype=torch.complex128)
+            tensor([0.7071+0.j, 0.0000+0.j, 0.0000+0.j, 0.7071+0.j],
+                   dtype=torch.complex128)
         """
 
-        basis = set(basis)
+        # nqubits = len(next(iter(amplitudes.keys())))
+        nqubits = cls._validate_amplitudes(amplitudes=amplitudes, eigenstates=eigenstates)
+        basis = set(eigenstates)
         if basis == {"r", "g"}:
             one = "r"
         elif basis == {"0", "1"}:
-            one = "1"
+            raise NotImplementedError(
+                "{'0','1'} basis is related to XY Hamiltonian, which is not implemented"
+            )
         else:
             raise ValueError("Unsupported basis provided")
 
-        accum_state = StateVector.zero(num_sites=nqubits, **kwargs)
+        accum_state = StateVector.zero(num_sites=nqubits)
 
-        for state, amplitude in strings.items():
+        for state, amplitude in amplitudes.items():
             bin_to_int = int(
                 state.replace(one, "1").replace("g", "0"), 2
             )  # "0" basis is already in "0"
@@ -231,7 +252,10 @@ class StateVector(State):
 
         accum_state._normalize()
 
-        return accum_state
+        return accum_state, amplitudes
+
+    def overlap(self, other: StateVector, /) -> torch.Tensor:
+        return self.inner(other)
 
 
 def inner(left: StateVector, right: StateVector) -> torch.Tensor:
@@ -248,21 +272,27 @@ def inner(left: StateVector, right: StateVector) -> torch.Tensor:
     Examples:
         >>> factor = math.sqrt(2.0)
         >>> basis = ("r","g")
-        >>> nqubits = 2
         >>> string_state1 = {"gg":1.0,"rr":1.0}
         >>> state1 = StateVector.from_state_string(basis=basis,
-        >>>     nqubits=nqubits,strings=string_state1)
+            ... nqubits=nqubits,strings=string_state1)
         >>> string_state2 = {"gr":1.0/factor,"rr":1.0/factor}
         >>> state2 = StateVector.from_state_string(basis=basis,
-        >>>     nqubits=nqubits,strings=string_state2)
+            ... nqubits=nqubits,strings=string_state2)
+
+        >>> state1 = StateVector.from_state_amplitudes(eigenstates=basis,
+        ...     amplitudes=string_state1)
+        >>> string_state2 = {"gr":1.0/factor,"rr":1.0/factor}
+        >>> state2 = StateVector.from_state_amplitudes(eigenstates=basis,
+        ...     amplitudes=string_state2)
         >>> inner(state1,state2).item()
-        (0.4999999999999999+0j)
+        (0.49999999144286444+0j)
     """
 
-    assert (left.vector.shape == right.vector.shape) and (
-        left.vector.dim() == 1
-    ), "Shape of a and b should be the same and both needs to be 1D tesnor"
-    return torch.inner(left.vector, right.vector)
+    assert (left.vector.shape == right.vector.shape) and (left.vector.dim() == 1), (
+        "Shape of left.vector and right.vector should be",
+        " the same and both need to be 1D tesnor",
+    )
+    return left.inner(right)
 
 
 if __name__ == "__main__":