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
« 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
5import xarray as xr
6import xugrid as xu
7from fastcore.dispatch import typedispatch
9from imod.mf6.simulation import Modflow6Simulation
10from imod.typing import GridDataArray
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")
27 flowmodel = list(gwf_models.values())[0]
28 idomain = flowmodel.domain
29 idomain_top = copy.deepcopy(idomain.isel(layer=0))
31 return _partition_idomain(idomain_top, npartitions)
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
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 """
55 # get axis sizes
56 x_axis_size = idomain_grid.shape[0]
57 y_axis_size = idomain_grid.shape[1]
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
65 x_partition = _partition_1d(nr_partitions_x, x_axis_size)
66 y_partition = _partition_1d(nr_partitions_y, y_axis_size)
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
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 """
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 )
100 # quick exit if 1 partition is requested
101 if nr_partitions == 1:
102 return [(0, axis_size - 1)]
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 )
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 ]
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
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 """
130 factor = int(sqrt(number_partitions))
131 while factor > 0:
132 if number_partitions % factor == 0:
133 break
134 else:
135 factor -= 1
137 return factor, int(number_partitions / factor)