ENH: add spidr rollover

requirements.txt

@@ -1,12 +1,6 @@
 # List required packages in this file, one per line.
 numpy
 numba
-os
 pandas
-scipy.spatial
-concurrent.futures
-multiprocessing
-time
-tdqm
-gc
-pyCHX
+scipy
+# pyCHX

tpx3awkward/__init__.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 from . import _version
+from ._utils import ingest_from_files
 
 
 def _get_version():

tpx3awkward/_utils.py

@@ -2,16 +2,18 @@
 import numpy as np
 from pathlib import Path
 from numpy.typing import NDArray
-from typing import TypeVar, Union, Dict, Set, List
+from typing import TypeVar, Union, Dict, Set, List, Iterable
 import numba
 import pandas as pd
 from scipy.spatial import KDTree
-import concurrent.futures
 import multiprocessing
-import time
 from tqdm import tqdm
-from pyCHX.chx_packages import db, get_sid_filenames
+import warnings
 import gc
+try:
+    from pyCHX.chx_packages import db, get_sid_filenames
+except ImportError:
+    warnings.warn("Could not import pyCHX package. Proceeding without it...")
 
 IA = NDArray[np.uint64]
 UnSigned = TypeVar("UnSigned", IA, np.uint64)
@@ -27,17 +29,24 @@ def raw_as_numpy(fpath: Union[str, Path]) -> IA:
     ----------
 
     """
-    with open(fpath, "rb") as fin:
-        return np.frombuffer(fin.read(), dtype="<u8")
+
+    return np.fromfile(fpath, dtype="<u8")
 
 
 @numba.jit(nopython=True)
 def get_block(v: UnSigned, width: int, shift: int) -> UnSigned:
+
     return v >> np.uint64(shift) & np.uint64(2**width - 1)
 
 
+@numba.jit(nopython=True)
+def matches_nibble(data, nibble) -> numba.boolean:
+    return (int(data) >> 60) == nibble
+
+
 @numba.jit(nopython=True)
 def is_packet_header(v: UnSigned) -> UnSigned:
+    """Identify packet headers by magic number (TPX3 as ascii on lowest 8 bytes]"""
     return get_block(v, 32, 0) == 861425748
 
 
@@ -55,7 +64,6 @@ def classify_array(data: IA) -> NDArray[np.uint8]:
     6: frame driven data (id'd via 0xA upper nibble) (??)
     """
     output = np.zeros_like(data, dtype="<u1")
-    # identify packet headers by magic number (TPX3 as ascii on lowest 8 bytes]
     is_header = is_packet_header(data)
     output[is_header] = 1
     # get the highest nibble
@@ -92,10 +100,170 @@ def _shift_xy(chip, row, col):
     return out
 
 
+@numba.jit(nopython=True)
+def decode_xy(msg, chip):
+    # these names and math are adapted from c++ code
+    l_pix_addr = (msg >> np.uint(44)) & np.uint(0xFFFF)
+    # This is laid out 16ibts which are 2 interleaved 8 bit unsigned ints
+    #  CCCCCCCRRRRRRCRR
+    #  |dcol ||spix|^||
+    #  | 7   || 6  |1|2
+    #
+    # The high 7 bits of the column
+    # '1111111000000000'
+    dcol = (l_pix_addr & np.uint(0xFE00)) >> np.uint(8)
+    # The high 6 bits of the row
+    # '0000000111111000'
+    spix = (l_pix_addr & np.uint(0x01F8)) >> np.uint(1)
+    rowcol = _shift_xy(
+        chip,
+        # add the low 2 bits of the row
+        # '0000000000000011'
+        spix + (l_pix_addr & np.uint(0x3)),
+        # add the low 1 bit of the column
+        # '0000000000000100'
+        dcol + ((l_pix_addr & np.uint(0x4)) >> np.uint(2)),
+    )
+    return rowcol[1], rowcol[0]
+
+
+@numba.jit(nopython=True)
+def get_spidr(msg):
+    return msg & np.uint(0xFFFF)
+
+
+@numba.jit(nopython=True)
+def decode_message(msg, chip, spidr_epoch=0, spidr_cutoff=0):
+    """Decode TPX3 packages of the second type corresponding to photon events (id'd via 0xB upper nibble)
+
+    Parameters
+    ----------
+        msg (uint64): tpx3 binary package
+        chip (uint8): chip ID, 0..3
+        spidr_epoch: SPIDR epoch at the start of the file; the counter increments every time the SPIDR clock resets
+        spidr_cutoff: if a file spans two SPIDR epochs, this parameter defines a midpoint between the lowest SPIDR
+            id in the previous epoch and the highest SPIDR id in the following epoch; i.e. all packages with
+            SPIDR > spidr_cutoff are guaranteed to belong to the next epoch, and vice versa.
+
+    Returns
+    ----------
+        Arrays of pixel coordinates and timestamps.
+    """
+    x, y = decode_xy(msg, chip)  # or use x1, y1 = calculateXY(msg, chip) from the Vendor's code
+    # ToA is 14 bits
+    ToA = (msg >> np.uint(30)) & np.uint(0x3FFF)
+    # ToT is 10 bits; report in ns
+    ToT = ((msg >> np.uint(20)) & np.uint(0x3FF)) * 25
+    # FToA is 4 bits
+    FToA = (msg >> np.uint(16)) & np.uint(0xF)
+    # SPIDR time is 16 bits
+    SPIDR = np.uint64(get_spidr(msg))
+    # Apply the spidr_epoch rollover
+    SPIDR += np.uint64((spidr_epoch + (SPIDR < spidr_cutoff)) * (2**16))
+    # Counter value, in multiples of 1.5625 ns
+    count = (((SPIDR << np.uint(14)) + ToA) << np.uint(4)) - FToA
+    # Phase correction
+    phase = np.uint((x / 2) % 16) or np.uint(16)
+    # Calculate the timestamp
+    ts = np.uint64(count + phase)
+
+    return x, y, ToT, ts
+
+
+@numba.jit(nopython=True)
+def check_spidr_overflow(data, last_spidr=0):
+    """Detect if the SPIDR counter has overflowed in the current data chunk (file)
+
+    Assumes that only single SPIDR overflow can occur within one file, i.e. the duration of the file is <27 sec.
+
+    Arguments:
+    ----------
+    data : NDArray[np.unint64]
+        The stream of raw data from the timepix3 detector
+    last_spidr : int
+        The last SPIDR value at the tail of the previous tpx3 file
+
+    Returns:
+    -------
+    spidr_midpoint : int
+        A value of the spidr counter such that any datapoints in the file with SPIDR < midpoint should be assigned
+        to the next SPIDR epoch (kept track of and incremented outside the function), and vice versa.
+        If spidr_midpoint=0, all data points in the file belong to the previous SPIDR epoch.
+    last_spidr : int
+        The highest SPIDR counter value from the new epoch observed at the end of file.
+    """
+
+    head_min, head_max = 65535, 0
+    tail_min, tail_max = 65535, 0
+    for i in range(64):
+        if not is_packet_header(data[i]) and matches_nibble(data[i], 0xB):
+            head_min = min(get_spidr(data[i]), head_min)
+            head_max = max(get_spidr(data[i]), head_max)
+        if not is_packet_header(data[-i - 1]) and matches_nibble(data[-i - 1], 0xB):
+            tail_min = min(get_spidr(data[-i - 1]), tail_min)
+            tail_max = max(get_spidr(data[-i - 1]), tail_max)
+    if head_min > tail_max:
+        # Transition somewhere in the middle of the file
+        midpoint = (np.uint32(head_min) + np.uint32(tail_max)) // 2
+    elif (head_max - head_min) > 32768:
+        # Large drop in the head of the file
+        midpoint = (65535 + np.uint32(tail_max)) // 2  # Assume the smallest SPIDR from the last epoch at the head is 65535
+    elif (tail_max - tail_min) > 32768:
+        # Large drop in the tail of the file; find tail_max from the new epoch
+        tail_max = 0
+        for i in range(64):
+            if not is_packet_header(data[-i - 1]) and matches_nibble(data[-i - 1], 0xB):
+                tail_spidr = get_spidr(data[-i - 1])
+                if tail_spidr < 32768:
+                    tail_max = max(tail_spidr, tail_max)
+        midpoint = (np.uint32(head_min) + np.uint32(tail_max)) // 2
+    elif last_spidr > head_min:
+        # Transition occured between the files
+        midpoint = 65535  # Everything in this file is in the new SPIDR epoch
+    else:
+        midpoint = 0
+
+    return np.uint16(midpoint), np.uint16(tail_max)
+
+
+@numba.jit(nopython=True)
+def _ingest_raw_data_rollover(data: IA, spidr_epoch=0, last_spidr=0):
+
+    chips = np.zeros_like(data, dtype=np.uint8)
+    x = np.zeros_like(data, dtype="u2")
+    y = np.zeros_like(data, dtype="u2")
+    tot = np.zeros_like(data, dtype="u4")
+    ts = np.zeros_like(data, dtype="u8")
+
+    chip_indx = 0
+    spidr_cutoff, last_spidr = check_spidr_overflow(data, last_spidr=last_spidr)
+    i = 0
+    for msg in data:
+        if is_packet_header(msg):
+            chip_indx = np.uint8(get_block(msg, 8, 32))
+        elif matches_nibble(msg, 0xB):
+            chips[i] = chip_indx
+            _x, _y, _tot, _ts = decode_message(msg, chip_indx, spidr_epoch=spidr_epoch, spidr_cutoff=spidr_cutoff)
+            x[i] = _x
+            y[i] = _y
+            tot[i] = _tot
+            ts[i] = _ts
+            i += 1
+
+    # Sort the timestamps
+    indx = np.argsort(ts[:i], kind="mergesort")
+    chips = chips[indx]
+    x, y, tot, ts = x[indx], y[indx], tot[indx], ts[indx]
+
+    if spidr_cutoff > 0:
+        spidr_epoch += 1
+
+    return x, y, tot, ts, chips, spidr_epoch, last_spidr
+
+
 @numba.jit(nopython=True)
 def _ingest_raw_data(data: IA):
     types = np.zeros_like(data, dtype="<u1")
-    # identify packet headers by magic number (TPX3 as ascii on lowest 8 bytes]
     is_header = is_packet_header(data)
     types[is_header] = 1
     # get the highest nibble
@@ -241,6 +409,32 @@ def _ingest_raw_data(data: IA):
     return x, y, pix_addr, ToA, ToT, FToA, SPIDR, chip_number, basetime, timestamp
 
 
+def ingest_from_files(fpaths: List[Union[str, Path]]) -> Iterable[Dict[str, NDArray]]:
+    """Parse values out of a sequence of timepix3 files with rollover of SPIDR counter.
+
+    Parameters
+    ----------
+    fpaths : A sorted sequence of tpx3 filepaths.
+
+    Returns
+    -------
+    An iterable over the parsing results, each encapsulated in a dictionary
+    """
+
+    spidr_epoch, last_spidr = 0, 0
+    for fpath in fpaths:
+        data = raw_as_numpy(fpath)
+        x, y, tot, ts, chips, spidr_epoch, last_spidr = _ingest_raw_data_rollover(data, spidr_epoch, last_spidr=last_spidr)
+
+        yield {
+            k.strip(): v
+            for k, v in zip(
+                "x, y, tot, timestamp, chips".split(","),
+                (x, y, tot, ts, chips),
+            )
+        }
+
+
 def ingest_raw_data(data: IA) -> Dict[str, NDArray]:
     """
     Parse values out of raw timepix3 data stream.

tpx3awkward/tests/test_toa.py

@@ -0,0 +1,67 @@
+from tpx3awkward._utils import get_block, matches_nibble, get_spidr, check_spidr_overflow
+
+import numpy as np
+import pytest
+
+
+@pytest.fixture(scope="function")
+def data(n=10):
+    data = np.zeros(n, dtype=np.uint64)
+    return data
+
+
+@pytest.fixture(scope="function")
+def header_msg(chip=3):
+    return (np.uint8(chip) << np.uint(32)) + np.uint64(861425748)
+
+
+@pytest.fixture(scope="function")
+def empty_msg():
+    return np.uint64(0xb) << np.uint(60)
+
+
+def test_get_block(header_msg):
+    assert np.uint8(get_block(header_msg, 8, 32)) == 3
+
+
+def test_matches_nibble():
+    msg = np.uint64(0xb) << np.uint(60)
+    assert matches_nibble(msg, 0xb)
+
+
+def test_get_spidr(empty_msg):
+    msg1 = empty_msg + np.uint32(32768)
+    msg2 = empty_msg + np.uint32(65535)
+    assert get_spidr(msg1) == 32768
+    assert get_spidr(msg2) == 65535
+
+def test_check_spidr_overflow(empty_msg):
+    spidr_arr = [0, 0, 0, 1, 1, 1, 2, 2, 3, 4, 5, 6, 6, 7, 8, 9, 10]
+    spidr_arr.extend(range(10, 65535, 3))
+    spidr_arr.extend([65534, 65534, 65535, 65534, 65534, 65535, 65535, 65534])
+    data = [empty_msg + np.uint32(s) for s in spidr_arr]
+    midpoint, last_spidr = check_spidr_overflow(data, 0)
+    assert midpoint == 0
+    assert last_spidr == 65535
+
+    spidr_arr = list(range(20000, 65535))
+    spidr_arr.extend(range(0, 10000))
+    data = [empty_msg + np.uint32(s) for s in spidr_arr]
+    midpoint, last_spidr = check_spidr_overflow(data, 0)
+    assert midpoint == 14999
+    assert last_spidr == 9999
+
+    spidr_arr = [65534, 65534, 65535, 65534, 65534, 65535, 0, 65535, 1, 65534, 0, 0, 0, 1, 1, 2, 3, 4, 5]
+    spidr_arr.extend(range(5, 2000))
+    data = [empty_msg + np.uint32(s) for s in spidr_arr]
+    midpoint, last_spidr = check_spidr_overflow(data, 0)
+    assert midpoint == 33767
+    assert last_spidr == 1999
+
+    spidr_arr = list(range(20000, 65535))
+    spidr_arr.extend([65534, 65534, 65535, 65534, 65534, 65535, 0, 65535, 1, 65534, 0, 0, 1, 1, 2, 3, 4, 5])
+
+    data = [empty_msg + np.uint32(s) for s in spidr_arr]
+    midpoint, last_spidr = check_spidr_overflow(data, 0)
+    assert midpoint == 10002
+    assert last_spidr == 5