From 8c63960cb5555cd94f54a8f33e1f513d80e304af Mon Sep 17 00:00:00 2001
From: Philipp Schmidt <philipp.schmidt@xfel.eu>
Date: Thu, 26 Sep 2024 17:48:45 +0200
Subject: [PATCH] Fill pulse traces with NaN/-1 instead of throwing an
 exception when train trace is too short

---
 docs/changelog.md             |  1 +
 src/extra/components/_adq.pyx | 25 +++++++++++++++++++++----
 src/extra/components/adq.py   |  4 ++++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/docs/changelog.md b/docs/changelog.md
index a0d7a6e..c97d0d9 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -32,6 +32,7 @@ Changed:
   use by default (!221).
 - [`fit_gaussian()`][extra.utils.fit_gaussian] has a new `A_sign` parameter to
   specify the expected orientation of the peak (!222).
+- [`AdqRawChannel.pulse_data()`][extra.components.AdqRawChannel.pulse_data] no longer throws an exception if the pulse pattern refers to data beyond the acquired traces, but instead fills this with NaN or -1 depending on data type (!245).
 
 ## [2024.1.2]
 Added:
diff --git a/src/extra/components/_adq.pyx b/src/extra/components/_adq.pyx
index b1214ec..6743802 100644
--- a/src/extra/components/_adq.pyx
+++ b/src/extra/components/_adq.pyx
@@ -4,6 +4,7 @@
 from cython cimport numeric
 from libc.limits cimport INT_MAX
 from libc.string cimport memcpy
+from numpy.math cimport NAN
 
 
 def _reshape_flat_pulses(
@@ -15,12 +16,18 @@ def _reshape_flat_pulses(
     if pulse_ids.shape[0] > out.shape[0]:
         raise ValueError('pulse output buffer smaller than pulse count')
 
-    cdef int start, end, i, j = -1, \
+    cdef int start, end, i, j = -1, k, \
         cur_pid, prev_pid = INT_MAX, pid_offset = 0, \
         num_trains = traces.shape[0], trace_len = traces.shape[1]
 
     cdef size_t bytes_per_pulse = samples_per_pulse * sizeof(traces[0, 0])
 
+    # Prepare the placeholder value for float and int types, using the
+    # fact that NAN != NAN for the true float value while becoming equal
+    # after integer conversion.
+    cdef numeric placeholder_val = -1 if <numeric>NAN == <numeric>NAN \
+        else <numeric>NAN
+
     for i in range(pulse_ids.shape[0]):
         cur_pid = pulse_ids[i]
 
@@ -37,6 +44,16 @@ def _reshape_flat_pulses(
         end = start + samples_per_pulse
 
         if end > trace_len:
-            raise ValueError(f'trace too short for pulse at {start}:{end}')
-
-        memcpy(&out[i, 0], &traces[j, start], bytes_per_pulse)
+            # The end of this pulse is beyond the trace length, try to
+            # copy what remains and fill with placeholder as needed.
+
+            if start < trace_len:
+                # There is still some data to copy into this pulse.
+                memcpy(&out[i, 0], &traces[j, start],
+                       (trace_len - start) * sizeof(traces[0, 0]))
+
+            # Fill in the rest of this pulse with the placeholder value.
+            for k in range(max(trace_len - start, 0), samples_per_pulse):
+                out[i, k] = placeholder_val
+        else:
+            memcpy(&out[i, 0], &traces[j, start], bytes_per_pulse)
diff --git a/src/extra/components/adq.py b/src/extra/components/adq.py
index b7cb015..82ea7e1 100644
--- a/src/extra/components/adq.py
+++ b/src/extra/components/adq.py
@@ -1046,6 +1046,10 @@ def pulse_data(self, labelled=True, pulse_dim='pulseId', train_roi=(),
         This process depends on the `first_pulse_offset` and potentially
         `single_pulse_length` the component was initialized with.
 
+        If the pulse information refers to data beyond the acquired
+        traces, it is filled by np.nan for floating data types or
+        -1 for integer types.
+
         Args:
             labelled (bool, optional): Whether data is returned as a
                 labelled xarray (default) or unlabelled ndarray.