add test + change of node_size

skais/ais/ais_trajectory.py

@@ -224,7 +224,7 @@ class AISTrajectory(AISPoints):
                 result.append((row['ts_sec'], current_label))
         return result
 
-    def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', features=None):
+    def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', features=None, node_size=0):
         nb_channels = 1
         if features is not None:
             nb_channels = len(features)
@@ -234,10 +234,19 @@ class AISTrajectory(AISPoints):
         if bounding_box != 'fit':
             raise ValueError("feature not implemented")
         positions = self.df[['longitude', 'latitude']].to_numpy()
-        range_longitude = (min(positions[0, :]), max(positions[0, :]))
-        range_latitude = (min(positions[1, :]), max(positions[1, :]))
+        range_longitude = (min(positions[:, 0]), max(positions[:, 0]))
+        range_latitude = (min(positions[:, 1]), max(positions[:, 1]))
         for longitude, latitude in positions:
-            x_coord = width * (longitude - range_longitude[0]) / (range_longitude[1] - range_longitude[0])
-            y_coord = height * (longitude - range_latitude[0]) / (range_latitude[1] - range_latitude[0])
-            data[x_coord, y_coord, :] = 1
+            x_coord = max(min(height - int(height * (latitude - range_latitude[0]) / (range_latitude[1] - range_latitude[0])) - 1, height - 1), 0)
+            y_coord = max(min(int((width - 1) * (longitude - range_longitude[0]) / (range_longitude[1] - range_longitude[0])), width - 1), 0)
+
+            x_lower_bound = max(0, x_coord - node_size)
+            x_upper_bound = min(height - 1, x_coord + node_size)
+
+            y_lower_bound = max(0, y_coord - node_size)
+            y_upper_bound = min(width - 1, y_coord + node_size)
+
+            for x in range(x_lower_bound, x_upper_bound + 1):
+                for y in range(y_lower_bound, y_upper_bound + 1):
+                    data[x, y] = [1]
         return data

skais/tests/ais/test_ais_trajectory.py

@@ -379,3 +379,53 @@ class TestAISTrajectory(unittest.TestCase):
         expected = [(0, 1), (6600, 2), (12600, 1)]
 
         self.assertListEqual(result, expected)
+
+    def test_generate_array_from_positions(self):
+        trajectory = AISTrajectory(
+            pd.DataFrame(
+                {
+                    "latitude": [0, 10, 0, -10],
+                    "longitude": [0, 10, 10, -10],
+                    "ts_sec": [i for i in range(4)]
+                }
+            )
+        )
+
+        result = trajectory.generate_array_from_positions(height=9, width=9, link=False, bounding_box='fit',
+                                                          features=None, node_size=0).reshape((9, 9))
+        expected = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 1, 0, 0, 0, 1],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [1, 0, 0, 0, 0, 0, 0, 0, 0]])
+
+        np.testing.assert_array_equal(result, expected)
+
+    def test_generate_array_from_positions_node_size(self):
+        trajectory = AISTrajectory(
+            pd.DataFrame(
+                {
+                    "latitude": [0, 10, 0, -10],
+                    "longitude": [0, 10, 10, -10],
+                    "ts_sec": [i for i in range(4)]
+                }
+            )
+        )
+
+        result = trajectory.generate_array_from_positions(height=9, width=9, link=False, bounding_box='fit',
+                                                          features=None, node_size=1).reshape((9, 9))
+        expected = np.array([[0, 0, 0, 0, 0, 0, 0, 1, 1],
+                             [0, 0, 0, 0, 0, 0, 0, 1, 1],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [0, 0, 0, 1, 1, 1, 0, 1, 1],
+                             [0, 0, 0, 1, 1, 1, 0, 1, 1],
+                             [0, 0, 0, 1, 1, 1, 0, 1, 1],
+                             [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                             [1, 1, 0, 0, 0, 0, 0, 0, 0],
+                             [1, 1, 0, 0, 0, 0, 0, 0, 0]])
+
+        np.testing.assert_array_equal(result, expected)
\ No newline at end of file