set bounding box with 2 numbers

skais/ais/ais_trajectory.py

+import numbers
 import random
 
 import pandas as pd
@@ -232,7 +233,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', ref_index=-1, features=None,
                                       node_size=0):
         nb_channels = 1
 
@@ -241,7 +242,7 @@ class AISTrajectory(AISPoints):
             lower_lon, upper_lon = (min(positions[:, 0]), max(positions[:, 0]))
             lower_lat, upper_lat = (min(positions[:, 1]), max(positions[:, 1]))
         elif bounding_box == 'centered':
-            center_lon, center_lat = positions[-1]
+            center_lon, center_lat = positions[ref_index]
             min_lon, max_lon = (min(positions[:, 0]), max(positions[:, 0]))
             min_lat, max_lat = (min(positions[:, 1]), max(positions[:, 1]))
 
@@ -253,10 +254,19 @@ class AISTrajectory(AISPoints):
             upper_lon = center_lon + distance_to_center
             lower_lon = center_lon - distance_to_center
         elif type(bounding_box) is list:
+            if type(bounding_box[0]) is not numbers.Number:
                 upper_lon = bounding_box[1][0]
                 lower_lon = bounding_box[0][0]
                 upper_lat = bounding_box[1][1]
                 lower_lat = bounding_box[0][1]
+            else:
+                center_lon, center_lat = positions[ref_index]
+                distance_to_center_lon = bounding_box[0]
+                distance_to_center_lat = bounding_box[1]
+                upper_lat = center_lat + distance_to_center_lat
+                lower_lat = center_lat - distance_to_center_lat
+                upper_lon = center_lon + distance_to_center_lon
+                lower_lon = center_lon - distance_to_center_lon
         else:
             raise ValueError(f"Option not supported: {bounding_box}")