Coverage for src/abm_initialization_collection/sample/remove_edge_regions.py: 100%

13 statements  

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

1import pandas as pd 

2 

3 

4def remove_edge_regions( 

5 samples: pd.DataFrame, edge_threshold: int, edge_padding: float 

6) -> pd.DataFrame: 

7 """ 

8 Removes regions at edges. 

9 

10 Parameters 

11 ---------- 

12 samples 

13 Sample cell ids and coordinates. 

14 threshold 

15 Number of edge positions per axis needed to assign edge region. 

16 padding 

17 Distance from axis limits to assign edge positions. 

18 

19 Returns 

20 ------- 

21 : 

22 Samples with edge cells removed. 

23 """ 

24 

25 # Get ids of cell at edge. 

26 x_edge_ids = find_edge_ids("x", samples, edge_threshold, edge_padding) 

27 y_edge_ids = find_edge_ids("y", samples, edge_threshold, edge_padding) 

28 

29 # Filter samples for cells not at edge. 

30 all_edge_ids = set(x_edge_ids + y_edge_ids) 

31 samples_filtered = samples[~samples["id"].isin(all_edge_ids)] 

32 

33 return samples_filtered.reset_index(drop=True) 

34 

35 

36def find_edge_ids(axis: str, samples: pd.DataFrame, threshold: float, padding: float) -> list[int]: 

37 """ 

38 Finds ids of cells with voxels touching edges of given axis. 

39 

40 Parameters 

41 ---------- 

42 axis : {'x', 'y', 'z'} 

43 The name of axis to check. 

44 samples 

45 Sample cell ids and coordinates. 

46 threshold 

47 Number of edge positions per axis needed to assign edge region. 

48 padding 

49 Distance from axis limits to assign edge positions. 

50 

51 Returns 

52 ------- 

53 : 

54 List of edge cell ids. 

55 """ 

56 

57 # Get min and max coordinate for given axis. 

58 axis_min = samples[axis].min() + padding 

59 axis_max = samples[axis].max() - padding 

60 

61 # Check for cell ids located at edges. 

62 edges = samples.groupby("id").apply( 

63 lambda g: len(g[(g[axis] <= axis_min) | (g[axis] >= axis_max)]) 

64 ) 

65 edge_ids = edges[edges > threshold] 

66 

67 return list(edge_ids.index)