diff --git a/CHANGELOG.md b/CHANGELOG.md index e59e02665e..21e17cf91b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index 6f7492ebf3..97665150ea 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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, @@ -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: @@ -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 + 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) diff --git a/tests/conftest.py b/tests/conftest.py index 64974de460..88c6eb4764 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 @@ -381,7 +381,7 @@ 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 @@ -389,7 +389,7 @@ def regcoil_modular_coils(): 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 @@ -417,7 +417,7 @@ 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 @@ -425,7 +425,7 @@ def regcoil_windowpane_coils(): 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" @@ -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 diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 4adf058c81..4d004f06cd 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -4,9 +4,12 @@ import pytest import desc.examples +from desc.backend import jax, jit +from desc.compute import get_transforms from desc.equilibrium import Equilibrium from desc.examples import get from desc.geometry import FourierRZToroidalSurface, ZernikeRZToroidalSection +from desc.geometry.surface import _constant_offset_surface from desc.grid import LinearGrid from desc.utils import rpz2xyz @@ -154,7 +157,7 @@ def test_constant_offset_surface_circle(self): r_offset_surf = data["x_offset_surface"] r_surf = data["x"] dists = np.linalg.norm(r_surf - r_offset_surf, axis=1) - np.testing.assert_allclose(dists, 1, atol=1e-16) + np.testing.assert_allclose(dists, 1, atol=1e-14) R00_offset_ind = s_offset.R_basis.get_idx(M=0, N=0) R00_offset = s_offset.R_lmn[R00_offset_ind] R10_offset_ind = s_offset.R_basis.get_idx(M=1, N=0) @@ -171,7 +174,7 @@ def test_constant_offset_surface_circle(self): np.array([R00_offset_ind, R10_offset_ind]), ), 0, - atol=9e-15, + atol=1e-13, ) np.testing.assert_allclose( np.delete( @@ -179,7 +182,7 @@ def test_constant_offset_surface_circle(self): Zneg10_offset_ind, ), 0, - atol=9e-15, + atol=1e-13, ) grid_compute = LinearGrid(M=10, N=10) data = s.compute(["x", "e_theta", "e_zeta"], basis="rpz", grid=grid_compute) @@ -187,7 +190,7 @@ def test_constant_offset_surface_circle(self): ["x", "e_theta", "e_zeta"], basis="rpz", grid=grid_compute ) dists = np.linalg.norm(data["x"] - data_offset["x"], axis=1) - np.testing.assert_allclose(dists, 1, atol=1e-16) + np.testing.assert_allclose(dists, 1, atol=1e-14) correct_data_offset = { "e_theta": np.vstack( ( @@ -212,6 +215,45 @@ def test_constant_offset_surface_circle(self): err_msg=f"Failed test at comparison of {key}", ) + @pytest.mark.unit + def test_constant_offset_surface_circle_jax_transformable(self, capsys): + """Test constant offset algorithm is jax transformable.""" + s = FourierRZToroidalSurface() + grid = LinearGrid(M=3, N=2) + transforms = get_transforms(["R", "Z"], s, grid) + transforms["R"].build_pinv() + transforms["Z"].build_pinv() + offset = 1 + + def fun(params): + return _constant_offset_surface(s, offset, grid, transforms, params) + + # ensure is jitable + R_lmn, Z_lmn, data, _ = jit(fun)(s.params_dict) + + # make sure that the function is not recompiled + with jax.log_compiles(): + R_lmn, Z_lmn, data, _ = jit(fun)(s.params_dict) + + out = capsys.readouterr() + assert out.out == "" + + # check gradient is correct + # R00 of offset should change with R00 of original surface + grad_R00 = jax.grad(lambda params: fun(params)[0][s.R_basis.get_idx(M=0, N=0)])( + s.params_dict + ) + # check that the gradient is nonzero for the R00 component + assert np.any(np.abs(grad_R00["R_lmn"][s.R_basis.get_idx(M=0, N=0)]) > 1e-10) + # check that the gradient is zero otherwise + non_R00_indices = np.where(s.R_basis.modes.sum(axis=1) != 0)[0] + assert np.all(np.abs(grad_R00["R_lmn"][non_R00_indices]) < 1e-10) + + # check gradient is correct + np.testing.assert_allclose( + grad_R00["R_lmn"][s.R_basis.get_idx(M=0, N=0)], 1.0, atol=1e-10 + ) + @pytest.mark.slow @pytest.mark.unit def test_constant_offset_surface_rot_ellipse(self):