Make evaluation times uniform including final time (#330)

* Make evaluation times uniform including final time

* New strategy - repeat last sample

* Zero final sample.

* Black

* Make `_tot_duration` equal to seq duration.

* Change adapt method to full array input

* Edit simulation tests

* Simplify appending 0 and final time in eval times

* Modify tests to compare with simulation

* Black

* Remove unnecesary verification

* Remove unused import

* Black bis

* Ignore flexible eval_times instruction type.

* Update pulser/simulation/simulation.py

Assign value as np.ndarray

Co-authored-by: Henrique Silvério <henrique.silverio@tecnico.ulisboa.pt>

* Update pulser/simulation/simulation.py

Assign value as np.ndarray (bis)

Co-authored-by: Henrique Silvério <henrique.silverio@tecnico.ulisboa.pt>

* Passing all CI tests

* Changing sampler tests

* Add tests for empty and  small values of eval time

* Make test assert stronger increasing sampling rate

Co-authored-by: Henrique Silvério <henrique.silverio@tecnico.ulisboa.pt>

pulser/simulation/simulation.py

@@ -122,7 +122,9 @@ def __init__(
         self._collapse_ops: list[qutip.Qobj] = []
 
         self.sampling_times = self._adapt_to_sampling_rate(
-            np.arange(self._tot_duration, dtype=np.double) / 1000
+            # Include extra time step for final instruction from samples:
+            np.arange(self._tot_duration + 1, dtype=np.double)
+            / 1000
         )
         self.evaluation_times = evaluation_times  # type: ignore
 
@@ -319,13 +321,9 @@ def evaluation_times(self, value: Union[str, ArrayLike, float]) -> None:
         """Sets times at which the results of this simulation are returned."""
         if isinstance(value, str):
             if value == "Full":
-                self._eval_times_array = np.append(
-                    self.sampling_times, self._tot_duration / 1000
-                )
+                eval_times = np.copy(self.sampling_times)
             elif value == "Minimal":
-                self._eval_times_array = np.array(
-                    [self.sampling_times[0], self._tot_duration / 1000]
-                )
+                eval_times = np.array([])
             else:
                 raise ValueError(
                     "Wrong evaluation time label. It should "
@@ -337,43 +335,36 @@ def evaluation_times(self, value: Union[str, ArrayLike, float]) -> None:
                 raise ValueError(
                     "evaluation_times float must be between 0 " "and 1."
                 )
-            extended_times = np.append(
-                self.sampling_times, self._tot_duration / 1000
-            )
             indices = np.linspace(
                 0,
-                len(extended_times) - 1,
-                int(value * len(extended_times)),
+                len(self.sampling_times) - 1,
+                int(value * len(self.sampling_times)),
                 dtype=int,
             )
-            self._eval_times_array = extended_times[indices]
+            # Note: if `value` is very small `eval_times` is an empty list:
+            eval_times = self.sampling_times[indices]
         elif isinstance(value, (list, tuple, np.ndarray)):
-            t_max = np.max(value)
-            t_min = np.min(value)
-            if t_max > self._tot_duration / 1000:
+            if np.max(value, initial=0) > self._tot_duration / 1000:
                 raise ValueError(
                     "Provided evaluation-time list extends "
                     "further than sequence duration."
                 )
-            if t_min < 0:
+            if np.min(value, initial=0) < 0:
                 raise ValueError(
                     "Provided evaluation-time list contains "
                     "negative values."
                 )
-            # Ensure the list of times is sorted
-            eval_times = np.array(np.sort(value))
-            if t_min > 0:
-                eval_times = np.insert(eval_times, 0, 0.0)
-            if t_max < self._tot_duration / 1000:
-                eval_times = np.append(eval_times, self._tot_duration / 1000)
-            self._eval_times_array = eval_times
-            # always include initial and final times
+            eval_times = np.array(value)
         else:
             raise ValueError(
                 "Wrong evaluation time label. It should "
                 "be `Full`, `Minimal`, an array of times or a "
                 + "float between 0 and 1."
             )
+        # Ensure 0 and final time are included:
+        self._eval_times_array = np.union1d(
+            eval_times, [0.0, self._tot_duration / 1000]
+        )
         self._eval_times_instruction = value
 
     def draw(
@@ -430,10 +421,11 @@ def _extract_samples(self) -> None:
 
         def prepare_dict() -> dict[str, np.ndarray]:
             # Duration includes retargeting, delays, etc.
+            # Also adds extra time step for final instruction
             return {
-                "amp": np.zeros(self._tot_duration),
-                "det": np.zeros(self._tot_duration),
-                "phase": np.zeros(self._tot_duration),
+                "amp": np.zeros(self._tot_duration + 1),
+                "det": np.zeros(self._tot_duration + 1),
+                "phase": np.zeros(self._tot_duration + 1),
             }
 
         def write_samples(
@@ -461,6 +453,7 @@ def write_samples(
                 noise_amp = np.random.normal(1.0, 1.0e-3) * np.exp(
                     -((r / w0) ** 2)
                 )
+
             samples_dict["amp"][slot.ti : slot.tf] += (
                 _pulse.amplitude.samples * noise_amp
             )
@@ -579,8 +572,8 @@ def _adapt_to_sampling_rate(self, full_array: np.ndarray) -> np.ndarray:
         """Adapt list to correspond to sampling rate."""
         indices = np.linspace(
             0,
-            self._tot_duration - 1,
-            int(self._sampling_rate * self._tot_duration),
+            len(full_array) - 1,
+            int(self._sampling_rate * (self._tot_duration + 1)),
             dtype=int,
         )
         return cast(np.ndarray, full_array[indices])

pulser/tests/test_sequence_sampler.py

@@ -21,9 +21,18 @@
 def assert_same_samples_as_sim(seq: pulser.Sequence) -> None:
     """Check against the legacy sample extraction in the simulation module."""
     got = sample(seq)
-    want = pulser.Simulation(seq).samples
+    want = pulser.Simulation(seq).samples.copy()
 
-    assert_nested_dict_equality(got, want)
+    def truncate_samples(samples_dict):
+        for key, value in samples_dict.items():
+            if isinstance(value, dict):
+                if value:  # Dictionary is not empty
+                    samples_dict[key] = truncate_samples(value)
+            else:
+                samples_dict[key] = value[:-1]
+        return samples_dict
+
+    assert_nested_dict_equality(got, truncate_samples(want))
 
 
 def assert_nested_dict_equality(got, want: dict) -> None:
@@ -139,9 +148,9 @@ def z() -> np.ndarray:
             }
         },
     }
-    want["Local"]["ground-rydberg"]["batman"]["amp"][200:401] = a_samples
+    want["Local"]["ground-rydberg"]["batman"]["amp"][200:400] = a_samples
     want["Local"]["ground-rydberg"]["superman"]["amp"][0:200] = a_samples
-    want["Local"]["ground-rydberg"]["superman"]["amp"][200:401] = a_samples
+    want["Local"]["ground-rydberg"]["superman"]["amp"][200:400] = a_samples
 
     got = sample(seq_with_SLM)
     assert_nested_dict_equality(got, want)

pulser/tests/test_simulation.py

@@ -300,7 +300,9 @@ def test_get_hamiltonian():
         simple_sim.get_hamiltonian(-10)
     # Constant detuning, so |rr><rr| term is C_6/r^6 - 2*detuning for any time
     simple_ham = simple_sim.get_hamiltonian(143)
-    assert simple_ham[0, 0] == Chadoq2.interaction_coeff / 10**6 - 2 * detun
+    assert np.isclose(
+        simple_ham[0, 0], Chadoq2.interaction_coeff / 10**6 - 2 * detun
+    )
 
     np.random.seed(123)
     simple_sim_noise = Simulation(
@@ -455,7 +457,7 @@ def test_eval_times():
     assert sim._eval_times_instruction == "Full"
     np.testing.assert_almost_equal(
         sim._eval_times_array,
-        np.append(sim.sampling_times, sim._tot_duration / 1000),
+        sim.sampling_times,
     )
 
     sim = Simulation(seq, sampling_rate=1.0)
@@ -476,6 +478,18 @@ def test_eval_times():
         np.array([0, sim.sampling_times[-3], sim._tot_duration / 1000]),
     )
 
+    sim.evaluation_times = []
+    np.testing.assert_almost_equal(
+        sim._eval_times_array,
+        np.array([0, sim._tot_duration / 1000]),
+    )
+
+    sim.evaluation_times = 0.0001
+    np.testing.assert_almost_equal(
+        sim._eval_times_array,
+        np.array([0, sim._tot_duration / 1000]),
+    )
+
     sim = Simulation(seq, sampling_rate=1.0)
     sim.evaluation_times = [sim.sampling_times[-10], sim.sampling_times[-3]]
     np.testing.assert_almost_equal(
@@ -492,13 +506,12 @@ def test_eval_times():
 
     sim = Simulation(seq, sampling_rate=1.0)
     sim.evaluation_times = 0.4
-    extended_tlist = np.append(sim.sampling_times, sim._tot_duration / 1000)
     np.testing.assert_almost_equal(
-        extended_tlist[
+        sim.sampling_times[
             np.linspace(
                 0,
-                len(extended_tlist) - 1,
-                int(0.4 * len(extended_tlist)),
+                len(sim.sampling_times) - 1,
+                int(0.4 * len(sim.sampling_times)),
                 dtype=int,
             )
         ],
@@ -562,7 +575,7 @@ def test_dephasing():
     sim = Simulation(
         seq, sampling_rate=0.01, config=SimConfig(noise="dephasing")
     )
-    assert sim.run().sample_final_state() == Counter({"0": 482, "1": 518})
+    assert sim.run().sample_final_state() == Counter({"0": 595, "1": 405})
     assert len(sim._collapse_ops) != 0
     with pytest.warns(UserWarning, match="first-order"):
         reg = Register.from_coordinates([(0, 0), (0, 10)], prefix="q")
@@ -642,7 +655,7 @@ def test_get_xy_hamiltonian():
 
     assert np.isclose(np.linalg.norm(simple_seq.magnetic_field[0:2]), 1)
 
-    simple_sim = Simulation(simple_seq, sampling_rate=0.01)
+    simple_sim = Simulation(simple_seq, sampling_rate=0.03)
     with pytest.raises(
         ValueError, match="less than or equal to the sequence duration"
     ):