Add Klein-Gordon equation examples (data gen + FNO training)#713
Open
gpartin wants to merge 2 commits intoneuraloperator:mainfrom
Open
Add Klein-Gordon equation examples (data gen + FNO training)#713gpartin wants to merge 2 commits intoneuraloperator:mainfrom
gpartin wants to merge 2 commits intoneuraloperator:mainfrom
Conversation
Add two new Sphinx Gallery examples demonstrating the Klein-Gordon equation: 1. examples/data_gen/plot_klein_gordon_1d_solver.py - 1D Klein-Gordon solver using FiniteDiff and leapfrog integration - Compares wave behavior across mass parameters (m=0, 2, 10) - Spacetime diagrams and energy conservation verification 2. examples/models/plot_FNO_klein_gordon.py - Generates training data via leapfrog finite-difference solver - Trains 1D FNO to learn the KG time-evolution operator - Evaluates predictions on held-out test data - Demonstrates cross-mass generalization limits These examples fill a gap in the existing gallery: no hyperbolic PDE with a mass/dispersion term was previously covered. The Klein-Gordon equation (Klein 1926, Gordon 1926) is a standard benchmark in computational physics.
There was a problem hiding this comment.
Pull request overview
Adds two new Sphinx Gallery examples to the examples/ gallery demonstrating the 1D Klein–Gordon equation, covering both (1) finite-difference data generation/visualization and (2) training an FNO to learn the time-evolution operator and evaluating cross-mass generalization.
Changes:
- Add a 1D Klein–Gordon finite-difference solver example using
FiniteDiff, with visualization and an energy drift check. - Add an FNO training example that generates synthetic Klein–Gordon data, trains an FNO, evaluates test error, and sweeps mass for generalization.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
examples/data_gen/plot_klein_gordon_1d_solver.py |
New FD-based Klein–Gordon solver + plots + energy conservation check. |
examples/models/plot_FNO_klein_gordon.py |
New end-to-end example: synthetic Klein–Gordon data gen, FNO training, evaluation, and cross-mass sweep. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| import matplotlib.pyplot as plt | ||
| import sys | ||
| from neuralop.models import FNO | ||
| from neuralop import Trainer |
Comment on lines
+232
to
+234
| l2loss = LpLoss(d=1, p=2) | ||
| h1loss = H1Loss(d=1) | ||
|
|
Comment on lines
+110
to
+117
| u_xx = fd.dx(fd.dx(u)) | ||
| u_prev = u + 0.5 * dt**2 * (c**2 * u_xx - m**2 * u) | ||
|
|
||
| snapshots = [u.cpu().numpy().copy()] | ||
| energies = [] | ||
|
|
||
| for step in range(nt): | ||
| u_xx = fd.dx(fd.dx(u)) |
Comment on lines
+110
to
+117
| u_xx = fd.dx(fd.dx(u)) | ||
| u_prev = u + 0.5 * dt**2 * (c**2 * u_xx - m**2 * u) | ||
|
|
||
| snapshots = [u.cpu().numpy().copy()] | ||
| energies = [] | ||
|
|
||
| for step in range(nt): | ||
| u_xx = fd.dx(fd.dx(u)) |
Comment on lines
+300
to
+304
| x_grid = np.linspace(0, 1, nx) | ||
|
|
||
| fig, axes = plt.subplots(3, 1, figsize=(10, 8)) | ||
| for idx in range(3): | ||
| data_x = test_x[idx, 0].numpy() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds two Sphinx Gallery examples demonstrating the Klein-Gordon equation — a relativistic wave equation with a mass/dispersion term. This fills a gap in the existing example gallery, which covers elliptic (Darcy), parabolic (diffusion-advection), and incompressible fluid (Burgers, Navier-Stokes) PDEs, but no hyperbolic PDE with a mass term.
New files
1.
examples/data_gen/plot_klein_gordon_1d_solver.pyA finite-difference solver for the 1D Klein-Gordon equation using the library's
FiniteDiffclass and leapfrog (Verlet) time integration:add-klein-gordon-example\frac{\partial^2 u}{\partial t^2} = c^2 \frac{\partial^2 u}{\partial x^2} - m^2 uadd-klein-gordon-example
Demonstrates:
2.
examples/models/plot_FNO_klein_gordon.pyTrains a 1D FNO to learn the Klein-Gordon time-evolution operator:
Notes
# %%cells,.. math::blocks,.. raw:: htmldividers)blackFiniteDiff,FNO,AdamW,LpLoss,H1Loss)