Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dommaschk potentials implemented in JAX #483

Merged
merged 77 commits into from
Mar 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
77 commits
Select commit Hold shift + click to select a range
7d47f01
initial commit of dommaschk potentials implemented in JAX
dpanici Apr 8, 2023
27fd9d3
add doi to test
dpanici Apr 8, 2023
52a1f57
add test for assertion error
dpanici Apr 8, 2023
5fbc429
add tests and DommaschkMagneticField, need to fix vertical field test…
dpanici Apr 9, 2023
9ce422a
fix typo in args of V_m_l causing test to fail
dpanici Apr 11, 2023
c22b702
update docstring
dpanici Apr 11, 2023
67d3bbc
update arg in test
dpanici Apr 11, 2023
9a18671
add function to fit a magnetic field with Dommaschk potentials
dpanici Apr 20, 2023
060ac54
fix typo in Neumann term of Dommaschk potential. Add assertions to en…
dpanici Apr 25, 2023
797797d
add test for dommaschk fitting using MagneticField class
dpanici Apr 25, 2023
a3137c1
add documentation and update dommaschk function to more succinctly ca…
dpanici Apr 26, 2023
70b2ddc
Merge branch 'master' into dp/dommaschk
dpanici Apr 27, 2023
d0e4d13
implement generalized formulas from eqns 31 and 32 which are valid fo…
dpanici Apr 29, 2023
46c4503
lower test tols
dpanici Apr 29, 2023
a148759
change test tol
dpanici Apr 30, 2023
89e6d19
Merge branch 'master' into dp/dommaschk
dpanici Jun 3, 2023
7c8de1b
Merge branch 'master' into dp/dommaschk
dpanici Jun 7, 2023
e03fd6f
move checks to Dommaschk class, and make some ranges into jnp ranges
dpanici Jun 11, 2023
b8aaf1a
Merge branch 'master' into dp/dommaschk
dpanici Jun 20, 2023
496db5a
Merge branch 'master' into dp/dommaschk
dpanici Jun 20, 2023
3c845de
add 1d catch and check if B0 is 1d array with size 0 to avoid jnp.iss…
dpanici Jun 20, 2023
73a3c58
move read_mgrid to its own function
dpanici Jun 24, 2023
e64bea2
fix test, add check for ints since DommaschkField as written cannot h…
dpanici Jun 24, 2023
ea2d6fd
attempt to jaxify dommaschk
dpanici Jun 27, 2023
ed69632
tests pass but fitting is very slow now
dpanici Jun 28, 2023
3307c63
remove unneeded class creations in fitting, but still is very slow du…
dpanici Jun 28, 2023
cb23ebd
jaxify all domaschk potential things further so that it can be jit'd.…
dpanici Jun 28, 2023
3938852
Merge branch 'master' into dp/dommaschk
dpanici Jun 28, 2023
5c0e53e
add sym option to fitting
dpanici Jun 28, 2023
75b7cdf
jit compute inside of fit
dpanici Jun 28, 2023
c022586
change fitting to use jacfwd instead of for loops to find A,reduce te…
dpanici Jun 28, 2023
a5abda0
add jit to fit jacfwd call
dpanici Jun 29, 2023
c6c7ac9
Merge branch 'master' into dp/dommaschk
dpanici Jun 29, 2023
65ed97f
Merge branch 'master' into dp/dommaschk
f0uriest Oct 17, 2023
f2e9e8d
Merge branch 'master' into dp/dommaschk
dpanici Nov 7, 2023
eb29f6e
Reorganize tests to undo merge issues
f0uriest Nov 9, 2023
027c40a
Apply jit to dommaschk sub-functions
f0uriest Nov 9, 2023
f92cedb
Fix edge case of derivative of 0**0
f0uriest Nov 9, 2023
8c2a284
Use jax loop for dommaschk evaluation
f0uriest Nov 9, 2023
f9172ef
Fix order of return values in read_mgrid
f0uriest Nov 9, 2023
849711b
Adjust test tolerances
f0uriest Nov 9, 2023
84a8057
Merge branch 'master' into dp/dommaschk
dpanici Nov 13, 2023
e93a260
Merge branch 'dp/dommaschk' of github.com:PlasmaControl/DESC into dp/…
dpanici Nov 16, 2023
709eef5
Merge branch 'master' into dp/dommaschk
dpanici Nov 16, 2023
52fb9fb
fix test, fix array counting so that the num_modes = sum(len(a_arr)+l…
dpanici Nov 17, 2023
b973703
Merge branch 'master' into dp/dommaschk
dpanici Nov 17, 2023
c5eaa50
Merge branch 'master' into dp/dommaschk
f0uriest Nov 21, 2023
4c21350
Merge branch 'master' into dp/dommaschk
dpanici Nov 23, 2023
df3740b
address some comments
dpanici Nov 23, 2023
c27ac92
change to use lstsq from jnp.linalg for fit:
dpanici Nov 23, 2023
2382f16
correct use of residual from lstsq
dpanici Nov 24, 2023
3f6fd03
correct use of residual from lstsq, again
dpanici Nov 24, 2023
f2345da
Merge branch 'master' into dp/dommaschk
dpanici Nov 29, 2023
4ee2801
Merge branch 'master' into dp/dommaschk
dpanici Dec 2, 2023
13eaef8
Merge branch 'master' into dp/dommaschk
f0uriest Dec 10, 2023
124c8da
Merge branch 'master' into dp/dommaschk
dpanici Dec 10, 2023
7fb3d29
Merge branch 'master' into dp/dommaschk
dpanici Jan 20, 2024
4eb1c8f
remove unneeded jnp.any statements
dpanici Jan 20, 2024
602ec8d
Merge branch 'master' into dp/dommaschk
dpanici Feb 5, 2024
ae23a87
update docstring
dpanici Feb 5, 2024
74c0050
Merge branch 'master' into dp/dommaschk
f0uriest Mar 1, 2024
e5b2de4
Move magnetic fields to their own module
f0uriest Mar 1, 2024
05673fa
Fix imports
f0uriest Mar 1, 2024
8ef075f
Merge branch 'master' into dp/dommaschk
f0uriest Mar 2, 2024
29f948a
Merge branch 'master' into dp/dommaschk
f0uriest Mar 5, 2024
8c9ad2e
Merge branch 'master' into dp/dommaschk
dpanici Mar 6, 2024
612b45f
Merge branch 'master' into dp/dommaschk
dpanici Mar 6, 2024
6bb8f48
add comments
dpanici Mar 6, 2024
82edbe3
Merge branch 'master' into dp/dommaschk
dpanici Mar 6, 2024
98a9431
Merge branch 'dp/dommaschk' of github.com:PlasmaControl/DESC into dp/…
dpanici Mar 6, 2024
b4e3d37
update changelog
dpanici Mar 6, 2024
9118116
Merge branch 'master' into dp/dommaschk
f0uriest Mar 7, 2024
1cc8b0a
Cast lists to jax arrays
f0uriest Mar 7, 2024
813ff15
Merge branch 'master' into dp/dommaschk
dpanici Mar 11, 2024
b0422ff
Merge branch 'master' into dp/dommaschk
f0uriest Mar 12, 2024
acfbebe
add dommaschk to API
dpanici Mar 12, 2024
5830208
remove extra indent from docstring
dpanici Mar 13, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,14 @@ in the MAKEGRID format for use with other codes.
least squares fit is now weighted inversely with the distance from the axis to improve
the accuracy for low aspect ratio.
- Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z`
keyword arguments, which form a hollow cylindrical bounding box. If the field line
trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying
function of the distance to the box to stop the trajectory and prevent tracing the field
line out to infinity, which is both costly and unnecessary when making a Poincare plot,
keyword arguments, which form a hollow cylindrical bounding box. If the field line
trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying
function of the distance to the box to stop the trajectory and prevent tracing the field
line out to infinity, which is both costly and unnecessary when making a Poincare plot,
the principle purpose of the function.
- Adds a new class ``DommaschkPotentialField`` which allows creation of magnetic fields based
off of the vacuum potentials detailed in Representations for Vacuum Potentials in Stellarators
https://doi.org/10.1016/0010-4655(86)90109-8.


Speed Improvements
Expand Down
22 changes: 14 additions & 8 deletions desc/compute/_surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -626,7 +626,9 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
parameterization=[
"desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
],
)
def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi"] = (
Expand All @@ -649,7 +651,9 @@ def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwar
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
parameterization=[
"desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
],
)
def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi_t"] = (
Expand All @@ -670,7 +674,9 @@ def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kw
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
parameterization=[
"desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
],
)
def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi_z"] = (
Expand All @@ -691,7 +697,7 @@ def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kw
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.CurrentPotentialField",
parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
)
def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi"] = transforms["potential"](
Expand All @@ -714,7 +720,7 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.CurrentPotentialField",
parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
)
def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi_t"] = transforms["potential_dtheta"](
Expand All @@ -737,7 +743,7 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.magnetic_fields.CurrentPotentialField",
parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
)
def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
data["Phi_z"] = transforms["potential_dzeta"](
Expand All @@ -762,8 +768,8 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
coordinates="tz",
data=["Phi_t", "Phi_z", "e_theta", "e_zeta", "|e_theta x e_zeta|"],
parameterization=[
"desc.magnetic_fields.CurrentPotentialField",
"desc.magnetic_fields.FourierCurrentPotentialField",
"desc.magnetic_fields._current_potential.CurrentPotentialField",
"desc.magnetic_fields._current_potential.FourierCurrentPotentialField",
],
)
def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
Expand Down
8 changes: 4 additions & 4 deletions desc/compute/data_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,15 +211,15 @@ def _decorator(func):
"desc.geometry.curve.FourierPlanarCurve",
"desc.geometry.core.Curve",
],
"desc.magnetic_fields.CurrentPotentialField": [
"desc.magnetic_fields._current_potential.CurrentPotentialField": [
"desc.geometry.surface.FourierRZToroidalSurface",
"desc.geometry.core.Surface",
"desc.magnetic_fields.MagneticField",
"desc.magnetic_fields._core.MagneticField",
],
"desc.magnetic_fields.FourierCurrentPotentialField": [
"desc.magnetic_fields._current_potential.FourierCurrentPotentialField": [
"desc.geometry.surface.FourierRZToroidalSurface",
"desc.geometry.core.Surface",
"desc.magnetic_fields.MagneticField",
"desc.magnetic_fields._core.MagneticField",
],
"desc.coils.SplineXYZCoil": [
"desc.geometry.curve.SplineXYZCurve",
Expand Down
16 changes: 16 additions & 0 deletions desc/magnetic_fields/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
"""Classes for Magnetic Fields."""

from ._core import (
PoloidalMagneticField,
ScalarPotentialField,
ScaledMagneticField,
SplineMagneticField,
SumMagneticField,
ToroidalMagneticField,
VerticalMagneticField,
_MagneticField,
field_line_integrate,
read_BNORM_file,
)
from ._current_potential import CurrentPotentialField, FourierCurrentPotentialField
from ._dommaschk import DommaschkPotentialField, dommaschk_potential
Loading
Loading