fpfind 3.3.0__py3-none-any.whl

fpfind/lib/utils.py

@@ -0,0 +1,716 @@
+import argparse
+import enum
+import pathlib
+import re
+import warnings
+from dataclasses import dataclass
+from multiprocessing import Process, Queue
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import scipy.fft
+from numpy.typing import NDArray
+
+import fpfind.lib._logging as logging
+from fpfind import NP_PRECISEFLOAT
+from fpfind.lib.constants import TSRES, PeakFindingFailed
+from fpfind.lib.parse_epochs import date2epoch, epoch2int, int2epoch, read_T1, read_T2
+from fpfind.lib.parse_timestamps import read_a1_kth_timestamp
+from fpfind.lib.typing import (
+    Complex_,
+    DetectorArray,
+    Float,
+    Integer,
+    PathLike,
+    TimestampArray,
+)
+
+# Create alias to external dependency for fast histogramming
+try:
+    from S15lib.g2lib.g2lib import histogram  # noqa: F401 # type: ignore
+except ModuleNotFoundError:
+    pass
+
+logger, log = logging.get_logger("fpfind")
+
+
+def get_overlap(*arrays):
+    """Returns right-truncated arrays of largest possible common length.
+
+    Used for comparing timestamps of different length, e.g. raw timestamps
+    vs chopper-generated individual epochs.
+    """
+    overlap = min(map(len, arrays))
+    arrays = [a[:overlap] for a in arrays]
+    return overlap, arrays
+
+
+@dataclass
+class PeakStatistics:
+    signal: NDArray[np.number]
+    background: NDArray[np.number]
+
+    def has_signal(self):
+        return len(self.signal) != 0
+
+    def has_background(self):
+        return len(self.background) != 0
+
+    def has_data(self):
+        return self.has_signal() and self.has_background()
+
+    @property
+    def max(self):
+        if not self.has_signal():
+            return None
+        return np.max(self.signal)
+
+    @property
+    def mean(self):
+        if not self.has_background():
+            return None
+        return np.mean(self.background)  # type: ignore (not checked, TODO)
+
+    @property
+    def stdev(self):
+        if not self.has_background():
+            return None
+        return np.std(self.background)  # type: ignore (not checked, TODO)
+
+    @property
+    def total(self):
+        if not self.has_data():
+            return None
+        return sum(self.signal) - len(self.signal) * self.mean  # type: ignore (None accounted for)
+
+    @property
+    def significance(self):
+        if self.stdev == 0:
+            return 0
+        if not self.has_data():
+            return 0
+        return round((self.max - self.mean) / self.stdev, 3)  # type: ignore (None accounted for)
+
+    @property
+    def significance_raw(self):
+        if self.stdev == 0:
+            return None
+        full = np.hstack(self.signal, self.background)  # type: ignore (not checked)
+        return (np.max(full) - np.mean(full)) / np.std(full)
+
+    @property
+    def significance2(self):
+        if self.stdev == 0:
+            return None
+        # Estimate stdev after grouping in bins of 'len(signal)'
+        if not self.has_signal():
+            return None
+        length = (len(self.background) // len(self.signal)) * len(self.signal)
+        rebinned = np.sum(
+            self.background[:length].reshape(-1, len(self.signal)), axis=1
+        )
+        stdev = np.std(rebinned)
+        if stdev == 0:
+            return None
+        return round(self.total / stdev, 3)  # type: ignore (None accounted for)
+
+    @property
+    def g2(self):
+        if self.mean is None or self.max is None:
+            return None
+        if self.mean == 0:
+            return None
+        return self.max / self.mean
+
+
+def get_statistics(
+    hist: NDArray[np.number],
+    resolution: Optional[Float] = None,
+    center: Optional[Float] = None,
+    window: Float = 0.0,
+) -> PeakStatistics:
+    """Returns statistics of histogram, after performing cross-correlation.
+
+    Args:
+        hist: Timing histogram to analyze.
+        resolution: Resolution of the histogram.
+        center: Timing center of the peak, if known beforehand.
+        window: Desired timing window width to exclude from background mean calculation.
+    """
+    # Fallback to simple statistics, if not other arguments supplied
+    if resolution is None:
+        if center is None:
+            return PeakStatistics(hist, hist)
+        else:
+            raise ValueError("Resolution must be supplied if 'center' is supplied.")
+
+    # Guess non-negative center bin position, assuming aligned at zero
+    if center is None:
+        bin_center = np.argmax(hist)
+    else:
+        bin_center = np.abs(center) // resolution
+        if center < 0:
+            bin_center = len(hist) - bin_center
+
+    # Retrieve size of symmetrical window
+    num_windowbins_onesided = int(np.ceil(window / 2 / resolution))
+    bin_offset_left = max(0, bin_center - num_windowbins_onesided)
+    bin_offset_right = min(len(hist), bin_center + num_windowbins_onesided)
+
+    # Avoid tails of the cross-correlation by taking only half of the spectrum
+    # Use-case when same timestamp is used to obtain histogram, resulting in deadtime
+    bin_offset_left_bg = bin_offset_left // 2
+    bin_offset_right_bg = (len(hist) + bin_offset_right) // 2
+
+    # Retrieve signal
+    signal = hist[bin_offset_left : bin_offset_right + 1]
+    background = np.hstack(
+        (
+            hist[bin_offset_left_bg:bin_offset_left],
+            hist[bin_offset_right + 1 : bin_offset_right_bg + 1],
+        )
+    )
+    return PeakStatistics(signal, background)
+
+
+def generate_fft(
+    arr: TimestampArray,
+    num_bins: Integer,
+    time_res: Float,
+) -> NDArray[np.complex128]:
+    """Returns the FFT and frequency resolution for the set of timestamps.
+
+    Assumes the inputs are real-valued, i.e. the FFT output is symmetrical.
+
+    Args:
+        arr: The timestamp series.
+        num_bins: The number of bins in the time/frequency domain.
+        bin_size: The size of each timing bin, in ns.
+
+    Note:
+        This function is technically not cacheable due to the mutability of
+        np.ndarray.
+    """
+    if len(arr) == 0:
+        raise ValueError("Array is empty!")
+    bin_arr = np.bincount(np.int32((arr // time_res) % num_bins), minlength=num_bins)
+    return scipy.fft.rfft(bin_arr)  # type: ignore (dispatchable)
+
+
+def get_xcorr(
+    afft: NDArray[Complex_],
+    bfft: NDArray[Complex_],
+    filter: Optional[NDArray[np.number]] = None,
+) -> NDArray[np.float64]:
+    """Returns the cross-correlation.
+
+    Note:
+        The conjugation operation on an FFT is essentially a time-reversal
+        operation on the original time-series data.
+    """
+    fft = np.conjugate(afft) * bfft
+    if filter is not None:
+        fft = fft * filter
+    result = scipy.fft.irfft(fft)
+    return np.abs(result)  # type: ignore (dispatchable)
+
+
+def convert_histogram_fft(hist: NDArray[np.number], time_bins: NDArray[np.number]):
+    """Adjust axes to estimate position."""
+    hlen = len(hist) // 2
+    hist = np.hstack((hist[hlen:], hist[:hlen]))
+    time_bins = np.hstack((-np.flip(time_bins[1 : hlen + 1]), time_bins[:hlen]))
+    return hist, time_bins
+
+
+def get_timing_delay_fft(
+    hist: NDArray[np.number],
+    time_bins: NDArray[np.number],
+    include_negative: bool = False,
+) -> Tuple[np.signedinteger, np.signedinteger]:
+    """Returns the timing delay.
+
+    Args:
+        hist: Timing histogram
+        time_bins: Time bins
+
+    Example:
+        >>> get_timing_delay_fft([1,3,0,1], [2,4,6,8])
+        (4, -8)
+
+    """
+    # [0, 1, 2, 3]  -->  (0, -4)
+    ppos = np.argmax(hist)
+    ptime = time_bins[ppos]
+    if ppos == 0:
+        npos = 0
+    else:
+        npos = len(hist) - ppos
+    ntime = -time_bins[npos]
+    result = (ptime, ntime) if np.abs(ppos) < np.abs(npos) else (ntime, ptime)
+    return result
+
+
+def get_delay_at_index_fft(
+    time_bins: NDArray[np.number], pos: Integer
+) -> Tuple[np.number, np.number]:
+    assert pos >= 0
+    pos2 = len(time_bins) - pos if pos != 0 else 0  # corner case
+    ptime = time_bins[pos]
+    ntime = -time_bins[pos2]
+    result = (ptime, ntime) if np.abs(pos) < np.abs(pos2) else (ntime, ptime)
+    return result
+
+
+def get_top_k_delays_fft(
+    hist: NDArray[np.number], time_bins: NDArray[np.number], k: Integer
+):
+    assert k >= 1
+    if k == 1:
+        return [get_delay_at_index_fft(time_bins, np.argmax(hist))]
+
+    xs_raw = np.argpartition(hist, -k)[-k:]
+    ys_raw = hist[xs_raw]
+    sort = ys_raw.argsort()[::-1]  # descending order
+    result = [
+        (y, *get_delay_at_index_fft(time_bins, x))
+        for x, y in zip(xs_raw[sort], ys_raw[sort])
+    ]
+    return result
+
+
+def slice_timestamps(
+    ts: TimestampArray,
+    start: Union[Float, None] = None,
+    duration: Union[Float, None] = None,
+) -> TimestampArray:
+    dtype = ts.dtype
+    if start is not None:
+        ts = ts - start  # note: 'ts -= start' is in-place
+        ts = ts[ts >= 0]
+    else:
+        ts = ts - ts[0]
+    if duration is not None:
+        if len(ts) == 0:
+            warnings.warn("No data available.")
+            return np.array([], dtype=dtype)
+        if duration >= ts[-1]:
+            warnings.warn(
+                f"Desired duration ({duration * 1e-9:.3f} s) is longer "
+                f"than available data ({ts[-1] * 1e-9:.3f} s)"
+            )
+        ts = ts[ts < duration]
+    return ts
+
+
+def histogram_fft(  # noqa: PLR0913
+    alice: TimestampArray,
+    bob: TimestampArray,
+    num_bins: Integer,
+    resolution: Float = 1,
+    num_wraps: Integer = 1,
+    acq_start: Optional[Float] = None,
+    freq_corr: Float = 0.0,
+    filter: Optional[NDArray[np.number]] = None,
+    statistics: bool = False,
+    center: Optional[Float] = None,
+    window: Float = 0.0,
+):
+    """Returns the cross-correlation histogram.
+
+    Args:
+        acq_start: Starting timing, relative to first common timestamp.
+        filter: Optional filter in frequency-space.
+    """
+    if not isinstance(num_wraps, (int, np.integer)):
+        warnings.warn(
+            "Number of wraps is not an integer - "
+            "statistical significance will be lower."
+        )
+
+    duration = num_wraps * num_bins * resolution
+    first_timestamp = max(alice[0], bob[0])
+    last_timestamp = min(alice[-1], bob[-1])
+    if first_timestamp + duration > last_timestamp:
+        warnings.warn(
+            f"Desired duration of timestamps ({duration} ns) "
+            f"exceeds available data ({last_timestamp - first_timestamp} ns)."
+        )
+
+    # Normalize timestamps
+    acq_start = 0 if acq_start is None else acq_start
+    alice -= first_timestamp + acq_start
+    bob -= first_timestamp + acq_start
+    bob = bob * (1 + freq_corr)
+
+    # Generate FFT
+    afft, alen = generate_fft(alice, num_bins, resolution)  # TODO: Check if correct
+    bfft, blen = generate_fft(bob, num_bins, resolution)  # TODO: Check if correct
+    result = get_xcorr(afft, bfft, filter)
+    bins = np.arange(num_bins) * resolution
+    if statistics:
+        return (
+            result,
+            bins,
+            alen,
+            blen,
+            get_statistics(result, resolution, center, window),
+        )
+    return result, bins
+
+
+# https://stackoverflow.com/a/23941599
+class ArgparseCustomFormatter(argparse.RawDescriptionHelpFormatter):
+    RAW_INDICATOR = "rawtext|"
+
+    def _format_action_invocation(self, action):
+        if not action.option_strings:
+            _ = self._metavar_formatter(action, action.dest)(1)
+            print(action, _)
+            (metavar,) = _
+            return metavar
+        else:
+            parts = []
+            # if the Optional doesn't take a value, format is:
+            #    -s, --long
+            if action.nargs == 0:
+                parts.extend(action.option_strings)
+
+            # if the Optional takes a value, format is:
+            #    -s ARGS, --long ARGS
+            # change to
+            #    -s, --long ARGS
+            else:
+                default = action.dest.upper()
+                args_string = self._format_args(action, default)
+                for option_string in action.option_strings:
+                    # parts.append('%s %s' % (option_string, args_string))
+                    parts.append("%s" % option_string)
+                parts[-1] += " %s" % args_string
+            return ", ".join(parts)
+
+    def _split_lines(self, text, width):
+        marker = ArgparseCustomFormatter.RAW_INDICATOR
+        if text.startswith(marker):
+            return text[len(marker) :].splitlines()
+        return super()._split_lines(text, width)
+
+
+def get_first_overlapping_epoch(
+    dir1: PathLike,
+    dir2: PathLike,
+    first_epoch: Optional[Union[str, bytes]] = None,
+    return_length: bool = False,
+):
+    """Get epoch name of smallest overlapping epoch.
+
+    If 'return_length' is True, the return value is a tuple of the epoch name
+    and the number of continguous overlapping epochs starting from said epoch.
+    """
+    epochints1 = [epoch2int(fp.name) for fp in pathlib.Path(dir1).glob("*")]
+    epochints2 = [epoch2int(fp.name) for fp in pathlib.Path(dir2).glob("*")]
+    epochints = set(epochints1).intersection(epochints2)
+
+    # Exclude epochs smaller than 'first_epoch', if supplied
+    if first_epoch is not None:
+        epochint_first = epoch2int(first_epoch)
+        epochints = set([v for v in epochints if v >= epochint_first])
+
+    # Calculate number of overlapping epochs
+    if len(epochints) == 0:
+        min_epoch = None
+        num_epochs = 0
+    else:
+        min_epochint = min(epochints)
+        min_epoch = int2epoch(min_epochint)
+        num_epochs = 1
+        while (min_epochint + num_epochs) in epochints:
+            num_epochs += 1
+
+    # Return result
+    if return_length:
+        return min_epoch, num_epochs
+    return min_epoch
+
+
+def iterate_epochs(
+    epoch: Union[str, bytes], length: Optional[int] = None, step: int = 1
+):
+    """Stream incremental epoch names, starting from specified epoch.
+
+    Mainly as a convenience function. This is a generator, so it
+    should be fed into a consumer, e.g. `list(iterate_epochs(...))`.
+    If 'length' is None, this becomes an infinite stream of epochs.
+    The total number of epochs emitted does not change, if 'step' is specified.
+
+    Args:
+        epoch: Starting epoch, in hex.
+        length: Total number of epochs to emit.
+        step: Separation between consecutive epochs.
+
+    Example:
+        >>> for epoch in iterate_epochs("bbbbaaa0", length=3, step=2):
+        ...     print(epoch)
+        bbbbaaa0
+        bbbbaaa2
+        bbbbaaa4
+    """
+    epochint = epoch2int(epoch)
+    if length is not None:
+        for i in range(0, length * step, step):
+            yield int2epoch(epochint + i)
+    else:
+        while True:
+            yield int2epoch(epochint)
+            epochint += step
+
+
+def get_timestamp(
+    dirname: PathLike,
+    file_type: str,
+    first_epoch: Union[str, bytes],
+    skip_epoch: int,
+    num_of_epochs: int,
+):
+    epochdir = pathlib.Path(dirname)
+    timestamp = np.array([], dtype=NP_PRECISEFLOAT)
+    for i in range(num_of_epochs):
+        epoch_name = epochdir / int2epoch(epoch2int(first_epoch) + skip_epoch + i)
+        reader = read_T1 if file_type == "T1" else read_T2
+        timestamp = np.append(timestamp, reader(epoch_name)[0])
+    return timestamp
+
+
+def get_timestamp_pattern(
+    dirname: PathLike,
+    file_type: str,
+    first_epoch: Union[str, bytes],
+    skip_epoch: int,
+    num_of_epochs: int,
+) -> Tuple[TimestampArray, DetectorArray]:
+    epochdir = pathlib.Path(dirname)
+    timestamp = np.array([], dtype=NP_PRECISEFLOAT)
+    patterns = np.array([], dtype=np.uint32)
+    reader = read_T1 if file_type == "T1" else read_T2
+    for i in range(num_of_epochs):
+        epoch_name = epochdir / int2epoch(epoch2int(first_epoch) + skip_epoch + i)
+        ts, ps = reader(epoch_name)
+        timestamp = np.append(timestamp, ts)
+        patterns = np.append(patterns, ps)
+    return timestamp, patterns
+
+
+def normalize_timestamps(
+    *T: TimestampArray, skip: float = 0.0, preserve_relative: bool = True
+):
+    """Shifts timestamp arrays to reference zero.
+
+    Args:
+        T: List of timestamp arrays.
+        skip: Duration to skip, in seconds.
+        preserve_relative: Preserve relative durations between arrays.
+    """
+    if not preserve_relative:
+        T = tuple(slice_timestamps(ts) for ts in T)
+
+    start_time = max([ts[0] for ts in T]) + skip * 1e9  # units of ns
+    T = tuple(slice_timestamps(ts, start_time) for ts in T)
+    return T
+
+
+def parse_docstring_description(docstring: Optional[str]) -> str:
+    if docstring is None:
+        return ""
+
+    placeholder = "~~~PLACEHOLDER~~~"
+    # Remove all annotated sections, including changelog
+    d, *_ = re.split(r"\n[a-zA-Z0-9\s]+:\n", docstring)
+
+    # Replace all newlines except the first
+    d = re.sub(r"\n+", placeholder, d, count=1)
+    d = re.sub(r"\n+", " ", d)
+    d = re.sub(placeholder, "\n\n", d)
+    return d
+
+
+def timestamp2epoch(
+    filename: PathLike,
+    resolution: TSRES = TSRES.PS4,
+    legacy: bool = False,
+    full: bool = False,
+) -> str:
+    (t,), _ = read_a1_kth_timestamp(
+        filename,
+        [0],
+        legacy,
+        resolution=resolution,
+        fractional=False,
+    )
+    epoch = epoch2int(date2epoch())  # get current epoch (32-bit)
+    # epoch_header = epoch >> 17  # increments every ~20h (first 15 bits)
+
+    if full:
+        epochint = read_a1_kth_timestamp(
+            filename, [0], legacy=True, resolution=TSRES.PS125, fractional=False
+        )
+    epochint = (epoch & ~((1 << 17) - 1)) | (int(t) >> 37)  # 54 - 17
+    return int2epoch(epochint)
+
+
+def xcorr(abs: NDArray[np.number], bbs: NDArray[np.number]):
+    """Performs the main cross-correlation routine."""
+    afft: NDArray[Complex_] = scipy.fft.rfft(abs)  # type: ignore (dispatchable)
+    bfft: NDArray[Complex_] = scipy.fft.rfft(bbs)  # type: ignore (dispatchable)
+    ys = get_xcorr(afft, bfft)
+    return ys
+
+
+def xcorr_process(abs: NDArray[np.number], bbs: NDArray[np.number]):
+    """Runs 'xcorr()' in a separate process for interruptible FFT.
+
+    Note:
+        Activate this with caution, because spawning a separate process incurs a penalty
+        of pickling and piping inputs/outputs between processes. Rough estimates yield
+        20% increase in runtime, e.g. 17.6(2)s -> 20.2(2) for N = 2^27 bins.
+    """
+    q = Queue()
+    args = (q, abs, bbs)
+
+    def _xcorr(q, *args, **kwargs):
+        result = xcorr(*args, **kwargs)
+        q.put(result)
+
+    p = Process(target=_xcorr, args=args)
+    p.start()
+    try:
+        return q.get()
+    except KeyboardInterrupt:  # SIGINT
+        p.terminate()
+        raise
+
+
+def generate_fft_bins(
+    arr: TimestampArray,
+    num_bins: Integer,
+    time_res: Float,
+) -> NDArray[np.signedinteger]:
+    if len(arr) == 0:
+        raise ValueError("Array is empty!")
+    bin_arr = np.bincount(np.int32((arr // time_res) % num_bins), minlength=num_bins)
+    return bin_arr
+
+
+def histogram_fft2(
+    ats: TimestampArray,
+    bts: TimestampArray,
+    start: Optional[Float],
+    duration: Optional[Float],
+    N: Integer,
+    r: Float,
+    interruptible: bool = False,
+):
+    """Convenience function that wraps histogram routines."""
+    ats_early: TimestampArray = slice_timestamps(ats, start, duration)
+    bts_early: TimestampArray = slice_timestamps(bts, start, duration)
+    abs = generate_fft_bins(ats_early, N, r)
+    bbs = generate_fft_bins(bts_early, N, r)
+    _xcorr = xcorr_process if interruptible else xcorr
+    ys = _xcorr(abs, bbs)
+
+    _dtype = np.int32  # signed number needed for negative delays
+    if N * r > 2147483647:  # int32 max
+        _dtype = np.int64
+    xs = np.arange(N, dtype=_dtype) * r
+    return xs, ys
+
+
+def fold_histogram(
+    xs: NDArray[np.number], ys: NDArray[np.number], binning_factor: int = 2
+) -> Tuple[NDArray[np.number], NDArray[np.number]]:
+    xs = xs[::binning_factor]
+    ys = np.sum(ys.reshape(-1, binning_factor), axis=1)
+    return xs, ys
+
+
+class CoarseHistogramStrategy(enum.Enum):
+    RESOLUTION = enum.auto()
+    BINS = enum.auto()
+
+
+def histogram_fft3(
+    ats: TimestampArray,
+    bts: TimestampArray,
+    start: Optional[Float],
+    duration: Float,
+    num_bins: Integer,
+    resolution: Float,
+    max_duration: Float,
+    strategy: CoarseHistogramStrategy = CoarseHistogramStrategy.RESOLUTION,
+    interruptible: bool = False,
+):
+    """Histogram with max duration limit."""
+    ceil = lambda v: int(np.ceil(v))  # noqa: E731 (using lambda cleaner)
+    factor = ceil(resolution * num_bins / max_duration)
+    factor = 1 << ceil(np.log2(factor))  # better equal to 2^k for FFT/reshape
+    duration = min(duration, max_duration)
+    args = (ats, bts, start, duration, num_bins, resolution, interruptible)
+
+    if factor == 1:
+        xs, ys = histogram_fft2(*args)
+
+    elif strategy is CoarseHistogramStrategy.RESOLUTION:
+        resolution = resolution / factor
+        xs, ys = histogram_fft2(*args)
+        xs, ys = fold_histogram(xs, ys, factor)
+
+    elif strategy is CoarseHistogramStrategy.BINS:
+        num_bins = int(num_bins // factor)
+        xs, ys = histogram_fft2(*args)
+
+    else:
+        raise NotImplementedError
+
+    return xs, ys
+
+
+def match_dts(dt1s_early, dt1s_late, ddt_window):
+    dt1s_early = sorted(dt1s_early)
+    dt1s_late = sorted(dt1s_late)
+
+    i = j = 0  # two-pointer method
+    best = (None, (0, 0, 0), (0, 0, 0))  # min_ddt, early, late
+    while i < len(dt1s_early) and j < len(dt1s_late):
+        dt1_early = dt1s_early[i]
+        dt1_late = dt1s_late[j]
+        ddt = abs(dt1_early[0] - dt1_late[0])
+        if best[0] is None or (ddt < best[0] and ddt < ddt_window):
+            best = (ddt, dt1_early, dt1_late)
+
+        if dt1_early[0] < dt1_late[0]:
+            i += 1
+        else:
+            j += 1
+
+    if best[0] is None:
+        raise PeakFindingFailed("No coincident dt")
+
+    # Valid point found
+    # Approximate the resolvable resolution by finding the
+    # time differences found using the smallest resolution
+    def resolve(best_pt, pts) -> Tuple[int, int, float]:
+        left = best_pt[0] - best_pt[1]
+        right = best_pt[0] + best_pt[1]
+        assert len(pts) > 0
+        for pt in pts:
+            if left <= pt[0] <= right:
+                break
+        return pt  # pyright: ignore[reportPossiblyUnboundVariable]
+
+    dt1_early = resolve(best[1], dt1s_early)
+    dt1_late = resolve(best[2], dt1s_late)
+
+    # Pass control back with expected variables
+    r = min([dt1_early[1], dt1_late[1]])
+    dt1 = dt1_early[0]
+    _dt1 = dt1_late[0]
+    return dt1, _dt1, r