Closed
Conversation
Add a KleinGordon stepper for the relativistic wave equation uₜₜ = c² Δu - m² u, using the same spectral diagonalization approach as the Wave stepper. Key differences from Wave: - Uses Klein-Gordon dispersion: ω(k) = √(c²|k|² + m²) - Has mass gap: no modes with ω < m exist - DC mode (k=0) is diagonalizable when m > 0 (ω(0) = m) - Setting mass=0 recovers the standard wave equation Files: - exponax/stepper/_klein_gordon.py: KleinGordon stepper class - exponax/stepper/__init__.py: Updated exports - tests/test_builtin_solvers.py: Added to instantiation tests
- Add test_klein_gordon.py with dedicated solver tests: - Instantiation and output shape checks - mass=0 recovers Wave stepper (numerical equivalence) - Analytical standing-mode correctness - Mass gap test (k=0 oscillates at omega=m) - Energy conservation bounds - Update stepper list in __init__.py docstring to include KleinGordon
There was a problem hiding this comment.
Pull request overview
Adds a new Fourier-spectral, analytically-stepped timestepper for the d-dimensional Klein–Gordon equation, expanding exponax’s built-in PDE solver suite in the same style as the existing Wave stepper.
Changes:
- Introduces
KleinGordonstepper with Fourier-space diagonalization andmass=0compatibility withWave. - Adds a dedicated test suite covering instantiation, mass→wave equivalence, analytical mode evolution, mass gap, and energy behavior.
- Wires the new stepper into
exponax.stepperexports and the built-in instantiation test.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
exponax/stepper/_klein_gordon.py |
New Klein–Gordon spectral stepper implementation (diagonalization + optional DC correction for mass=0). |
exponax/stepper/__init__.py |
Exposes KleinGordon in imports, docs, and __all__. |
tests/test_klein_gordon.py |
New tests for Klein–Gordon behavior and compatibility with Wave. |
tests/test_builtin_solvers.py |
Ensures KleinGordon is included in the “all built-ins instantiate” test. |
💡 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.
exponax/stepper/__init__.py
Outdated
Comment on lines
99
to
101
| from ._klein_gordon import KleinGordon | ||
| from ._wave import Wave | ||
|
|
tests/test_klein_gordon.py
Outdated
| import jax.numpy as jnp | ||
| import pytest | ||
|
|
||
| import exponax as ex |
tests/test_klein_gordon.py
Outdated
| from exponax.stepper import KleinGordon, Wave | ||
|
|
||
| L = 2 * jnp.pi | ||
| PI = jnp.pi |
tests/test_klein_gordon.py
Outdated
| def test_energy_bounded(self): | ||
| """Total energy should be conserved (bounded) over many steps.""" | ||
| k0, c, m, N, dt = 3, 1.0, 2.0, 64, 0.005 | ||
| stepper, x, u0, omega = self._make_stepper_and_ic(k0, c, m, N, dt) |
tests/test_klein_gordon.py
Outdated
Comment on lines
+163
to
+164
| # KE + gradient PE + mass PE | ||
| return jnp.sum(v**2 + c**2 * jnp.abs(jnp.fft.rfft(h))**2 + m**2 * h**2) |
tests/test_klein_gordon.py
Outdated
| u = stepper(u) | ||
| e_final = energy(u) | ||
|
|
||
| # Spectral solver should conserve energy to machine precision |
…gy formula - Move KleinGordon import to alphabetical position (after _hyper_diffusion) - Remove unused 'import exponax as ex' (F401) - Remove unused PI constant (F841) - Decouple test_energy_bounded from helper to avoid unused variables (F841) - Fix energy formula: add |k|^2 weighting for gradient PE, proper dx scaling - Fix misleading 'machine precision' comment to match rel=1e-3 tolerance
Owner
|
Thanks again for your interest in extending Exponax. 😊 As I mentioned in #103, unfortunately, I do not have the resources to review AI-generated feature-extending Pull Requests. For contributions that extend Exponax's feature set, please first open an issue to discuss them. |
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 a spectral timestepper for the Klein-Gordon equation, extending exponax's PDE solver collection with the relativistic generalization of the wave equation.
What this PR adds
Stepper:
exponax/stepper/_klein_gordon.pyTests:
tests/test_klein_gordon.py(9 test cases)Integration:
exponax/stepper/__init__.pyexports and \all\tests/test_builtin_solvers.pyinstantiation test__init__.pyWhy Klein-Gordon?
The Klein-Gordon equation is the simplest relativistic wave equation and a cornerstone of quantum field theory. It generalizes the existing Wave stepper with a mass term that introduces a mass gap in the dispersion relation. exponax already has diffusion, advection, wave, Burgers, KdV, and Kuramoto-Sivashinsky steppers — Klein-Gordon is a natural addition.
Design decisions