PlasmaControl / DESC

Stellarator Equilibrium and Optimization Suite
MIT License
97 stars 26 forks source link

Use FFTs in poloidal direction where possible #641

Open f0uriest opened 1 year ago

f0uriest commented 1 year ago

Several objectives now use QuadratureGrid or LinearGrid for their default (to get accurate volume or surface averages), and because they are uniform in the poloidal direction we could possibly use FFTs to do that part of the transform. This would like see a decent speed bump and memory savings, similar to using FFTs for the poloidal direction. This would be added as another option to the Transform class.

rahulgaur104 commented 8 months ago

I am interested in this one. So, it would be like poloidal * Grid where poloidal is a FourierRZToroidalSurface object. I am assume the poloidal dependence is separable from the non-poloidal dependence.

f0uriest commented 8 months ago

yeah you can take a look at desc.transform.Transform to see how we do it for toroidal stuff. I think you'd first do a regular matmul for the radial part, which goes from the fourier-zernike basis to a double fourier series on each surface, then do an fft in both poloidal and toroidal directions (possibly simultaneously with fft2?). The main annoyance is that our fourier series are in the double angle form so you need to reshape and re-arange and pad stuff to be able to use a regular fft. See how we do it for the toroidal part, its probably similar though might also get complicated if the grid is symmetric (since that means poloidal nodes are ommitted).