valerian-chen
commented
1 year ago Looks excellent, very excited for this. I've been thinking of reducing the redundant calls to evaluate H, but your solution is much more elegant than the few ideas I've had.
Just a minor note g and b are not parallel, so the components of eps should either be in the (x,y,g) basis or the (perp_1, perp_2,parallel) basis. If I understand you, it's probably the latter, but I'll go through the code carefully and figure it out. Also, at first glance, perhaps DielectricTensor rather than PolarisationTensor? I'll read through the code first and make sure I've got you, though.
ZedThree
(WIP till #77 is merged)
To do:
Refactor the Hamiltonian code (
find_Hand thefind_dH_*functions). Deletes almost 4k lines and has a 30% speed-up.90% of Scotty's run time is in evaluating the Hamiltonian,
find_H. If we can make that faster, we can improve the performance of the whole code. This PR does that by roughly halving the number of times the Hamiltonian needs to be evaluated through reusing stencil points in the derivatives.Instead of calling
find_dH_*for each derivative we want,Hamiltonian.derivativescomputes all the derivatives at once so that any common points can be reused between different derivatives.For instance, $\partialR$ includes all the points for $\partial^2{RR}$, so we only need to evaluate $H$ at one more point and we can immediately calculate the second derivative. I played around with some fancier stencils for the mixed derivatives, but although they kept the error scaling at second order, they had a larger error coefficient, which resulted in
solve_ivptaking more steps, wiping out any performance gains. I did use a slightly different stencil for the forward x central difference derivatives, like $\partial^2_{RK_R}$, as this did improve things overall.How this works in practice:
{(-1,): -0.5, (1,): 0.5}defines the central first differenceHamiltonian.derivativesuses its internalapply_stencilto calculate the derivative for a tuple of directions such as("q_R", )for $\partial_R$ and("q_R", "q_Z")for $\partial^2_{RZ}$apply_stenciluses the stencil to create aCoordOffsetobject that holds the offsets for all directions for each point in the stencil that needs evaluatingdict) to see if $H$ has already been evaluated at this offset. If it has, it can be used immediately, otherwise $H$ is evaluated at the actual position and stored in the cacheapply_stencilis just sum the product of the weights and $H$ at each offset, and divide by the product of the grid spacingsThe other change to the Hamiltonian is to pull out a
PolarisationTensorclass that caches $\epsilon\Vert, \epsilon\perp, \epsilon_g$ to be reused in calculating $\alpha, \beta, \gamma$With these changes, we can then create a
Hamiltonianobject that contains some of the information that is reused, such as the field and density objects, and then we can pass much fewer arguments to the solver.