From 420c74773f35f876766f13f3a8a52349d57d0466 Mon Sep 17 00:00:00 2001 From: Misko Date: Sun, 7 Apr 2024 15:02:21 -0700 Subject: [PATCH] nuba optimize segmentation --- spf/dataset/spf_dataset.py | 2 + spf/sdrpluto/sdr_controller.py | 138 +++++++++++++++++++++++---------- 2 files changed, 97 insertions(+), 43 deletions(-) diff --git a/spf/dataset/spf_dataset.py b/spf/dataset/spf_dataset.py index 887e12a7..a82d0618 100644 --- a/spf/dataset/spf_dataset.py +++ b/spf/dataset/spf_dataset.py @@ -11,6 +11,7 @@ import yaml from compress_pickle import load from deepdiff import DeepDiff +from numba import njit from torch.utils.data import Dataset from spf.dataset.rover_idxs import ( # v3rx_column_names, @@ -49,6 +50,7 @@ from spf.rf import ULADetector +@njit def pi_norm(x): return ((x + np.pi) % (2 * np.pi)) - np.pi diff --git a/spf/sdrpluto/sdr_controller.py b/spf/sdrpluto/sdr_controller.py index c7a1baa8..0def84a1 100644 --- a/spf/sdrpluto/sdr_controller.py +++ b/spf/sdrpluto/sdr_controller.py @@ -9,6 +9,7 @@ import adi import matplotlib.pyplot as plt import numpy as np +from numba import jit, njit from tqdm import tqdm from spf.dataset.spf_dataset import pi_norm @@ -661,82 +662,130 @@ def setup_rxtx_and_phase_calibration( return None, None +@njit +def circular_diff_to_mean(angles, mean): + a = np.abs(mean - angles) % (2 * np.pi) + b = 2 * np.pi - a + dists = np.empty(a.shape[0]) + for i in range(a.shape[0]): + dists[i] = min(a[i], b[i]) + return dists + + # returns circular_stddev and trimmed cricular stddev +@njit def circular_stddev(v, u, trim=50.0): - diff_from_mean = np.abs( - np.vstack([v - (u - 2 * np.pi), v - u, v - (u + 2 * np.pi)]) - ).min(axis=0) - _diff_from_mean = diff_from_mean[ - diff_from_mean <= np.percentile(diff_from_mean, 100.0 - trim) - ] - stddev = np.sqrt((diff_from_mean**2).sum() / (diff_from_mean.shape[0] - 1)) + diff_from_mean = circular_diff_to_mean(angles=v, mean=u) + + diff_from_mean_squared = diff_from_mean**2 + stddev = np.sqrt(diff_from_mean_squared.sum() / (diff_from_mean.shape[0] - 1)) + + mask = diff_from_mean <= np.percentile(diff_from_mean, 100.0 - trim) + _diff_from_mean_squared = diff_from_mean_squared[mask] + trimmed_stddev = np.sqrt( - (_diff_from_mean**2).sum() / (_diff_from_mean.shape[0] - 1) + _diff_from_mean_squared.sum() / (_diff_from_mean_squared.shape[0] - 1) ) return stddev, trimmed_stddev # returns circular mean and trimmed circular mean +@njit def circular_mean(angles, trim=50.0): # n = angles.shape[0] # assert angles.ndim == 1 or np.prod(a.shape[1:]) == 1 - cm = np.arctan2(np.sin(angles).sum(), np.cos(angles).sum()) % (2 * np.pi) - dists = np.vstack([2 * np.pi - np.abs(cm - angles), np.abs(cm - angles)]).min( - axis=0 - ) - _angles = angles[dists < np.percentile(dists, 100.0 - trim)] - _cm = np.arctan2(np.sin(_angles).sum(), np.cos(_angles).sum()) % (2 * np.pi) + _sin_angles = np.sin(angles) + _cos_angles = np.cos(angles) + cm = np.arctan2(_sin_angles.sum(), _cos_angles.sum()) % (2 * np.pi) + + ##non JIT + # dists = np.vstack((2 * np.pi - np.abs(cm - angles), np.abs(cm - angles))).min( + # axis=0 + # ) + + # JIT version + dists = circular_diff_to_mean(angles=angles, mean=cm) + + mask = dists < np.percentile(dists, 100.0 - trim) + _cm = np.arctan2(_sin_angles[mask].sum(), _cos_angles[mask].sum()) % (2 * np.pi) return pi_norm(cm), pi_norm(_cm) -def chunkify_array_start_end_idxs(v, window_size, stride): - assert (v.shape[0] - window_size) % stride == 0 - steps = 1 + (v.shape[0] - window_size) // stride - for step in range(steps): - yield step * stride, step * stride + window_size +def segment_session( + z, + receiver, + session_idx, + window_size, + stride, + trim, + mean_diff_threshold, + max_stddev_threshold, +): + signal_matrix = z.receivers[receiver].signal_matrix[session_idx] + pd = get_phase_diff(signal_matrix) + return simple_segment( + pd, + window_size=window_size, + stride=stride, + trim=trim, + mean_diff_threshold=mean_diff_threshold, # + max_stddev_threshold=max_stddev_threshold, # just eyeballed this + ) +@njit def windowed_trimmed_circular_mean_and_stddev(v, window_size, stride, trim=50.0): - windows = [] - for start_idx, end_idx in chunkify_array_start_end_idxs( - v, window_size=window_size, stride=stride - ): + assert (v.shape[0] - window_size) % stride == 0 + n_steps = 1 + (v.shape[0] - window_size) // stride + # for step in range(steps): + # yield + step_stats = np.zeros((n_steps, 2), dtype=np.float64) + step_idxs = np.zeros((n_steps, 2), dtype=np.int64) + steps = np.arange(n_steps) + # start_idx, end_idx + step_idxs[:, 0] = steps * stride + step_idxs[:, 1] = step_idxs[:, 0] + window_size + for step in range(n_steps): + start_idx, end_idx = step_idxs[step] _v = v[start_idx:end_idx] - assert _v.shape[0] > 0 trimmed_cm = circular_mean(_v, trim=trim)[1] - windows.append( - { - "start_idx": start_idx, - "end_idx": end_idx, - "mean": trimmed_cm, - "stddev": circular_stddev(_v, trimmed_cm, trim=trim)[1], - } - ) - return windows + step_stats[step, 0] = trimmed_cm + step_stats[step, 1] = circular_stddev(_v, trimmed_cm, trim=trim)[1] + + return step_idxs, step_stats def simple_segment( v, window_size, stride, trim, mean_diff_threshold, max_stddev_threshold ): valid_windows = [] - for window in windowed_trimmed_circular_mean_and_stddev( + window_idxs, window_stats = windowed_trimmed_circular_mean_and_stddev( v, window_size=window_size, stride=stride, trim=trim - ): + ) + for step in range(window_idxs.shape[0]): + start_idx, end_idx = window_idxs[step] + mean, stddev = window_stats[step] # is this a valid region - if window["stddev"] < max_stddev_threshold: + if stddev < max_stddev_threshold: if ( len(valid_windows) > 0 # if this is the first window - and valid_windows[-1]["end_idx"] - >= window["start_idx"] # check for overlap - and abs(valid_windows[-1]["mean"] - window["mean"]) + and valid_windows[-1]["end_idx"] >= start_idx # check for overlap + and abs(valid_windows[-1]["mean"] - mean) <= mean_diff_threshold # if not within tolerance ): # append to previous window - valid_windows[-1]["end_idx"] = window["end_idx"] - valid_windows[-1]["mean"] = window["mean"] # recompute later + valid_windows[-1]["end_idx"] = end_idx + valid_windows[-1]["mean"] = mean # recompute later else: # add a new window - valid_windows.append(window.copy()) + valid_windows.append( + { + "start_idx": int(start_idx), + "end_idx": int(end_idx), + "mean": mean, + "stddev": stddev, + } + ) # re-compute final stats as they are off for window in valid_windows: _v = v[window["start_idx"] : window["end_idx"]] @@ -745,10 +794,13 @@ def simple_segment( return valid_windows +@njit def get_phase_diff(signal_matrix): - return pi_norm(np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) + diffs = (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])).astype(np.float64) + return pi_norm(diffs) +@njit def get_avg_phase(signal_matrix, trim=0.0): return circular_mean(get_phase_diff(signal_matrix=signal_matrix), trim=50.0)