DESC icon indicating copy to clipboard operation
DESC copied to clipboard

Add Ability to Write Boozmn.nc style output files

Open dpanici opened this issue 1 year ago • 33 comments

Merge in #1211 first, THEN merge in #1197 , then modify this PR to use the vectorized boozer transform from #1197, then it is ready to merge

Add make_boozmn_output utility function to vmec_utils, which allows DESC to write boozmn.nc style output files containing Boozer harmonics. I/O based heavily off of the hidden symmetries version of booz_xform

  • [x] Add test of DESC-computed quantities versus boozmn file written - very good comparison
  • [x] Add test of our boozmn.nc against an actual C++ version of booz_xform output file at same resolution/nsurfs and equilibrium - comparison not great, need to debug still (check order of m,n possibly)
  • [x] improve speed by computing all data at once
  • [x] make the prefactor and ["B"] basis evaluated at the necessary rho,theta_B,zeta_B nodes a name in data index (so dont need to recompute the sum of basis modes each time)
  • [x] add non-symmetric test and fix for non-symmetric
  • [ ] add documentation (in form of notebook or a page) on our Boozer transform setup and how to use and the meaning of our "|B|_mn_B" and how to get into the familiar double angle form
  • [ ] update algorithm to use vectorized boozer transform now that #1197 is merged

future work (or work if someone else wants to do on this PR go ahead)

  • [x] add capacity to save info at any surface, not just the half grid? though would need to update docstrings for everything saved... and unsure of who would use that as anyone using the .nc likely just wants it in the usual format with the half grid - pushing to a future PR
  • [ ] cannot precompute quantities that assume a single surface is passed in, i.e. theta_B, zeta_B which need nu which needs w_Boozer which needs B_theta_mn.... which itself is found with transforms["B"].fit("B_theta"), so that quantity cannot be computed using a volumetric grid right now (maybe, can FourierZernike fit then evaluate the fourier series? and pass that into the compute fxn)
  • [ ] precomputing transforms["B"], or possibly using same transform and modifying the basis matrices by multiplying by zernike_radial(rho_i)/zernike_radial(rho_{i-1}) (unsure if this second part is really possible, for now I jitted the transform build and it works reasonably well)
  • [x] #720 could help speed this up as a nontrivial amount of time is spent in zernike_to_fourier - pushing to a future PR

I will note that when comparing against the Cpp or Fortran version, DESC will always be more accurate as the other versions must interpolate the Fourier coefficients to get the Rmnc etc on the half grid, while we can exactly calculate what they are on the half grid because we know quantities for every rho in [0,1]

timing info

The time it takes to save the HELIOTRON example to a wout file with ns=100 is ~40s, then running the C++ version of booz_xform with M_Booz=N_Booz=20 on a single thread (remember it is multithreaded so this) is very quick (10s) and on 8 threads is 2s

That is compared to running the make_boozmn_output function in DESC, which takes 90s (50s on gpu though and the util function uses numpy so a lot of memory movement is happening). So there is definitely room for improvement here, on the CPU and on GPU, but at least on GPU it is comparable speed.

dpanici avatar Sep 19 '23 16:09 dpanici