drbx 2.0.0__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.
- drbx-2.0.0/LICENSE +21 -0
- drbx-2.0.0/PKG-INFO +395 -0
- drbx-2.0.0/README.md +366 -0
- drbx-2.0.0/pyproject.toml +57 -0
- drbx-2.0.0/setup.cfg +4 -0
- drbx-2.0.0/src/drbx/__init__.py +20 -0
- drbx-2.0.0/src/drbx/__main__.py +5 -0
- drbx-2.0.0/src/drbx/cli.py +586 -0
- drbx-2.0.0/src/drbx/config/__init__.py +20 -0
- drbx-2.0.0/src/drbx/config/boutinp.py +556 -0
- drbx-2.0.0/src/drbx/config/model.py +34 -0
- drbx-2.0.0/src/drbx/config/normalization.py +66 -0
- drbx-2.0.0/src/drbx/data/atomic_rates/__init__.py +1 -0
- drbx-2.0.0/src/drbx/data/atomic_rates/iz_AMJUEL_H.x_2.1.5.json +213 -0
- drbx-2.0.0/src/drbx/data/atomic_rates/iz_AMJUEL_H.x_2.3.9a.json +213 -0
- drbx-2.0.0/src/drbx/data/atomic_rates/rec_AMJUEL_H.x_2.1.8.json +213 -0
- drbx-2.0.0/src/drbx/data/atomic_rates/rec_AMJUEL_H.x_2.3.13a.json +213 -0
- drbx-2.0.0/src/drbx/geometry/__init__.py +209 -0
- drbx-2.0.0/src/drbx/geometry/embedding.py +56 -0
- drbx-2.0.0/src/drbx/geometry/essos_import.py +1650 -0
- drbx-2.0.0/src/drbx/geometry/fci_geometry.py +4622 -0
- drbx-2.0.0/src/drbx/geometry/fci_maps.py +85 -0
- drbx-2.0.0/src/drbx/geometry/island_divertor.py +291 -0
- drbx-2.0.0/src/drbx/geometry/metric_tensor.py +99 -0
- drbx-2.0.0/src/drbx/geometry/open_slab.py +150 -0
- drbx-2.0.0/src/drbx/geometry/rotating_ellipse.py +253 -0
- drbx-2.0.0/src/drbx/geometry/shifted_torus.py +225 -0
- drbx-2.0.0/src/drbx/geometry/stellarator.py +287 -0
- drbx-2.0.0/src/drbx/geometry/vmec_extender_import.py +499 -0
- drbx-2.0.0/src/drbx/geometry/vmex_import.py +321 -0
- drbx-2.0.0/src/drbx/linear/__init__.py +37 -0
- drbx-2.0.0/src/drbx/linear/dispersion.py +138 -0
- drbx-2.0.0/src/drbx/linear/eigen.py +91 -0
- drbx-2.0.0/src/drbx/native/__init__.py +224 -0
- drbx-2.0.0/src/drbx/native/array_backend.py +64 -0
- drbx-2.0.0/src/drbx/native/deck_runner.py +779 -0
- drbx-2.0.0/src/drbx/native/electromagnetic.py +250 -0
- drbx-2.0.0/src/drbx/native/expression.py +173 -0
- drbx-2.0.0/src/drbx/native/fci.py +295 -0
- drbx-2.0.0/src/drbx/native/fci_2_field_rhs.py +182 -0
- drbx-2.0.0/src/drbx/native/fci_4_field_rhs.py +1267 -0
- drbx-2.0.0/src/drbx/native/fci_boundaries.py +2494 -0
- drbx-2.0.0/src/drbx/native/fci_differentiable_case.py +304 -0
- drbx-2.0.0/src/drbx/native/fci_drb_EB_rhs.py +1243 -0
- drbx-2.0.0/src/drbx/native/fci_drb_rhs.py +190 -0
- drbx-2.0.0/src/drbx/native/fci_halo.py +1575 -0
- drbx-2.0.0/src/drbx/native/fci_helpers.py +350 -0
- drbx-2.0.0/src/drbx/native/fci_model.py +294 -0
- drbx-2.0.0/src/drbx/native/fci_neutral.py +139 -0
- drbx-2.0.0/src/drbx/native/fci_operators.py +4081 -0
- drbx-2.0.0/src/drbx/native/fci_sharding.py +597 -0
- drbx-2.0.0/src/drbx/native/fci_sheath_recycling.py +206 -0
- drbx-2.0.0/src/drbx/native/fci_time_integrator.py +96 -0
- drbx-2.0.0/src/drbx/native/fci_vorticity.py +198 -0
- drbx-2.0.0/src/drbx/native/fluid_1d.py +330 -0
- drbx-2.0.0/src/drbx/native/hasegawa_wakatani.py +196 -0
- drbx-2.0.0/src/drbx/native/limiters.py +57 -0
- drbx-2.0.0/src/drbx/native/mesh.py +238 -0
- drbx-2.0.0/src/drbx/native/metrics.py +234 -0
- drbx-2.0.0/src/drbx/native/neutrals/__init__.py +58 -0
- drbx-2.0.0/src/drbx/native/neutrals/atomic_rates.py +134 -0
- drbx-2.0.0/src/drbx/native/neutrals/detachment_sol_model.py +221 -0
- drbx-2.0.0/src/drbx/native/neutrals/reactions.py +164 -0
- drbx-2.0.0/src/drbx/native/neutrals/recycling_sol_model.py +197 -0
- drbx-2.0.0/src/drbx/native/sol_flux_tube.py +133 -0
- drbx-2.0.0/src/drbx/native/stellarator_turbulence.py +343 -0
- drbx-2.0.0/src/drbx/native/transport.py +134 -0
- drbx-2.0.0/src/drbx/native/units.py +32 -0
- drbx-2.0.0/src/drbx/native/vorticity.py +252 -0
- drbx-2.0.0/src/drbx/runtime/__init__.py +53 -0
- drbx-2.0.0/src/drbx/runtime/artifacts.py +161 -0
- drbx-2.0.0/src/drbx/runtime/memory.py +144 -0
- drbx-2.0.0/src/drbx/runtime/output.py +374 -0
- drbx-2.0.0/src/drbx/runtime/paths.py +9 -0
- drbx-2.0.0/src/drbx/runtime/performance.py +161 -0
- drbx-2.0.0/src/drbx/runtime/run_config.py +184 -0
- drbx-2.0.0/src/drbx/runtime/scheduler.py +99 -0
- drbx-2.0.0/src/drbx/runtime/state.py +40 -0
- drbx-2.0.0/src/drbx/validation/__init__.py +424 -0
- drbx-2.0.0/src/drbx/validation/autodiff_diffusion.py +329 -0
- drbx-2.0.0/src/drbx/validation/autodiff_diffusion_uncertainty.py +235 -0
- drbx-2.0.0/src/drbx/validation/diverted_tokamak_movie.py +558 -0
- drbx-2.0.0/src/drbx/validation/essos_fieldline_import_campaign.py +181 -0
- drbx-2.0.0/src/drbx/validation/essos_imported_artifact_audit.py +535 -0
- drbx-2.0.0/src/drbx/validation/essos_imported_drb_movie_campaign.py +2826 -0
- drbx-2.0.0/src/drbx/validation/essos_imported_fci_campaign.py +5241 -0
- drbx-2.0.0/src/drbx/validation/essos_imported_pytree_campaign.py +406 -0
- drbx-2.0.0/src/drbx/validation/essos_vmec_closed_field_campaign.py +314 -0
- drbx-2.0.0/src/drbx/validation/essos_vmec_closed_field_transient_campaign.py +629 -0
- drbx-2.0.0/src/drbx/validation/essos_vmec_fieldline_surface_campaign.py +620 -0
- drbx-2.0.0/src/drbx/validation/fluid_1d_mms_convergence.py +250 -0
- drbx-2.0.0/src/drbx/validation/geometry_lineouts.py +136 -0
- drbx-2.0.0/src/drbx/validation/geometry_slices.py +178 -0
- drbx-2.0.0/src/drbx/validation/publication_plotting.py +91 -0
- drbx-2.0.0/src/drbx/validation/stellarator_drb_pytree_campaign.py +621 -0
- drbx-2.0.0/src/drbx/validation/stellarator_fci_geometry_campaign.py +200 -0
- drbx-2.0.0/src/drbx/validation/stellarator_fci_operator_campaign.py +304 -0
- drbx-2.0.0/src/drbx/validation/stellarator_fci_suite_campaign.py +264 -0
- drbx-2.0.0/src/drbx/validation/stellarator_metric_mms_campaign.py +289 -0
- drbx-2.0.0/src/drbx/validation/stellarator_neutral_physics_campaign.py +255 -0
- drbx-2.0.0/src/drbx/validation/stellarator_sheath_recycling_campaign.py +331 -0
- drbx-2.0.0/src/drbx/validation/stellarator_sol_showcase.py +628 -0
- drbx-2.0.0/src/drbx/validation/stellarator_vorticity_campaign.py +304 -0
- drbx-2.0.0/src/drbx/validation/vmec_extender_edge_field_campaign.py +260 -0
- drbx-2.0.0/src/drbx/validation/vmec_extender_sol_smoke_campaign.py +365 -0
- drbx-2.0.0/src/drbx.egg-info/PKG-INFO +395 -0
- drbx-2.0.0/src/drbx.egg-info/SOURCES.txt +174 -0
- drbx-2.0.0/src/drbx.egg-info/dependency_links.txt +1 -0
- drbx-2.0.0/src/drbx.egg-info/entry_points.txt +2 -0
- drbx-2.0.0/src/drbx.egg-info/requires.txt +19 -0
- drbx-2.0.0/src/drbx.egg-info/top_level.txt +1 -0
- drbx-2.0.0/tests/test_autodiff_methods.py +54 -0
- drbx-2.0.0/tests/test_boutinp.py +186 -0
- drbx-2.0.0/tests/test_cli_command_surface.py +73 -0
- drbx-2.0.0/tests/test_cli_run.py +541 -0
- drbx-2.0.0/tests/test_detachment_control.py +58 -0
- drbx-2.0.0/tests/test_diffusion_precision_benchmark.py +39 -0
- drbx-2.0.0/tests/test_essos_fieldline_import.py +2360 -0
- drbx-2.0.0/tests/test_fast_research_checks.py +99 -0
- drbx-2.0.0/tests/test_fci_differentiable.py +98 -0
- drbx-2.0.0/tests/test_fci_halo.py +649 -0
- drbx-2.0.0/tests/test_fci_model_state.py +180 -0
- drbx-2.0.0/tests/test_fci_neutrals_3d.py +98 -0
- drbx-2.0.0/tests/test_fci_operators.py +2755 -0
- drbx-2.0.0/tests/test_fci_operators_domain_decomp.py +1236 -0
- drbx-2.0.0/tests/test_fci_sharded_2field.py +64 -0
- drbx-2.0.0/tests/test_geometry_fci_maps.py +336 -0
- drbx-2.0.0/tests/test_hasegawa_wakatani.py +359 -0
- drbx-2.0.0/tests/test_island_divertor.py +95 -0
- drbx-2.0.0/tests/test_landreman_paul_turbulence.py +129 -0
- drbx-2.0.0/tests/test_linear_dispersion.py +246 -0
- drbx-2.0.0/tests/test_mms_convergence.py +81 -0
- drbx-2.0.0/tests/test_mms_shifted_torus_2_field.py +764 -0
- drbx-2.0.0/tests/test_mms_shifted_torus_4_field.py +2050 -0
- drbx-2.0.0/tests/test_mms_shifted_torus_EB.py +1238 -0
- drbx-2.0.0/tests/test_mms_slab_2_field.py +916 -0
- drbx-2.0.0/tests/test_multigrid_preconditioner.py +537 -0
- drbx-2.0.0/tests/test_native_array_backend.py +26 -0
- drbx-2.0.0/tests/test_native_atomic_rates.py +84 -0
- drbx-2.0.0/tests/test_native_detachment_sol.py +83 -0
- drbx-2.0.0/tests/test_native_electromagnetic.py +286 -0
- drbx-2.0.0/tests/test_native_expression.py +78 -0
- drbx-2.0.0/tests/test_native_fluid_1d.py +120 -0
- drbx-2.0.0/tests/test_native_limiters.py +71 -0
- drbx-2.0.0/tests/test_native_mesh.py +333 -0
- drbx-2.0.0/tests/test_native_metrics.py +251 -0
- drbx-2.0.0/tests/test_native_reactions.py +93 -0
- drbx-2.0.0/tests/test_native_recycling_sol.py +99 -0
- drbx-2.0.0/tests/test_native_transport.py +83 -0
- drbx-2.0.0/tests/test_native_vorticity.py +122 -0
- drbx-2.0.0/tests/test_normalization.py +37 -0
- drbx-2.0.0/tests/test_open_field_line_sol.py +108 -0
- drbx-2.0.0/tests/test_packaging_metadata.py +78 -0
- drbx-2.0.0/tests/test_restartable_diffusion_tutorial.py +51 -0
- drbx-2.0.0/tests/test_rotating_ellipse_fci.py +245 -0
- drbx-2.0.0/tests/test_rotating_ellipse_filament.py +98 -0
- drbx-2.0.0/tests/test_run_config.py +69 -0
- drbx-2.0.0/tests/test_runtime_artifacts.py +152 -0
- drbx-2.0.0/tests/test_runtime_memory.py +34 -0
- drbx-2.0.0/tests/test_runtime_performance.py +63 -0
- drbx-2.0.0/tests/test_runtime_precision.py +40 -0
- drbx-2.0.0/tests/test_scheduler.py +65 -0
- drbx-2.0.0/tests/test_shifted_torus_4_field_blob.py +435 -0
- drbx-2.0.0/tests/test_shifted_torus_4_field_free_decay.py +1099 -0
- drbx-2.0.0/tests/test_shifted_torus_EB_blob.py +2962 -0
- drbx-2.0.0/tests/test_state.py +15 -0
- drbx-2.0.0/tests/test_stellarator_turbulence.py +64 -0
- drbx-2.0.0/tests/test_validation_autodiff_diffusion.py +61 -0
- drbx-2.0.0/tests/test_validation_autodiff_diffusion_uncertainty.py +27 -0
- drbx-2.0.0/tests/test_validation_diverted_tokamak_movie.py +165 -0
- drbx-2.0.0/tests/test_validation_fluid_1d_mms_convergence.py +64 -0
- drbx-2.0.0/tests/test_validation_stellarator_fci_campaigns.py +1170 -0
- drbx-2.0.0/tests/test_validation_vmec_extender_edge_field_campaign.py +81 -0
- drbx-2.0.0/tests/test_validation_vmec_extender_sol_smoke_campaign.py +81 -0
- drbx-2.0.0/tests/test_vmec_extender_import.py +252 -0
- drbx-2.0.0/tests/test_vmex_import.py +99 -0
drbx-2.0.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Rogerio Jorge
|
|
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.
|
drbx-2.0.0/PKG-INFO
ADDED
|
@@ -0,0 +1,395 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: drbx
|
|
3
|
+
Version: 2.0.0
|
|
4
|
+
Summary: JAX-first edge and scrape-off-layer plasma code with differentiable, hardware-agnostic solver infrastructure.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Classifier: Programming Language :: Python :: 3
|
|
7
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
8
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: jax
|
|
14
|
+
Requires-Dist: scipy
|
|
15
|
+
Requires-Dist: matplotlib
|
|
16
|
+
Requires-Dist: netCDF4
|
|
17
|
+
Requires-Dist: rich
|
|
18
|
+
Requires-Dist: pillow
|
|
19
|
+
Requires-Dist: solvax>=0.8.6
|
|
20
|
+
Requires-Dist: tomli; python_version < "3.11"
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: pytest; extra == "dev"
|
|
23
|
+
Requires-Dist: pytest-cov; extra == "dev"
|
|
24
|
+
Requires-Dist: pytest-xdist; extra == "dev"
|
|
25
|
+
Provides-Extra: docs
|
|
26
|
+
Requires-Dist: mkdocs; extra == "docs"
|
|
27
|
+
Requires-Dist: pymdown-extensions; extra == "docs"
|
|
28
|
+
Dynamic: license-file
|
|
29
|
+
|
|
30
|
+
# DRBX
|
|
31
|
+
|
|
32
|
+
[](https://github.com/uwplasma/drbx/actions/workflows/test.yml)
|
|
33
|
+
[](https://github.com/uwplasma/drbx/actions/workflows/docs.yml)
|
|
34
|
+
[](https://github.com/uwplasma/drbx/actions/workflows/coverage.yml)
|
|
35
|
+
[](https://pypi.org/project/drbx/)
|
|
36
|
+
[](https://pypi.org/project/drbx/)
|
|
37
|
+
[](LICENSE)
|
|
38
|
+
[](https://jax.readthedocs.io/)
|
|
39
|
+
[](https://drbx.readthedocs.io/)
|
|
40
|
+
|
|
41
|
+
**DRBX is a JAX-based, end-to-end differentiable drift-reduced Braginskii
|
|
42
|
+
(DRB) code for edge and scrape-off-layer (SOL) plasma turbulence** — on both
|
|
43
|
+
closed and open field lines, in axisymmetric (tokamak) and non-axisymmetric
|
|
44
|
+
(stellarator) geometry via the flux-coordinate-independent (FCI) approach.
|
|
45
|
+
|
|
46
|
+
Because the whole model is written in JAX, every simulation is `jit`-compiled,
|
|
47
|
+
runs on CPU or GPU unchanged, and is differentiable: you can take gradients of
|
|
48
|
+
any output (a saturated fluctuation energy, a transport level) with respect to
|
|
49
|
+
any input (a density gradient, an adiabaticity, a diffusivity) through the
|
|
50
|
+
solver. To our knowledge no other published DRB SOL turbulence code is
|
|
51
|
+
differentiable, and none combines differentiability with FCI stellarator
|
|
52
|
+
geometry.
|
|
53
|
+
|
|
54
|
+
Documentation: [drbx.readthedocs.io](https://drbx.readthedocs.io/en/latest/).
|
|
55
|
+
|
|
56
|
+
## Stellarator turbulence in three dimensions
|
|
57
|
+
|
|
58
|
+
Four-field drift-reduced turbulence on a rotating-ellipse stellarator — a torus
|
|
59
|
+
whose elliptical cross-section rotates with the toroidal angle. The cutaway
|
|
60
|
+
shows the density fluctuations on a flux surface and through the interior;
|
|
61
|
+
every frame is a `jit`-compiled, differentiable JAX step:
|
|
62
|
+
|
|
63
|
+

|
|
64
|
+
|
|
65
|
+
*Reproduce with [`examples/stellarator/stellarator_3d_render.py`](examples/stellarator/stellarator_3d_render.py).*
|
|
66
|
+
|
|
67
|
+
The same geometry supports closed and open field lines: core field lines (blue)
|
|
68
|
+
stay on flux surfaces, while beyond a toroidal limiter the scrape-off-layer
|
|
69
|
+
field lines (red) end on the limiter plate, where a Bohm sheath drains the
|
|
70
|
+
plasma:
|
|
71
|
+
|
|
72
|
+

|
|
73
|
+
|
|
74
|
+
*Reproduce with [`examples/stellarator/stellarator_3d_render.py`](examples/stellarator/stellarator_3d_render.py).*
|
|
75
|
+
|
|
76
|
+
## Install
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install drbx # from PyPI
|
|
80
|
+
# or, from source:
|
|
81
|
+
git clone https://github.com/uwplasma/drbx && cd drbx && pip install -e .
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
Runtime dependencies are `jax`, `scipy`, `matplotlib`, `netCDF4`, `rich`,
|
|
85
|
+
`pillow`, and [`solvax`](https://github.com/uwplasma/SOLVAX). Python 3.10-3.12.
|
|
86
|
+
|
|
87
|
+
## Quick start
|
|
88
|
+
|
|
89
|
+
Run a simulation from a TOML deck, or inspect one without running it:
|
|
90
|
+
|
|
91
|
+
```bash
|
|
92
|
+
drbx inspect examples/inputs/restartable_diffusion.toml # resolve and print the plan
|
|
93
|
+
drbx run examples/inputs/restartable_diffusion.toml # run and write artifacts
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
From Python, a differentiable turbulence run is a few lines:
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
import jax.numpy as jnp
|
|
100
|
+
import numpy as np
|
|
101
|
+
from drbx.native.hasegawa_wakatani import HasegawaWakataniParameters, hw_grid, hw_run
|
|
102
|
+
|
|
103
|
+
grid = hw_grid(64, 2 * jnp.pi * 8)
|
|
104
|
+
params = HasegawaWakataniParameters(adiabaticity=1.0, gradient=1.0)
|
|
105
|
+
rng = np.random.default_rng(0)
|
|
106
|
+
zeta0 = jnp.fft.fft2(jnp.asarray(1e-2 * rng.standard_normal((64, 64))))
|
|
107
|
+
n0 = jnp.fft.fft2(jnp.asarray(1e-2 * rng.standard_normal((64, 64))))
|
|
108
|
+
zeta, n = hw_run(zeta0, n0, grid, params, dt=5e-3, steps=500) # jit-compiled, differentiable
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
Every example below is a flat script: parameters at the top, run, plot.
|
|
112
|
+
|
|
113
|
+
## Highlights
|
|
114
|
+
|
|
115
|
+
**Turbulence on closed and open field lines.** The same multi-mode seed on the
|
|
116
|
+
rotating-ellipse stellarator, with all field lines closed (top) and with a
|
|
117
|
+
limiter opening the outer flux surfaces into a sheath-drained scrape-off layer
|
|
118
|
+
(bottom). Four toroidal cross-sections; the mode pattern differs plane by plane
|
|
119
|
+
because the flux surfaces rotate:
|
|
120
|
+
|
|
121
|
+

|
|
122
|
+
|
|
123
|
+

|
|
124
|
+
|
|
125
|
+
*Reproduce with [`examples/stellarator/stellarator_turbulence.py`](examples/stellarator/stellarator_turbulence.py).*
|
|
126
|
+
|
|
127
|
+
**Island divertor.** A sheared rotational transform with resonant perturbations
|
|
128
|
+
forms island chains and a stochastic edge. The open scrape-off layer emerges
|
|
129
|
+
from the field itself: multi-transit field-line tracing marks the finite
|
|
130
|
+
connection-length region, and the turbulence drains through it:
|
|
131
|
+
|
|
132
|
+

|
|
133
|
+
|
|
134
|
+
*Reproduce with [`examples/stellarator/island_divertor.py`](examples/stellarator/island_divertor.py).*
|
|
135
|
+
|
|
136
|
+
**Imported fields: real coils and VMEC equilibria.** The same closed/open
|
|
137
|
+
field-line machinery runs on imported fields: the vacuum Biot-Savart field of
|
|
138
|
+
the Landreman-Paul quasi-axisymmetric coil set (via ESSOS) shows nested closed
|
|
139
|
+
surfaces inside a chaotic open edge, and a VMEC equilibrium (via VMEX)
|
|
140
|
+
provides closed surfaces whose traced rotational transform matches the
|
|
141
|
+
equilibrium's `iotaf` profile to ~1e-6 — with the open field lines confined to
|
|
142
|
+
the thin island/stochastic layer just beyond the last closed vacuum surface,
|
|
143
|
+
in the island-divertor sense:
|
|
144
|
+
|
|
145
|
+
| Coil field (closed core, open edge) | VMEC equilibrium + coil-field SOL |
|
|
146
|
+
|---|---|
|
|
147
|
+
|  |  |
|
|
148
|
+
|
|
149
|
+
*Reproduce with [`examples/geometry-3D/essos-field-lines/closed_open_vacuum_poincare.py`](examples/geometry-3D/essos-field-lines/closed_open_vacuum_poincare.py) (left) and [`examples/geometry-3D/vmex/closed_open_field_lines.py`](examples/geometry-3D/vmex/closed_open_field_lines.py) (right).*
|
|
150
|
+
|
|
151
|
+
And the turbulence runs on that same imported field: a four-field interchange
|
|
152
|
+
simulation on the Landreman-Paul VMEC equilibrium (real metric, `|B|`, and
|
|
153
|
+
surface-preserving parallel maps; recovered iota ~ 0.42) with a limiter opening
|
|
154
|
+
a scrape-off layer outside the last closed surface — the closed core is
|
|
155
|
+
retained while the open SOL drains through the Bohm sheath:
|
|
156
|
+
|
|
157
|
+

|
|
158
|
+
|
|
159
|
+
*Reproduce with [`examples/stellarator/landreman_paul_turbulence.py`](examples/stellarator/landreman_paul_turbulence.py).*
|
|
160
|
+
|
|
161
|
+
**Neutrals and detachment.** The open SOL flux tube reaches the two-point Bohm
|
|
162
|
+
steady state; the hermes-3 neutral model (packaged AMJUEL atomic rates, target
|
|
163
|
+
recycling) builds the neutral cushion; and with an evolved temperature the SOL
|
|
164
|
+
detaches — the target cools through 1 eV and the target ion flux rolls over
|
|
165
|
+
(the SD1D benchmark):
|
|
166
|
+
|
|
167
|
+
| Open SOL flux tube | Recycling SOL with neutrals |
|
|
168
|
+
|---|---|
|
|
169
|
+
|  |  |
|
|
170
|
+
|
|
171
|
+
*Reproduce with [`examples/sol/open_sol_flux_tube.py`](examples/sol/open_sol_flux_tube.py) (left) and [`examples/sol/recycling_sol.py`](examples/sol/recycling_sol.py) (right).*
|
|
172
|
+
|
|
173
|
+

|
|
174
|
+
|
|
175
|
+
*Reproduce with [`examples/benchmarks/b6_detachment_rollover.py`](examples/benchmarks/b6_detachment_rollover.py).*
|
|
176
|
+
|
|
177
|
+
## Gradient-based optimization through the physics
|
|
178
|
+
|
|
179
|
+
Because the solver is differentiable end to end, control and design problems
|
|
180
|
+
become gradient computations. Example: **detachment control** — find the
|
|
181
|
+
upstream density that places the divertor target exactly at the 1 eV
|
|
182
|
+
detachment threshold (Dudson et al., PPCF 61, 065008; Body et al., NME 41,
|
|
183
|
+
101819). The sensitivity `dTe_target/dn_up` is computed by forward-mode
|
|
184
|
+
autodiff **through the entire 20,000-step stiff SOL solve**, and a trust-region
|
|
185
|
+
Newton iteration walks down the detachment cliff to the threshold:
|
|
186
|
+
|
|
187
|
+

|
|
188
|
+
|
|
189
|
+
*Reproduce with [`examples/autodiff/detachment_control.py`](examples/autodiff/detachment_control.py).*
|
|
190
|
+
|
|
191
|
+
A second example: **turbulence optimization** — find the adiabaticity (the
|
|
192
|
+
parallel electron conductivity) at which saturated drift-wave transport drops
|
|
193
|
+
to a quarter of its hydrodynamic-regime level, the classic
|
|
194
|
+
hydrodynamic-to-adiabatic transition of Hasegawa-Wakatani turbulence (Camargo,
|
|
195
|
+
Biskamp & Scott, *Phys. Plasmas* 2, 48 (1995)). A damped Newton iteration on
|
|
196
|
+
the adiabaticity, with forward-mode gradients through the saturated turbulence,
|
|
197
|
+
converges in 7 iterations; an independent long run verifies a 3.96x flux
|
|
198
|
+
reduction against the 4x target. Left column: the initial hydrodynamic state;
|
|
199
|
+
right column: the optimized adiabatic state:
|
|
200
|
+
|
|
201
|
+

|
|
202
|
+
|
|
203
|
+
*Reproduce with [`examples/tokamak/hasegawa_wakatani_optimization.py`](examples/tokamak/hasegawa_wakatani_optimization.py).*
|
|
204
|
+
|
|
205
|
+
The same machinery recovers a transport-drive parameter by gradient descent
|
|
206
|
+
through nonlinear drift-wave turbulence
|
|
207
|
+
([inverse design](examples/tokamak/drift_wave_inverse_design.py)), and the
|
|
208
|
+
[differentiation-methods example](examples/autodiff/differentiation_methods.py)
|
|
209
|
+
measures which method is cheapest (forward mode for a few parameters, ~2x a
|
|
210
|
+
forward run; reverse mode for parameter fields; checkpointing when memory-bound).
|
|
211
|
+
|
|
212
|
+
## Performance and parallel execution
|
|
213
|
+
|
|
214
|
+
Single-CPU turbulence throughput is about 2 million cell-updates per second in
|
|
215
|
+
float64, and one gradient through a 200-step rollout costs 2-3x a forward run.
|
|
216
|
+
The full four-field FCI step — four RHS evaluations, each with a GMRES
|
|
217
|
+
potential inversion — compiles as a single `jit` program with no host
|
|
218
|
+
synchronization inside, which roughly halves the stellarator-turbulence
|
|
219
|
+
step time on one CPU
|
|
220
|
+
([details](docs/performance_and_differentiability.md)):
|
|
221
|
+
|
|
222
|
+
| Turbulence performance | Cost of each differentiation method |
|
|
223
|
+
|---|---|
|
|
224
|
+
|  |  |
|
|
225
|
+
|
|
226
|
+
*Reproduce with [`examples/benchmarks/performance_benchmark.py`](examples/benchmarks/performance_benchmark.py) (left) and [`examples/autodiff/differentiation_methods.py`](examples/autodiff/differentiation_methods.py) (right).*
|
|
227
|
+
|
|
228
|
+
The FCI stack runs across devices with `shard_map`: the sharded step is
|
|
229
|
+
bit-exact against single-device execution (checksums identical at every device
|
|
230
|
+
count), and on a 36-core host with one core per shard a 1.05M-cell step reaches
|
|
231
|
+
a 4.5x speedup at 16 shards, while one NVIDIA A4000 GPU runs the same step ~21x
|
|
232
|
+
faster than a single CPU shard. Both ratios are modest precisely because the
|
|
233
|
+
single-device kernel is now fast — the solvax GMRES potential solve and the
|
|
234
|
+
sync-free RHS cut the per-step baseline sharply, so there is simply less work
|
|
235
|
+
left to parallelize
|
|
236
|
+
([demo](examples/benchmarks/fci_sharded_strong_scaling.py)):
|
|
237
|
+
|
|
238
|
+

|
|
239
|
+
|
|
240
|
+
*Reproduce with [`examples/benchmarks/fci_sharded_strong_scaling.py`](examples/benchmarks/fci_sharded_strong_scaling.py).*
|
|
241
|
+
|
|
242
|
+
## Hasegawa-Wakatani benchmark
|
|
243
|
+
|
|
244
|
+
The standard two-field drift-wave turbulence benchmark: grown from noise
|
|
245
|
+
through the linear instability (growth rate verified against the analytic
|
|
246
|
+
dispersion relation to ~1e-14) into nonlinear E×B transport. Detailed
|
|
247
|
+
verification figures — dispersion scans, MMS convergence orders, and
|
|
248
|
+
gradient-vs-finite-difference checks — are in the
|
|
249
|
+
[documentation](https://drbx.readthedocs.io/en/latest/):
|
|
250
|
+
|
|
251
|
+

|
|
252
|
+
|
|
253
|
+
*Reproduce with [`examples/tokamak/drift_wave_turbulence.py`](examples/tokamak/drift_wave_turbulence.py).*
|
|
254
|
+
|
|
255
|
+
## Reproducing the figures and movies
|
|
256
|
+
|
|
257
|
+
Every figure and movie above states the script that generates it. Each is a
|
|
258
|
+
flat pedagogical file in [`examples/`](examples/): all parameters at the top
|
|
259
|
+
with comments, explicit model/geometry/boundary-condition setup through the
|
|
260
|
+
public API, progress printed while it runs, plot written at the end. The
|
|
261
|
+
embedded copies live compressed in `docs/media/`; the scripts regenerate
|
|
262
|
+
full-quality versions under `output/`.
|
|
263
|
+
|
|
264
|
+
## What it does
|
|
265
|
+
|
|
266
|
+
| Capability | What ships |
|
|
267
|
+
|---|---|
|
|
268
|
+
| **Turbulence models** | Hasegawa-Wakatani drift-wave (pseudo-spectral, differentiable); FCI 2-field, 4-field interchange (density/vorticity/parallel flows), and electromagnetic drift-reduced stacks with curvature and vorticity/potential closures |
|
|
269
|
+
| **Geometry** | Rotating-ellipse stellarator (closed core + optional limiter SOL), island-divertor field (emergent stochastic SOL), shifted-torus helical flux tube, open slab SOL, imported ESSOS coil / VMEC / hybrid equilibria — metrics by autodiff of analytic embeddings |
|
|
270
|
+
| **Field-line topology** | Closed and open field lines; FCI traced field-line maps; multi-transit connection-length tracing; Bohm sheath + target recycling closure on open endpoints |
|
|
271
|
+
| **Neutrals (hermes-3 model)** | Packaged AMJUEL ionization/recombination + charge-exchange rates (no external database); Galilean-invariant plasma-neutral coupling; recycling SOL and a self-consistent detaching SOL (implicit Spitzer conduction, self-limiting radiation, SD1D rollover) |
|
|
272
|
+
| **Linear solver** | `drbx.linear` linearizes any model about an equilibrium; drift-wave, shear-Alfven, and interchange dispersion reproduced to machine precision |
|
|
273
|
+
| **Differentiability** | `jit`/`grad`/`vmap` through every model — sensitivity, uncertainty propagation, inverse design, detachment control; forward/reverse/checkpointed methods measured and gated to agree |
|
|
274
|
+
| **Parallelism** | Multi-device `shard_map` FCI stepping (bit-exact vs single device) with halo exchange; CPU strong scaling demonstrated, GPU-ready |
|
|
275
|
+
| **Solvers** | Structured solves via [`solvax`](https://github.com/uwplasma/SOLVAX) (spectral Fourier-Helmholtz elliptic, tridiagonal, Krylov, preconditioners) |
|
|
276
|
+
| **Runtime** | TOML-deck CLI (`drbx inspect` / `run`) and a small Python API; restartable runs; portable JSON/NPZ artifacts |
|
|
277
|
+
|
|
278
|
+
## Validation
|
|
279
|
+
|
|
280
|
+
`drbx` is validated against a ladder of literature-anchored benchmarks.
|
|
281
|
+
Each rung has a test (or a documented gate) and an example that regenerates
|
|
282
|
+
its figure.
|
|
283
|
+
|
|
284
|
+
Verified today (each with a passing test):
|
|
285
|
+
|
|
286
|
+
| Case | Anchor | What is checked |
|
|
287
|
+
|------|--------|-----------------|
|
|
288
|
+
| Method of manufactured solutions | Riva et al., *Phys. Plasmas* 21, 062301 (2014); Dudson et al. 23, 062303 (2016) | operator / 1D-fluid / FCI convergence order → 2 |
|
|
289
|
+
| Resistive drift-wave dispersion | Dudson et al., *Comput. Phys. Commun.* 180, 1467 (2009) | growth rate and frequency vs analytic dispersion |
|
|
290
|
+
| Shear-Alfvén wave dispersion | Stegmeir et al., *Phys. Plasmas* 26, 052517 (2019) | phase velocity vs analytic (with electron inertia) |
|
|
291
|
+
| Interchange / Rayleigh-Taylor | curvature-driven flute dispersion | growth rate vs `√(gκ)·k_y/k` analytic |
|
|
292
|
+
| FCI on non-axisymmetric geometry | Shanahan et al., *PPCF* 61, 025007 (2019, BSTING) | parallel-operator MMS; differentiable rollout (grad vs FD 6e-11) |
|
|
293
|
+
| Rotating-ellipse (`l = 2`) FCI | Stegmeir et al., *Comput. Phys. Commun.* 198, 139 (2016, GRILLIX) | direct & traced-field-line parallel gradient converge at order 2 on a genuinely non-axisymmetric metric; shape-differentiable; a seeded four-field filament generates interchange vorticity on the rotating surfaces |
|
|
294
|
+
| Island-divertor field (B8) | Shanahan et al., *J. Plasma Phys.* 90 (2024, BSTING); GBS island-divertor studies | sheared-iota island chains + stochastic edge; closed core and finite-connection-length open SOL emerge from multi-transit tracing; turbulence drains through the emergent divertor masks |
|
|
295
|
+
| Open-field-line SOL flux tube | two-point / Bohm-sheath SOL theory (Stangeby, *The Plasma Boundary of Magnetic Fusion Devices*, 2000) | parallel flow reaches Mach 1 at the targets; target density = half upstream; exact Bohm particle balance and sheath-recycling accounting |
|
|
296
|
+
| Neutrals and recycling (hermes-3 model) | hermes-3: Dudson et al., *Comput. Phys. Commun.* 296, 108991 (2024); AMJUEL atomic rates | physically-correct ionization/recombination/CX rates; exact plasma↔neutral particle & momentum conservation; neutrals conserve on the 3D closed rotating ellipse and recycle on the open slab |
|
|
297
|
+
| SD1D detachment rollover (B6) | SD1D: Dudson et al., *Plasma Phys. Control. Fusion* 61, 065008 (2019) | self-consistent SOL (evolved temperature, implicit Spitzer conduction, self-limiting radiation): the target cools through 1 eV into the recombining regime and the target ion flux rolls over as upstream density rises; differentiable |
|
|
298
|
+
| Differentiable inverse design | — | gradient descent through turbulence recovers a drive parameter |
|
|
299
|
+
|
|
300
|
+
Planned rungs (seeded-blob inertial scaling and others) are
|
|
301
|
+
tracked in the project planning notes; benchmark reports live under
|
|
302
|
+
[docs/](docs/linear_dispersion_benchmark.md) and
|
|
303
|
+
[docs/validation_gallery.md](docs/validation_gallery.md).
|
|
304
|
+
|
|
305
|
+
## Examples
|
|
306
|
+
|
|
307
|
+
Flagship simulations, by geometry:
|
|
308
|
+
|
|
309
|
+
| | Turbulence flagship | Geometry |
|
|
310
|
+
|---|---|---|
|
|
311
|
+
| **Tokamak** | [drift-wave turbulence](examples/tokamak/drift_wave_turbulence.py) (Hasegawa-Wakatani; linear phase B2-verified, differentiable) + [inverse design](examples/tokamak/drift_wave_inverse_design.py) | periodic flux tube |
|
|
312
|
+
| **Stellarator** | [turbulence on closed + open field lines](examples/stellarator/stellarator_turbulence.py) (four-field, limiter SOL, movies) + [Landreman-Paul turbulence](examples/stellarator/landreman_paul_turbulence.py) (four-field on the imported LP VMEC equilibrium, closed core + sheath-drained SOL) + [3D renders](examples/stellarator/stellarator_3d_render.py) (cutaway turbulence movie, field-line topology) + [island divertor](examples/stellarator/island_divertor.py) (B8: Poincare, connection lengths, emergent open SOL) + [rotating-ellipse FCI](examples/stellarator/rotating_ellipse_fci.py) (parallel-operator convergence) + [seeded filament](examples/stellarator/rotating_ellipse_filament.py) + [differentiable FCI drift-reduced model](examples/stellarator/fci_differentiable.py) | rotating ellipse (closed core + limiter SOL) + shifted-torus helical + imported [ESSOS/VMEC](examples/geometry-3D/) |
|
|
313
|
+
| **Coils (vacuum)** | [Landreman-Paul closed + open field lines](examples/geometry-3D/essos-field-lines/closed_open_vacuum_poincare.py) (ESSOS Biot-Savart, Poincare classification) | imported coil field |
|
|
314
|
+
| **VMEC equilibria** | [closed field lines from a wout file](examples/geometry-3D/vmex/closed_field_lines.py) (VMEX import; traced rotational transform matches the equilibrium `iotaf` profile to ~1e-6) + [closed + open field lines](examples/geometry-3D/vmex/closed_open_field_lines.py) (coil field with the VMEC last closed flux surface overlaid) | imported VMEC equilibrium (Landreman-Paul precise QA) |
|
|
315
|
+
| **SOL (open)** | [open SOL flux tube](examples/sol/open_sol_flux_tube.py) (parallel transport to Bohm-sheath targets; two-point steady state) + [recycling SOL](examples/sol/recycling_sol.py) (neutrals, ionization/recombination, detachment onset) | open slab flux tube |
|
|
316
|
+
|
|
317
|
+
Open-field-line SOL:
|
|
318
|
+
[open slab flux tube](examples/sol/open_sol_flux_tube.py) — parallel
|
|
319
|
+
transport to Bohm-sheath-bounded targets, relaxing to the classic two-point
|
|
320
|
+
steady state (Mach 1 at the targets, target density half the upstream density),
|
|
321
|
+
with the FCI sheath/recycling closure on the target plates.
|
|
322
|
+
|
|
323
|
+
Benchmarks, differentiable, and geometry examples:
|
|
324
|
+
|
|
325
|
+
- Linear dispersion (B2/B3):
|
|
326
|
+
[examples/benchmarks/linear_dispersion.py](examples/benchmarks/linear_dispersion.py)
|
|
327
|
+
reproduces the drift-wave and shear-Alfvén dispersion relations from the
|
|
328
|
+
linear solver.
|
|
329
|
+
- Autodiff: [gradient-based detachment control](examples/autodiff/detachment_control.py)
|
|
330
|
+
(forward-mode sensitivity through the stiff SOL solve, trust-region Newton onto
|
|
331
|
+
the 1 eV threshold), [inverse design through turbulence](examples/tokamak/drift_wave_inverse_design.py)
|
|
332
|
+
(recover a parameter by gradient descent through a nonlinear drift-wave run),
|
|
333
|
+
[choosing the most efficient differentiation method](examples/autodiff/differentiation_methods.py)
|
|
334
|
+
(forward vs reverse vs checkpointed reverse — same gradient, different cost),
|
|
335
|
+
plus [sensitivity](examples/autodiff_diffusion_sensitivity.py),
|
|
336
|
+
[uncertainty](examples/autodiff_diffusion_uncertainty.py), and reduced
|
|
337
|
+
[inverse design](examples/autodiff_diffusion_inverse_design.py).
|
|
338
|
+
- Stellarator FCI and imported geometry:
|
|
339
|
+
[examples/geometry-3D/](examples/geometry-3D/).
|
|
340
|
+
- Start with [examples/model_selection_guide.py](examples/model_selection_guide.py)
|
|
341
|
+
to choose a model family, dimension, and boundary conditions.
|
|
342
|
+
|
|
343
|
+
The examples are self-contained — no external plasma code is needed to run
|
|
344
|
+
them. Large figures and movies are hosted in GitHub releases so the checkout
|
|
345
|
+
stays small.
|
|
346
|
+
|
|
347
|
+
## Geometry and parallelization
|
|
348
|
+
|
|
349
|
+
The FCI operator and domain-decomposition stack (`FciGeometry3D`,
|
|
350
|
+
`fci_operators`, halo exchange) was contributed by **Aiken Xie** in
|
|
351
|
+
[PR #3](https://github.com/uwplasma/drbx/pull/3) and is incorporated here.
|
|
352
|
+
Built on it, the drift-reduced two-field step runs across multiple devices with
|
|
353
|
+
`shard_map`: the domain is decomposed into halo-exchanged shards and the sharded
|
|
354
|
+
RK4 step is **bit-exact** against the single-device step (checked to ~1e-16 for
|
|
355
|
+
single-device and forced-four-device runs in
|
|
356
|
+
[`tests/test_fci_sharded_2field.py`](tests/test_fci_sharded_2field.py)). On a
|
|
357
|
+
36-core Linux host with one core bound per shard, a 1.05M-cell step reaches a
|
|
358
|
+
**4.5x speedup at 16 shards** (1.18 s → 0.27 s), and one NVIDIA A4000 GPU runs
|
|
359
|
+
the same step **~21x faster than a single CPU shard** (checksums identical)
|
|
360
|
+
([strong-scaling script](examples/benchmarks/fci_sharded_strong_scaling.py),
|
|
361
|
+
[docs](docs/performance_and_differentiability.md)).
|
|
362
|
+
|
|
363
|
+
## Documentation
|
|
364
|
+
|
|
365
|
+
- Physics and numerics: [physics_models.md](docs/physics_models.md),
|
|
366
|
+
[equation_to_code_map.md](docs/equation_to_code_map.md),
|
|
367
|
+
[code_structure.md](docs/code_structure.md).
|
|
368
|
+
- Performance and differentiability:
|
|
369
|
+
[performance_and_differentiability.md](docs/performance_and_differentiability.md),
|
|
370
|
+
[profiling_runtime.md](docs/profiling_runtime.md).
|
|
371
|
+
- Validation: [validation_gallery.md](docs/validation_gallery.md).
|
|
372
|
+
- Testing policy: [testing_strategy.md](docs/testing_strategy.md).
|
|
373
|
+
|
|
374
|
+
## Testing
|
|
375
|
+
|
|
376
|
+
```bash
|
|
377
|
+
pytest -q -m "not slow" # full fast suite
|
|
378
|
+
pytest -q -m "not slow" --cov=drbx --cov-branch # with coverage
|
|
379
|
+
```
|
|
380
|
+
|
|
381
|
+
CI runs the full fast suite on Python 3.10–3.12.
|
|
382
|
+
|
|
383
|
+
## Releases
|
|
384
|
+
|
|
385
|
+
The current development series is described in
|
|
386
|
+
[docs/release_notes_2_0_0.md](docs/release_notes_2_0_0.md).
|
|
387
|
+
|
|
388
|
+
## Citing
|
|
389
|
+
|
|
390
|
+
If you use DRBX in published work, please cite this repository
|
|
391
|
+
(https://github.com/uwplasma/DRBX).
|
|
392
|
+
|
|
393
|
+
## License
|
|
394
|
+
|
|
395
|
+
MIT — see [LICENSE](LICENSE).
|