fix energy observables at time 0. (#218)

Author: a-quelle-pasqal

docs/emu_sv/notebooks/getting_started.ipynb

@@ -192,7 +192,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -211,7 +211,7 @@
    ],
    "source": [
     "last_elem = -1\n",
-    "results.state[last_elem].vector[:10]"
+    "results.state[last_elem].data[:10]"
    ]
   },
   {

emu_mps/mps_backend.py

@@ -71,7 +71,6 @@ def run(self) -> Results:
 
     @staticmethod
     def _run(impl: MPSBackendImpl) -> Results:
-        impl.fill_results()  # at t == 0 for pulser compatibility
         while not impl.is_finished():
             impl.progress()
 

emu_mps/mps_backend_impl.py

@@ -260,7 +260,7 @@ def init_initial_state(self, initial_state: State | None = None) -> None:
         self.state = initial_state
         self.state.orthogonalize(0)
 
-    def init_hamiltonian(self) -> None:
+    def init_noiseless_hamiltonian(self) -> None:
         """
         Must be called AFTER init_dark_qubits otherwise,
         too many factors are put in the Hamiltonian
@@ -275,7 +275,9 @@ def init_hamiltonian(self) -> None:
             num_gpus_to_use=self.resolved_num_gpus,
             dim=self.dim,
         )
+        self.update_H_no_noise()
 
+    def update_H(self) -> None:
         update_H(
             hamiltonian=self.hamiltonian,
             omega=self.omega[self.timestep_index, :],
@@ -284,6 +286,15 @@ def init_hamiltonian(self) -> None:
             noise=self.lindblad_noise,
         )
 
+    def update_H_no_noise(self) -> None:
+        update_H(
+            hamiltonian=self.hamiltonian,
+            omega=self.omega[self.timestep_index, :],
+            delta=self.delta[self.timestep_index, :],
+            phi=self.phi[self.timestep_index, :],
+            noise=torch.zeros(self.dim, self.dim, dtype=dtype),  # no noise
+        )
+
     def init_baths(self) -> None:
         self.left_baths = [
             torch.ones(1, 1, 1, dtype=dtype, device=self.state.factors[0].device)
@@ -300,7 +311,9 @@ def get_current_left_bath(self) -> torch.Tensor:
     def init(self) -> None:
         self.init_dark_qubits()
         self.init_initial_state(self.config.initial_state)
-        self.init_hamiltonian()
+        self.init_noiseless_hamiltonian()
+        self.fill_results()  # at t == 0 for pulser compatibility
+        self.update_H()
         self.init_baths()
 
     def is_finished(self) -> bool:
@@ -466,13 +479,7 @@ def timestep_complete(self) -> None:
 
         if not self.is_finished():
             self.target_time = self.target_times[self.timestep_index + 1]
-            update_H(
-                hamiltonian=self.hamiltonian,
-                omega=self.omega[self.timestep_index, :],
-                delta=self.delta[self.timestep_index, :],
-                phi=self.phi[self.timestep_index, :],
-                noise=self.lindblad_noise,
-            )
+            self.update_H()
             self.init_baths()
 
         self.statistics.data.append(time.time() - self.time)
@@ -709,19 +716,8 @@ def do_random_quantum_jump(self) -> None:
         assert math.isclose(norm_after_normalizing, 1, abs_tol=1e-10)
         self.set_jump_threshold(norm_after_normalizing**2)
 
-    def remove_noise_from_hamiltonian(self) -> None:
-        # Remove the noise from self.hamiltonian for the callbacks.
-        # Since update_H is called at the start of do_time_step this is safe.
-        update_H(
-            hamiltonian=self.hamiltonian,
-            omega=self.omega[self.timestep_index - 1, :],  # Meh
-            delta=self.delta[self.timestep_index - 1, :],
-            phi=self.phi[self.timestep_index - 1, :],
-            noise=torch.zeros(self.dim, self.dim, dtype=dtype),  # no noise
-        )
-
     def timestep_complete(self) -> None:
-        self.remove_noise_from_hamiltonian()
+        self.update_H_no_noise()
         super().timestep_complete()
 
 

emu_sv/custom_callback_implementations.py

@@ -27,9 +27,9 @@ 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=dtype, device=state.vector.device)
+    occupation = torch.zeros(nqubits, dtype=dtype, device=state.data.device)
     for i in range(nqubits):
-        state_tensor = state.vector.view(2**i, 2, -1)
+        state_tensor = state.data.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
@@ -55,8 +55,8 @@ def qubit_occupation_sv_den_mat_impl(
     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()
+    occupation = torch.zeros(nqubits, dtype=dtype, device=state.data.device)
+    diag_state_tensor = state.data.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
@@ -76,10 +76,10 @@ def correlation_matrix_sv_impl(
     TODO: extend to arbitrary two-point correlation ❬ψ|AᵢBⱼ|ψ❭
     """
     nqubits = state.n_qudits
-    correlation = torch.zeros(nqubits, nqubits, dtype=dtype, device=state.vector.device)
+    correlation = torch.zeros(nqubits, nqubits, dtype=dtype, device=state.data.device)
 
     for i in range(nqubits):
-        select_i = state.vector.view(2**i, 2, -1)
+        select_i = state.data.view(2**i, 2, -1)
         select_i = select_i[:, 1]
         correlation[i, i] = torch.linalg.vector_norm(select_i) ** 2
         for j in range(i + 1, nqubits):  # select the upper triangle
@@ -104,7 +104,7 @@ def correlation_matrix_sv_den_mat_impl(
     """
     nqubits = state.n_qudits
     correlation = torch.zeros(nqubits, nqubits, dtype=dtype)
-    state_diag_matrix = state.matrix.diagonal()
+    state_diag_matrix = state.data.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, :]
@@ -127,9 +127,9 @@ def energy_variance_sv_impl(
     """
     Custom implementation of the energy variance ❬ψ|H²|ψ❭-❬ψ|H|ψ❭² for the state vector solver.
     """
-    hstate = hamiltonian * state.vector
+    hstate = hamiltonian * state.data
     h_squared = torch.vdot(hstate, hstate).real
-    energy = torch.vdot(state.vector, hstate).real
+    energy = torch.vdot(state.data, hstate).real
     en_var: torch.Tensor = h_squared - energy**2
     return en_var.cpu()
 
@@ -145,8 +145,8 @@ def energy_variance_sv_den_mat_impl(
     Custom implementation of the energy variance tr(ρH²)-tr(ρH)² for the
     lindblad equation solver.
     """
-    h_dense_matrix = hamiltonian.h_eff(state.matrix)  # Hρ
-    gpu = state.matrix.is_cuda
+    h_dense_matrix = hamiltonian.h_eff(state.data)  # Hρ
+    gpu = state.data.is_cuda
     h_squared_dense_mat = hamiltonian.expect(
         DensityMatrix(h_dense_matrix, gpu=gpu)
     )  # tr(ρH²)
@@ -165,7 +165,7 @@ def energy_second_moment_sv_impl(
     Custom implementation of the second moment of energy ❬ψ|H²|ψ❭
     for the state vector solver.
     """
-    hstate = hamiltonian * state.vector
+    hstate = hamiltonian * state.data
     en_2_mom: torch.Tensor = torch.vdot(hstate, hstate).real
     return en_2_mom.cpu()
 
@@ -181,8 +181,8 @@ def energy_second_moment_den_mat_impl(
     Custom implementation of the second moment of energy tr(ρH²) for the
     lindblad equation solver.
     """
-    h_dense_matrix = hamiltonian.h_eff(state.matrix)  # Hρ
-    gpu = state.matrix.is_cuda
+    h_dense_matrix = hamiltonian.h_eff(state.data)  # Hρ
+    gpu = state.data.is_cuda
 
     return hamiltonian.expect(
         DensityMatrix(h_dense_matrix, gpu=gpu)

emu_sv/dense_operator.py

@@ -45,10 +45,10 @@ def __init__(
         gpu: bool = True,
     ):
         device = "cuda" if gpu and DEVICE_COUNT > 0 else "cpu"
-        self.matrix = matrix.to(dtype=dtype, device=device)
+        self.data = matrix.to(dtype=dtype, device=device)
 
     def __repr__(self) -> str:
-        return repr(self.matrix)
+        return repr(self.data)
 
     def __matmul__(self, other: Operator) -> DenseOperator:
         """
@@ -63,7 +63,7 @@ def __matmul__(self, other: Operator) -> DenseOperator:
         assert isinstance(
             other, DenseOperator
         ), "DenseOperator can only be multiplied with a DenseOperator."
-        return DenseOperator(self.matrix @ other.matrix)
+        return DenseOperator(self.data @ other.data)
 
     def __add__(self, other: Operator) -> DenseOperator:
         """
@@ -78,7 +78,7 @@ def __add__(self, other: Operator) -> DenseOperator:
         assert isinstance(
             other, DenseOperator
         ), "DenseOperator can only be added to another DenseOperator."
-        return DenseOperator(self.matrix + other.matrix)
+        return DenseOperator(self.data + other.data)
 
     def __rmul__(self, scalar: complex) -> DenseOperator:
         """
@@ -91,7 +91,7 @@ def __rmul__(self, scalar: complex) -> DenseOperator:
             A new DenseOperator scaled by the given scalar.
         """
 
-        return DenseOperator(scalar * self.matrix)
+        return DenseOperator(scalar * self.data)
 
     def apply_to(self, other: State) -> StateVector:
         """
@@ -107,7 +107,7 @@ def apply_to(self, other: State) -> StateVector:
             other, StateVector
         ), "DenseOperator can only be applied to a StateVector."
 
-        return StateVector(self.matrix @ other.vector)
+        return StateVector(self.data @ other.data)
 
     def expect(self, state: State) -> torch.Tensor:
         """
@@ -123,7 +123,7 @@ def expect(self, state: State) -> torch.Tensor:
             state, StateVector
         ), "Only expectation values of StateVectors are supported."
 
-        return torch.vdot(state.vector, self.apply_to(state).vector).cpu()
+        return torch.vdot(state.data, self.apply_to(state).data).cpu()
 
     @classmethod
     def _from_operator_repr(

emu_sv/density_matrix_state.py

@@ -48,12 +48,12 @@ def __init__(
         # NOTE: this accepts also zero matrices.
 
         device = "cuda" if gpu and DEVICE_COUNT > 0 else "cpu"
-        self.matrix = matrix.to(dtype=dtype, device=device)
+        self.data = 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])
+        nqudits = math.log2(self.data.shape[0])
         return int(nqudits)
 
     @classmethod
@@ -72,9 +72,9 @@ def __rmul__(self, scalar: complex) -> DensityMatrix:
     def _normalize(self) -> None:
         # NOTE: use this in the callbacks
         """Normalize the density matrix state"""
-        matrix_trace = torch.trace(self.matrix)
+        matrix_trace = torch.trace(self.data)
         if not torch.allclose(matrix_trace, torch.tensor(1.0, dtype=torch.float64)):
-            self.matrix = self.matrix / matrix_trace
+            self.data = self.data / matrix_trace
 
     def overlap(self, other: State) -> torch.Tensor:
         """
@@ -104,12 +104,10 @@ def overlap(self, other: State) -> torch.Tensor:
             other, DensityMatrix
         ), "Other state also needs to be a DensityMatrix"
         assert (
-            self.matrix.shape == other.matrix.shape
+            self.data.shape == other.data.shape
         ), "States do not have the same number of sites"
 
-        return torch.vdot(
-            self.matrix.flatten(), other.matrix.to(self.matrix.device).flatten()
-        )
+        return torch.vdot(self.data.flatten(), other.data.to(self.data.device).flatten())
 
     @classmethod
     def from_state_vector(cls, state: StateVector) -> DensityMatrix:
@@ -124,7 +122,7 @@ def from_state_vector(cls, state: StateVector) -> DensityMatrix:
             dtype=torch.complex128)
             bell_state = StateVector(bell_state_vec, gpu=False)
             density = DensityMatrix.from_state_vector(bell_state)
-            print(density.matrix)
+            print(density.data)
             ```
 
             Output:
@@ -137,9 +135,7 @@ def from_state_vector(cls, state: StateVector) -> DensityMatrix:
             ```
         """
 
-        return cls(
-            torch.outer(state.vector, state.vector.conj()), gpu=state.vector.is_cuda
-        )
+        return cls(torch.outer(state.data, state.data.conj()), gpu=state.data.is_cuda)
 
     @classmethod
     def _from_state_amplitudes(
@@ -170,7 +166,7 @@ def _from_state_amplitudes(
             n = 2
             dense_mat=DensityMatrix.from_state_amplitudes(eigenstates=eigenstates,
             amplitudes={"rr":1.0,"gg":1.0})
-            print(dense_mat.matrix)
+            print(dense_mat.data)
             ```
 
             Output:
@@ -224,7 +220,7 @@ def sample(
             ```
         """
 
-        probabilities = torch.abs(self.matrix.diagonal())
+        probabilities = torch.abs(self.data.diagonal())
 
         outcomes = torch.multinomial(probabilities, num_shots, replacement=True)
 

emu_sv/hamiltonian.py

@@ -150,6 +150,6 @@ def expect(self, state: StateVector) -> torch.Tensor:
         assert isinstance(
             state, StateVector
         ), "Currently, only expectation values of StateVectors are supported"
-        en = torch.vdot(state.vector, self * state.vector)
+        en = torch.vdot(state.data, self * state.data)
         assert torch.allclose(en.imag, torch.zeros_like(en.imag), atol=1e-8)
         return en.real

emu_sv/lindblad_operator.py

@@ -28,7 +28,7 @@ class RydbergLindbladian:
         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)
+        pulser_lindblads (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
@@ -52,7 +52,7 @@ def __init__(
         omegas: torch.Tensor,
         deltas: torch.Tensor,
         phis: torch.Tensor,
-        pulser_linblads: list[torch.Tensor],
+        pulser_lindblads: list[torch.Tensor],
         interaction_matrix: torch.Tensor,
         device: torch.device,
     ):
@@ -61,7 +61,7 @@ def __init__(
         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.pulser_lindblads: list[torch.Tensor] = pulser_lindblads
         self.device: torch.device = device
         self.complex = self.phis.any()
 
@@ -185,7 +185,7 @@ def __matmul__(self, density_matrix: torch.Tensor) -> torch.Tensor:
         """
 
         # compute -0.5i ∑ₖ Lₖ† Lₖ
-        sum_lindblad_local = compute_noise_from_lindbladians(self.pulser_linblads).to(
+        sum_lindblad_local = compute_noise_from_lindbladians(self.pulser_lindblads).to(
             self.device
         )
         # Heff = Hρ  -0.5i ∑ₖ Lₖ† Lₖ ρ
@@ -204,14 +204,14 @@ def __matmul__(self, density_matrix: torch.Tensor) -> torch.Tensor:
                 qubit,
             )
             for qubit in range(self.nqubits)
-            for L in self.pulser_linblads
+            for L in self.pulser_lindblads
         )
 
         return H_den_matrix + 1.0j * L_den_matrix_Ldag
 
     def expect(self, state: DensityMatrix) -> torch.Tensor:
         """Return the energy expectation value E=tr(H𝜌)"""
-        en = (self.h_eff(state.matrix)).trace()
+        en = (self.h_eff(state.data)).trace()
 
         assert torch.allclose(en.imag, torch.zeros_like(en.imag), atol=1e-8)
         return en.real