Operator Learning with Domain Decomposition for Geometry Generalization in PDE Solving

Imagine you are trying to solve a massive, complex puzzle, like figuring out how heat spreads through a strangely shaped metal plate or how water flows through a winding underground cave system. In the world of physics and engineering, these puzzles are called Partial Differential Equations (PDEs).

Traditionally, solving these puzzles is like trying to build a house by hand, brick by brick. It's accurate, but it takes forever. In recent years, scientists have tried using AI (Neural Operators) to act like a "super-architect" that can instantly predict the solution. However, there's a big catch: this AI is like a student who only studied in a classroom with square desks. If you put them in a room with round tables or triangular desks (a new, weird shape), they get confused and fail. They need to be retrained from scratch for every new shape, which requires massive amounts of data.

This paper introduces a clever new method called Schwarz Neural Inference (SNI) to fix this. Here is how it works, explained with simple analogies:

1. The Problem: The "One-Size-Fits-None" AI

Current AI solvers try to learn the solution for the entire shape at once. If the shape is a weird, jagged coastline, the AI has never seen anything like it before. It's like asking a chef who only knows how to cook a perfect square pizza to suddenly cook a perfect pizza shaped like a dragon. They don't have the right tools or experience.

2. The Solution: The "Lego" Strategy (Domain Decomposition)

The authors propose a "Local-to-Global" approach. Instead of trying to solve the whole dragon-shaped pizza at once, they break the problem down into tiny, manageable Lego bricks.

Step 1: The Training Phase (Learning the Bricks)
The AI is trained only on simple, basic shapes (like triangles, squares, and simple polygons). It learns how to solve the physics problem perfectly on these small, simple blocks. Think of this as teaching the chef how to bake a perfect square cookie. Once they master the square, they understand the fundamental rules of baking.
Step 2: The Inference Phase (Building the Dragon)
When you give the AI a new, crazy shape (like the dragon), the system automatically chops that shape into those same simple Lego bricks.
- The AI solves the problem for each small brick individually (using the knowledge it gained in Step 1).
- Then, it uses a special "stitching" algorithm (called the Schwarz method) to glue these small solutions together.
- It does this iteratively: "Okay, this brick says the heat is high here, so I'll tell the neighbor brick to adjust. Now the neighbor brick adjusts, and I adjust again." They keep talking to each other until the whole picture makes sense.

3. The Magic Trick: "Symmetry" and "Translation"

One of the paper's cleverest points is how it handles the fact that a "square" in the training data might be a "rectangle" in the real world.

The Analogy: Imagine you learned to ride a bike on a flat, straight road. Now you need to ride on a hill.
The Fix: The system uses mathematical "symmetries" (like rotation, scaling, or shifting) to translate the problem. It tells the AI: "Hey, this hill is just a flat road that's been tilted and stretched. Apply the same rules you know, just adjusted for the tilt." This allows the AI to handle shapes it has never seen before without needing new training data.

4. Why This Matters

Data Efficiency: You don't need millions of examples of every possible weird shape. You just need examples of simple shapes. It's like learning the alphabet instead of memorizing every book in the library.
Geometry Generalization: It works on any shape, whether it's a simple circle or a complex, multi-holed industrial part.
Speed: Because the AI only solves small, simple pieces, it's much faster and more accurate than trying to guess the whole complex shape at once.

Summary

Think of this paper as teaching an AI to be a master builder rather than a memorizer.

Old Way: Memorize the blueprint for every possible house shape. (Impossible, takes too much data).
New Way (SNI): Learn how to build a perfect wall, a perfect window, and a perfect roof. Then, when you need a weird house, just assemble those perfect parts together and let them "talk" to each other until the house is finished.

This approach bridges the gap between the rigid world of traditional math and the flexible world of AI, making it possible to solve complex engineering problems on any shape, anywhere, with much less data.

1. Problem Statement

Neural operators (e.g., DeepONet, FNO, GNOT) have shown promise in solving Partial Differential Equations (PDEs) by learning mappings between function spaces. However, they face two critical bottlenecks:

Data Hunger: Training requires massive amounts of paired data (input functions to solutions), which is expensive to generate via traditional numerical solvers.
Geometry Generalization: Existing methods struggle to generalize to unseen geometries. While some approaches handle variations in boundary conditions or coefficients, they often fail when the domain shape differs significantly from the training distribution. Current methods relying on geometry parametrization or coordinate representation cannot easily adapt to entirely novel, complex topologies without retraining or generating new data.

The core challenge is how to solve PDEs on arbitrary, complex geometries using a neural operator trained only on simple, basic shapes, without requiring massive datasets for every new shape.

2. Methodology: Operator Learning with Domain Decomposition

The authors propose a local-to-global framework that integrates Neural Operator learning with classical Domain Decomposition Methods (DDMs). The core idea is to decompose a complex, arbitrary domain into smaller, manageable subdomains (basic shapes) where a neural operator can generalize effectively, solve the local problems, and then stitch the solutions together.

The framework consists of three main stages:

A. Training Data Generation (Local Operator Training)

Instead of training on complex domains, the framework trains a neural operator on a library of basic shapes.

Basic Shapes: The authors use simple polygons (with up to $n$ vertices) bounded within a compact region. These are chosen because they are flexible enough to approximate any 2D domain and allow for tractable sampling.
Boundary Conditions: To handle mixed boundary conditions (Dirichlet and Neumann), the boundaries of these basic shapes are randomly partitioned.
Data Augmentation: To enhance generalization, Lie Point Symmetry Data Augmentation (LPSDA) is applied. Transformations like rotation, scaling, and value shifting are applied to the input functions and boundary conditions, leveraging the inherent symmetries of the PDEs to create an infinite supply of training data from a finite set.
Model: A GNOT (General Neural Operator Transformer) is used as the local solver due to its flexibility in handling irregular meshes and diverse inputs.

B. Inference: Schwarz Neural Inference (SNI)

For a new, arbitrary target domain, the framework employs an iterative algorithm called Schwarz Neural Inference (SNI), inspired by the Additive Schwarz Method (ASM).

Decomposition: The target domain is partitioned into overlapping subdomains using a graph partitioning algorithm (e.g., METIS) followed by an iterative extension to create overlaps.
Normalization: Since the local subdomains may have different scales or coordinate systems than the training data, pre-processing transforms ( $T_k$ ) are applied to map local inputs (geometry, boundary conditions, coefficients) back into the training distribution range.
Local Inference: The trained neural operator ( $G^\dagger$ ) solves the PDE on each subdomain independently using the normalized inputs.
Stitching & Update: The local solutions are extended to the global domain. An Additive Schwarz-Richardson iteration updates the global solution by combining the local corrections.
Post-processing: The final solution is mapped back via inverse transforms ( $\tilde{T}_k$ ).

The process repeats until a convergence criterion is met.

3. Theoretical Analysis

The paper provides a theoretical convergence analysis for the SNI algorithm:

Assumptions: The PDE operator is self-adjoint and coercive elliptic. The classical Schwarz iteration is contractive, and the learned local operator has a uniform error bound.
Result: The authors prove that SNI converges to a fixed point. Furthermore, they derive an error bound showing that the error between the SNI solution and the true solution is proportional to the local neural operator's approximation error, scaled by the overlap factor and step size. This theoretically guarantees that if the local operator is accurate, the global solution will be accurate.

4. Key Contributions

Local-to-Global Framework: A novel paradigm that decouples geometry generalization from operator learning by solving PDEs on basic shapes and stitching them via DDMs.
Schwarz Neural Inference (SNI): An iterative inference algorithm that adapts classical DDMs to neural operators, enabling solutions on arbitrary geometries without retraining.
Symmetry-Based Data Generation: A scheme combining random shape generation with Lie point symmetry augmentation to efficiently train local operators on diverse basic shapes.
Theoretical Guarantees: A proof of convergence and error bounds for the proposed iterative scheme, linking the global error to the local operator's generalization capability.

5. Experimental Results

The framework was evaluated on various linear and nonlinear PDEs (Laplace, Darcy, Heat, Nonlinear Laplace) with Dirichlet, Neumann, and mixed boundary conditions.

Geometry Generalization: SNI significantly outperformed direct inference baselines (GNOT) on unseen complex domains (including multiply connected domains with holes and corners).
- Error Reduction: Achieved 34.8% to 96.8% reduction in $L_2$ relative error compared to baselines across stationary problems.
- Consistency: The method maintained consistent accuracy across different domain complexities (A, B, C), whereas baselines failed significantly on complex domains.
Data Efficiency:
- SNI achieved comparable or better accuracy than direct inference with significantly less training data.
- Even with small datasets, SNI's error was lower than the validation error of the direct inference model trained on larger datasets.
Time-Dependent Problems: The framework was successfully extended to the 2D Heat equation using space-time decomposition, reducing errors by 54.1% to 74.2%.
Ablation Studies:
- Hyperparameters: The number of partitions ( $K$ ) and step size ( $\tau$ ) primarily affect convergence speed, while the extension depth ( $d$ ) had minimal impact on final accuracy in the tested range.
- Architecture: The framework is architecture-agnostic; using Geo-FNO instead of GNOT yielded similar results, though GNOT showed slightly better generalization.
- Augmentation: Symmetry-based augmentation (rotation + scaling) was crucial for reducing performance variance across different domains.

6. Significance and Impact

Bridging the Gap: This work addresses the "chicken-and-egg" problem in operator learning (need data to train, need solvers to get data) by allowing training on simple, synthetic shapes while solving complex real-world geometries.
Industrial Applicability: By enabling the use of pre-trained local operators on arbitrary industrial geometries without retraining, it offers a practical path for deploying neural operators in engineering simulations (e.g., fluid dynamics, heat transfer).
Theoretical Foundation: It provides one of the first rigorous convergence analyses for neural operator-based domain decomposition, validating the "local-to-global" approach mathematically.
Future Directions: The authors suggest this framework can be extended to 3D problems, non-overlapping decompositions, and more complex PDE types (e.g., Navier-Stokes), potentially revolutionizing how PDE solvers are designed in the era of AI.

In summary, the paper presents a robust, theoretically grounded, and empirically superior method for solving PDEs on arbitrary geometries by leveraging the modularity of domain decomposition and the expressiveness of neural operators.