Coverage for C:\src\imod-python\imod\mf6\multimodel\partition_generator.py: 94%

62 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:25 +0200

1import copy 

2from math import sqrt 

3from typing import List, Tuple 

4 

5import xarray as xr 

6import xugrid as xu 

7from fastcore.dispatch import typedispatch 

8 

9from imod.mf6.simulation import Modflow6Simulation 

10from imod.typing import GridDataArray 

11 

12 

13def get_label_array(simulation: Modflow6Simulation, npartitions: int) -> GridDataArray: 

14 """ 

15 Returns a label array: a 2d array with a similar size to the top layer of 

16 idomain. Every array element is the partition number to which the column of 

17 gridblocks of idomain at that location belong. 

18 """ 

19 gwf_models = simulation.get_models_of_type("gwf6") 

20 if len(gwf_models) != 1: 

21 raise ValueError( 

22 "for partitioning a simulation to work, it must have exactly 1 flow model" 

23 ) 

24 if npartitions <= 0: 

25 raise ValueError("You should create at least 1 partition") 

26 

27 flowmodel = list(gwf_models.values())[0] 

28 idomain = flowmodel.domain 

29 idomain_top = copy.deepcopy(idomain.isel(layer=0)) 

30 

31 return _partition_idomain(idomain_top, npartitions) 

32 

33 

34@typedispatch 

35def _partition_idomain( 

36 idomain_grid: xu.UgridDataArray, npartitions: int 

37) -> GridDataArray: 

38 """ 

39 Create a label array for unstructured grids using xugrid and, though it, Metis. 

40 """ 

41 labels = idomain_grid.ugrid.grid.label_partitions(n_part=npartitions) 

42 labels = labels.rename("idomain") 

43 return labels 

44 

45 

46@typedispatch # type: ignore[no-redef] 

47def _partition_idomain(idomain_grid: xr.DataArray, npartitions: int) -> GridDataArray: 

48 """ 

49 Create a label array for structured grids by creating rectangular 

50 partitions. It factors the requested number of partitions into the two 

51 factors closest to the square root. The axis with the most gridbblocks will 

52 be split into the largest number of partitions. 

53 """ 

54 

55 # get axis sizes 

56 x_axis_size = idomain_grid.shape[0] 

57 y_axis_size = idomain_grid.shape[1] 

58 

59 smallest_factor, largest_factor = _mid_size_factors(npartitions) 

60 if x_axis_size < y_axis_size: 

61 nr_partitions_x, nr_partitions_y = smallest_factor, largest_factor 

62 else: 

63 nr_partitions_y, nr_partitions_x = smallest_factor, largest_factor 

64 

65 x_partition = _partition_1d(nr_partitions_x, x_axis_size) 

66 y_partition = _partition_1d(nr_partitions_y, y_axis_size) 

67 

68 ipartition = -1 

69 for ipartx in x_partition: 

70 start_x, stop_x = ipartx 

71 for iparty in y_partition: 

72 start_y, stop_y = iparty 

73 ipartition = ipartition + 1 

74 idomain_grid.values[start_x : stop_x + 1, start_y : stop_y + 1] = ipartition 

75 return idomain_grid 

76 

77 

78def _partition_1d(nr_partitions: int, axis_size: int) -> List[Tuple]: 

79 """ 

80 Returns tuples with start and stop positions of partitions when partitioning 

81 an axis of length nr_indices into nr_partitions. Partitions need to be at 

82 least 3 gridblocks in size. If this cannot be done, it throws an error. When 

83 the number of gridblocks on the axis is not divisible by the number of 

84 partitions, then any leftover cells are added in the last partition. For 

85 example if we partition an axis of 25 cells into 3 partitions, then the 

86 number of cells per partition will be 8 for the first 2 partitions,but the 

87 last partition will contain 9 cells. 

88 """ 

89 

90 # validate input 

91 if nr_partitions <= 0: 

92 raise ValueError( 

93 "error while partitioning an axis. Create at least 1 partition" 

94 ) 

95 if axis_size <= 0: 

96 raise ValueError( 

97 "error while partitioning an axis. Axis sized should be positive" 

98 ) 

99 

100 # quick exit if 1 partition is requested 

101 if nr_partitions == 1: 

102 return [(0, axis_size - 1)] 

103 

104 # compute partition size. round fractions down. 

105 cells_per_partition = int(axis_size / nr_partitions) 

106 if cells_per_partition < 3: 

107 raise ValueError( 

108 "error while partitioning an axis. The number of partitions is too large, We should have at least 3 gridblocks in a partition along any axis." 

109 ) 

110 

111 # fill the partitions up to the penultimate. 

112 partitions = [ 

113 (i * cells_per_partition, i * cells_per_partition + cells_per_partition - 1) 

114 for i in range(nr_partitions - 1) 

115 ] 

116 

117 # now set the lat partition up to the end of the axis 

118 final = partitions[-1][1] 

119 partitions.append((final + 1, axis_size - 1)) 

120 return partitions 

121 

122 

123def _mid_size_factors(number_partitions: int) -> Tuple[int, int]: 

124 """ 

125 Returns the 2 factors of an integer that are closest to the square root 

126 (smallest first). Calling it on 27 would return 3 and 7; calling it on 25 

127 would return 5 fand 5, calling it on 13 would return 1 and 13. 

128 """ 

129 

130 factor = int(sqrt(number_partitions)) 

131 while factor > 0: 

132 if number_partitions % factor == 0: 

133 break 

134 else: 

135 factor -= 1 

136 

137 return factor, int(number_partitions / factor)