From 4527cbfecedca4cf3f9f6b55f4d448b8a6948f12 Mon Sep 17 00:00:00 2001
From: HGSilveri <henrique.silverio@tecnico.ulisboa.pt>
Date: Tue, 19 Nov 2024 13:04:58 +0100
Subject: [PATCH] Include total sequence duration in Results

---
 pulser-core/pulser/backend/observable.py          | 11 ++++-------
 pulser-core/pulser/backend/results.py             |  9 ++++++++-
 pulser-core/pulser/result.py                      |  4 +++-
 pulser-simulation/pulser_simulation/simulation.py |  4 ++--
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/pulser-core/pulser/backend/observable.py b/pulser-core/pulser/backend/observable.py
index 2f7030e3..cb264df1 100644
--- a/pulser-core/pulser/backend/observable.py
+++ b/pulser-core/pulser/backend/observable.py
@@ -37,7 +37,6 @@ def __call__(
         state: State,
         hamiltonian: Operator,
         result: Results,
-        time_tol: float,
     ) -> None:
         """Specifies a call to the callback at a specific time.
 
@@ -53,8 +52,6 @@ def __call__(
             state: The current state.
             hamiltonian: The Hamiltonian at this time.
             result: The Results object to store the result in.
-            time_tol: Tolerance below which two time values are considered
-                equal.
         """
         pass
 
@@ -79,7 +76,6 @@ def __call__(
         state: State,
         hamiltonian: Operator,
         result: Results,
-        time_tol: float,
     ) -> None:
         """Specifies a call to the observable at a specific time.
 
@@ -98,6 +94,9 @@ def __call__(
             time_tol: Tolerance below which two time values are considered
                 equal.
         """
+        time_tol = (
+            (0.5 / result.total_duration) if result.total_duration else 1e-6
+        )
         if (
             self.evaluation_times is not None
             and config.is_time_in_evaluation_times(
@@ -105,7 +104,7 @@ def __call__(
             )
         ) or config.is_evaluation_time(t, tol=time_tol):
             value_to_store = self.apply(
-                config=config, t=t, state=state, hamiltonian=hamiltonian
+                config=config, state=state, hamiltonian=hamiltonian
             )
             result._store(
                 observable_name=self.name(), time=t, value=value_to_store
@@ -129,7 +128,6 @@ def apply(
         self,
         *,
         config: EmulationConfig,
-        t: float,
         state: State,
         hamiltonian: Operator,
     ) -> Any:
@@ -137,7 +135,6 @@ def apply(
 
         Args:
             config: The config object passed to the backend.
-            t: The relative time as a float between 0 and 1.
             state: The current state.
             hamiltonian: The Hamiltonian at this time.
 
diff --git a/pulser-core/pulser/backend/results.py b/pulser-core/pulser/backend/results.py
index eb8c7690..e2812fd3 100644
--- a/pulser-core/pulser/backend/results.py
+++ b/pulser-core/pulser/backend/results.py
@@ -24,10 +24,17 @@
 
 @dataclass
 class Results:
-    """A collection of results."""
+    """A collection of results.
+
+    Args:
+        atoms_order: The order of the atoms/qudits in the results.
+        meas_basis: The measurement basis.
+        total_duration: The total duration of the sequence, in ns.
+    """
 
     atom_order: tuple[QubitId, ...]
     meas_basis: str
+    total_duration: int
     _results: dict[str, dict[float, Any]] = field(init=False)
 
     def __post_init__(self) -> None:
diff --git a/pulser-core/pulser/result.py b/pulser-core/pulser/result.py
index b2cbba6b..b078cae5 100644
--- a/pulser-core/pulser/result.py
+++ b/pulser-core/pulser/result.py
@@ -17,7 +17,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections import Counter
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import matplotlib.pyplot as plt
@@ -47,6 +47,8 @@ def __getattr__(name: str) -> Any:
 class Result(ABC, backend_results.Results):
     """Base class for storing the result of an observable at specific time."""
 
+    total_duration: int = field(default=0, init=False)
+
     @property
     def sampling_dist(self) -> dict[str, float]:
         """Sampling distribution of the measured bitstring.
diff --git a/pulser-simulation/pulser_simulation/simulation.py b/pulser-simulation/pulser_simulation/simulation.py
index d017f86a..dd9d4320 100644
--- a/pulser-simulation/pulser_simulation/simulation.py
+++ b/pulser-simulation/pulser_simulation/simulation.py
@@ -597,7 +597,7 @@ def _run_solver() -> CoherentResults:
                     self._meas_basis,
                     state,
                     self._meas_basis in self.basis_name,
-                    evaluation_time=t / self._tot_duration * 1e-3,
+                    evaluation_time=t / self._tot_duration * 1e3,
                 )
                 for state, t in zip(result.states, self._eval_times_array)
             ]
@@ -685,7 +685,7 @@ def _run_solver() -> CoherentResults:
                 tuple(self._hamiltonian._qdict),
                 self._meas_basis,
                 total_count[ind],
-                evaluation_time=t / self._tot_duration * 1e-3,
+                evaluation_time=t / self._tot_duration * 1e3,
             )
             for ind, t in enumerate(self._eval_times_array)
         ]