zoomy_jax.mesh.partition module

zoomy_jax.mesh.partition module#

Mesh partitioning for MPI-parallel FVM solves.

Provides graph-based partitioning (via pymetis when available) and local sub-mesh extraction with ghost-cell bookkeeping. All routines operate on the NumPy-level mesh classes (BaseMesh / LSQMesh) so that they run before the JAX conversion step. If pymetis is not installed, a simple contiguous splitter is used instead.

class zoomy_jax.mesh.partition.PartitionInfo(rank, owned_cells, ghost_cells, local_faces, send_map=<factory>, recv_map=<factory>, global_to_local=<factory>)#

Bases: object

Partition bookkeeping for a single MPI rank.

Parameters:

  • rank (int) –

  • owned_cells (ndarray) –

  • ghost_cells (ndarray) –

  • local_faces (ndarray) –

  • send_map (Dict[int, ndarray]) –

  • recv_map (Dict[int, ndarray]) –

  • global_to_local (ndarray) –

rank#

The MPI rank this partition belongs to.

Type:

int

owned_cells#

Global cell indices owned by this rank (inner cells only).

Type:

np.ndarray[int]

ghost_cells#

Global cell indices needed from other ranks for face-stencil computations. These appear after owned cells in the local numbering.

Type:

np.ndarray[int]

local_faces#

Global face indices where at least one adjacent cell is owned.

Type:

np.ndarray[int]

send_map#

{neighbor_rank: local cell indices to send}. These are owned cells that are ghosts on a neighbouring rank.

Type:

dict[int, np.ndarray[int]]

recv_map#

{neighbor_rank: local indices where received data goes}. Local indices point into the ghost region of the local array (i.e. offset by len(owned_cells)).

Type:

dict[int, np.ndarray[int]]

global_to_local#

Mapping from global cell index to local cell index. Only entries for owned_cells and ghost_cells are meaningful; all others are set to -1.

Type:

np.ndarray[int]

rank: int#
owned_cells: ndarray#
ghost_cells: ndarray#
local_faces: ndarray#
send_map: Dict[int, ndarray]#
recv_map: Dict[int, ndarray]#
global_to_local: ndarray#
zoomy_jax.mesh.partition.partition_mesh(mesh, n_parts)#

Partition mesh into n_parts sub-domains.

Parameters:

  • mesh (BaseMesh (or FVMMesh / LSQMesh)) – The full (global) Zoomy mesh. Only topology arrays (face_cells, n_inner_cells, n_cells, n_faces) are used.

  • n_parts (int) – Number of partitions (typically MPI comm size).

Returns:

One entry per rank.

Return type:

list[PartitionInfo]

zoomy_jax.mesh.partition.extract_local_mesh(mesh, partition)#

Build a local mesh for one rank from the global mesh and its PartitionInfo.

  • Owned cells come first (indices 0 .. n_owned-1)

  • Ghost cells follow (indices n_owned .. n_owned+n_ghost-1)

  • All topology arrays are remapped to local numbering.

  • Boundary-face bookkeeping is restricted to faces that touch an owned cell.

The returned mesh is an LSQMesh (if available) or FVMMesh with precomputed geometry, ready for convert_mesh_to_jax.

Parameters:

partition (PartitionInfo) –