MAGNET: an open-source library for mesh agglomeration by Graph Neural Networks

TL;DR

MAGNET addresses mesh agglomeration by reframing it as a graph partitioning problem on the dual mesh and solving it with Graph Neural Networks, including SAGE-based models and reinforcement-learning–driven partitioners/refiners. The approach is benchmarked against METIS and k-means across 2D and 3D test cases (including brain geometries and heterogeneous domains), showing competitive partition quality and computational efficiency, especially when GPU acceleration is available. MAGNET provides a modular, open-source framework with generation, I/O, and quality-metrics tools and seamlessly interfaces with the lymph PolyDG library to validate PDE discretizations on agglomerated meshes. The work highlights extensibility, potential integration with multigrid hierarchies, and avenues for richer RL models to further improve performance on complex polytopal meshes.

Abstract

We introduce MAGNET, an open-source Python library designed for mesh agglomeration in both two- and three-dimensions, based on employing Graph Neural Networks (GNN). MAGNET serves as a comprehensive solution for training a variety of GNN models, integrating deep learning and other advanced algorithms such as METIS and k-means to facilitate mesh agglomeration and quality metric computation. The library's introduction is outlined through its code structure and primary features. The GNN framework adopts a graph bisection methodology that capitalizes on connectivity and geometric mesh information via SAGE convolutional layers, in line with the methodology proposed by Antonietti et al. (2024). Additionally, the proposed MAGNET library incorporates reinforcement learning to enhance the accuracy and robustness of the model for predicting coarse partitions within a multilevel framework. A detailed tutorial is provided to guide the user through the process of mesh agglomeration and the training of a GNN bisection model. We present several examples of mesh agglomeration conducted by MAGNET, demonstrating the library's applicability across various scenarios. Furthermore, the performance of the newly introduced models is contrasted with that of METIS and k-means, illustrating that the proposed GNN models are competitive regarding partition quality and computational efficiency. Finally, we exhibit the versatility of MAGNET's interface through its integration with Lymph, an open-source library implementing discontinuous Galerkin methods on polytopal grids for the numerical discretization of multiphysics differential problems.

Figures (20)

• Figure 1: Example of a mesh obtained from the centroidal Voronoi tessellation of a cube (left) and a representation of its computational graph $\mathcal{G}$: in red the nodes $\mathcal{V}$, in blue the edges $\mathcal{E}$ (right).
• Figure 2: Representation of how the reinforcement learning agent bisects the graph. First, all nodes are in subset $\mathcal{S}_1$ (blue) except for one node with minimum degree (the orange one). At each step, the agent picks a blue node and flips it to an orange one until the two sets have the same number of nodes.
• Figure 3: Code structure of MAGNET, highlighting the agglomeration models class hierarchy. Arrows denote import relations.
• Figure 4: Examples of the meshes that can be generated using the Sub-module of MAGNET.
• Figure 5: Training and validation losses history plot of the two-dimensional and three-dimensional versions of the SAGE-Base model of Section \ref{['sec:sagebase']}.