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

24 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2024-09-25 19:25 +0000

1from math import floor, sqrt 

2 

3import numpy as np 

4from hexalattice import hexalattice 

5 

6 

7def make_grid_coordinates( 

8 grid: str, 

9 bounds: tuple[int, int, int], 

10 increment_xy: float, 

11 increment_z: float, 

12) -> list: 

13 """ 

14 Get all coordinates within given bounding box for selected grid type. 

15 

16 Parameters 

17 ---------- 

18 grid : {'rect', 'hex'} 

19 Type of grid. 

20 bounds 

21 Bounds in the x, y, and z directions. 

22 increment_xy 

23 Increment size in x/y directions. 

24 increment_z 

25 Increment size in z direction. 

26 

27 Returns 

28 ------- 

29 : 

30 List of grid coordinates. 

31 """ 

32 

33 if grid == "rect": 

34 return make_rect_grid_coordinates(bounds, increment_xy, increment_z) 

35 

36 if grid == "hex": 

37 return make_hex_grid_coordinates(bounds, increment_xy, increment_z) 

38 

39 message = f"invalid grid type {grid}" 

40 raise ValueError(message) 

41 

42 

43def make_rect_grid_coordinates( 

44 bounds: tuple[int, int, int], 

45 increment_xy: float, 

46 increment_z: float, 

47) -> list: 

48 """ 

49 Get list of bounded (x, y, z) coordinates for rect grid. 

50 

51 Parameters 

52 ---------- 

53 bounds 

54 Bounds in the x, y, and z directions. 

55 increment_xy 

56 Increment size in x/y directions. 

57 increment_z 

58 Increment size in z direction. 

59 

60 Returns 

61 ------- 

62 : 

63 List of grid coordinates. 

64 """ 

65 

66 x_bound, y_bound, z_bound = bounds 

67 

68 z_indices = np.arange(0, z_bound, increment_z) 

69 x_indices = np.arange(0, x_bound, increment_xy) 

70 y_indices = np.arange(0, y_bound, increment_xy) 

71 

72 return [(x, y, z) for z in z_indices for x in x_indices for y in y_indices] 

73 

74 

75def make_hex_grid_coordinates( 

76 bounds: tuple[int, int, int], 

77 increment_xy: float, 

78 increment_z: float, 

79) -> list: 

80 """ 

81 Get list of bounded (x, y, z) coordinates for hex grid. 

82 

83 Coordinates are offset in sets of three z slices to form a face-centered 

84 cubic (FCC) packing. 

85 

86 Parameters 

87 ---------- 

88 bounds 

89 Bounds in the x, y, and z directions. 

90 increment_xy 

91 Increment size in x/y directions. 

92 increment_z 

93 Increment size in z direction. 

94 

95 Returns 

96 ------- 

97 : 

98 List of grid coordinates. 

99 """ 

100 

101 x_bound, y_bound, z_bound = bounds 

102 

103 z_indices = np.arange(0, z_bound, increment_z) 

104 z_offsets = [(i % 3) for i in range(len(z_indices))] 

105 

106 xy_indices, _ = hexalattice.create_hex_grid( 

107 nx=floor(x_bound / increment_xy), 

108 ny=floor(y_bound / increment_xy * sqrt(3)), 

109 min_diam=increment_xy, 

110 align_to_origin=False, 

111 do_plot=False, 

112 ) 

113 

114 x_offsets = [(increment_xy / 2) if z_offset == 1 else 0 for z_offset in z_offsets] 

115 y_offsets = [(increment_xy / 2) * sqrt(3) / 3 * z_offset for z_offset in z_offsets] 

116 

117 return [ 

118 (x + x_offset, y + y_offset, z) 

119 for z, x_offset, y_offset in zip(z_indices, x_offsets, y_offsets) 

120 for x, y in xy_indices 

121 if round(x + x_offset) < x_bound and round(y + y_offset) < y_bound 

122 ]