diff --git a/skais/ais/ais_trajectory.py b/skais/ais/ais_trajectory.py index 8ac8a11726c6ca62773e46fffcbf28399792ee8c..e22b445cb3b2b5067afdc6607ca0c2f1cc319a44 100644 --- a/skais/ais/ais_trajectory.py +++ b/skais/ais/ais_trajectory.py @@ -53,6 +53,19 @@ def apply_func_on_window(dat, func, radius, on_edge='copy'): raise ValueError +def apply_func_on_window_2(dat, func, radius): + result = np.zeros(dat.shape[0]) + for i in range(0, dat.shape[0]): + ts = dat[i][0] + lower_bound = max(0, np.argmin(dat > (ts - radius))) + upper_bound = np.argmax(dat < ts + radius) + if upper_bound == 0: + upper_bound = dat.shape[0] + data = dat[lower_bound:upper_bound] + result[i] = func(data) + return result + + def apply_time_sequence(dat, time, func): result = np.empty(dat.shape[0]) result[0] = func(dat[0], dat[1], time[0], time[1]) @@ -252,6 +265,19 @@ class AISTrajectory(AISPoints): else: self.df[new_column] = result + def apply_func_on_time_window_2(self, func, radius, column, new_column=None): + if type(column) is not list: + df_column = ['ts_sec', column] + else: + df_column = column + ['ts_sec'] + dat = self.df[df_column].to_numpy() + result = apply_func_on_window_2(dat, func, radius) + + if new_column is None: + self.df[column] = result + else: + self.df[new_column] = result + # TODO rename function/simplify def apply_func_on_time_sequence(self, func, column, new_column=None): dat = self.df[column].to_numpy()