test for custom observables

emu_sv/sv_config.py

@@ -74,7 +74,7 @@ def __init__(
         self.gpu = gpu
         self.krylov_tolerance = krylov_tolerance
 
-        for num,obs in enumerate(self.callbacks):
+        for num,obs in enumerate(self.callbacks): # monkey patch 
             obs_copy = copy.deepcopy(obs)
             if isinstance(obs, QubitDensity): 
                 obs_copy.apply = MethodType(custom_qubit_density, obs)  # type: ignore[method-assign]

test/emu_sv/test_custom_callbacks.py

@@ -6,6 +6,8 @@
     custom_qubit_density,
     custom_correlation_matrix,
     custom_energy,
+    custom_energy_variance,
+    custom_second_momentum_energy
 )
 from emu_base.base_classes.default_callbacks import (
     QubitDensity,
@@ -81,10 +83,10 @@ def test_custom_correlation():
             assert col == approx(expected[i][j], abs=1e-8)
 
 
-def test_custom_energy():
+def test_custom_energy_and_variance_and_second():
 
     torch.manual_seed(1337)
-    dtype = torch.complex128
+    dtype = torch.float64
 
     basis = ("r", "g")
     num_qubits = 4
@@ -105,8 +107,30 @@ def test_custom_energy():
     energy_mock = MockEnergy.return_value
     t = 1
     energy = custom_energy(energy_mock, config, t, state, h_rydberg)
-    print(energy)
-    expected = 0.438415110
+    expected_energy = 0.73826361936
+
+    assert energy == approx(expected_energy, abs=1e-8)
+
+    MockEnergy = MagicMock(spec=EnergyVariance)
+    energy_variance_mock = MockEnergy.return_value
+
+    energy_variance = custom_energy_variance(energy_variance_mock, config, t, state, h_rydberg)
+    expected_varaince = 3.67378968943955
+    assert energy_variance == approx(expected_varaince,abs=1e-8)
+
+    MockEnergy = MagicMock(spec=SecondMomentOfEnergy)
+    second_momentum_mock = MockEnergy.return_value
+
+    second_momentum = custom_second_momentum_energy(second_momentum_mock, config, t, state, h_rydberg)
+    expected_second = 4.2188228611101
+
+    assert second_momentum == approx(expected_second, abs=1e-8)
+
+
+
+
+
+
+
 
-    assert energy == approx(expected, abs=1e-8)
 

test/emu_sv/test_end_to_end.py

@@ -162,4 +162,3 @@ def test_end_to_end_afm_ring():
     assert approx(second_moment_energy,1e-8) == 13350.5053421
 
     correlation = result["correlation_matrix"][final_time]
-    print(correlation)