pyflwdir.FlwdirRaster.partition#

FlwdirRaster.partition(level: Literal['basin', 'subbasin'] = 'subbasin', n_processes: int = 4, *, seed: int = 42, refine: bool = True, min_subtree_size: int = 100000, imbalance_target: float = 1.005, upa: ndarray | None = None) → PartitionPlan[source]#

Divide the raster into parts, one for each process of an upstream accumulation.

Pass the returned plan to upstream_area() or accuflux(), with parallel=True, to accumulate each part in its own process. One plan serves any number of accumulations on the same flow directions:

plan = flw.partition(n_processes=4, level="subbasin")
uparea = flw.upstream_area(parallel=True, partition_plan=plan)
accu = flw.accuflux(data, parallel=True, partition_plan=plan)

Make a new plan when the flow directions change, for example after add_pits(): the accumulation checks the plan and refuses one that no longer fits.

level="basin" (FlowTopo Method 1) keeps every basin whole and groups the basins into n_processes contiguous, balanced parts, using METIS, a graph partitioning library, and a refinement of the boundaries. level="subbasin" (FlowTopo Method 2) opens the basins that keep the parts from balancing, along their mainstems (going up from the pit through the tributary with the most cells at each confluence), into tributary subtrees, and divides those and the other basins into n_processes parts. P_min is the most upstream mainstem cell where a tributary of another part enters: from there down to the pit, the mainstem is a section accumulated after the parts, in the parent process.

METIS is an optional dependency; install it with pip install pyflwdir[partition].

Parameters:

  • level ({'basin', 'subbasin'}, optional) – How to divide the raster, as above, by default "subbasin".

  • n_processes (int, optional) – Number of processes, one part each, by default 4. A part without cells, as when there are fewer basins than processes at basin level, starts no process. A tributary subtree is never cut, so the heaviest part is at least as heavy as the largest subtree of an opened basin, and more processes do not balance a region beyond that.

  • seed (int, optional) – Seed of the first METIS run, by default 42.

  • refine (bool, optional) – Move boundary nodes between parts after METIS to balance them, by default True.

  • min_subtree_size (int, optional) – At subbasin level, tributary subtrees smaller than this are grouped with a larger neighbour before METIS divides them, by default 100,000 cells as in the continental FlowTopo C workflow; the balance refinement still moves single subtrees. Regardless of this value, at least 8 * n_processes groups are kept, or every subtree if there are fewer.

  • imbalance_target (float, optional) – Largest part load divided by the mean load that the partition aims for, by default 1.005. At subbasin level, further basins open only while it is not reached.

  • upa (ndarray, optional) – Upstream area, one value a cell, finite and positive on the network. It decides between two tributaries with as many cells when a mainstem is traced, as in the FlowTopo C code; without it the one with the smaller cell index is taken.

Returns:

plan.parts holds the part of each cell, shaped like the raster, from 0 to n_processes - 1. The value n_processes marks the mainstem sections, accumulated after the parts, and -1 marks the cells that take no part: off the network, the cell of a basin of a single cell (as in the FlowTopo C code), and cells in, or draining into, a loop of the flow directions, which never reach a pit. plan.loads holds the cell count of each part, without the mainstem sections, and plan.stems how each mainstem section connects to the parts.

Return type:

PartitionPlan