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

12 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2024-07-26 20:12 +0000

1import numpy as np 

2import pandas as pd 

3 

4 

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

6 """ 

7 Filters 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 

16 Returns 

17 ------- 

18 : 

19 Filtered list of coordinates. 

20 """ 

21 

22 filtered_coordinates = [] 

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

24 

25 for x, y, z in coordinates: 

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

27 if coordinate_radius <= radius**2: 

28 if center: 

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

30 else: 

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

32 

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