-
Notifications
You must be signed in to change notification settings - Fork 22
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
Bounce averaging #854
base: master
Are you sure you want to change the base?
Bounce averaging #854
Conversation
| benchmark_name | dt(%) | dt(s) | t_new(s) | t_old(s) |
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
test_build_transform_fft_lowres | +9.64 +/- 6.35 | +4.95e-02 +/- 3.26e-02 | 5.63e-01 +/- 2.7e-02 | 5.14e-01 +/- 1.8e-02 |
test_build_transform_fft_midres | +3.15 +/- 7.76 | +1.89e-02 +/- 4.65e-02 | 6.19e-01 +/- 4.3e-02 | 6.00e-01 +/- 1.7e-02 |
test_build_transform_fft_highres | -1.45 +/- 2.46 | -1.47e-02 +/- 2.49e-02 | 9.98e-01 +/- 8.0e-03 | 1.01e+00 +/- 2.4e-02 |
test_equilibrium_init_lowres | -0.35 +/- 4.77 | -1.30e-02 +/- 1.80e-01 | 3.77e+00 +/- 1.2e-01 | 3.78e+00 +/- 1.4e-01 |
test_equilibrium_init_medres | +1.73 +/- 5.48 | +7.35e-02 +/- 2.33e-01 | 4.33e+00 +/- 1.8e-01 | 4.26e+00 +/- 1.5e-01 |
test_equilibrium_init_highres | -3.70 +/- 4.69 | -2.24e-01 +/- 2.85e-01 | 5.84e+00 +/- 2.1e-01 | 6.07e+00 +/- 1.9e-01 |
test_objective_compile_dshape_current | -5.37 +/- 4.94 | -2.17e-01 +/- 2.00e-01 | 3.82e+00 +/- 2.7e-02 | 4.04e+00 +/- 2.0e-01 |
test_objective_compile_atf | -1.29 +/- 2.46 | -1.09e-01 +/- 2.08e-01 | 8.34e+00 +/- 1.1e-01 | 8.45e+00 +/- 1.8e-01 |
test_objective_compute_dshape_current | +2.68 +/- 3.78 | +3.36e-05 +/- 4.73e-05 | 1.29e-03 +/- 2.5e-05 | 1.25e-03 +/- 4.0e-05 |
test_objective_compute_atf | +1.68 +/- 3.86 | +7.19e-05 +/- 1.65e-04 | 4.34e-03 +/- 1.2e-04 | 4.27e-03 +/- 1.1e-04 |
test_objective_jac_dshape_current | -2.08 +/- 5.23 | -8.05e-04 +/- 2.03e-03 | 3.79e-02 +/- 1.5e-03 | 3.87e-02 +/- 1.4e-03 |
test_objective_jac_atf | +0.24 +/- 3.52 | +4.57e-03 +/- 6.65e-02 | 1.89e+00 +/- 3.3e-02 | 1.89e+00 +/- 5.8e-02 |
test_perturb_1 | -0.85 +/- 1.71 | -1.16e-01 +/- 2.32e-01 | 1.34e+01 +/- 1.2e-01 | 1.35e+01 +/- 2.0e-01 |
test_perturb_2 | +0.93 +/- 2.35 | +1.78e-01 +/- 4.49e-01 | 1.93e+01 +/- 2.9e-01 | 1.91e+01 +/- 3.4e-01 |
test_proximal_jac_atf | -0.57 +/- 1.36 | -4.63e-02 +/- 1.10e-01 | 8.08e+00 +/- 4.8e-02 | 8.12e+00 +/- 9.9e-02 |
test_proximal_freeb_compute | +1.96 +/- 0.80 | +3.50e-03 +/- 1.42e-03 | 1.82e-01 +/- 1.1e-03 | 1.79e-01 +/- 8.9e-04 |
test_proximal_freeb_jac | +2.70 +/- 1.67 | +1.98e-01 +/- 1.22e-01 | 7.53e+00 +/- 1.1e-01 | 7.33e+00 +/- 5.0e-02 |
test_solve_fixed_iter | -2.03 +/- 12.76 | -3.78e-01 +/- 2.38e+00 | 1.83e+01 +/- 1.5e+00 | 1.86e+01 +/- 1.9e+00 | |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #854 +/- ##
==========================================
- Coverage 95.43% 95.40% -0.03%
==========================================
Files 87 88 +1
Lines 22313 22679 +366
==========================================
+ Hits 21294 21637 +343
- Misses 1019 1042 +23
|
See also #719 for some other related ideas. |
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.
Added a few suggestions. I'll add the analytical test soon. It's somewhere in my old laptop.
|
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.
Spline num/den separately, evaluate integrals with quadratures
…repeat... and add backend.take
I would like to use a Hermite spline for |B|, but it looks like interpax doesn't support that. Should I add that? |
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.
I have an idea for a simple tokamak test that we can use to test this PR. I'll add the details of that test and more comments soon.
could add one-liner jnp.where statement to zero out quadrature nodes at middle peak of W shaped wells. this option would resolve numerical instability for strong singularities, which could potentially cause a spike in profile |
|
||
import numpy as np | ||
from interpax import CubicHermiteSpline, PPoly, interp1d | ||
from jax.nn import softmax |
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.
should import this in backend
from interpax import CubicHermiteSpline, PPoly, interp1d | ||
from jax.nn import softmax | ||
from matplotlib import pyplot as plt | ||
from orthax.legendre import leggauss |
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.
pretty sure we could use scipy for this to avoid an extra dependency (we use scipy gauss legendre for the bootstrap stuff). Also orthax is still very WIP.
# Compute from analytic formula to avoid the issue of complex roots with small | ||
# imaginary parts and to avoid nan in gradient. | ||
r = func[c.shape[0]](*c[:-1], c[-1] - k, sentinel, eps, distinct) | ||
distinct = distinct and c.shape[0] > 3 |
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.
can't you have 2 duplicate real roots for a quadratic?
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.
yes, but i removed the duplicate in _root_quadratic
First axis enumerates the coefficients of power series. Second axis | ||
enumerates the splines along the field lines. Last axis enumerates the | ||
polynomials that compose the spline along a particular field line. | ||
B_z_ra_c : jnp.ndarray |
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.
isnt this just polyder(B_c)
?
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.
yes, but I don't wan't to recompute that for every integral of a particular pitch when i loop over bounce_integrate
return result | ||
|
||
|
||
def bounce_integral( |
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.
is this the only "entry point" to this new code? I assume all the _*
functions are only needed locally, but are the other non _*
functions meant to be semi-public (ie used by developers?) Or is this the only one?
return _filter_distinct(r, sentinel, eps) if distinct else r | ||
|
||
|
||
def _poly_der(c): |
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.
isn't this the same as jnp.polyder
?
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.
I could replace with
def jax_poly_der(c):
return jnp.vectorize(jnp.polyder, signature="(m)->(n)")(c.T).T
Under jit, they are the same speed and memory. However, this orders of magnitude slower on numpy backend.
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.
Actually testing again today looks like the bumpy backend does fine too. I only tested cpu though
def _poly_val(x, c): | ||
"""Evaluate the set of polynomials ``c`` at the points ``x``. | ||
|
||
Note this function is not the same as ``np.polynomial.polynomial.polyval(x,c)``. |
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.
np.polyval
assumes c
is in decreasing order of powers, while np.polynomial.polynomial.polyval
assumes c
is increasing powers of x. In theory would jnp.polyval
work here? (apart from it being possibly slower as it uses horners method which isn't great on gpu)
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.
Yes
Each of these are equivalent. Speed and memory under jit on cpu shows all are similar. Method 1 is a little faster on cpu, and likely faster still on gpu. Method 1 is my preference
# method 1
val = jnp.linalg.vecdot(
polyvander(x, c.shape[0] - 1), jnp.moveaxis(jnp.flipud(c), 0, -1)
)
# method 2
val = jnp.vectorize(jnp.polyval, signature="(m),(n)->(n)")(
jnp.moveaxis(c, 0, -1), x[..., np.newaxis]
).squeeze(axis=-1)
# method 3
val = jnp.einsum(
"...i,i...", x[..., np.newaxis] ** jnp.arange(c.shape[0] - 1, -1, -1), c
)
Main question I forgot to include in the actual review comment: There are a few different automorphisms (sin, arcsin) and different quadrature methods (gauss-legendre, gauss-chebyshev, tanh-sinh). Is there a clear winning combination? Or is it still problem dependent? |
This PR is completed. An alternative PR with more intelligent math (see #1045 ) will likely replace it. We should keep/merge both for comparison though.
After this pull request: