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

1 import copy

2 from math import sqrt

3 from typing import List, Tuple

4

5 import xarray as xr

6 import xugrid as xu

7 from fastcore.dispatch import typedispatch

8

9 from imod.mf6.simulation import Modflow6Simulation

10 from imod.typing import GridDataArray

11

12

13 def 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

35 def _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

47 def _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

78 def _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

123 def _mid_size_factors(number_partitions: int) -> (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 wouldd return 1 and 13.

128 """

129

130 factor = int(sqrt(number_partitions))

131 while factor > 0:

132 if number_partitions % factor == 0:

133 return (factor, int(number_partitions / factor))

134 else:

135 factor -= 1