Skip to content

Make a differentiable version of FourierRZToroidalSurface.constant_offset_surface #2016

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.

Sign up for GitHub

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

Open
dpanici wants to merge 24 commits into
base: master
Choose a base branch
from
+214 −63
Open

Make a differentiable version of FourierRZToroidalSurface.constant_offset_surface #2016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
5078e65
make a differentiable version of _constant_offset_surface
dpanici Nov 15, 2025
9c0c65d
add offset kwarg
dpanici Nov 15, 2025
36e4beb
Merge branch 'master' into dp/offset-differentiable
dpanici Nov 15, 2025
7623fe9
Merge branch 'master' into dp/offset-differentiable
dpanici Feb 2, 2026
3919bb3
change constant offset surface to reuse the jitable version's code, a…
dpanici Feb 2, 2026
9896d54
remove VolumeOffset
dpanici Feb 2, 2026
af3c12a
allow transform to be built and passed to constant offset surface
dpanici Feb 3, 2026
747eae7
reduce rootfind tol
dpanici Feb 3, 2026
f70b544
add test
dpanici Feb 3, 2026
573ad0f
Update desc/geometry/surface.py
dpanici Feb 3, 2026
7abd21d
fix incorrect n calc, though did not affect any subsequent calculations
dpanici Feb 3, 2026
08e5c95
Merge branch 'dp/offset-differentiable' of github.com:PlasmaControl/D…
dpanici Feb 3, 2026
156c89c
use arctan2 and correct angle span inside of rootfind, this should ma…
dpanici Feb 3, 2026
7ce498c
update docs
dpanici Feb 3, 2026
32b2879
make grid sym in tests with constant offset surface be true
dpanici Feb 3, 2026
16bd93e
fix test
dpanici Feb 3, 2026
3ca222c
update changelog
dpanici Feb 3, 2026
2e06b38
update tests again
dpanici Feb 3, 2026
d1f04a8
remove redundant part of test, and adjust tols
dpanici Feb 3, 2026
e811204
adjust tols
dpanici Feb 4, 2026
6079bdd
attempt to fix docs
dpanici Feb 4, 2026
0101568
another doc fix attempt
dpanici Feb 4, 2026
66f74d5
adjust maxiter
dpanici Feb 5, 2026
df3f909
Merge branch 'master' into dp/offset-differentiable
dpanici Feb 5, 2026
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
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,14 @@ alternative to fourier continuation methods.
- Adds ``"scipy-l-bfgs-b"`` optimizer option as a wrapper to scipy's ``"l-bfgs-b"`` method.
- Adds ``check_intersection`` flag to ``desc.magnetic_fields.FourierCurrentPotentialField.to_Coilset``, to allow the choice of checking the resulting coilset for intersections or not.
- Changes the import paths for ``desc.external`` to require reference to the sub-modules.
- Adds a differentiable utility for finding constant offset toroidal surfaces inside of optimizations. See [PR](https://github.com/PlasmaControl/DESC/pull/2016) for more details.


Bug Fixes

- No longer uses the full Hessian to compute the scale when ``x_scale="auto"`` and using a scipy optimizer that approximates the hessian (e.g. if using ``"scipy-bfgs"``, no longer attempts the Hessian computation to get the x_scale).
- ``SplineMagneticField.from_field()`` correctly uses the ``NFP`` input when given. Also adds this as a similar input option to ``MagneticField.save_mgrid()``.
- Fixes some bugs that hampered robustness of ``desc.geometry.FourierRZToroidalSurface.constant_offset_surface``, particularly when the given grid had stellarator symmetry or when NFP=1.

Performance Improvements

Expand Down
212 changes: 159 additions & 53 deletions desc/geometry/surface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
vmap,
)
from desc.basis import DoubleFourierSeries, ZernikePolynomial
from desc.compute import get_transforms
from desc.grid import Grid, LinearGrid
from desc.io import InputReader
from desc.optimizable import optimizable_parameter
Expand Down Expand Up @@ -666,7 +667,7 @@ def from_shape_parameters(
def constant_offset_surface(
self, offset, grid=None, M=None, N=None, full_output=False
):
"""Create a FourierRZSurface with constant offset from the base surface (self).
"""Create a new FourierRZToroidalSurface with constant offset from self.

Implementation of algorithm described in Appendix B of
"An improved current potential method for fast computation of
Expand All @@ -676,6 +677,10 @@ def constant_offset_surface(
NOTE: Must have the toroidal angle as the cylindrical toroidal angle
in order for this algorithm to work properly

NOTE: if one wants to use this inside of an optimization, one should
use the private method _constant_offset_surface directly, and refer to
the documentation in PR #2016 for more details.

Parameters
----------
base_surface : FourierRZToroidalSurface
Expand All @@ -684,11 +689,11 @@ def constant_offset_surface(
constant offset (in m) of the desired surface from the input surface
offset will be in the normal direction to the surface.
grid : Grid, optional
Grid object of the points on the given surface to evaluate the
Grid object of the points on the offset surface to evaluate the
offset points at, from which the offset surface will be created by fitting
offset points with the basis defined by the given M and N.
If None, defaults to a LinearGrid with M and N and NFP equal to the
base_surface.M and base_surface.N and base_surface.NFP
If None, defaults to a LinearGrid with M and N and NFP equal to twice the
base_surface.M and base_surface.N and NFP equal to base_surface.NFP
M : int, optional
Poloidal resolution of the basis used to fit the offset points
to create the resulting constant offset surface, by default equal
Expand Down Expand Up @@ -717,6 +722,8 @@ def constant_offset_surface(
coordinates on the offset surface, corresponding to the
``x`` points on the base_surface (i.e. the points to which the
offset surface was fit)
``transforms`` : dict containing the Transform objects used to
fit R and Z of the
info : tuple
2 element tuple containing residuals and number of iterations
for each point. Only returned if ``full_output`` is True
Expand All @@ -725,7 +732,7 @@ def constant_offset_surface(
M = check_nonnegint(M, "M")
N = check_nonnegint(N, "N")

base_surface = self
base_surface = self.copy()
if grid is None:
grid = LinearGrid(
M=base_surface.M * 2,
Expand All @@ -738,57 +745,23 @@ def constant_offset_surface(
), "base_surface must be a FourierRZToroidalSurface!"
M = base_surface.M if M is None else int(M)
N = base_surface.N if N is None else int(N)
base_surface.change_resolution(M=M, N=N)

def n_and_r_jax(nodes):
data = base_surface.compute(
["X", "Y", "Z", "n_rho"],
grid=Grid(nodes, jitable=True, sort=False),
method="jitable",
)

phi = nodes[:, 2]
re = jnp.vstack([data["X"], data["Y"], data["Z"]]).T
n = data["n_rho"]
n = rpz2xyz_vec(n, phi=phi)
r_offset = re + offset * n
return n, re, r_offset

def fun_jax(zeta_hat, theta, zeta):
nodes = jnp.vstack((jnp.ones_like(theta), theta, zeta_hat)).T
n, r, r_offset = n_and_r_jax(nodes)
return jnp.arctan(r_offset[0, 1] / r_offset[0, 0]) - zeta

vecroot = jit(
vmap(
lambda x0, *p: root_scalar(
fun_jax, x0, jac=None, args=p, full_output=full_output
)
)
R_lmn, Z_lmn, data, (res, niter) = _constant_offset_surface(
base_surface,
offset,
grid=grid,
)
if full_output:
zetas, (res, niter) = vecroot(
grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2]
)
else:
zetas = vecroot(grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2])

zetas = np.asarray(zetas)
nodes = np.vstack((np.ones_like(grid.nodes[:, 1]), grid.nodes[:, 1], zetas)).T
n, x, x_offsets = n_and_r_jax(nodes)

data = {}
data["n"] = xyz2rpz_vec(n, phi=nodes[:, 1])
data["x"] = xyz2rpz(x)
data["x_offset_surface"] = xyz2rpz(x_offsets)

offset_surface = FourierRZToroidalSurface.from_values(
data["x_offset_surface"],
theta=nodes[:, 1],
M=M,
N=N,
NFP=base_surface.NFP,
sym=base_surface.sym,

offset_surface = FourierRZToroidalSurface(
R_lmn,
Z_lmn,
data["transforms"]["R"].basis.modes[:, 1:],
data["transforms"]["Z"].basis.modes[:, 1:],
base_surface.NFP,
base_surface.sym,
)

if full_output:
return offset_surface, data, (res, niter)
else:
Expand Down Expand Up @@ -1205,3 +1178,136 @@ def _get_ess_scale(self, alpha=1.2, order=np.inf, min_value=1e-7):
scales.update(get_ess_scale(modes, alpha, order, min_value))

return scales


def _constant_offset_surface(
base_surface,
offset,
grid,
transforms=None,
params=None,
):
"""Create a FourierRZToroidalSurface with constant offset from the base surface.

Implementation of algorithm described in Appendix B of
"An improved current potential method for fast computation of
stellarator coil shapes", Landreman (2017)
https://iopscience.iop.org/article/10.1088/1741-4326/aa57d4

NOTE: Must have the toroidal angle as the cylindrical toroidal angle
in order for this algorithm to work properly

NOTE: this function lacks the checks of the constant_offset_surface
Copy link
Collaborator Author

@dpanici dpanici Nov 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

only is differentiable though if params is passed, otherwise the end will not be differentiatedcorrectlt

so that it is jittable/differentiable

Parameters
----------
base_surface : FourierRZToroidalSurface
Surface from which the constant offset surface will be found.
offset : float
constant offset (in m) of the desired surface from the input surface
offset will be in the normal direction to the surface.
grid : Grid, optional
Grid object of the points on the offset surface to evaluate the
offset points at, from which the offset surface will be created by fitting
offset points with the basis defined by the given M and N.
If None, defaults to a LinearGrid with M and N and NFP equal to twice the
base_surface.M and base_surface.N and NFP equal to base_surface.NFP
transforms: dict, optional
Transforms to use to fit the offset surface's R and Z, respectively. If None,
new transforms will be created using the given surface's M and N.
If given, should contain the keys ["R"] and ["Z"], with the pinv matrices
already built, and the corresponding grid should match the input grid.
params : dict, optional
dictionary of parameters to use when computing data from the base_surface.
If None, uses base_surface.params_dict, however the resulting computation
will not be differentiable with respect to the base_surface parameters
(since the JAX AD inside of an objective traces the params dictionaries
that are passed to their compute methods)

Returns
-------
R_lmn, Z_lmn : array-like
coefficients describing the offset surface geometry
data : dict
dictionary containing the following data, in the cylindrical basis:
``n`` : (``grid.num_nodes`` x 3) array of the unit surface normal on
the base_surface evaluated at the input ``grid``
``x`` : (``grid.num_nodes`` x 3) array of coordinates on
the base_surface evaluated at the input ``grid``
``x_offset_surface`` : (``grid.num_nodes`` x 3) array of the
coordinates on the offset surface, corresponding to the
``x`` points on the base_surface (i.e. the points to which the
offset surface was fit)
as well as the transforms used to fit R and Z.
info : tuple
2 element tuple containing residuals and number of iterations
for each point.

"""
if params is None:
params = base_surface.params_dict

def n_and_r_jax(nodes):
data = base_surface.compute(
["X", "Y", "Z", "n_rho"],
grid=Grid(nodes, jitable=True, sort=False),
method="jitable",
params=params,
)

phi = nodes[:, 2]
re = jnp.vstack([data["X"], data["Y"], data["Z"]]).T
n = data["n_rho"]
n = rpz2xyz_vec(n, phi=phi)
r_offset = re + offset * n
return n, re, r_offset

def fun_jax(zeta_hat, theta, zeta):
nodes = jnp.vstack((jnp.ones_like(theta), theta, zeta_hat)).T
n, r, r_offset = n_and_r_jax(nodes)
# add 2pi to the arctan2<0 so it matches our convention of
# zeta being btwn 0 and 2pi
zeta_offset = jnp.arctan2(r_offset[0, 1], r_offset[0, 0])
return jnp.where(zeta_offset < 0, zeta_offset + 2 * np.pi, zeta_offset) - zeta

vecroot = jit(
vmap(
lambda x0, *p: root_scalar(
fun_jax, x0, jac=None, args=p, full_output=True, tol=1e-12, maxiter=100
)
)
)
zetas, (res, niter) = vecroot(grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2])

zetas = jnp.asarray(zetas)
nodes = jnp.vstack((jnp.ones_like(grid.nodes[:, 1]), grid.nodes[:, 1], zetas)).T
n, x, x_offsets = n_and_r_jax(nodes)

data = {}
data["n"] = xyz2rpz_vec(n, phi=nodes[:, 2])
data["x"] = xyz2rpz(x)
data["x_offset_surface"] = xyz2rpz(x_offsets)

if transforms is None:
# NOTE: we are assuming here that the rootfind was successful for every point,
# so that the zeta=arctan(y/x) of the offset surface point are the same as
# the grid nodes' zeta values. If this is not the case, the fitting
# will be incorrect.
transforms = get_transforms(
obj=base_surface,
keys=["R", "Z"],
grid=grid,
# this is more robust than letting method become fft if
# jitable=False, and will also work within jitted functions
jitable=True,
)
transforms["R"].build_pinv()
transforms["Z"].build_pinv()

R_lmn = transforms["R"].fit(data["x_offset_surface"][:, 0])
Z_lmn = transforms["Z"].fit(data["x_offset_surface"][:, 2])

data["transforms"] = transforms

return R_lmn, Z_lmn, data, (res, niter)
12 changes: 6 additions & 6 deletions tests/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,7 @@ def regcoil_helical_coils_scan():
offset=0.2, # desired offset
M=16, # Poloidal resolution of desired offset surface
N=12, # Toroidal resolution of desired offset surface
grid=LinearGrid(M=32, N=16, NFP=eq.NFP),
grid=LinearGrid(M=32, N=16, NFP=eq.NFP, sym=eq.sym),
)
surface_current_field = FourierCurrentPotentialField.from_surface(
surf_winding, M_Phi=8, N_Phi=8
Expand Down Expand Up @@ -381,15 +381,15 @@ def regcoil_modular_coils():
offset=0.2, # desired offset
M=16, # Poloidal resolution of desired offset surface
N=12, # Toroidal resolution of desired offset surface
grid=LinearGrid(M=32, N=16, NFP=eq.NFP),
grid=LinearGrid(M=32, N=16, NFP=eq.NFP, sym=eq.sym),
)
M_Phi = 10
N_Phi = 10
M_egrid = 20
N_egrid = 20
M_sgrid = 40
N_sgrid = 40
lambda_regularization = 1e-18
lambda_regularization = 1e-20

surface_current_field = FourierCurrentPotentialField.from_surface(
surf_winding, M_Phi=M_Phi, N_Phi=N_Phi
Expand Down Expand Up @@ -417,15 +417,15 @@ def regcoil_windowpane_coils():
offset=0.2, # desired offset
M=16, # Poloidal resolution of desired offset surface
N=12, # Toroidal resolution of desired offset surface
grid=LinearGrid(M=32, N=16, NFP=eq.NFP),
grid=LinearGrid(M=32, N=16, NFP=eq.NFP, sym=eq.sym),
)
M_Phi = 10
N_Phi = 10
M_egrid = 20
N_egrid = 20
M_sgrid = 20
N_sgrid = 20
lambda_regularization = 1e-18
lambda_regularization = 1e-20

surface_current_field = FourierCurrentPotentialField.from_surface(
surf_winding, M_Phi=M_Phi, N_Phi=N_Phi, sym_Phi="sin"
Expand Down Expand Up @@ -457,7 +457,7 @@ def regcoil_PF_coils():
offset=0.2, # desired offset
M=16, # Poloidal resolution of desired offset surface
N=12, # Toroidal resolution of desired offset surface
grid=LinearGrid(M=32, N=16, NFP=eq.NFP),
grid=LinearGrid(M=32, N=16, NFP=eq.NFP, sym=eq.sym),
)
M_Phi = 10
N_Phi = 10
Expand Down
Loading
Loading