fixed engine tests

skais/process/data_augmentation/augmentation_engine.py

@@ -28,7 +28,7 @@ class AugmentationEngine:
                     self.pipelines.append(Pipeline([translator, flip]))
 
     def transform(self, x, verbose=0):
-        results = x.copy()
+        results = []
 
         iterator = self.pipelines
         if verbose > 0:

skais/process/data_augmentation/flip.py

@@ -14,6 +14,4 @@ class Flip(DataTransformer):
                 df = trajectory.df.copy()
                 df['latitude'] = -trajectory.df['latitude']
                 result.append(AISTrajectory(df))
-        else:
-            result += x.copy()
         return result

skais/tests/process/data_augmentation/test_engine.py

@@ -5,7 +5,8 @@ from skais.process.data_augmentation.augmentation_engine import AugmentationEngi
 
 import pandas as pd
 
-class Test_Engine(unittest.TestCase):
+
+class TestEngine(unittest.TestCase):
     def setUp(self):
         t1 = AISTrajectory(
             pd.DataFrame(
@@ -28,7 +29,6 @@ class Test_Engine(unittest.TestCase):
 
         self.trajectories = [t1, t2]
 
-
     def test_transform_simple_translation(self):
         engine = AugmentationEngine(translation_values=[(10, 0), (20, 0)], keep_original=False)
 
@@ -74,20 +74,116 @@ class Test_Engine(unittest.TestCase):
         )
         expected = [t1, t2, t3, t4]
 
-        self.assertEqual(len(expected), len(result))
+        self.assertEqual(len(result), len(expected))
         for t1, t2 in zip(result, expected):
             pd.testing.assert_frame_equal(t1.df, t2.df)
 
     def test_transform_simple_flip(self):
-        engine = AugmentationEngine(flip_values=[None, 0], keep_original=False)
+        engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=False)
+
+        result = engine.transform(self.trajectories)
+
+        t1 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+        t2 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [12 - i for i in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+        expected = [t1, t2]
+
+        self.assertEqual(len(expected), len(result))
+        for t1, t2 in zip(result, expected):
+            pd.testing.assert_frame_equal(t1.df, t2.df)
+
+    def test_transform_flip_and_translate(self):
+        engine = AugmentationEngine(translation_values=[(10, 0)], flip_values=[(None, 0)], keep_original=False)
 
         result = engine.transform(self.trajectories)
+        t1 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
+                    'longitude': [22 + i for i in range(10)]
+                }
+            )
+        )
+        t2 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [-12 + i for i in range(10)],
+                    'longitude': [22 + i for i in range(10)]
+                }
+            )
+        )
+
+        t3 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+        t4 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [12 - i for i in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+
+        t5 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
+                    'longitude': [22 + i for i in range(10)]
+                }
+            )
+        )
+        t6 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [12 - i for i in range(10)],
+                    'longitude': [22 + i for i in range(10)]
+                }
+            )
+        )
+
+        expected = [t1, t2, t3, t4, t5, t6]
+
+        self.assertEqual(len(expected), len(result))
+        for t1, t2 in zip(result, expected):
+            pd.testing.assert_frame_equal(t1.df, t2.df)
+
+    def test_engine_verbose(self):
+        engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=False)
+
+        result = engine.transform(self.trajectories, verbose=1)
 
         t1 = AISTrajectory(
             pd.DataFrame(
                 {
                     'ts_sec': [i for i in range(10)],
-                    'latitude': [0 - i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
                     'longitude': [12 + i for i in range(10)]
                 }
             )
@@ -107,5 +203,35 @@ class Test_Engine(unittest.TestCase):
         for t1, t2 in zip(result, expected):
             pd.testing.assert_frame_equal(t1.df, t2.df)
 
+    def test_engine_keep_original(self):
+        engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=True)
+
+        result = engine.transform(self.trajectories, verbose=1)
+
+        t1 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [0 for _ in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+        t2 = AISTrajectory(
+            pd.DataFrame(
+                {
+                    'ts_sec': [i for i in range(10)],
+                    'latitude': [12 - i for i in range(10)],
+                    'longitude': [12 + i for i in range(10)]
+                }
+            )
+        )
+        expected = self.trajectories + [t1, t2]
+
+        self.assertEqual(len(expected), len(result))
+        for t1, t2 in zip(result, expected):
+            pd.testing.assert_frame_equal(t1.df, t2.df)
+
+
 if __name__ == '__main__':
     unittest.main()