diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py index ba1cdf24..e04e5bf1 100644 --- a/tests/test_trajectory.py +++ b/tests/test_trajectory.py @@ -232,8 +232,8 @@ def test_get_length_euclidiean(self): def test_get_direction(self): df = pd.DataFrame([ {'geometry':Point(0,0), 't':datetime(2018,1,1,12,0,0)}, - {'geometry':Point(-6,10), 't':datetime(2018,1,1,12,0,0)}, - {'geometry':Point(6,6), 't':datetime(2018,1,1,12,0,1)} + {'geometry':Point(-6,10), 't':datetime(2018,1,1,12,0,1)}, + {'geometry':Point(6,6), 't':datetime(2018,1,1,12,0,2)} ]).set_index('t') geo_df = GeoDataFrame(df, crs={'init': '31256'}) traj = Trajectory(1,geo_df) @@ -241,6 +241,21 @@ def test_get_direction(self): expected_result = 45 self.assertAlmostEqual(result, expected_result, 1) + def test_split_by_daybreak(self): + df = pd.DataFrame([ + {'geometry':Point(0,0), 't':datetime(2018,1,1,12,0,0)}, + {'geometry':Point(-6,10), 't':datetime(2018,1,1,12,1,0)}, + {'geometry':Point(6,6), 't':datetime(2018,1,3,12,0,1)}, + {'geometry':Point(6,16), 't':datetime(2018,1,3,12,5,1)} + ]).set_index('t') + geo_df = GeoDataFrame(df, crs={'init': '31256'}) + traj = Trajectory(1,geo_df) + split = traj.split('daybreak') + result = len(split) + expected_result = 2 + self.assertEqual(result, expected_result) + + if __name__ == '__main__': unittest.main() diff --git a/trajectory.py b/trajectory.py index bac0bb23..4b11f89d 100644 --- a/trajectory.py +++ b/trajectory.py @@ -206,3 +206,15 @@ def clip(self, polygon): def intersection(self, feature): return overlay.intersection(self, feature) + + def split(self, mode='daybreak'): + result = [] + if mode == 'daybreak': + dfs = [group[1] for group in self.df.groupby(self.df.index.day)] + for i, df in enumerate(dfs): + result.append(Trajectory('{}_{}'.format(self.id, i), df)) + else: + raise ValueError('Invalid split mode {}. Must be one of [daybreak]'.format(mode)) + return result + +