-
Notifications
You must be signed in to change notification settings - Fork 24
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
HermiteSplineProfile to improve accuracy of coordinate mapping #1199
Conversation
Otherwise you'll get this error message in jitted functions: > g = obj.grad(obj.x()) tests/test_neoclassical.py:84: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ desc/objectives/objective_funs.py:415: in grad return jnp.atleast_1d(self._grad(x, constants).squeeze()) desc/derivatives.py:89: in __call__ return self.compute(*args, **kwargs) desc/derivatives.py:150: in compute return self._compute(*args, **kwargs) desc/objectives/objective_funs.py:335: in compute_scalar f = jnp.sum(self.compute_scaled_error(x, constants=constants) ** 2) / 2 desc/objectives/objective_funs.py:312: in compute_scaled_error [ desc/objectives/objective_funs.py:313: in <listcomp> obj.compute_scaled_error(*par, constants=const) desc/objectives/objective_funs.py:952: in compute_scaled_error f = self.compute(*args, **kwargs) desc/objectives/_neoclassical.py:235: in compute iota=SplineProfile(iota, df=iota_r, knots=self._rho, name="iota", jnp=jnp), desc/profiles.py:827: in __init__ values = np.atleast_1d(values) ../../../miniconda3/envs/desc-env/lib/python3.10/site-packages/numpy/core/shape_base.py:65: in atleast_1d ary = asanyarray(ary) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ self = TracerArrayConversionError(Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=5/0)>) tracer = Traced<ShapedArray(float64[2])>with<DynamicJaxprTrace(level=5/0)> def __init__(self, tracer: core.Tracer): super().__init__( "The numpy.ndarray conversion method __array__() was called on " > f"{tracer._error_repr()}{tracer._origin_msg()}") E IndexError: list index out of range
| 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 | |
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 |
Codecov ReportAttention: Patch coverage is
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
|
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
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) |
This might be a good idea regardless |
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.
one tiny comment on the docstring
HermiteSplineProfile
.Profile
class with orthogonal polynomials like Jacobi or Chebyshev #1201.