[FEATURE] - Recursive bisection for partioning

iMOD-WQ supports a recursive bisection cut. This has the nice property of resulting in rectangular partitions, which aligns reasonably well with structured topologies.

Here is a basic implementation that results in a label array:

 
# %%
import numpy as np
from typing import NamedTuple
class Partition2d(NamedTuple):
    weights: np.ndarray
    row_range: tuple[int, int]
    column_range: tuple[int, int]
# %%

def switch_axis(axis: int):
    return int(not bool(axis))

def bisection_cut(partitions, axis: int):
    out = []
    for p in partitions:
        weight_along_axis = p.weights.sum(axis=switch_axis(axis))
        
        # Calculate the cumulative weights and find the cut point
        cut_weight = weight_along_axis.sum() / 2
        cumulative_weight = weight_along_axis.cumsum()
        cut_index = int(np.searchsorted(cumulative_weight, cut_weight)) + 1
        
        # Ensure cut_index is valid
        if cut_index == len(weight_along_axis):
            cut_index -= 1
        elif cut_index == 0:
            cut_index = 1
        
        if axis == 0:  # Cut along rows
            p0 = Partition2d(
                p.weights[:cut_index, :], 
                (p.row_range[0], p.row_range[0] + cut_index), 
                p.column_range
            )
            p1 = Partition2d(
                p.weights[cut_index:, :], 
                (p.row_range[0] + cut_index, p.row_range[1]), 
                p.column_range
            )
        else:  # Cut along columns
            p0 = Partition2d(
                p.weights[:, :cut_index], 
                p.row_range, 
                (p.column_range[0], p.column_range[0] + cut_index)
            )
            p1 = Partition2d(
                p.weights[:, cut_index:], 
                p.row_range, 
                (p.column_range[0] + cut_index, p.column_range[1])
            )
        
        out.extend([p0, p1])
    
    return out

def weightkey(p: Partition2d):
    return p.weights.sum()

def ratiokey(p: Partition2d):
    # Probably a pretty bad measure
    nrow = p.row_range[1] - p.row_range[0]
    ncol = p.column_range[1] - p.column_range[0]
    return abs(nrow - ncol)

def recursive_bisection_cut(weights: np.ndarray, n_partition: int, partition_sortkey):
    nrow, ncol = weights.shape
    partitions = [Partition2d(weights, (0, nrow), (0, ncol))]
    axis = 0
    while len(partitions) < n_partition:
        remainder = n_partition - len(partitions)
        # Each cut results in +1 partitions
        n_cuts = min(len(partitions), remainder)
        partitions = bisection_cut(partitions[:n_cuts], axis) + partitions[n_cuts:]
        # Sort partitions by some property, e.g. sum of weights
        partitions = sorted(partitions, key=partition_sortkey, reverse=True)
        axis=switch_axis(axis)
    return partitions

def paint(weights, partitions):
    labels = np.zeros_like(weights)
    for i, p in enumerate(partitions):
        labels[p.row_range[0]: p.row_range[1], p.column_range[0]: p.column_range[1]] = i
    return labels
# %%
weights = np.random.rand(10, 10)
# %%
partitions = recursive_bisection_cut(weights, 13, weightkey)
labels = paint(weights, partitions)
import matplotlib.pyplot as plt
plt.imshow(labels)
# %%
 

But anyway, this scheme also works for unstructured grids; instead of cutting by row and column, we'd cut by x and y bounds. It might make more sense to add this xugrid instead as an alternative scheme to METIS partioning.

The most controversial subject is what to do in case the number of partitions isn't a power of 2. An easy way around is to sort by some measure and then bisect some partitions first, e.g. those with the largest weight.

I would expect this to be noticable inferior to METIS, since METIS doesn't need this "rounding off" when the desired number of partitions isn't a power of two.

Start by examining the existing METIS partitioning entry point and how alternative partitioning schemes are represented. Compare it with the proposed recursive_bisection_cut implementation, then define behavior for non-power-of-two partition counts and what integration with xugrid should look like. Done means the alternative scheme is integrated with a clear partition-count policy and produces the expected labels.