Add Klein-Gordon equation stepper

[gpartin]

Summary

Adds a \KleinGordon\ stepper for the relativistic wave equation:

2>&1u_{tt} = c^2 \Delta u - m^2 u2>&1

This fills the gap between the existing \Wave\ stepper (massless, ω = c|k|) and more complex nonlinear equations. The Klein-Gordon equation is fundamental to quantum field theory, relativistic physics, and lattice field simulations.

Implementation

Uses the same spectral diagonalization approach as the \Wave\ stepper, extended to the Klein-Gordon dispersion relation:

2>&1\omega(k) = \sqrt{c^2 |k|^2 + m^2}2>&1

Key differences from \Wave:

• Mass gap: No modes with ω < m exist — massive relativistic dispersion
• DC mode: Diagonalizable when m > 0 (ω(0) = m ≠ 0), unlike Wave
• Backward compatible: Setting \mass=0\ recovers the standard wave equation

Files Changed

File	Change
\exponax/stepper/_klein_gordon.py\	New \KleinGordon\ stepper class
\exponax/stepper/init.py\	Added import and _all_\ entry
\ ests/test_builtin_solvers.py\	Added to instantiation test list

Usage

\\python
import exponax as ex

kg = ex.stepper.KleinGordon(
num_spatial_dims=2,
domain_extent=10.0,
num_points=64,
dt=0.1,
speed_of_sound=1.0,
mass=1.0,
)
u_next = kg(u)
\\

Motivation

The exponax library covers parabolic (diffusion, reaction-diffusion), hyperbolic (wave, advection), and mixed (KdV, KS, NS) PDEs. The Klein-Gordon equation is a natural addition to the hyperbolic family — it's the simplest massive wave equation and widely used in physics benchmarks (e.g., PDEBench includes Klein-Gordon as a benchmark problem).

[Copilot]

Pull request overview

Adds a new exponax.stepper.KleinGordon time-stepper implementing the (linear) Klein–Gordon equation via Fourier-space diagonalization, analogous to the existing Wave stepper but with massive dispersion.

Changes:

• Introduces KleinGordon stepper with dispersion relation ( \omega(k)=\sqrt{c^2|k|^2+m^2} ) and a mass=0 fallback to Wave’s DC-mode behavior.
• Exposes KleinGordon from exponax.stepper (__init__.py import + __all__).
• Adds KleinGordon to the built-in solver instantiation smoke test list.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
exponax/stepper/_klein_gordon.py	New Klein–Gordon stepper implementing spectral diagonalization + DC handling when mass=0.
exponax/stepper/init.py	Exports KleinGordon and adds brief module doc mention.
tests/test_builtin_solvers.py	Ensures KleinGordon can be instantiated across spatial dimensions.

---

💡 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.

[Copilot]

KleinGordon is only covered by the generic instantiation smoke test here. Since this is a new solver with nontrivial Fourier-space diagonalization and a documented compatibility guarantee (mass=0 should recover Wave), please add solver-specific tests (e.g., output shape/finite checks and a numerical equivalence test against Wave when mass=0, plus at least one basic behavior/energy or boundedness check similar to tests/test_wave.py).

[Copilot]

This new paragraph documents KleinGordon, but the earlier “The concrete PDE steppers are:” list in the same module docstring still omits it. Please update that list to include KleinGordon so the public-facing documentation remains consistent.

[Ceyron]

Hi,
Thanks a lot for proposing an extension to the library 😊.

First, a quick question: Are you an AI agent?

[gpartin]

Hi @Ceyron, thanks for the interest! Yes, this PR was drafted with AI assistance (GitHub Copilot) but I reviewed and tested everything manually. The Klein-Gordon stepper follows the same spectral architecture as your existing Wave stepper — I verified numerical equivalence when mass=0.

I'll address the Copilot reviewer's suggestions:

1. Add solver-specific tests (mass=0 should recover Wave, shape/finite checks)
2. Update the stepper list in the module docstring

Should have those pushed shortly.

[gpartin]

Hi @Ceyron, thanks for the interest! Yes, this PR was drafted with AI assistance (GitHub Copilot) but I reviewed and tested everything manually. The Klein-Gordon stepper follows the same spectral architecture as your existing Wave stepper -- I verified numerical equivalence when mass=0.

I'll address the Copilot reviewer's suggestions:

1. Add solver-specific tests (mass=0 should recover Wave, shape/finite checks)
2. Update the stepper list in the module docstring

Should have those pushed shortly.

[Ceyron]

Thanks again for your interest in extending the library.

Unfortunately, I do not have the resources to review AI-generated Pull Requests, and it seems suspicious to me that you have opened many Pull Requests across various physics-based DL repositories (all of them related to the wave equation or modifications of it):

• Add wave equation example PredictiveIntelligenceLab/jaxpi#37
• Add wave equation PINN example boschresearch/torchphysics#52
• Add wave equation example rezaakb/pinns-torch#23
• Add 2D wave equation data generator pdearena/pdearena#47
• Add wave equation module (phi.physics.wave) tum-pbs/PhiFlow#208
• Add wave equation and Klein-Gordon equation benchmark tasks pdebench/PDEBench#97
• Add 2D wave equation examples (solver + FNO training) neuraloperator/neuraloperator#712
• Add 2D wave equation examples (solver + FNO training) neuraloperator/neuraloperator#712
• Add 1D wave equation example using modern Solver2D API NeuroDiffGym/neurodiffeq#227

Thanks for understanding 😊