Make a differentiable version of FourierRZToroidalSurface.constant_offset_surface
#2016
Conversation
| 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 |
There was a problem hiding this comment.
only is differentiable though if params is passed, otherwise the end will not be differentiatedcorrectlt
Memory benchmark result| Test Name | %Δ | Master (MB) | PR (MB) | Δ (MB) | Time PR (s) | Time Master (s) |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
test_objective_jac_w7x | 5.55 % | 3.831e+03 | 4.044e+03 | 212.55 | 38.55 | 36.15 |
test_proximal_jac_w7x_with_eq_update | 1.87 % | 6.486e+03 | 6.608e+03 | 121.59 | 159.86 | 159.77 |
test_proximal_freeb_jac | -0.26 % | 1.321e+04 | 1.318e+04 | -33.90 | 83.13 | 81.62 |
test_proximal_freeb_jac_blocked | 1.07 % | 7.456e+03 | 7.536e+03 | 80.13 | 72.29 | 72.20 |
test_proximal_freeb_jac_batched | -0.82 % | 7.489e+03 | 7.428e+03 | -61.05 | 72.29 | 72.70 |
test_proximal_jac_ripple | -2.18 % | 3.540e+03 | 3.463e+03 | -77.25 | 64.68 | 64.59 |
test_proximal_jac_ripple_bounce1d | 0.17 % | 3.510e+03 | 3.515e+03 | 5.80 | 75.34 | 75.43 |
test_eq_solve | 1.53 % | 1.975e+03 | 2.006e+03 | 30.28 | 93.17 | 93.26 |For the memory plots, go to the summary of |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #2016 +/- ##
==========================================
- Coverage 95.77% 95.52% -0.26%
==========================================
Files 101 101
Lines 27796 27881 +85
==========================================
+ Hits 26622 26632 +10
- Misses 1174 1249 +75
🚀 New features to boost your workflow:
|
…nd to make jitable version also perform the fit
dpanici
requested review from
a team,
YigitElma,
ddudt,
f0uriest,
rahulgaur104 and
unalmis
and removed request for
a team
| params = base_surface.params_dict | ||
|
|
||
| def n_and_r_jax(nodes): | ||
| data = base_surface.compute( |
There was a problem hiding this comment.
Are we sure this is differentiable?
| to create the resulting constant offset surface, by default equal | ||
| to base_surface.N. If basis is given, this is ignored. | ||
| R_basis, Z_basis: DoubleFourierSeries, optional | ||
| Basis to use to fit the offset surface's R and Z, respectivelu. If None, |
There was a problem hiding this comment.
| Basis to use to fit the offset surface's R and Z, respectivelu. If None, | |
| Basis to use to fit the offset surface's R and Z, respectively. If None, |
| 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 passedto their compute methods) |
There was a problem hiding this comment.
| that are passedto their compute methods) | |
| that are passed to their compute methods) |
| data["x"] = xyz2rpz(x) | ||
| data["x_offset_surface"] = xyz2rpz(x_offsets) | ||
|
|
||
| offset_zetas = data["x_offset_surface"][:, 1] |
There was a problem hiding this comment.
I don't think I follow the code. I know this was added in a previous PR but I wanted to ask to see if this code is correct or not. So, the above rootfind basically tries to find some zeta where if you start from (theta, zeta) and go in the normal direction by offset amount, your zeta is equal to the original grid's zeta. In short, find some points on the original surface with uniform spacing in theta but non-uniform spacing in zeta that give points in offset surface that are uniformly spaced in zeta (and assign a different theta definition that is uniform).
Why do you find the offset zetas here? They should be already equal to grid.nodes[:, 2]?
| n, x, x_offsets = n_and_r_jax(nodes) | ||
|
|
||
| data = {} | ||
| data["n"] = xyz2rpz_vec(n, phi=nodes[:, 1]) |
There was a problem hiding this comment.
nodes[:, 1] is theta, isn't it?
| Toroidal resolution of the basis used to fit the offset points | ||
| to create the resulting constant offset surface, by default equal | ||
| to base_surface.N. If basis is given, this is ignored. | ||
| R_basis, Z_basis: DoubleFourierSeries, optional |
There was a problem hiding this comment.
Giving the basis doesn't seem efficient, why not give the transforms directly? Instead of rebuilding them everytime, they will be reused, and they store the information on basis too if you need to access that.
| data["x_offset_surface"] = xyz2rpz(x_offsets) | ||
|
|
||
| offset_zetas = data["x_offset_surface"][:, 1] | ||
| offset_rtz_nodes = jnp.vstack( |
There was a problem hiding this comment.
I guess since offset_zeta==grid.nodes[:,2] and you choose some new theta definition for the offset surface that matches grid.nodes[:,1], you can just use the original grid here. And since that will stay constant, you can pass the pre-built transforms directly
| ``x`` points on the base_surface (i.e. the points to which the | ||
| offset surface was fit) | ||
| as well as the DoubleFourierSeries bases used to fit R and Z. | ||
| Only returned if ``full_output`` is True |
There was a problem hiding this comment.
Docs of the grid can also be updated. The default grid resolution is twice the M,N
Adds differentiable version of offset surface algorithm. Potentially useful if one desires an objective that concerns a conformal vessel to a changing eq during optimization.
notebook showing example use:
test_differentiable_offset.ipynb