Coverage for src/abm_initialization_collection/coordinate/filter_coordinate_bounds.py: 100%

1import numpy as np 

2import pandas as pd 

3 

4 

5def filter_coordinate_bounds(coordinates: list, radius: float, *, center: bool) -> pd.DataFrame: 

6 """ 

7 Filter list for coordinates with given radius. 

8 

9 Parameters 

10 ---------- 

11 coordinates 

12 List of (x, y, z) coordinates. 

13 radius 

14 Maximum valid radius of coordinate. 

15 center 

16 True if coordinates should be centered, False otherwise. 

17 

18 Returns 

19 ------- 

20 : 

21 Filtered list of coordinates. 

22 """ 

23 

24 filtered_coordinates = [] 

25 x_center, y_center, _ = np.array(coordinates).mean(axis=0) 

26 

27 for x, y, z in coordinates: 

28 coordinate_radius = (x - x_center) ** 2 + (y - y_center) ** 2 

29 if coordinate_radius <= radius**2: 

30 if center: 

31 filtered_coordinates.append((x - x_center, y - y_center, z)) 

32 else: 

33 filtered_coordinates.append((x, y, z)) 

34 

35 return pd.DataFrame(filtered_coordinates, columns=["x", "y", "z"])