gkx 1.7.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- gkx-1.7.1/LICENSE +21 -0
- gkx-1.7.1/PKG-INFO +356 -0
- gkx-1.7.1/README.md +298 -0
- gkx-1.7.1/pyproject.toml +125 -0
- gkx-1.7.1/setup.cfg +4 -0
- gkx-1.7.1/src/gkx/__init__.py +31 -0
- gkx-1.7.1/src/gkx/_version.py +1 -0
- gkx-1.7.1/src/gkx/api/__init__.py +393 -0
- gkx-1.7.1/src/gkx/artifacts/__init__.py +29 -0
- gkx-1.7.1/src/gkx/artifacts/io.py +1088 -0
- gkx-1.7.1/src/gkx/artifacts/nonlinear_netcdf.py +1069 -0
- gkx-1.7.1/src/gkx/artifacts/plotting.py +812 -0
- gkx-1.7.1/src/gkx/artifacts/spectral_layout.py +310 -0
- gkx-1.7.1/src/gkx/artifacts/zonal_plots.py +248 -0
- gkx-1.7.1/src/gkx/benchmarking/__init__.py +2 -0
- gkx-1.7.1/src/gkx/benchmarking/shared.py +898 -0
- gkx-1.7.1/src/gkx/cli.py +388 -0
- gkx-1.7.1/src/gkx/config.py +368 -0
- gkx-1.7.1/src/gkx/core/__init__.py +1 -0
- gkx-1.7.1/src/gkx/core/grid.py +232 -0
- gkx-1.7.1/src/gkx/core/velocity.py +485 -0
- gkx-1.7.1/src/gkx/data/advanced_collision_six_moment.json +27 -0
- gkx-1.7.1/src/gkx/data/advanced_collision_six_moment.npy +0 -0
- gkx-1.7.1/src/gkx/data/cyclone_reference_adiabatic.csv +12 -0
- gkx-1.7.1/src/gkx/data/cyclone_reference_kinetic.csv +8 -0
- gkx-1.7.1/src/gkx/data/etg_reference.csv +4 -0
- gkx-1.7.1/src/gkx/data/kbm_reference.csv +6 -0
- gkx-1.7.1/src/gkx/data/tem_reference.csv +7 -0
- gkx-1.7.1/src/gkx/diagnostics/__init__.py +34 -0
- gkx-1.7.1/src/gkx/diagnostics/analysis.py +664 -0
- gkx-1.7.1/src/gkx/diagnostics/growth_rates.py +673 -0
- gkx-1.7.1/src/gkx/diagnostics/growth_windows.py +617 -0
- gkx-1.7.1/src/gkx/diagnostics/metadata.py +898 -0
- gkx-1.7.1/src/gkx/diagnostics/modes.py +345 -0
- gkx-1.7.1/src/gkx/diagnostics/moments.py +903 -0
- gkx-1.7.1/src/gkx/diagnostics/nonlinear_gradient_evidence.py +916 -0
- gkx-1.7.1/src/gkx/diagnostics/nonlinear_gradient_statistics.py +952 -0
- gkx-1.7.1/src/gkx/diagnostics/nonlinear_replicates.py +886 -0
- gkx-1.7.1/src/gkx/diagnostics/nonlinear_transport_optimization.py +932 -0
- gkx-1.7.1/src/gkx/diagnostics/normalization.py +120 -0
- gkx-1.7.1/src/gkx/diagnostics/quasilinear_calibration.py +885 -0
- gkx-1.7.1/src/gkx/diagnostics/quasilinear_model_selection.py +728 -0
- gkx-1.7.1/src/gkx/diagnostics/quasilinear_transport.py +625 -0
- gkx-1.7.1/src/gkx/diagnostics/stellarator_transport_reports.py +912 -0
- gkx-1.7.1/src/gkx/diagnostics/transport.py +698 -0
- gkx-1.7.1/src/gkx/diagnostics/transport_windows.py +911 -0
- gkx-1.7.1/src/gkx/diagnostics/validation_gates.py +644 -0
- gkx-1.7.1/src/gkx/diagnostics/zonal_validation.py +733 -0
- gkx-1.7.1/src/gkx/geometry/__init__.py +48 -0
- gkx-1.7.1/src/gkx/geometry/analytic.py +415 -0
- gkx-1.7.1/src/gkx/geometry/autodiff_checks.py +888 -0
- gkx-1.7.1/src/gkx/geometry/backend_discovery.py +319 -0
- gkx-1.7.1/src/gkx/geometry/booz_xform_bridge.py +712 -0
- gkx-1.7.1/src/gkx/geometry/core.py +214 -0
- gkx-1.7.1/src/gkx/geometry/differentiable.py +342 -0
- gkx-1.7.1/src/gkx/geometry/flux_tube.py +540 -0
- gkx-1.7.1/src/gkx/geometry/flux_tube_contract.py +302 -0
- gkx-1.7.1/src/gkx/geometry/imported_miller.py +941 -0
- gkx-1.7.1/src/gkx/geometry/imported_vmec.py +699 -0
- gkx-1.7.1/src/gkx/geometry/kernels.py +417 -0
- gkx-1.7.1/src/gkx/geometry/miller_eik.py +135 -0
- gkx-1.7.1/src/gkx/geometry/numerics.py +242 -0
- gkx-1.7.1/src/gkx/geometry/sensitivity.py +329 -0
- gkx-1.7.1/src/gkx/geometry/vmec_boozer_constants.py +91 -0
- gkx-1.7.1/src/gkx/geometry/vmec_boozer_core.py +997 -0
- gkx-1.7.1/src/gkx/geometry/vmec_boozer_derivatives.py +810 -0
- gkx-1.7.1/src/gkx/geometry/vmec_boundary_chain.py +730 -0
- gkx-1.7.1/src/gkx/geometry/vmec_eik.py +232 -0
- gkx-1.7.1/src/gkx/geometry/vmec_field_line_sampling.py +910 -0
- gkx-1.7.1/src/gkx/geometry/vmec_flux_tube_reports.py +938 -0
- gkx-1.7.1/src/gkx/geometry/vmec_state_controls.py +546 -0
- gkx-1.7.1/src/gkx/geometry/vmec_state_sensitivity.py +841 -0
- gkx-1.7.1/src/gkx/geometry/vmec_tensor_mapping.py +84 -0
- gkx-1.7.1/src/gkx/objectives/__init__.py +1 -0
- gkx-1.7.1/src/gkx/objectives/autodiff_validation.py +839 -0
- gkx-1.7.1/src/gkx/objectives/core.py +375 -0
- gkx-1.7.1/src/gkx/objectives/eigen.py +399 -0
- gkx-1.7.1/src/gkx/objectives/geometry.py +158 -0
- gkx-1.7.1/src/gkx/objectives/gradient_gates.py +658 -0
- gkx-1.7.1/src/gkx/objectives/portfolio.py +641 -0
- gkx-1.7.1/src/gkx/objectives/portfolio_guard.py +530 -0
- gkx-1.7.1/src/gkx/objectives/sampling.py +197 -0
- gkx-1.7.1/src/gkx/objectives/solver_gradient_reports.py +230 -0
- gkx-1.7.1/src/gkx/objectives/solver_vmec.py +442 -0
- gkx-1.7.1/src/gkx/objectives/stellarator.py +469 -0
- gkx-1.7.1/src/gkx/objectives/stellarator_contracts.py +180 -0
- gkx-1.7.1/src/gkx/objectives/stellarator_reduced.py +637 -0
- gkx-1.7.1/src/gkx/objectives/stellarator_tables.py +423 -0
- gkx-1.7.1/src/gkx/objectives/vmec_boozer.py +342 -0
- gkx-1.7.1/src/gkx/objectives/vmec_boozer_context.py +307 -0
- gkx-1.7.1/src/gkx/objectives/vmec_boozer_fd.py +822 -0
- gkx-1.7.1/src/gkx/objectives/vmec_boozer_gradients.py +701 -0
- gkx-1.7.1/src/gkx/objectives/vmec_boozer_line_search.py +882 -0
- gkx-1.7.1/src/gkx/objectives/vmec_candidate_admission.py +562 -0
- gkx-1.7.1/src/gkx/objectives/vmec_transport.py +636 -0
- gkx-1.7.1/src/gkx/objectives/vmec_transport_admission.py +542 -0
- gkx-1.7.1/src/gkx/objectives/vmec_transport_branch.py +405 -0
- gkx-1.7.1/src/gkx/objectives/vmec_transport_optimization.py +562 -0
- gkx-1.7.1/src/gkx/objectives/zonal.py +851 -0
- gkx-1.7.1/src/gkx/operators/__init__.py +16 -0
- gkx-1.7.1/src/gkx/operators/collision.py +34 -0
- gkx-1.7.1/src/gkx/operators/linear/__init__.py +83 -0
- gkx-1.7.1/src/gkx/operators/linear/cache_arrays.py +219 -0
- gkx-1.7.1/src/gkx/operators/linear/cache_builder.py +947 -0
- gkx-1.7.1/src/gkx/operators/linear/cache_model.py +175 -0
- gkx-1.7.1/src/gkx/operators/linear/collision_tables.py +195 -0
- gkx-1.7.1/src/gkx/operators/linear/collisions.py +994 -0
- gkx-1.7.1/src/gkx/operators/linear/dissipation.py +921 -0
- gkx-1.7.1/src/gkx/operators/linear/linked.py +277 -0
- gkx-1.7.1/src/gkx/operators/linear/moments.py +227 -0
- gkx-1.7.1/src/gkx/operators/linear/params.py +360 -0
- gkx-1.7.1/src/gkx/operators/linear/rhs.py +108 -0
- gkx-1.7.1/src/gkx/operators/linear/streaming.py +555 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/__init__.py +46 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/brackets.py +312 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/collisions.py +221 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/device_z.py +687 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/device_z_reports.py +290 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/diagnostic_state.py +767 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/diagnostics.py +349 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/domain_decomposition.py +577 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/parallel.py +157 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/parallel_contracts_domain.py +231 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/parallel_contracts_spectral.py +356 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/parallel_contracts_strategy.py +160 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/policies.py +595 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/projection.py +223 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/rhs.py +204 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/spectral_core.py +533 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/spectral_identity_integrator.py +692 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/spectral_identity_reports.py +481 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/spectral_identity_rhs.py +287 -0
- gkx-1.7.1/src/gkx/operators/nonlinear/spectral_layout.py +214 -0
- gkx-1.7.1/src/gkx/parallel/__init__.py +89 -0
- gkx-1.7.1/src/gkx/parallel/batch.py +159 -0
- gkx-1.7.1/src/gkx/parallel/decomposition.py +399 -0
- gkx-1.7.1/src/gkx/parallel/identity.py +108 -0
- gkx-1.7.1/src/gkx/parallel/independent.py +623 -0
- gkx-1.7.1/src/gkx/parallel/integrators.py +312 -0
- gkx-1.7.1/src/gkx/parallel/state.py +90 -0
- gkx-1.7.1/src/gkx/parallel/velocity.py +119 -0
- gkx-1.7.1/src/gkx/parallel/velocity_drive.py +516 -0
- gkx-1.7.1/src/gkx/parallel/velocity_hermite.py +289 -0
- gkx-1.7.1/src/gkx/parallel/velocity_plan.py +179 -0
- gkx-1.7.1/src/gkx/parallel/velocity_streaming.py +328 -0
- gkx-1.7.1/src/gkx/runtime.py +472 -0
- gkx-1.7.1/src/gkx/solvers/__init__.py +1 -0
- gkx-1.7.1/src/gkx/solvers/linear/__init__.py +38 -0
- gkx-1.7.1/src/gkx/solvers/linear/implicit.py +781 -0
- gkx-1.7.1/src/gkx/solvers/linear/integrator_diagnostics.py +367 -0
- gkx-1.7.1/src/gkx/solvers/linear/integrators.py +803 -0
- gkx-1.7.1/src/gkx/solvers/linear/krylov.py +627 -0
- gkx-1.7.1/src/gkx/solvers/linear/krylov_algorithms.py +909 -0
- gkx-1.7.1/src/gkx/solvers/linear/parallel.py +334 -0
- gkx-1.7.1/src/gkx/solvers/linear/parallel_common.py +76 -0
- gkx-1.7.1/src/gkx/solvers/linear/parallel_electrostatic.py +945 -0
- gkx-1.7.1/src/gkx/solvers/linear/parallel_streaming.py +653 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/__init__.py +71 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/diagnostic_integration.py +477 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/diagnostics.py +901 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/explicit.py +571 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/imex.py +454 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/imex_diagnostics.py +816 -0
- gkx-1.7.1/src/gkx/solvers/nonlinear/state_integration.py +831 -0
- gkx-1.7.1/src/gkx/solvers/time/__init__.py +43 -0
- gkx-1.7.1/src/gkx/solvers/time/diffrax_core.py +168 -0
- gkx-1.7.1/src/gkx/solvers/time/diffrax_linear.py +476 -0
- gkx-1.7.1/src/gkx/solvers/time/diffrax_nonlinear.py +452 -0
- gkx-1.7.1/src/gkx/solvers/time/diffrax_streaming.py +580 -0
- gkx-1.7.1/src/gkx/solvers/time/explicit.py +603 -0
- gkx-1.7.1/src/gkx/solvers/time/explicit_cfl.py +405 -0
- gkx-1.7.1/src/gkx/solvers/time/explicit_diagnostics.py +455 -0
- gkx-1.7.1/src/gkx/solvers/time/explicit_steps.py +320 -0
- gkx-1.7.1/src/gkx/solvers/time/runners.py +287 -0
- gkx-1.7.1/src/gkx/terms/__init__.py +37 -0
- gkx-1.7.1/src/gkx/terms/assembly.py +879 -0
- gkx-1.7.1/src/gkx/terms/config.py +71 -0
- gkx-1.7.1/src/gkx/terms/fields.py +474 -0
- gkx-1.7.1/src/gkx/terms/linear_terms.py +498 -0
- gkx-1.7.1/src/gkx/terms/nonlinear.py +887 -0
- gkx-1.7.1/src/gkx/utils/__init__.py +1 -0
- gkx-1.7.1/src/gkx/utils/callbacks.py +143 -0
- gkx-1.7.1/src/gkx/workflows/__init__.py +3 -0
- gkx-1.7.1/src/gkx/workflows/demo.py +235 -0
- gkx-1.7.1/src/gkx/workflows/linear.py +937 -0
- gkx-1.7.1/src/gkx/workflows/nonlinear.py +709 -0
- gkx-1.7.1/src/gkx/workflows/runtime/__init__.py +5 -0
- gkx-1.7.1/src/gkx/workflows/runtime/artifacts.py +150 -0
- gkx-1.7.1/src/gkx/workflows/runtime/chunks.py +222 -0
- gkx-1.7.1/src/gkx/workflows/runtime/commands.py +900 -0
- gkx-1.7.1/src/gkx/workflows/runtime/config.py +273 -0
- gkx-1.7.1/src/gkx/workflows/runtime/diagnostic_arrays.py +310 -0
- gkx-1.7.1/src/gkx/workflows/runtime/diagnostics.py +506 -0
- gkx-1.7.1/src/gkx/workflows/runtime/execution.py +292 -0
- gkx-1.7.1/src/gkx/workflows/runtime/initial_conditions.py +505 -0
- gkx-1.7.1/src/gkx/workflows/runtime/initial_phi.py +214 -0
- gkx-1.7.1/src/gkx/workflows/runtime/orchestration_artifacts.py +683 -0
- gkx-1.7.1/src/gkx/workflows/runtime/orchestration_scan.py +783 -0
- gkx-1.7.1/src/gkx/workflows/runtime/policies.py +338 -0
- gkx-1.7.1/src/gkx/workflows/runtime/results.py +106 -0
- gkx-1.7.1/src/gkx/workflows/runtime/startup.py +354 -0
- gkx-1.7.1/src/gkx/workflows/runtime/toml.py +269 -0
- gkx-1.7.1/src/gkx.egg-info/PKG-INFO +356 -0
- gkx-1.7.1/src/gkx.egg-info/SOURCES.txt +206 -0
- gkx-1.7.1/src/gkx.egg-info/dependency_links.txt +1 -0
- gkx-1.7.1/src/gkx.egg-info/entry_points.txt +2 -0
- gkx-1.7.1/src/gkx.egg-info/requires.txt +41 -0
- gkx-1.7.1/src/gkx.egg-info/top_level.txt +1 -0
gkx-1.7.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 UW Plasma
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
gkx-1.7.1/PKG-INFO
ADDED
|
@@ -0,0 +1,356 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: gkx
|
|
3
|
+
Version: 1.7.1
|
|
4
|
+
Summary: JAX gyrokinetic solver with Hermite-Laguerre velocity space
|
|
5
|
+
Author: GKX team
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/uwplasma/GKX
|
|
8
|
+
Project-URL: Documentation, https://github.com/uwplasma/GKX/tree/main/docs
|
|
9
|
+
Project-URL: Repository, https://github.com/uwplasma/GKX
|
|
10
|
+
Project-URL: Issues, https://github.com/uwplasma/GKX/issues
|
|
11
|
+
Keywords: gyrokinetics,jax,plasma-physics,tokamak,stellarator,turbulence
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
19
|
+
Requires-Python: >=3.10
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
License-File: LICENSE
|
|
22
|
+
Requires-Dist: jax
|
|
23
|
+
Requires-Dist: jaxlib
|
|
24
|
+
Requires-Dist: numpy
|
|
25
|
+
Requires-Dist: matplotlib
|
|
26
|
+
Requires-Dist: scipy
|
|
27
|
+
Requires-Dist: netCDF4
|
|
28
|
+
Requires-Dist: diffrax
|
|
29
|
+
Requires-Dist: equinox
|
|
30
|
+
Requires-Dist: solvax
|
|
31
|
+
Requires-Dist: tomli; python_version < "3.11"
|
|
32
|
+
Requires-Dist: tqdm
|
|
33
|
+
Provides-Extra: docs
|
|
34
|
+
Requires-Dist: sphinx; extra == "docs"
|
|
35
|
+
Requires-Dist: sphinx-rtd-theme; extra == "docs"
|
|
36
|
+
Requires-Dist: matplotlib; extra == "docs"
|
|
37
|
+
Provides-Extra: release
|
|
38
|
+
Requires-Dist: build; extra == "release"
|
|
39
|
+
Requires-Dist: twine; extra == "release"
|
|
40
|
+
Provides-Extra: dev
|
|
41
|
+
Requires-Dist: pytest; extra == "dev"
|
|
42
|
+
Requires-Dist: pytest-cov; extra == "dev"
|
|
43
|
+
Requires-Dist: ruff; extra == "dev"
|
|
44
|
+
Requires-Dist: mypy; extra == "dev"
|
|
45
|
+
Requires-Dist: mpmath; extra == "dev"
|
|
46
|
+
Requires-Dist: gmpy2; extra == "dev"
|
|
47
|
+
Requires-Dist: pandas; extra == "dev"
|
|
48
|
+
Requires-Dist: Pillow; extra == "dev"
|
|
49
|
+
Requires-Dist: mkdocs; extra == "dev"
|
|
50
|
+
Requires-Dist: sphinx; extra == "dev"
|
|
51
|
+
Requires-Dist: sphinx-rtd-theme; extra == "dev"
|
|
52
|
+
Requires-Dist: matplotlib; extra == "dev"
|
|
53
|
+
Requires-Dist: build; extra == "dev"
|
|
54
|
+
Requires-Dist: twine; extra == "dev"
|
|
55
|
+
Provides-Extra: assets
|
|
56
|
+
Requires-Dist: Pillow; extra == "assets"
|
|
57
|
+
Dynamic: license-file
|
|
58
|
+
|
|
59
|
+
# GKX
|
|
60
|
+
|
|
61
|
+
[](https://github.com/uwplasma/GKX/releases)
|
|
62
|
+
[](https://pypi.org/project/gkx/)
|
|
63
|
+
[](https://github.com/uwplasma/GKX/actions/workflows/ci.yml)
|
|
64
|
+
[](https://codecov.io/gh/uwplasma/GKX)
|
|
65
|
+
[](LICENSE)
|
|
66
|
+
[](pyproject.toml)
|
|
67
|
+
[](https://gkx.readthedocs.io)
|
|
68
|
+
|
|
69
|
+
GKX is a JAX-native gyrokinetic solver for linear stability,
|
|
70
|
+
nonlinear turbulence, differentiable analysis, and stellarator design. It uses
|
|
71
|
+
Fourier perpendicular coordinates, a Hermite-Laguerre velocity basis, and
|
|
72
|
+
field-aligned analytic, Miller, or VMEC geometry. The package runs on CPUs and
|
|
73
|
+
GPUs, exposes a Python API for autodiff and optimization, and provides a simple
|
|
74
|
+
executable for routine simulations.
|
|
75
|
+
|
|
76
|
+
## Installation
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install gkx
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
For development:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
git clone https://github.com/uwplasma/GKX
|
|
86
|
+
cd GKX
|
|
87
|
+
pip install -e .
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Quickstart
|
|
91
|
+
|
|
92
|
+
Run the built-in linear initial-value example:
|
|
93
|
+
|
|
94
|
+
```bash
|
|
95
|
+
gkx
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
The equivalent `gkx` entry point is also installed. The default run
|
|
99
|
+
prints setup, progress, elapsed time, and ETA, then writes its input, summary,
|
|
100
|
+
time series, eigenfunction, and a two-panel plot in the current directory.
|
|
101
|
+
|
|
102
|
+
Run a checked-in case or plot an existing result:
|
|
103
|
+
|
|
104
|
+
```bash
|
|
105
|
+
gkx examples/linear/axisymmetric/cyclone.toml
|
|
106
|
+
gkx run-runtime-nonlinear \
|
|
107
|
+
--config examples/nonlinear/axisymmetric/runtime_cyclone_nonlinear.toml \
|
|
108
|
+
--steps 200 --out cyclone.out.nc
|
|
109
|
+
gkx --plot cyclone.out.nc
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
Generate the small VMEC equilibria used by the self-contained examples:
|
|
113
|
+
|
|
114
|
+
```bash
|
|
115
|
+
pip install vmex
|
|
116
|
+
cd examples/vmec
|
|
117
|
+
./generate_wouts.sh
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
Full documentation is hosted at **[gkx.readthedocs.io](https://gkx.readthedocs.io)**.
|
|
121
|
+
Start with the [quickstart](https://gkx.readthedocs.io/en/latest/quickstart.html) and
|
|
122
|
+
[input reference](https://gkx.readthedocs.io/en/latest/inputs.html) for linear,
|
|
123
|
+
nonlinear, Miller, VMEC, restart, quasilinear, and plotting workflows.
|
|
124
|
+
|
|
125
|
+
## Highlights
|
|
126
|
+
|
|
127
|
+
- Electrostatic and electromagnetic gyrokinetics with kinetic or Boltzmann species.
|
|
128
|
+
- Linear initial-value, dominant-eigenmode, and nonlinear turbulence solvers.
|
|
129
|
+
- Analytic s-alpha, Miller, imported VMEC, and differentiable VMEC/Boozer geometry.
|
|
130
|
+
- JAX JIT, forward/reverse autodiff, implicit eigenvalue derivatives, and UQ tools.
|
|
131
|
+
- Quasilinear transport diagnostics with explicit saturation-rule metadata.
|
|
132
|
+
- CPU/GPU execution and production parallelization for independent scans and ensembles.
|
|
133
|
+
- Restartable NetCDF output and `gkx --plot` publication-style figures.
|
|
134
|
+
- A limited conserving Lenard-Bernstein/Dougherty-like collision model, with
|
|
135
|
+
advanced multispecies and linearized Landau operators remaining research lanes.
|
|
136
|
+
|
|
137
|
+
## Main Validation Results
|
|
138
|
+
|
|
139
|
+
The release atlas compares growth rates, frequencies, eigenfunctions, and
|
|
140
|
+
nonlinear transport windows with established gyrokinetic reference results.
|
|
141
|
+
Promoted cases include Cyclone ITG, Cyclone Miller, KBM, W7-X, and HSX, with
|
|
142
|
+
ETG and kinetic-electron stress cases kept at their documented claim level.
|
|
143
|
+
|
|
144
|
+

|
|
145
|
+
|
|
146
|
+
The exact equations, normalization, grids, boundary conditions, diagnostic
|
|
147
|
+
windows, tolerances, and artifact provenance are in the
|
|
148
|
+
[benchmark documentation](docs/benchmarks.rst) and
|
|
149
|
+
[verification matrix](docs/verification_matrix.rst). A visual overlay alone is
|
|
150
|
+
not treated as parity evidence.
|
|
151
|
+
|
|
152
|
+
The advanced-collision research lane now generates the complete retained
|
|
153
|
+
finite-Larmor Coulomb moment algebra and checks it against independent
|
|
154
|
+
velocity-space projection, spectral convergence, drift-kinetic conservation,
|
|
155
|
+
finite-wavelength gyro-diffusion, and the H-theorem. The panel below passes all
|
|
156
|
+
operator-level gates, but does not claim production Landau transport yet;
|
|
157
|
+
the finite-Larmor collisional-ITG gate and drift-kinetic collisional-zonal gate
|
|
158
|
+
are closed, and the paper-resolution finite-wavelength zonal traces and
|
|
159
|
+
velocity sections now pass their independent literature gate.
|
|
160
|
+
The differentiable driven-current solve is equation-tested, and the direct
|
|
161
|
+
Coulomb hierarchy is now converged through ``(P,J)=(20,5)`` with a maximum
|
|
162
|
+
nested current change of 0.017%. Its collision-frequency and electric-field
|
|
163
|
+
normalizations are now closed against the stationary Spitzer problem: all
|
|
164
|
+
three operators saturate by ``t nu_ee=50``, remain linear over a 100x field
|
|
165
|
+
scan, and Coulomb reaches the high-charge analytic limit within 7.46%. The
|
|
166
|
+
arbitrary-order original model reproduces its published 11% low-charge current
|
|
167
|
+
deficit, while the converged improved model is within 0.31% of Coulomb over the
|
|
168
|
+
complete ion-charge scan. The converged finite-Larmor ITG artifact is accepted;
|
|
169
|
+
the generated equal-species operator remains a Python research path because a
|
|
170
|
+
production input-file selector would imply broader multispecies coefficients
|
|
171
|
+
that are not yet implemented.
|
|
172
|
+
|
|
173
|
+

|
|
174
|
+
|
|
175
|
+
Equations, thresholds, machine-readable results, literature links, and the
|
|
176
|
+
one-command reproduction recipe are in the [collision-operator
|
|
177
|
+
documentation](docs/operators.rst).
|
|
178
|
+
|
|
179
|
+

|
|
180
|
+
|
|
181
|
+
At ``(P,J)=(24,10)``, the drift-kinetic traces approach the Xiao residual and
|
|
182
|
+
the finite-wavelength tails reproduce the published original < improved <
|
|
183
|
+
Coulomb ordering at both ``kx rho_i=0.1`` and ``0.2``. The improved model is
|
|
184
|
+
also closer to Coulomb over ``t nu_ii <= 10`` at both wavenumbers. Equations,
|
|
185
|
+
velocity-space convergence, compact replay data, and the Figure 12--14 gate
|
|
186
|
+
are documented in [Operators and Terms](docs/operators.rst).
|
|
187
|
+
|
|
188
|
+
## Runtime and Memory
|
|
189
|
+
|
|
190
|
+

|
|
191
|
+
|
|
192
|
+
The panel reports measured cold wall time and peak memory for the tracked CPU,
|
|
193
|
+
GPU, and comparison-code runs. Cold JAX rows include startup and compilation.
|
|
194
|
+
Prepared Python simulations avoid recompiling a fixed geometry and numerical
|
|
195
|
+
policy, but their CPU/GPU throughput depends on the software stack and GPU
|
|
196
|
+
operating state. See
|
|
197
|
+
[performance](docs/performance.rst) for profiler artifacts, memory accounting,
|
|
198
|
+
current reproducibility notes, and the distinction between executable,
|
|
199
|
+
prepared, and distributed runs.
|
|
200
|
+
|
|
201
|
+
## Differentiable Python API
|
|
202
|
+
|
|
203
|
+
```python
|
|
204
|
+
import jax.numpy as jnp
|
|
205
|
+
|
|
206
|
+
from gkx import CycloneBaseCase, LinearParams, integrate_linear_from_config
|
|
207
|
+
from gkx.core.grid import build_spectral_grid
|
|
208
|
+
from gkx.geometry import SAlphaGeometry
|
|
209
|
+
|
|
210
|
+
cfg = CycloneBaseCase()
|
|
211
|
+
grid = build_spectral_grid(cfg.grid)
|
|
212
|
+
geometry = SAlphaGeometry.from_config(cfg.geometry)
|
|
213
|
+
parameters = LinearParams()
|
|
214
|
+
state = jnp.zeros((2, 2, grid.ky.size, grid.kx.size, grid.z.size), dtype=jnp.complex64)
|
|
215
|
+
state = state.at[0, 0, 0, 0, :].set(1.0e-3)
|
|
216
|
+
trajectory, potential = integrate_linear_from_config(
|
|
217
|
+
state, grid, geometry, parameters, cfg.time
|
|
218
|
+
)
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
For repeated nonlinear calls with fixed geometry and numerical policy, prepare
|
|
222
|
+
the compiled simulation once:
|
|
223
|
+
|
|
224
|
+
```python
|
|
225
|
+
from gkx.solvers.nonlinear.diagnostic_integration import prepare_nonlinear_explicit_diagnostics
|
|
226
|
+
|
|
227
|
+
simulation = prepare_nonlinear_explicit_diagnostics(
|
|
228
|
+
initial_state,
|
|
229
|
+
grid,
|
|
230
|
+
geometry,
|
|
231
|
+
parameters,
|
|
232
|
+
dt=0.02,
|
|
233
|
+
steps=400,
|
|
234
|
+
resolved_diagnostics=False,
|
|
235
|
+
)
|
|
236
|
+
time, diagnostics, final_state, fields = simulation.run()
|
|
237
|
+
```
|
|
238
|
+
|
|
239
|
+
The prepared object accepts another same-shape initial state without rebuilding
|
|
240
|
+
the scan. A matched rebuilt cache/parameter PyTree can also remain dynamic for
|
|
241
|
+
autodiff; geometry layout is fixed, and dynamic-geometry compile reuse remains
|
|
242
|
+
an active differentiability lane.
|
|
243
|
+
|
|
244
|
+
The planted two-mode inverse problem below recovers two gradient parameters and
|
|
245
|
+
checks the autodiff Jacobian against finite differences. The single-mode demo in
|
|
246
|
+
the docs intentionally demonstrates non-identifiability rather than exact
|
|
247
|
+
parameter recovery.
|
|
248
|
+
|
|
249
|
+

|
|
250
|
+
|
|
251
|
+
See [differentiable geometry](docs/geometry.rst),
|
|
252
|
+
[algorithms](docs/algorithms.rst), and [stellarator optimization](docs/stellarator_optimization.rst)
|
|
253
|
+
for JVP, VJP, implicit differentiation, conditioning, covariance, and finite-difference gates.
|
|
254
|
+
|
|
255
|
+
## Quasilinear Modeling
|
|
256
|
+
|
|
257
|
+

|
|
258
|
+
|
|
259
|
+
The current quasilinear implementation is a scoped model-development and
|
|
260
|
+
optimization-screening result. It supports ranking and correlation studies but
|
|
261
|
+
is not a runtime/TOML absolute-flux predictor. Absolute-flux
|
|
262
|
+
promotion remains rejected when the declared Solovev and shaped-pressure stress
|
|
263
|
+
outliers are retained. Model definitions, derivations, calibration splits,
|
|
264
|
+
uncertainty, residual anatomy, and holdout gates are in the
|
|
265
|
+
[quasilinear documentation](docs/quasilinear.rst).
|
|
266
|
+
|
|
267
|
+
## QA ITG Optimization
|
|
268
|
+
|
|
269
|
+
The VMEX-style examples append a GKX growth-rate, quasilinear, or
|
|
270
|
+
nonlinear-window residual to the aspect-ratio, mean-iota, and quasisymmetry
|
|
271
|
+
objective tuples. The baseline follows the max-mode-5 QA workflow; all
|
|
272
|
+
transport comparisons use solved VMEC equilibria.
|
|
273
|
+
|
|
274
|
+

|
|
275
|
+
|
|
276
|
+
These rows are not promoted turbulent-flux designs. Their matched long
|
|
277
|
+
post-transient nonlinear audits use converged post-transient heat-flux windows
|
|
278
|
+
and do not show a statistically significant reduction relative to the strict
|
|
279
|
+
QA baseline. They are useful negative transfer evidence for improving objective
|
|
280
|
+
conditioning and optimizer choice.
|
|
281
|
+
|
|
282
|
+
The RBC(1,1) scan is a landscape and noise/convergence diagnostic, not a source
|
|
283
|
+
of admitted optimized candidates. It compares linear growth, all shipped
|
|
284
|
+
quasilinear rules, and replicated long-window nonlinear transport.
|
|
285
|
+
|
|
286
|
+

|
|
287
|
+
|
|
288
|
+
Reproducible scripts are in [examples/optimization](examples/optimization), and
|
|
289
|
+
full objective equations, optimizer policies, comparison fingerprints, and
|
|
290
|
+
long-window audits are in the [optimization documentation](docs/stellarator_optimization.rst).
|
|
291
|
+
|
|
292
|
+
## Parallelization
|
|
293
|
+
|
|
294
|
+
Production parallelization currently covers independent `k_y` scans,
|
|
295
|
+
quasilinear/UQ ensembles, and file-backed independent tasks with deterministic
|
|
296
|
+
ordering and serial identity gates. Sensitivity sweeps can use the same deterministic independent-work reconstruction, but they need a dedicated
|
|
297
|
+
matched scaling artifact before any speedup claim is promoted.
|
|
298
|
+
|
|
299
|
+
Nonlinear whole-state and domain decomposition remain diagnostic. Species-first
|
|
300
|
+
and Hermite-second decomposition, explicit Hermite halo exchange, field-moment
|
|
301
|
+
collectives, and physical transport-window identity must pass before a
|
|
302
|
+
nonlinear parallelization speedup is claimed. See [parallelization](docs/parallelization.rst).
|
|
303
|
+
|
|
304
|
+
## Current Claim Scope
|
|
305
|
+
|
|
306
|
+
Validated release claims are bounded by the [release scope](docs/release_scope.rst):
|
|
307
|
+
|
|
308
|
+
- Standard electrostatic/electromagnetic full gyrokinetics is validated only on
|
|
309
|
+
the promoted cases and observables in the verification matrix.
|
|
310
|
+
- Quasilinear outputs are diagnostics and screening models, not universal
|
|
311
|
+
absolute nonlinear heat-flux predictions.
|
|
312
|
+
- Nonlinear optimization evidence requires matched, replicated, long
|
|
313
|
+
post-transient windows; startup or reduced envelopes are not production evidence.
|
|
314
|
+
- W7-X zonal long-window recurrence/damping and W7-X TEM / kinetic-electron extensions are deferred.
|
|
315
|
+
- Production nonlinear domain decomposition, equilibrium ExB flow shear,
|
|
316
|
+
species-coupled collisions, and linearized Landau/Sugama operators remain open.
|
|
317
|
+
|
|
318
|
+
## Examples and Documentation
|
|
319
|
+
|
|
320
|
+
The repository keeps small runnable examples under:
|
|
321
|
+
|
|
322
|
+
- [`examples/linear`](examples/linear): axisymmetric and stellarator linear runs.
|
|
323
|
+
- [`examples/nonlinear`](examples/nonlinear): nonlinear turbulence and restarts.
|
|
324
|
+
- [`examples/optimization`](examples/optimization): differentiable QA workflows.
|
|
325
|
+
- [`examples/theory_and_demos`](examples/theory_and_demos): numerical and autodiff demonstrations.
|
|
326
|
+
- [`benchmarks`](benchmarks): comparison inputs, drivers, and compact result indexes.
|
|
327
|
+
|
|
328
|
+
Detailed user and developer documentation:
|
|
329
|
+
|
|
330
|
+
- [Physics and equations](docs/theory.rst)
|
|
331
|
+
- [Operators and models](docs/operators.rst)
|
|
332
|
+
- [Numerics and solvers](docs/numerics.rst)
|
|
333
|
+
- [Geometry](docs/geometry.rst)
|
|
334
|
+
- [Outputs and plotting](docs/outputs.rst)
|
|
335
|
+
- [Testing and validation](docs/testing.rst)
|
|
336
|
+
- [Code structure](docs/code_structure.rst)
|
|
337
|
+
- [Release and research scope](docs/release_scope.rst)
|
|
338
|
+
|
|
339
|
+
## Testing
|
|
340
|
+
|
|
341
|
+
```bash
|
|
342
|
+
pytest
|
|
343
|
+
python tools/release/run_test_gates.py fast
|
|
344
|
+
python tools/release/run_test_gates.py wide-coverage \
|
|
345
|
+
--shards 48 --timeout 300 --fail-under 95 \
|
|
346
|
+
--pytest-arg=-o --pytest-arg=addopts= --pytest-arg=-m --pytest-arg="not slow"
|
|
347
|
+
python -m sphinx -W -b html docs docs/_build/html
|
|
348
|
+
```
|
|
349
|
+
|
|
350
|
+
The package-wide CI coverage gate is at least 95%. Physics, convergence,
|
|
351
|
+
comparison, differentiability, and performance gates are required in addition
|
|
352
|
+
to line coverage.
|
|
353
|
+
|
|
354
|
+
## License
|
|
355
|
+
|
|
356
|
+
GKX is distributed under the [MIT License](LICENSE).
|
gkx-1.7.1/README.md
ADDED
|
@@ -0,0 +1,298 @@
|
|
|
1
|
+
# GKX
|
|
2
|
+
|
|
3
|
+
[](https://github.com/uwplasma/GKX/releases)
|
|
4
|
+
[](https://pypi.org/project/gkx/)
|
|
5
|
+
[](https://github.com/uwplasma/GKX/actions/workflows/ci.yml)
|
|
6
|
+
[](https://codecov.io/gh/uwplasma/GKX)
|
|
7
|
+
[](LICENSE)
|
|
8
|
+
[](pyproject.toml)
|
|
9
|
+
[](https://gkx.readthedocs.io)
|
|
10
|
+
|
|
11
|
+
GKX is a JAX-native gyrokinetic solver for linear stability,
|
|
12
|
+
nonlinear turbulence, differentiable analysis, and stellarator design. It uses
|
|
13
|
+
Fourier perpendicular coordinates, a Hermite-Laguerre velocity basis, and
|
|
14
|
+
field-aligned analytic, Miller, or VMEC geometry. The package runs on CPUs and
|
|
15
|
+
GPUs, exposes a Python API for autodiff and optimization, and provides a simple
|
|
16
|
+
executable for routine simulations.
|
|
17
|
+
|
|
18
|
+
## Installation
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
pip install gkx
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
For development:
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
git clone https://github.com/uwplasma/GKX
|
|
28
|
+
cd GKX
|
|
29
|
+
pip install -e .
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
## Quickstart
|
|
33
|
+
|
|
34
|
+
Run the built-in linear initial-value example:
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
gkx
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
The equivalent `gkx` entry point is also installed. The default run
|
|
41
|
+
prints setup, progress, elapsed time, and ETA, then writes its input, summary,
|
|
42
|
+
time series, eigenfunction, and a two-panel plot in the current directory.
|
|
43
|
+
|
|
44
|
+
Run a checked-in case or plot an existing result:
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
gkx examples/linear/axisymmetric/cyclone.toml
|
|
48
|
+
gkx run-runtime-nonlinear \
|
|
49
|
+
--config examples/nonlinear/axisymmetric/runtime_cyclone_nonlinear.toml \
|
|
50
|
+
--steps 200 --out cyclone.out.nc
|
|
51
|
+
gkx --plot cyclone.out.nc
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
Generate the small VMEC equilibria used by the self-contained examples:
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
pip install vmex
|
|
58
|
+
cd examples/vmec
|
|
59
|
+
./generate_wouts.sh
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
Full documentation is hosted at **[gkx.readthedocs.io](https://gkx.readthedocs.io)**.
|
|
63
|
+
Start with the [quickstart](https://gkx.readthedocs.io/en/latest/quickstart.html) and
|
|
64
|
+
[input reference](https://gkx.readthedocs.io/en/latest/inputs.html) for linear,
|
|
65
|
+
nonlinear, Miller, VMEC, restart, quasilinear, and plotting workflows.
|
|
66
|
+
|
|
67
|
+
## Highlights
|
|
68
|
+
|
|
69
|
+
- Electrostatic and electromagnetic gyrokinetics with kinetic or Boltzmann species.
|
|
70
|
+
- Linear initial-value, dominant-eigenmode, and nonlinear turbulence solvers.
|
|
71
|
+
- Analytic s-alpha, Miller, imported VMEC, and differentiable VMEC/Boozer geometry.
|
|
72
|
+
- JAX JIT, forward/reverse autodiff, implicit eigenvalue derivatives, and UQ tools.
|
|
73
|
+
- Quasilinear transport diagnostics with explicit saturation-rule metadata.
|
|
74
|
+
- CPU/GPU execution and production parallelization for independent scans and ensembles.
|
|
75
|
+
- Restartable NetCDF output and `gkx --plot` publication-style figures.
|
|
76
|
+
- A limited conserving Lenard-Bernstein/Dougherty-like collision model, with
|
|
77
|
+
advanced multispecies and linearized Landau operators remaining research lanes.
|
|
78
|
+
|
|
79
|
+
## Main Validation Results
|
|
80
|
+
|
|
81
|
+
The release atlas compares growth rates, frequencies, eigenfunctions, and
|
|
82
|
+
nonlinear transport windows with established gyrokinetic reference results.
|
|
83
|
+
Promoted cases include Cyclone ITG, Cyclone Miller, KBM, W7-X, and HSX, with
|
|
84
|
+
ETG and kinetic-electron stress cases kept at their documented claim level.
|
|
85
|
+
|
|
86
|
+

|
|
87
|
+
|
|
88
|
+
The exact equations, normalization, grids, boundary conditions, diagnostic
|
|
89
|
+
windows, tolerances, and artifact provenance are in the
|
|
90
|
+
[benchmark documentation](docs/benchmarks.rst) and
|
|
91
|
+
[verification matrix](docs/verification_matrix.rst). A visual overlay alone is
|
|
92
|
+
not treated as parity evidence.
|
|
93
|
+
|
|
94
|
+
The advanced-collision research lane now generates the complete retained
|
|
95
|
+
finite-Larmor Coulomb moment algebra and checks it against independent
|
|
96
|
+
velocity-space projection, spectral convergence, drift-kinetic conservation,
|
|
97
|
+
finite-wavelength gyro-diffusion, and the H-theorem. The panel below passes all
|
|
98
|
+
operator-level gates, but does not claim production Landau transport yet;
|
|
99
|
+
the finite-Larmor collisional-ITG gate and drift-kinetic collisional-zonal gate
|
|
100
|
+
are closed, and the paper-resolution finite-wavelength zonal traces and
|
|
101
|
+
velocity sections now pass their independent literature gate.
|
|
102
|
+
The differentiable driven-current solve is equation-tested, and the direct
|
|
103
|
+
Coulomb hierarchy is now converged through ``(P,J)=(20,5)`` with a maximum
|
|
104
|
+
nested current change of 0.017%. Its collision-frequency and electric-field
|
|
105
|
+
normalizations are now closed against the stationary Spitzer problem: all
|
|
106
|
+
three operators saturate by ``t nu_ee=50``, remain linear over a 100x field
|
|
107
|
+
scan, and Coulomb reaches the high-charge analytic limit within 7.46%. The
|
|
108
|
+
arbitrary-order original model reproduces its published 11% low-charge current
|
|
109
|
+
deficit, while the converged improved model is within 0.31% of Coulomb over the
|
|
110
|
+
complete ion-charge scan. The converged finite-Larmor ITG artifact is accepted;
|
|
111
|
+
the generated equal-species operator remains a Python research path because a
|
|
112
|
+
production input-file selector would imply broader multispecies coefficients
|
|
113
|
+
that are not yet implemented.
|
|
114
|
+
|
|
115
|
+

|
|
116
|
+
|
|
117
|
+
Equations, thresholds, machine-readable results, literature links, and the
|
|
118
|
+
one-command reproduction recipe are in the [collision-operator
|
|
119
|
+
documentation](docs/operators.rst).
|
|
120
|
+
|
|
121
|
+

|
|
122
|
+
|
|
123
|
+
At ``(P,J)=(24,10)``, the drift-kinetic traces approach the Xiao residual and
|
|
124
|
+
the finite-wavelength tails reproduce the published original < improved <
|
|
125
|
+
Coulomb ordering at both ``kx rho_i=0.1`` and ``0.2``. The improved model is
|
|
126
|
+
also closer to Coulomb over ``t nu_ii <= 10`` at both wavenumbers. Equations,
|
|
127
|
+
velocity-space convergence, compact replay data, and the Figure 12--14 gate
|
|
128
|
+
are documented in [Operators and Terms](docs/operators.rst).
|
|
129
|
+
|
|
130
|
+
## Runtime and Memory
|
|
131
|
+
|
|
132
|
+

|
|
133
|
+
|
|
134
|
+
The panel reports measured cold wall time and peak memory for the tracked CPU,
|
|
135
|
+
GPU, and comparison-code runs. Cold JAX rows include startup and compilation.
|
|
136
|
+
Prepared Python simulations avoid recompiling a fixed geometry and numerical
|
|
137
|
+
policy, but their CPU/GPU throughput depends on the software stack and GPU
|
|
138
|
+
operating state. See
|
|
139
|
+
[performance](docs/performance.rst) for profiler artifacts, memory accounting,
|
|
140
|
+
current reproducibility notes, and the distinction between executable,
|
|
141
|
+
prepared, and distributed runs.
|
|
142
|
+
|
|
143
|
+
## Differentiable Python API
|
|
144
|
+
|
|
145
|
+
```python
|
|
146
|
+
import jax.numpy as jnp
|
|
147
|
+
|
|
148
|
+
from gkx import CycloneBaseCase, LinearParams, integrate_linear_from_config
|
|
149
|
+
from gkx.core.grid import build_spectral_grid
|
|
150
|
+
from gkx.geometry import SAlphaGeometry
|
|
151
|
+
|
|
152
|
+
cfg = CycloneBaseCase()
|
|
153
|
+
grid = build_spectral_grid(cfg.grid)
|
|
154
|
+
geometry = SAlphaGeometry.from_config(cfg.geometry)
|
|
155
|
+
parameters = LinearParams()
|
|
156
|
+
state = jnp.zeros((2, 2, grid.ky.size, grid.kx.size, grid.z.size), dtype=jnp.complex64)
|
|
157
|
+
state = state.at[0, 0, 0, 0, :].set(1.0e-3)
|
|
158
|
+
trajectory, potential = integrate_linear_from_config(
|
|
159
|
+
state, grid, geometry, parameters, cfg.time
|
|
160
|
+
)
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
For repeated nonlinear calls with fixed geometry and numerical policy, prepare
|
|
164
|
+
the compiled simulation once:
|
|
165
|
+
|
|
166
|
+
```python
|
|
167
|
+
from gkx.solvers.nonlinear.diagnostic_integration import prepare_nonlinear_explicit_diagnostics
|
|
168
|
+
|
|
169
|
+
simulation = prepare_nonlinear_explicit_diagnostics(
|
|
170
|
+
initial_state,
|
|
171
|
+
grid,
|
|
172
|
+
geometry,
|
|
173
|
+
parameters,
|
|
174
|
+
dt=0.02,
|
|
175
|
+
steps=400,
|
|
176
|
+
resolved_diagnostics=False,
|
|
177
|
+
)
|
|
178
|
+
time, diagnostics, final_state, fields = simulation.run()
|
|
179
|
+
```
|
|
180
|
+
|
|
181
|
+
The prepared object accepts another same-shape initial state without rebuilding
|
|
182
|
+
the scan. A matched rebuilt cache/parameter PyTree can also remain dynamic for
|
|
183
|
+
autodiff; geometry layout is fixed, and dynamic-geometry compile reuse remains
|
|
184
|
+
an active differentiability lane.
|
|
185
|
+
|
|
186
|
+
The planted two-mode inverse problem below recovers two gradient parameters and
|
|
187
|
+
checks the autodiff Jacobian against finite differences. The single-mode demo in
|
|
188
|
+
the docs intentionally demonstrates non-identifiability rather than exact
|
|
189
|
+
parameter recovery.
|
|
190
|
+
|
|
191
|
+

|
|
192
|
+
|
|
193
|
+
See [differentiable geometry](docs/geometry.rst),
|
|
194
|
+
[algorithms](docs/algorithms.rst), and [stellarator optimization](docs/stellarator_optimization.rst)
|
|
195
|
+
for JVP, VJP, implicit differentiation, conditioning, covariance, and finite-difference gates.
|
|
196
|
+
|
|
197
|
+
## Quasilinear Modeling
|
|
198
|
+
|
|
199
|
+

|
|
200
|
+
|
|
201
|
+
The current quasilinear implementation is a scoped model-development and
|
|
202
|
+
optimization-screening result. It supports ranking and correlation studies but
|
|
203
|
+
is not a runtime/TOML absolute-flux predictor. Absolute-flux
|
|
204
|
+
promotion remains rejected when the declared Solovev and shaped-pressure stress
|
|
205
|
+
outliers are retained. Model definitions, derivations, calibration splits,
|
|
206
|
+
uncertainty, residual anatomy, and holdout gates are in the
|
|
207
|
+
[quasilinear documentation](docs/quasilinear.rst).
|
|
208
|
+
|
|
209
|
+
## QA ITG Optimization
|
|
210
|
+
|
|
211
|
+
The VMEX-style examples append a GKX growth-rate, quasilinear, or
|
|
212
|
+
nonlinear-window residual to the aspect-ratio, mean-iota, and quasisymmetry
|
|
213
|
+
objective tuples. The baseline follows the max-mode-5 QA workflow; all
|
|
214
|
+
transport comparisons use solved VMEC equilibria.
|
|
215
|
+
|
|
216
|
+

|
|
217
|
+
|
|
218
|
+
These rows are not promoted turbulent-flux designs. Their matched long
|
|
219
|
+
post-transient nonlinear audits use converged post-transient heat-flux windows
|
|
220
|
+
and do not show a statistically significant reduction relative to the strict
|
|
221
|
+
QA baseline. They are useful negative transfer evidence for improving objective
|
|
222
|
+
conditioning and optimizer choice.
|
|
223
|
+
|
|
224
|
+
The RBC(1,1) scan is a landscape and noise/convergence diagnostic, not a source
|
|
225
|
+
of admitted optimized candidates. It compares linear growth, all shipped
|
|
226
|
+
quasilinear rules, and replicated long-window nonlinear transport.
|
|
227
|
+
|
|
228
|
+

|
|
229
|
+
|
|
230
|
+
Reproducible scripts are in [examples/optimization](examples/optimization), and
|
|
231
|
+
full objective equations, optimizer policies, comparison fingerprints, and
|
|
232
|
+
long-window audits are in the [optimization documentation](docs/stellarator_optimization.rst).
|
|
233
|
+
|
|
234
|
+
## Parallelization
|
|
235
|
+
|
|
236
|
+
Production parallelization currently covers independent `k_y` scans,
|
|
237
|
+
quasilinear/UQ ensembles, and file-backed independent tasks with deterministic
|
|
238
|
+
ordering and serial identity gates. Sensitivity sweeps can use the same deterministic independent-work reconstruction, but they need a dedicated
|
|
239
|
+
matched scaling artifact before any speedup claim is promoted.
|
|
240
|
+
|
|
241
|
+
Nonlinear whole-state and domain decomposition remain diagnostic. Species-first
|
|
242
|
+
and Hermite-second decomposition, explicit Hermite halo exchange, field-moment
|
|
243
|
+
collectives, and physical transport-window identity must pass before a
|
|
244
|
+
nonlinear parallelization speedup is claimed. See [parallelization](docs/parallelization.rst).
|
|
245
|
+
|
|
246
|
+
## Current Claim Scope
|
|
247
|
+
|
|
248
|
+
Validated release claims are bounded by the [release scope](docs/release_scope.rst):
|
|
249
|
+
|
|
250
|
+
- Standard electrostatic/electromagnetic full gyrokinetics is validated only on
|
|
251
|
+
the promoted cases and observables in the verification matrix.
|
|
252
|
+
- Quasilinear outputs are diagnostics and screening models, not universal
|
|
253
|
+
absolute nonlinear heat-flux predictions.
|
|
254
|
+
- Nonlinear optimization evidence requires matched, replicated, long
|
|
255
|
+
post-transient windows; startup or reduced envelopes are not production evidence.
|
|
256
|
+
- W7-X zonal long-window recurrence/damping and W7-X TEM / kinetic-electron extensions are deferred.
|
|
257
|
+
- Production nonlinear domain decomposition, equilibrium ExB flow shear,
|
|
258
|
+
species-coupled collisions, and linearized Landau/Sugama operators remain open.
|
|
259
|
+
|
|
260
|
+
## Examples and Documentation
|
|
261
|
+
|
|
262
|
+
The repository keeps small runnable examples under:
|
|
263
|
+
|
|
264
|
+
- [`examples/linear`](examples/linear): axisymmetric and stellarator linear runs.
|
|
265
|
+
- [`examples/nonlinear`](examples/nonlinear): nonlinear turbulence and restarts.
|
|
266
|
+
- [`examples/optimization`](examples/optimization): differentiable QA workflows.
|
|
267
|
+
- [`examples/theory_and_demos`](examples/theory_and_demos): numerical and autodiff demonstrations.
|
|
268
|
+
- [`benchmarks`](benchmarks): comparison inputs, drivers, and compact result indexes.
|
|
269
|
+
|
|
270
|
+
Detailed user and developer documentation:
|
|
271
|
+
|
|
272
|
+
- [Physics and equations](docs/theory.rst)
|
|
273
|
+
- [Operators and models](docs/operators.rst)
|
|
274
|
+
- [Numerics and solvers](docs/numerics.rst)
|
|
275
|
+
- [Geometry](docs/geometry.rst)
|
|
276
|
+
- [Outputs and plotting](docs/outputs.rst)
|
|
277
|
+
- [Testing and validation](docs/testing.rst)
|
|
278
|
+
- [Code structure](docs/code_structure.rst)
|
|
279
|
+
- [Release and research scope](docs/release_scope.rst)
|
|
280
|
+
|
|
281
|
+
## Testing
|
|
282
|
+
|
|
283
|
+
```bash
|
|
284
|
+
pytest
|
|
285
|
+
python tools/release/run_test_gates.py fast
|
|
286
|
+
python tools/release/run_test_gates.py wide-coverage \
|
|
287
|
+
--shards 48 --timeout 300 --fail-under 95 \
|
|
288
|
+
--pytest-arg=-o --pytest-arg=addopts= --pytest-arg=-m --pytest-arg="not slow"
|
|
289
|
+
python -m sphinx -W -b html docs docs/_build/html
|
|
290
|
+
```
|
|
291
|
+
|
|
292
|
+
The package-wide CI coverage gate is at least 95%. Physics, convergence,
|
|
293
|
+
comparison, differentiability, and performance gates are required in addition
|
|
294
|
+
to line coverage.
|
|
295
|
+
|
|
296
|
+
## License
|
|
297
|
+
|
|
298
|
+
GKX is distributed under the [MIT License](LICENSE).
|