HermiteSplineProfile to improve accuracy of coordinate mapping

[unalmis]

• Adds HermiteSplineProfile.
  • Relevant for Implement 1D Profile class with orthogonal polynomials like Jacobi or Chebyshev #1201.
• Resolves Improve root finding accuracy on low res grids #1200.
  • Another solution would be to use the radial basis polynomials of the quadrature grid for spectral spline of iota.
• Resolves Issue when map coordinates is used under JIT #1203

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.78 +/- 5.00     | +4.04e-03 +/- 2.59e-02 |  5.22e-01 +/- 2.1e-02  |  5.18e-01 +/- 1.5e-02  |
 test_build_transform_fft_midres         |     +0.79 +/- 2.64     | +4.73e-03 +/- 1.59e-02 |  6.05e-01 +/- 1.3e-02  |  6.00e-01 +/- 9.4e-03  |
 test_build_transform_fft_highres        |     +0.32 +/- 2.02     | +3.17e-03 +/- 2.01e-02 |  9.96e-01 +/- 1.7e-02  |  9.93e-01 +/- 1.1e-02  |
 test_equilibrium_init_lowres            |     +1.07 +/- 3.79     | +3.98e-02 +/- 1.41e-01 |  3.75e+00 +/- 1.2e-01  |  3.71e+00 +/- 8.0e-02  |
 test_equilibrium_init_medres            |     +0.11 +/- 3.58     | +4.63e-03 +/- 1.51e-01 |  4.22e+00 +/- 1.0e-01  |  4.22e+00 +/- 1.1e-01  |
 test_equilibrium_init_highres           |     +0.28 +/- 1.99     | +1.59e-02 +/- 1.12e-01 |  5.62e+00 +/- 1.1e-01  |  5.61e+00 +/- 2.9e-02  |
 test_objective_compile_dshape_current   |     +0.58 +/- 1.06     | +2.22e-02 +/- 4.02e-02 |  3.82e+00 +/- 3.7e-02  |  3.80e+00 +/- 1.6e-02  |
 test_objective_compile_atf              |     -0.37 +/- 1.68     | -3.07e-02 +/- 1.38e-01 |  8.23e+00 +/- 1.3e-01  |  8.26e+00 +/- 5.0e-02  |
 test_objective_compute_dshape_current   |     -3.12 +/- 6.94     | -4.05e-05 +/- 9.03e-05 |  1.26e-03 +/- 4.1e-05  |  1.30e-03 +/- 8.0e-05  |
 test_objective_compute_atf              |     +1.38 +/- 6.33     | +5.88e-05 +/- 2.69e-04 |  4.31e-03 +/- 1.7e-04  |  4.25e-03 +/- 2.0e-04  |
 test_objective_jac_dshape_current       |     +4.21 +/- 8.53     | +1.58e-03 +/- 3.20e-03 |  3.91e-02 +/- 1.3e-03  |  3.76e-02 +/- 2.9e-03  |
 test_objective_jac_atf                  |     -1.87 +/- 3.59     | -3.52e-02 +/- 6.76e-02 |  1.85e+00 +/- 4.8e-02  |  1.89e+00 +/- 4.7e-02  |
 test_perturb_1                          |     +0.67 +/- 1.41     | +8.94e-02 +/- 1.87e-01 |  1.34e+01 +/- 1.5e-01  |  1.33e+01 +/- 1.1e-01  |
 test_perturb_2                          |     +0.05 +/- 1.45     | +9.70e-03 +/- 2.69e-01 |  1.85e+01 +/- 2.3e-01  |  1.85e+01 +/- 1.3e-01  |
 test_proximal_jac_atf                   |     -0.58 +/- 0.97     | -4.66e-02 +/- 7.83e-02 |  8.00e+00 +/- 5.4e-02  |  8.05e+00 +/- 5.7e-02  |
 test_proximal_freeb_compute             |     +0.55 +/- 0.84     | +9.82e-04 +/- 1.49e-03 |  1.79e-01 +/- 1.1e-03  |  1.78e-01 +/- 9.8e-04  |
 test_proximal_freeb_jac                 |     -0.55 +/- 0.96     | -4.07e-02 +/- 7.10e-02 |  7.32e+00 +/- 4.9e-02  |  7.36e+00 +/- 5.2e-02  |
 test_solve_fixed_iter                   |     +0.58 +/- 12.52    | +1.03e-01 +/- 2.22e+00 |  1.79e+01 +/- 1.5e+00  |  1.78e+01 +/- 1.6e+00  |

[unalmis]

After merging #1188 locally , the error in the map_coordinates_test that fails on that branch is reduced by a factor of 20. If that error is not acceptable, we'll either have to use more than eq.L_grid when making the quadrature grid in eq.compute_profile or change the jacfun in map_coordinates to replace the radial derivatives with zero when the radial coordinate of in basis and outbasis matches, so that the Jacobian is effectively 2d.

[codecov]

Codecov Report

Attention: Patch coverage is 84.12698% with 10 lines in your changes missing coverage. Please review.

Project coverage is 95.42%. Comparing base (dd3f472) to head (236af07).
Report is 13 commits behind head on master.

Files	Patch %	Lines
desc/profiles.py	84.09%	7 Missing ⚠️
desc/equilibrium/coords.py	70.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1199      +/-   ##
==========================================
- Coverage   95.47%   95.42%   -0.05%     
==========================================
  Files          87       87              
  Lines       22277    22313      +36     
==========================================
+ Hits        21268    21293      +25     
- Misses       1009     1020      +11     

Files	Coverage Δ
desc/compute/utils.py	96.46% <ø> (ø)
desc/equilibrium/equilibrium.py	95.71% <100.00%> (+0.01%)	⬆️
desc/equilibrium/coords.py	91.80% <70.00%> (-0.67%)	⬇️
desc/profiles.py	95.94% <84.09%> (-1.27%)	⬇️

... and 1 file with indirect coverage changes

[dpanici]

Can you add a simple check trying to solve with this? I get an error when I try, I can try to debug today/tmrw, I think tho it is due to eq.params_dict["p_l"] here being a 2D array instead of a 1D array like it is for the usual SplineProfile, and I think we expect everything inside of params_dict to be a 1D array. To fix, either

• making the params of HermiteSplineProfile be twice the length and have the 2nd half be the deriv values?
• or we could have a p_r_l (and etc for i and c) that is for parameters describing the derivs of the profiles. This could lead nicely into allowing for the derivaitves of profiles to be given as an input (like the I'(r) profile which is often given in other codes). This is more work to do though and could be made as an issue for future work

import numpy as np
from desc.examples import get
from desc.profiles import HermiteSplineProfile

eq = get("DSHAPE")
rho = np.linspace(0,1.0,20,endpoint=True)
eq.pressure = HermiteSplineProfile(r=rho, f = eq.pressure(rho), dfdr = eq.pressure(rho,dr=1))
eq.solve()

yields

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 8
      6 rho = np.linspace(0,1.0,20,endpoint=True)
      7 eq.pressure = HermiteSplineProfile(r=rho, f = eq.pressure(rho), dfdr = eq.pressure(rho,dr=1))
----> 8 eq.solve()

File [~/Research/DESC/desc/equilibrium/equilibrium.py:2058](http://localhost:8888/lab/tree/~/Research/DESC/desc/equilibrium/equilibrium.py#line=2057), in Equilibrium.solve(self, objective, constraints, optimizer, ftol, xtol, gtol, maxiter, x_scale, options, verbose, copy)
   2043 warnif(
   2044     self.N > self.N_grid or self.M > self.M_grid or self.L > self.L_grid,
   2045     msg="Equilibrium has one or more spectral resolutions "
   (...)
   2049     + "to avoid this warning.",
   2050 )
   2051 errorif(
   2052     self.bdry_mode == "poincare",
   2053     NotImplementedError,
   2054     "Solving equilibrium with poincare XS as BC is not supported yet "
   2055     + "on master branch.",
   2056 )
-> 2058 things, result = optimizer.optimize(
   2059     self,
   2060     objective,
   2061     constraints,
   2062     ftol=ftol,
   2063     xtol=xtol,
   2064     gtol=gtol,
   2065     x_scale=x_scale,
   2066     verbose=verbose,
   2067     maxiter=maxiter,
   2068     options=options,
   2069     copy=copy,
   2070 )
   2072 return things[0], result

File [~/Research/DESC/desc/optimize/optimizer.py:220](http://localhost:8888/lab/tree/~/Research/DESC/desc/optimize/optimizer.py#line=219), in Optimizer.optimize(self, things, objective, constraints, ftol, xtol, gtol, ctol, x_scale, verbose, maxiter, options, copy)
    218     objective.build(verbose=verbose)
    219 if linear_constraint is not None and not linear_constraint.built:
--> 220     linear_constraint.build(verbose=verbose)
    221 if nonlinear_constraint is not None and not nonlinear_constraint.built:
    222     nonlinear_constraint.build(verbose=verbose)

File [~/Research/DESC/desc/backend.py:136](http://localhost:8888/lab/tree/~/Research/DESC/desc/backend.py#line=135), in execute_on_cpu.<locals>.wrapper(*args, **kwargs)
    133 @functools.wraps(func)
    134 def wrapper(*args, **kwargs):
    135     with jax.default_device(jax.devices("cpu")[0]):
--> 136         return func(*args, **kwargs)

File [~/Research/DESC/desc/objectives/objective_funs.py:174](http://localhost:8888/lab/tree/~/Research/DESC/desc/objectives/objective_funs.py#line=173), in ObjectiveFunction.build(self, use_jit, verbose)
    172         if verbose > 0:
    173             print("Building objective: " + objective.name)
--> 174         objective.build(use_jit=self.use_jit, verbose=verbose)
    175     self._dim_f += objective.dim_f
    176 if self._dim_f == 1:

File [~/Research/DESC/desc/objectives/linear_objectives.py:691](http://localhost:8888/lab/tree/~/Research/DESC/desc/objectives/linear_objectives.py#line=690), in FixBoundaryR.build(self, use_jit, verbose)
    689     scales = compute_scaling_factors(eq)
    690     self._normalization = scales["a"]
--> 691 super().build(use_jit=use_jit, verbose=verbose)

File [~/Research/DESC/desc/objectives/linear_objectives.py:192](http://localhost:8888/lab/tree/~/Research/DESC/desc/objectives/linear_objectives.py#line=191), in FixParameters.build(self, use_jit, verbose)
    190 # default target
    191 if self.target is None and self.bounds is None:
--> 192     self.target = np.concatenate(
    193         [
    194             np.atleast_1d(param[idx])
    195             for param, idx in zip(tree_leaves(thing.params_dict), self._indices)
    196         ]
    197     )
    199 super().build(use_jit=use_jit, verbose=verbose)

ValueError: all the input arrays must have same number of dimensions, but the array at index 0 has 1 dimension(s) and the array at index 18 has 2 dimension(s)

[f0uriest]

or change the jacfun in map_coordinates to replace the radial derivatives with zero when the radial coordinate of in basis and outbasis matches, so that the Jacobian is effectively 2d.

This might be a good idea regardless

[dpanici]

one tiny comment on the docstring