💻 computer science

Geometry-aware LegONet for PDE Learning on Arbitrary Domains

This paper introduces Geometry-aware LegONet (gLegONet), a modular operator learning framework that decouples physical mechanisms from specific geometries by pretraining variational blocks on an ambient spectral domain and adapting them to arbitrary target domains through boundary-defined affine manifolds, thereby enabling reusable PDE solving and law discovery without geometry-specific retraining.

Original authors: Jiahao Zhang, Yueqi Wang, Guang Lin

Published 2026-09-01

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiahao Zhang, Yueqi Wang, Guang Lin

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of physics, many natural phenomena are described by a specific type of mathematical rule known as a partial differential equation. These rules act like a universal instruction manual for how things change over time and space, governing everything from the flow of water around a rock to the spread of heat through a metal plate. For decades, scientists have relied on classical computers to solve these rules, breaking the physical world into a grid of tiny points and calculating the changes step by step. More recently, artificial intelligence has entered the field, learning to predict these changes by studying vast amounts of data. However, a significant hurdle remains: most of these AI models are trained on very specific shapes, like perfect squares or circles. If a scientist wants to use a model trained on a square to predict what happens on a twisted, irregular shape, the model often fails. It is as if a map drawn for a city grid cannot be used to navigate a winding mountain path without being completely redrawn. This limitation forces researchers to retrain their AI from scratch for every new shape, wasting time and computational power, and preventing the discovery of universal physical laws that should work anywhere.

A team of researchers at Purdue University has developed a new approach to solve this problem, creating a system they call Geometry-aware LegONet. Their work, published in a recent study, introduces a way to separate the core physics of a problem from the specific shape of the container it lives in. Instead of training a single, massive AI model to handle every possible shape, they trained a small library of specialized "mechanism blocks" on a simple, standard square. These blocks learn the fundamental behaviors of physics, such as how heat diffuses or how fluids transport energy, without worrying about the boundaries of the container. Once these blocks are trained, they are frozen, meaning their internal settings are locked and never changed again.

The innovation lies in how these frozen blocks are applied to new, complex shapes. When the researchers want to simulate a physical process on a strange, irregular domain, they do not retrain the AI. Instead, they use a mathematical technique to construct a special set of coordinates that fit the new shape perfectly. This process involves sampling points along the boundary of the new shape and using linear algebra to create a "boundary-adapted" framework. This framework acts as a bridge, translating the frozen, shape-agnostic physics blocks into the specific language of the new geometry. The result is that the same pre-trained blocks can be assembled to solve problems on a bunny-shaped domain, a pinwheel, or a heart-shaped region without any additional neural network training. The physics remains the same; only the algebraic interface changes to fit the new container.

The researchers tested this method on a wide variety of challenging scenarios. They successfully simulated complex fluid flows around obstacles, heat diffusion in twisted shapes, and wave patterns on irregular surfaces. In one test, they modeled a fluid flow around a cylinder embedded in a channel, a scenario that requires the simulation to respect the boundaries of both the channel walls and the cylinder simultaneously. The system maintained the correct physical constraints at the boundaries with extreme precision, keeping errors near the limits of computer numerical precision. In another test involving a volume-constrained chemical reaction, the method preserved the total amount of material in the system, a crucial physical law that many other AI methods struggle to maintain over long periods. The system achieved these results in seconds, whereas traditional AI methods required hundreds or thousands of seconds to reach similar accuracy, and only after being trained specifically for that shape.

Beyond just predicting how a system evolves, the researchers demonstrated that their method could also work backward to discover the underlying physical laws. By feeding the system short, noisy observations of a physical process on an unseen shape, the system could identify the correct mathematical terms governing the behavior. It successfully recovered the coefficients of reaction and diffusion terms, even when the data was sparse and collected from only a few sensors scattered across the domain. This capability suggests that the system does not just memorize patterns but truly understands the modular components of the physics, allowing it to reconstruct the governing equations from limited data.

The significance of this work is that it changes how reusable scientific software can be built. Previously, the reusable unit in machine learning for physics was often a full solution map for a specific set of shapes. This new approach shifts the focus to the physical mechanisms themselves. By decoupling the learning of physical laws from the geometry of the domain, the researchers have created a library of universal building blocks. These blocks can be deployed on any new geometry through a deterministic algebraic assembly process, rather than requiring a new, expensive training phase. This separation means that as long as the physics remains the same, the same learned blocks can be used to simulate phenomena on any shape, from the microscopic to the macroscopic, provided the geometry is sampled correctly.

The study confirms that this approach works across diverse boundary conditions, including those where the value of the field is fixed, where the flow across the boundary is zero, or where a mix of both applies. It also handles complex interactions, such as when a physical process requires solving an auxiliary equation to determine the flow field, or when higher-order derivatives are involved. The researchers found that the accuracy of the method depends primarily on the resolution of the initial training space and the number of coordinates retained for the new shape, rather than on the number of boundary samples, once a sufficient threshold is met. This indicates that the bottleneck for precision is the capacity of the reduced space to represent the physics, not the enforcement of the boundary itself.

In the end, this research offers a practical path forward for scientific machine learning. It moves the field away from the idea that every new problem requires a new, monolithic AI model. Instead, it proposes a system where the core knowledge of physics is learned once and stored in a reusable library, ready to be assembled into a solver for any new geometry. This makes the discovery of physical laws and the simulation of complex systems more efficient and robust, ensuring that the laws of physics are not tied to the specific shapes used to train the computer. The method does not claim to be the fastest predictor for a single, fixed geometry, but rather provides a flexible, interpretable, and reusable framework for tackling the vast diversity of shapes found in the natural world.

Technical Summary: Geometry-aware LegONet for PDE Learning on Arbitrary Domains

Problem Statement
Current learned Partial Differential Equation (PDE) solvers, including Physics-Informed Neural Networks (PINNs), Neural Operators (e.g., FNO, DeepONet), and learned simulators, often entangle the learned physical operators with the specific geometry, boundary conditions, and discretization used during training. This coupling limits the reusability of these models; when the same physical laws are applied to a new domain or boundary configuration, the models typically require retraining, architectural modifications, or additional penalty-based boundary treatments. While recent "geometry-aware" neural operators attempt to condition solutions on domain shapes, they often still rely on training data specific to families of geometries or require complex attention mechanisms that do not strictly separate interior physics from boundary realization.

Methodology: Geometry-aware LegONet (gLegONet)
The authors propose gLegONet, a framework that decouples the learning of physical mechanisms from their realization on specific geometries. The method extends the "Lego-like" variational operator learning concept by introducing a deterministic, algebraic boundary interface.

Ambient Pretraining (Offline):
- Physical mechanisms (e.g., diffusion, transport, reaction) are pretrained once as modular neural blocks on a fixed, regular "ambient" spectral domain (a square $Q = [-1, 1]^2$ ).
- These blocks, denoted as $\mathcal{B}^Q_\theta = \{F^\theta_i\}$ , learn to map coefficient vectors of the ambient field to the coefficients of the resulting physical mechanism (e.g., mapping Fourier coefficients of $u$ to those of $\Delta u$ ).
- Once trained, these neural parameters are frozen. They represent reusable "Lego" blocks of physics independent of any specific target geometry.
Boundary-Adapted Coordinates (Online):
- For a new target domain $\Omega \subset Q$ , the method constructs a coordinate system that inherently satisfies the boundary conditions.
- Boundary constraints (Dirichlet, Neumann, Robin, or mixed) are sampled at discrete points on $\partial \Omega$ to form a linear system $Ca = d(t)$ , where $a$ represents the ambient coefficients.
- The solution space is decomposed into an affine lift $a_{bc}(t)$ (satisfying non-homogeneous data) and a null-space component $N_\Omega z(t)$ (satisfying homogeneous constraints).
- The matrix $N_\Omega$ is computed via Singular Value Decomposition (SVD) of the constraint matrix $C$ , followed by mass-orthonormalization with respect to the $L^2(\Omega)$ inner product. This ensures that any state $a(t) = a_{bc}(t) + N_\Omega z(t)$ satisfies the boundary conditions by construction.
Galerkin Realization and Rollout:
- The frozen ambient blocks are realized on the target domain by projecting their outputs onto the boundary-adapted tangent space defined by $N_\Omega$ .
- The dynamics evolve in the reduced coordinate space $z(t)$ via the equation:
  $\dot{z}(t) = \sum_{i} c_i q^\Omega_i(z(t), t) - N_\Omega^\top M_\Omega \dot{a}_{bc}(t)$
  where $q^\Omega_i$ is the reduced response of the $i$ -th mechanism and $M_\Omega$ is the mass matrix.
- This process requires no neural retraining for new domains; only the algebraic assembly of $C$ , $N_\Omega$ , and $M_\Omega$ changes.
Sparse Law Identification:
- The framework supports inverse problems by using the same boundary-adapted coordinates to regress sparse, noisy observations against a dictionary of realized mechanisms (both frozen neural blocks and analytic terms).
- Coefficients for the governing law are identified via regularized least squares in the reduced space, enabling long-time prediction on unseen domains.

Key Contributions

Separation of Physics and Geometry: The paper introduces a method where the "interior physics" (learned operators) and "boundary realization" (algebraic constraints) are strictly separated. Changing the domain alters only the coordinate interface, not the learned physics.
Zero-Shot Domain Transfer: The method demonstrates that a library of mechanisms pretrained on a simple square can be deployed on arbitrary, unseen domains (e.g., rosettes, bunny shapes, pinwheels) with complex boundary conditions without any geometry-specific neural retraining.
Hard Boundary Enforcement: Unlike penalty-based methods (e.g., standard PINNs) where boundary residuals are minimized but not guaranteed, gLegONet enforces boundary conditions to the tolerance of the null-space construction throughout the entire time rollout.
Modular Operator Composition: The framework supports the assembly of complex systems (e.g., fourth-order Swift-Hohenberg equations, Navier-Stokes vorticity transport) by composing lower-order frozen blocks (e.g., Laplacian) within the boundary-adapted space, avoiding the need to train high-order specific blocks.

Results
The authors evaluate gLegONet on several benchmarks involving unseen geometries and diverse boundary conditions:

Forward Simulations: On five manufactured solution benchmarks (reaction-diffusion, transport, coupled systems), gLegONet achieves final errors in the range of $10^{-4}$ to $10^{-3}$ with runtimes of 2–20 seconds. This compares favorably to baselines like PINN, FNO, and UNO, which require significantly longer training/inference times and often exhibit higher errors or boundary leakage.
Constrained Dynamics: In a volume-constrained Allen-Cahn simulation on a disk, the method preserves global mass conservation and boundary constraints with high precision, outperforming PINN in mass drift by two orders of magnitude.
Complex Operators: The method successfully simulates flow past an embedded cylinder (Navier-Stokes) and clamped Swift-Hohenberg dynamics on a heart-shaped domain, demonstrating the ability to handle multiply connected domains and higher-order operator compositions without retraining.
Inverse Discovery: In sparse identification tests on "peanut" and "channel" domains, the method recovers governing equations from short-time, noisy, sparse observations. It successfully identifies active coefficients and maintains stable long-time rollouts, whereas baselines using ambient Fourier bases or sparse finite differences fail to recover stable laws under the same sensor budgets.

Significance and Claims
The paper claims that gLegONet advances the state of scientific machine learning by changing the "unit of reuse." Instead of retraining a monolithic solver for every new geometry, the method allows for the deployment of a pretrained library of physical mechanisms onto new domains via an algebraic boundary assembly.

The authors emphasize that this approach:

Converts arbitrary-domain PDE learning from a geometry-specific retraining problem into an algebraic assembly problem.
Provides a natural space for physical law discovery where candidate features are geometry-compatible by construction.
Clarifies error sources: since boundary leakage is eliminated by the coordinate construction, remaining errors are attributable to the expressiveness of the ambient basis, the accuracy of the frozen blocks, or time discretization, rather than weak boundary enforcement.

The work is presented as a step toward "scientific machine-learning libraries" where physical mechanisms are learned once and deployed on new configurations through boundary-adapted coordinates, bridging the gap between modular neural learning and classical numerical principles of separating interior operators from boundary constraints.

Technical Summary: Geometry-aware LegONet for PDE Learning on Arbitrary Domains

More like this