From a505c432c823f1b325f161af592f2d4597f477e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mauro=20Mendiz=C3=A1bal?= <murogrande@yahoo.com>
Date: Wed, 22 Jan 2025 19:54:36 +0100
Subject: [PATCH] [Emu-SV]: observables - improve performance (#17)

* monkey patch: rydbergHamiltonian and custom energyVariance apply
* monkey patch: rydbergHamiltonian and custom CorrelationMatrix and Second apply
*bug fixes: gpu and monkey patch issue with emu-mps
*test for the new updates in observables
---------
Co-authored-by: MauroMendizabal <mauro.mendizabal-pico@pasqal.com>
---
 emu_sv/custom_callback_implementations.py     |  65 +++++-
 emu_sv/hamiltonian.py                         |  12 +-
 emu_sv/state_vector.py                        |   2 +-
 emu_sv/sv_backend.py                          |   8 +-
 emu_sv/sv_config.py                           |  44 ++--
 emu_sv/time_evolution.py                      |  15 +-
 .../examples/basic_algeb_notebook.ipynb       | 216 +++++++++++++++---
 pyproject.toml                                |   2 +-
 test/emu_mps/test_endtoend.py                 |  20 ++
 test/emu_sv/test_custom_callbacks.py          |  96 +++++++-
 test/emu_sv/test_end_to_end.py                |  26 ++-
 test/emu_sv/test_hamiltonian.py               |   2 +-
 test/emu_sv/test_time_evolution.py            |   2 +-
 13 files changed, 435 insertions(+), 75 deletions(-)

diff --git a/emu_sv/custom_callback_implementations.py b/emu_sv/custom_callback_implementations.py
index 6542558..958648c 100644
--- a/emu_sv/custom_callback_implementations.py
+++ b/emu_sv/custom_callback_implementations.py
@@ -1,17 +1,78 @@
 import math
 from typing import Any
 
+import torch
+
 from emu_base.base_classes.config import BackendConfig
-from emu_base.base_classes.default_callbacks import QubitDensity
+from emu_base.base_classes.default_callbacks import (
+    QubitDensity,
+    EnergyVariance,
+    SecondMomentOfEnergy,
+    CorrelationMatrix,
+)
 from emu_base.base_classes.operator import Operator
 
 from emu_sv import StateVector
+from emu_sv.hamiltonian import RydbergHamiltonian
 
 
-def custom_qubit_density(
+def qubit_density_sv_impl(
     self: QubitDensity, config: BackendConfig, t: int, state: StateVector, H: Operator
 ) -> Any:
 
     num_qubits = int(math.log2(len(state.vector)))
     state_tensor = state.vector.reshape((2,) * num_qubits)
     return [(state_tensor.select(i, 1).norm() ** 2).item() for i in range(num_qubits)]
+
+
+def correlation_matrix_sv_impl(
+    self: CorrelationMatrix,
+    config: BackendConfig,
+    t: int,
+    state: StateVector,
+    H: Operator,
+) -> Any:
+    """'Sparse' implementation of <𝜓| nᵢ nⱼ | 𝜓 >"""
+    num_qubits = int(math.log2(len(state.vector)))
+    state_tensor = state.vector.reshape((2,) * num_qubits)
+
+    correlation_matrix = []
+    for numi in range(num_qubits):
+        one_correlation = []
+        select_i = state_tensor.select(numi, 1)
+        for numj in range(num_qubits):
+            if numj < numi:
+                one_correlation.append((select_i.select(numj, 1).norm() ** 2).item())
+            elif numj > numi:  # the selected atom is deleted
+                one_correlation.append((select_i.select(numj - 1, 1).norm() ** 2).item())
+            else:
+                one_correlation.append((select_i.norm() ** 2).item())
+
+        correlation_matrix.append(one_correlation)
+    return correlation_matrix
+
+
+def energy_variance_sv_impl(
+    self: EnergyVariance,
+    config: BackendConfig,
+    t: int,
+    state: StateVector,
+    H: RydbergHamiltonian,
+) -> Any:
+    hstate = H * state.vector
+    h_squared = torch.vdot(hstate, hstate)
+    h_state = torch.vdot(state.vector, hstate)
+    return (h_squared.real - h_state.real**2).item()
+
+
+def second_momentum_sv_impl(
+    self: SecondMomentOfEnergy,
+    config: BackendConfig,
+    t: int,
+    state: StateVector,
+    H: RydbergHamiltonian,
+) -> Any:
+
+    hstate = H * state.vector
+    h_squared = torch.vdot(hstate, hstate)
+    return (h_squared.real).item()
diff --git a/emu_sv/hamiltonian.py b/emu_sv/hamiltonian.py
index c4ab70a..65b01ab 100644
--- a/emu_sv/hamiltonian.py
+++ b/emu_sv/hamiltonian.py
@@ -5,6 +5,8 @@
 
 import torch
 
+from emu_sv.state_vector import StateVector
+
 
 class RydbergHamiltonian:
     """
@@ -34,7 +36,7 @@ class RydbergHamiltonian:
             strengths between each pair of qubits.
 
     Methods:
-        __matmul__(vec): Performs matrix-vector multiplication with a vector.
+        __mul__(vec): Performs matrix-vector multiplication with a vector.
         _diag_elemts(): Constructs the diagonal elements of the Hamiltonian
             based on `deltas` and `interaction_matrix`.
         _size(): Calculates the memory size of the `RydbergHamiltonian` object in MiB.
@@ -54,7 +56,7 @@ def __init__(
         self.diag: torch.Tensor = self._create_diagonal().to(device=device)
         self.inds = torch.tensor([1, 0], device=device)  # flips the state, for 𝜎ₓ
 
-    def __matmul__(self, vec: torch.Tensor) -> torch.Tensor:
+    def __mul__(self, vec: torch.Tensor) -> torch.Tensor:
         """
         Performs a matrix-vector multiplication between the `RydbergHamiltonian` form and
         a torch vector
@@ -142,3 +144,9 @@ def _create_diagonal(self) -> torch.Tensor:
                 )  # note the j-1 since i was already removed
                 i_j_fixed += self.interaction_matrix[i, j]
         return diag
+
+    def expect(self, state: StateVector) -> float | complex:
+        assert isinstance(
+            state, StateVector
+        ), "currently, only expectation values of StateVectors are supported"
+        return torch.vdot(state.vector, self * state.vector).item()  # type: ignore [no-any-return]
diff --git a/emu_sv/state_vector.py b/emu_sv/state_vector.py
index fd0db97..5b9e8b4 100644
--- a/emu_sv/state_vector.py
+++ b/emu_sv/state_vector.py
@@ -109,7 +109,7 @@ def make(cls, num_sites: int, gpu: bool = False) -> StateVector:
         tensor([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=torch.complex128)
         """
 
-        device = "gpu" if gpu else "cpu"
+        device = "cuda" if gpu else "cpu"
         ground_state = torch.zeros(2**num_sites, dtype=dtype, device=device)
         ground_state[0] = torch.tensor(1.0, dtype=dtype, device=device)
         return cls(ground_state)
diff --git a/emu_sv/sv_backend.py b/emu_sv/sv_backend.py
index ddd9dc3..161651f 100644
--- a/emu_sv/sv_backend.py
+++ b/emu_sv/sv_backend.py
@@ -4,7 +4,7 @@
 from pulser import Sequence
 from emu_base.pulser_adapter import PulserData
 from emu_sv.time_evolution import do_time_step
-from emu_sv import StateVector, DenseOperator
+from emu_sv import StateVector
 import torch
 from time import time
 from resource import RUSAGE_SELF, getrusage
@@ -53,7 +53,7 @@ def run(self, sequence: Sequence, sv_config: BackendConfig) -> Results:
 
             start = time()
 
-            state.vector = do_time_step(
+            state.vector, H = do_time_step(
                 dt,
                 omega[step],
                 delta[step],
@@ -62,10 +62,6 @@ def run(self, sequence: Sequence, sv_config: BackendConfig) -> Results:
                 sv_config.krylov_tolerance,
             )
 
-            # TODO: Rydberg Mamiltonian should be a dense operator
-
-            H = DenseOperator  # Energy, SecondMomentun... and Variance are not implemented
-
             for callback in sv_config.callbacks:
                 callback(sv_config, (step + 1) * sv_config.dt, state, H, results)
 
diff --git a/emu_sv/sv_config.py b/emu_sv/sv_config.py
index 424c250..b7280a2 100644
--- a/emu_sv/sv_config.py
+++ b/emu_sv/sv_config.py
@@ -1,21 +1,26 @@
 from emu_base.base_classes import (
-    BitStrings,
-    StateResult,
     CorrelationMatrix,
     QubitDensity,
-    Fidelity,
+    EnergyVariance,
+    SecondMomentOfEnergy,
 )
 
+import copy
+
 
 from emu_base import BackendConfig
 from emu_sv import StateVector
 from typing import Any
 
+from emu_sv.custom_callback_implementations import (
+    qubit_density_sv_impl,
+    energy_variance_sv_impl,
+    second_momentum_sv_impl,
+    correlation_matrix_sv_impl,
+)
 
 from types import MethodType
 
-from emu_sv.custom_callback_implementations import custom_qubit_density
-
 
 class SVConfig(BackendConfig):
     """
@@ -55,26 +60,23 @@ def __init__(
     ):
         super().__init__(**kwargs)
 
-        observables = set(map(type, self.callbacks))
-
-        supported_observables = {
-            BitStrings,
-            StateResult,
-            CorrelationMatrix,
-            QubitDensity,
-            Fidelity,
-        }
-
-        unsupported_observables = observables - supported_observables
-        if unsupported_observables:
-            raise ValueError(f"{unsupported_observables} are not supported in emu-sv")
         self.initial_state = initial_state
         self.dt = dt
         self.max_krylov_dim = max_krylov_dim
         self.gpu = gpu
         self.krylov_tolerance = krylov_tolerance
 
-        for obs in self.callbacks:
+        for num, obs in enumerate(self.callbacks):  # monkey patch
+            obs_copy = copy.deepcopy(obs)
             if isinstance(obs, QubitDensity):
-                # mypy: ignoring dynamically replacing method
-                obs.apply = MethodType(custom_qubit_density, obs)  # type: ignore[method-assign]
+                obs_copy.apply = MethodType(qubit_density_sv_impl, obs)  # type: ignore[method-assign]
+                self.callbacks[num] = obs_copy
+            elif isinstance(obs, EnergyVariance):
+                obs_copy.apply = MethodType(energy_variance_sv_impl, obs)  # type: ignore[method-assign]
+                self.callbacks[num] = obs_copy
+            elif isinstance(obs, SecondMomentOfEnergy):
+                obs_copy.apply = MethodType(second_momentum_sv_impl, obs)  # type: ignore[method-assign]
+                self.callbacks[num] = obs_copy
+            elif isinstance(obs, CorrelationMatrix):
+                obs_copy.apply = MethodType(correlation_matrix_sv_impl, obs)  # type: ignore[method-assign]
+                self.callbacks[num] = obs_copy
diff --git a/emu_sv/time_evolution.py b/emu_sv/time_evolution.py
index d4e8386..fc3f91d 100644
--- a/emu_sv/time_evolution.py
+++ b/emu_sv/time_evolution.py
@@ -11,15 +11,20 @@ def do_time_step(
     full_interaction_matrix: torch.Tensor,
     state_vector: torch.Tensor,
     krylov_tolerance: float,
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, RydbergHamiltonian]:
     ham = RydbergHamiltonian(
         omegas=omega,
         deltas=delta,
         interaction_matrix=full_interaction_matrix,
         device=state_vector.device,
     )
-    op = lambda x: -1j * dt * (ham @ x)
-
-    return krylov_exp(
-        op, state_vector, norm_tolerance=krylov_tolerance, exp_tolerance=krylov_tolerance
+    op = lambda x: -1j * dt * (ham * x)
+    return (
+        krylov_exp(
+            op,
+            state_vector,
+            norm_tolerance=krylov_tolerance,
+            exp_tolerance=krylov_tolerance,
+        ),
+        ham,
     )
diff --git a/examples/emu_sv_examples/examples/basic_algeb_notebook.ipynb b/examples/emu_sv_examples/examples/basic_algeb_notebook.ipynb
index ef6baaf..a7a564c 100644
--- a/examples/emu_sv_examples/examples/basic_algeb_notebook.ipynb
+++ b/examples/emu_sv_examples/examples/basic_algeb_notebook.ipynb
@@ -107,7 +107,7 @@
     {
      "data": {
       "text/plain": [
-       "Counter({'00': 503, '11': 497})"
+       "Counter({'00': 519, '11': 481})"
       ]
      },
      "execution_count": 4,
@@ -189,7 +189,7 @@
       "        -0.7071+3.5853e-06j], dtype=torch.complex128)\n",
       "\n",
       "sampling the resulting state\n",
-      "Counter({'01': 660, '00': 174, '11': 166})\n"
+      "Counter({'01': 685, '11': 178, '00': 137})\n"
      ]
     }
    ],
@@ -215,13 +215,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Counter({'11': 516, '01': 484})\n"
+      "Counter({'01': 502, '11': 498})\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "516"
+       "498"
       ]
      },
      "execution_count": 6,
@@ -245,7 +245,7 @@
     {
      "data": {
       "text/plain": [
-       "Counter({'111': 513, '000': 487})"
+       "Counter({'000': 504, '111': 496})"
       ]
      },
      "execution_count": 7,
@@ -270,7 +270,7 @@
     {
      "data": {
       "text/plain": [
-       "Counter({'111': 515, '000': 485})"
+       "Counter({'000': 516, '111': 484})"
       ]
      },
      "execution_count": 8,
@@ -333,9 +333,9 @@
      "output_type": "stream",
      "text": [
       "\n",
-      "Custom multiplication: sparse matrix with random vector tensor([-1.9323-1.3664j,  0.2566+1.8306j, -0.3390+1.0845j, -0.2470-0.7743j,\n",
-      "         0.4235-0.8698j,  0.4089+0.2061j, -0.3048-0.3941j,  0.8510-4.2901j,\n",
-      "         0.9998-1.7926j,  3.1959+1.0241j], dtype=torch.complex128)\n"
+      "Custom multiplication: sparse matrix with random vector tensor([-0.4097-0.4712j,  0.3177-0.1494j,  0.3953+1.5042j, -0.6309-0.5509j,\n",
+      "         0.4014+0.3531j, -0.3201+0.1923j, -0.4379+1.5370j, -0.9880+0.3821j,\n",
+      "         0.5249+0.2123j, -0.4576-0.4910j], dtype=torch.complex128)\n"
      ]
     }
    ],
@@ -350,7 +350,7 @@
     "h_custom = RydbergHamiltonian(\n",
     "    omegas=omega, deltas=delta, interaction_matrix=interaction_matrix, device=device\n",
     ")\n",
-    "res_sparse = h_custom @ v\n",
+    "res_sparse = h_custom * v\n",
     "print(\"\\nCustom multiplication: sparse matrix with random vector\", res_sparse[0:10])"
    ]
   },
@@ -379,7 +379,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "id": "2792a87c",
    "metadata": {},
    "outputs": [
@@ -390,7 +390,7 @@
        "       dtype=torch.complex128)"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -421,10 +421,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "id": "78a06bcd",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  12.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j, -12.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j, -12.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,  12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [ 12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j, -12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j, -12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j]],\n",
+       "       dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "x = {\"g\"+\"r\":1.0, \"r\"+\"g\":1.0}\n",
     "z = {\"g\"+\"g\":1.0, \"r\"+\"r\":-1.0}\n",
@@ -448,7 +467,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 14,
    "id": "73712109",
    "metadata": {},
    "outputs": [],
@@ -468,10 +487,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "id": "fe7a2eb8",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+24.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+24.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.-24.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.-24.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+24.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+24.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.-24.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.-24.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]],\n",
+       "       dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "x = {\"g\"+\"r\":1.0, \"r\"+\"g\":1.0}\n",
     "y = {\"g\"+\"r\":-1.0j, \"r\"+\"g\":1.0j}\n",
@@ -494,10 +532,36 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "id": "86129cc9",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  12.+0.j,\n",
+       "           0.+0.j,   0.+24.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  12.+0.j,   0.+0.j,\n",
+       "           0.+24.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.-24.j,\n",
+       "           0.+0.j, -12.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.-24.j,   0.+0.j,\n",
+       "         -12.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,  12.+0.j,   0.+0.j,   0.+24.j,   0.+0.j,   0.+0.j,\n",
+       "           0.+0.j,   0.+0.j],\n",
+       "        [ 12.+0.j,   0.+0.j,   0.+24.j,   0.+0.j,   0.+0.j,   0.+0.j,\n",
+       "           0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.-24.j,   0.+0.j, -12.+0.j,   0.+0.j,   0.+0.j,\n",
+       "           0.+0.j,   0.+0.j],\n",
+       "        [  0.-24.j,   0.+0.j, -12.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,\n",
+       "           0.+0.j,   0.+0.j]], dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "#summing 2 operators\n",
     "oper_a+oper_b"
@@ -505,10 +569,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 17,
    "id": "c241eff2",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  60.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,  60.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j, -60.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j, -60.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,  60.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [ 60.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j,   0.+0.j, -60.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j],\n",
+       "        [  0.+0.j,   0.+0.j, -60.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j,   0.+0.j]],\n",
+       "       dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "#scalar multiplication\n",
     "5.0*oper_a"
@@ -516,10 +599,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 18,
    "id": "bc8bc349",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([ 0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j, 12.+0.j,  0.+0.j,  0.+0.j],\n",
+       "       dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "#operator applied to a StateVector\n",
     "state = StateVector.make(3) #|000>\n",
@@ -528,10 +623,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "id": "75d2c16e",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor(0.+0.j, dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "# expectation value\n",
     "expectation_000 = oper_a.expect(state)\n",
@@ -540,10 +646,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "id": "f679f77d",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([ 0.0000+0.j,  0.0000+0.j, -8.4853+0.j,  0.0000+0.j,  0.0000+0.j,  8.4853+0.j,\n",
+       "         0.0000+0.j,  0.0000+0.j], dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "# aplication of an operator to a StateVector\n",
     "\n",
@@ -560,10 +678,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "id": "b74a1316",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor(0.+0.j, dtype=torch.complex128)\n",
+      "tensor(0.+0.j, dtype=torch.complex128)\n"
+     ]
+    }
+   ],
    "source": [
     "#expectation value\n",
     "print(oper_a.expect(from_string))\n",
@@ -572,15 +699,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 22,
    "id": "bd3cb100",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[0.+0.j, 0.+0.j, 0.+288.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+288.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+288.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+288.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+288.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+288.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+288.j, 0.+0.j, 0.+0.j, 0.+0.j],\n",
+       "        [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+288.j, 0.+0.j, 0.+0.j]],\n",
+       "       dtype=torch.complex128)"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "## operator - operator multipliation\n",
     "\n",
     "oper_a @ oper_b"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e18a7fcf",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
diff --git a/pyproject.toml b/pyproject.toml
index adeefdc..94907d9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,7 +41,7 @@ exclude_lines = [
 
 [tool.flake8]
 max-complexity = 18
-max-line-length = 100
+max-line-length = 110
 exclude = """
     .git,
     .venv,
diff --git a/test/emu_mps/test_endtoend.py b/test/emu_mps/test_endtoend.py
index b43cd2e..7c31bef 100644
--- a/test/emu_mps/test_endtoend.py
+++ b/test/emu_mps/test_endtoend.py
@@ -18,6 +18,10 @@
     MPSConfig,
     StateResult,
     QubitDensity,
+    Energy,
+    EnergyVariance,
+    SecondMomentOfEnergy,
+    CorrelationMatrix,
 )
 
 import pulser.noise_model
@@ -86,6 +90,10 @@ def simulate(
             BitStrings(evaluation_times=times, num_shots=1000),
             Fidelity(evaluation_times=times, state=fidelity_state),
             QubitDensity(evaluation_times=times, basis={"r", "g"}, nqubits=nqubits),
+            Energy(evaluation_times=times),
+            EnergyVariance(evaluation_times=times),
+            SecondMomentOfEnergy(evaluation_times=times),
+            CorrelationMatrix(basis={"r", "g"}, evaluation_times=times, nqubits=nqubits),
         ],
         noise_model=noise_model,
         interaction_cutoff=interaction_cutoff,
@@ -233,6 +241,18 @@ def test_end_to_end_afm_ring():
     q_density = result["qubit_density"][final_time]
     assert approx(q_density, 1e-3) == [0.578] * 10
 
+    energy = result["energy"][final_time]
+    assert approx(energy, 1e-8) == -115.3437071169735
+
+    energy_variance = result["energy_variance"][final_time]
+    assert approx(energy_variance, 1e-6) == 45.90602999801922
+
+    second_moment_energy = result["second_moment_of_energy"][final_time]
+    assert approx(second_moment_energy, 1e-6) == 13350.07680148
+
+    correlation_matrix = result["correlation_matrix"][final_time]
+    print(correlation_matrix)
+
 
 def test_end_to_end_afm_line_with_state_preparation_errors():
     torch.manual_seed(seed)
diff --git a/test/emu_sv/test_custom_callbacks.py b/test/emu_sv/test_custom_callbacks.py
index a1287d8..6f7e824 100644
--- a/test/emu_sv/test_custom_callbacks.py
+++ b/test/emu_sv/test_custom_callbacks.py
@@ -1,13 +1,29 @@
+import torch
 from emu_sv.state_vector import StateVector
 from emu_sv.dense_operator import DenseOperator
 from emu_sv.sv_config import SVConfig
-from emu_sv.custom_callback_implementations import custom_qubit_density
-from emu_base.base_classes.default_callbacks import QubitDensity
+from emu_sv.custom_callback_implementations import (
+    qubit_density_sv_impl,
+    correlation_matrix_sv_impl,
+    energy_variance_sv_impl,
+    second_momentum_sv_impl,
+)
+from emu_base.base_classes.default_callbacks import (
+    QubitDensity,
+    CorrelationMatrix,
+    EnergyVariance,
+    SecondMomentOfEnergy,
+)
 
 from pytest import approx
 
 from unittest.mock import MagicMock
 
+from emu_sv.hamiltonian import RydbergHamiltonian
+
+# device = "cuda"
+device = "cpu"
+
 
 def test_custom_qubit_density():
     # set up for state
@@ -30,6 +46,80 @@ def test_custom_qubit_density():
 
     t = 1
 
-    qubit_density = custom_qubit_density(qubit_density_mock, config, t, state, H_mock)
+    qubit_density = qubit_density_sv_impl(qubit_density_mock, config, t, state, H_mock)
     expected = [0.5] * num_qubits
     assert qubit_density == approx(expected, abs=1e-8)
+
+
+def test_custom_correlation():
+    # set up for state
+    basis = ("r", "g")
+    num_qubits = 4
+    strings = {"rgrg": 1.0, "grgr": 1.0}
+    state = StateVector.from_state_string(
+        basis=basis, nqubits=num_qubits, strings=strings
+    )
+    config = SVConfig()
+    operator_mock = MagicMock(spec=DenseOperator)
+    H_mock = operator_mock.return_value
+    correlation_matrix_mock = MagicMock(spec=CorrelationMatrix)
+    correlation_mock = correlation_matrix_mock.return_value
+    t = 1
+    correlation = correlation_matrix_sv_impl(correlation_mock, config, t, state, H_mock)
+
+    expected = []
+    for qubiti in range(num_qubits):
+        correlation_one = []
+        for qubitj in range(num_qubits):
+            if (qubiti + qubitj) % 2 == 0:
+                correlation_one.append(0.5)
+            else:
+                correlation_one.append(0.0)
+        expected.append(correlation_one)
+
+    for i, row in enumerate(correlation):
+        for j, col in enumerate(row):
+            assert col == approx(expected[i][j], abs=1e-8)
+
+
+def test_custom_energy_and_variance_and_second():
+
+    torch.manual_seed(1337)
+    dtype = torch.float64
+
+    basis = ("r", "g")
+    num_qubits = 4
+    strings = {"rgrg": 1.0, "grgr": 1.0}
+    state = StateVector.from_state_string(
+        basis=basis, nqubits=num_qubits, strings=strings
+    )
+    config = SVConfig()
+
+    omega = torch.randn(num_qubits, dtype=dtype, device=device)
+    delta = torch.randn(num_qubits, dtype=dtype, device=device)
+    interaction_matrix = torch.randn((num_qubits, num_qubits))
+    h_rydberg = RydbergHamiltonian(
+        omegas=omega, deltas=delta, interaction_matrix=interaction_matrix, device=device
+    )
+
+    t = 1
+
+    energy_mock = MagicMock(spec=EnergyVariance)
+    energy_variance_mock = energy_mock.return_value
+
+    energy_variance = energy_variance_sv_impl(
+        energy_variance_mock, config, t, state, h_rydberg
+    )
+    expected_varaince = 3.67378968943955
+
+    assert energy_variance == approx(expected_varaince, abs=1e-8)
+
+    second_momentum_energy_mock = MagicMock(spec=SecondMomentOfEnergy)
+    second_momentum_mock = second_momentum_energy_mock.return_value
+
+    second_momentum = second_momentum_sv_impl(
+        second_momentum_mock, config, t, state, h_rydberg
+    )
+    expected_second = 4.2188228611101
+
+    assert second_momentum == approx(expected_second, abs=1e-8)
diff --git a/test/emu_sv/test_end_to_end.py b/test/emu_sv/test_end_to_end.py
index 049052b..4068004 100644
--- a/test/emu_sv/test_end_to_end.py
+++ b/test/emu_sv/test_end_to_end.py
@@ -5,7 +5,16 @@
 import emu_base.base_classes
 import emu_base.base_classes.default_callbacks
 
-from emu_base.base_classes import BitStrings, Fidelity, StateResult, QubitDensity
+from emu_base.base_classes import (
+    BitStrings,
+    Fidelity,
+    StateResult,
+    QubitDensity,
+    Energy,
+    EnergyVariance,
+    SecondMomentOfEnergy,
+    CorrelationMatrix,
+)
 
 from emu_sv.sv_config import SVConfig, StateVector
 from emu_sv.sv_backend import SVBackend
@@ -93,6 +102,10 @@ def simulate(
             BitStrings(evaluation_times=times, num_shots=1000),
             Fidelity(evaluation_times=times, state=fidelity_state),
             QubitDensity(evaluation_times=times, basis={"r", "g"}, nqubits=nqubits),
+            Energy(evaluation_times=times),
+            EnergyVariance(evaluation_times=times),
+            SecondMomentOfEnergy(evaluation_times=times),
+            CorrelationMatrix(evaluation_times=times, basis={"r", "g"}, nqubits=nqubits),
         ],
         noise_model=noise_model,
         interaction_cutoff=interaction_cutoff,
@@ -138,3 +151,14 @@ def test_end_to_end_afm_ring():
     assert torch.allclose(
         torch.tensor([0.578] * 10, dtype=torch.float64), q_density, atol=1e-3
     )
+
+    energy = result["energy"][final_time]  # (-115.34554274708604-2.1316282072803006e-14j)
+    assert approx(energy, 1e-8) == -115.34554274708604
+
+    energy_variance = result["energy_variance"][final_time]  # 45.911110563993134
+    assert approx(energy_variance, 1e-3) == 45.91111056399
+
+    second_moment_energy = result["second_moment_of_energy"][
+        final_time
+    ]  # 13350.505342183847
+    assert approx(second_moment_energy, 1e-6) == 13350.5053421
diff --git a/test/emu_sv/test_hamiltonian.py b/test/emu_sv/test_hamiltonian.py
index 4f87678..e26abf4 100644
--- a/test/emu_sv/test_hamiltonian.py
+++ b/test/emu_sv/test_hamiltonian.py
@@ -66,7 +66,7 @@ def test_dense_vs_sparse():
         omegas=omega, deltas=delta, interaction_matrix=interaction_matrix, device=device
     )
 
-    res_sparse = (h_custom @ v).reshape(-1)
+    res_sparse = (h_custom * v).reshape(-1)
 
     assert torch.allclose(res_sparse, res_dense)
 
diff --git a/test/emu_sv/test_time_evolution.py b/test/emu_sv/test_time_evolution.py
index 0a088d5..44f394e 100644
--- a/test/emu_sv/test_time_evolution.py
+++ b/test/emu_sv/test_time_evolution.py
@@ -73,7 +73,7 @@ def test_forward():
     h = sv_hamiltonian(interactions, omega, delta).to(device)
     dt = 1.0
     ed = torch.linalg.matrix_exp(-1j * dt * h) @ state
-    krylov = do_time_step(
+    krylov, _ = do_time_step(
         dt,
         omega,
         delta,