Make a differentiable version of FourierRZToroidalSurface.constant_offset_surface
#2016
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 constant_offset_surface( | ||
| self, offset, grid=None, M=None, N=None, full_output=False | ||
| ): | ||
| """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 @@ | |
| 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 @@ | |
| 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 | ||
| 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 @@ | |
| 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 @@ | |
| M = check_nonnegint(M, "M") | ||
| N = check_nonnegint(N, "N") | ||
|
|
||
| base_surface = self.copy() | ||
| if grid is None: | ||
| grid = LinearGrid( | ||
| M=base_surface.M * 2, | ||
|
|
@@ -738,57 +745,21 @@ | |
| ), "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) | ||
|
|
||
| R_lmn, Z_lmn, data, (res, niter) = _constant_offset_surface( | ||
| base_surface, offset, grid=grid | ||
| ) | ||
|
|
||
| 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 +1176,136 @@ | |
| 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( | ||
|
dpanici marked this conversation as resolved.
|
||
| ["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 | ||
|
Collaborator
Author
There was a problem hiding this comment. Once #1904 is in we can allow this to be fft, though I did notice robustness issues that I did not fully understand when I used fft instead of direct here. |
||
| 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) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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