mikegrudic / pytreegrav

Fast N-body gravitational force and potential in Python
MIT License
63 stars 9 forks source link

Added an option to use quadrupole expansions of mass #8

Closed bwkeller closed 3 years ago

bwkeller commented 3 years ago

I've been playing around with some old classic treecode papers, and pytreegrav has been really handy for that! Thanks @mikegrudic and everyone else who helped write this code. I added an option in the Octree class & methods to use a quadrupole expansion, rather than just the n=0 monopole. Naturally, this gives slightly more accurate potentials and accelerations for the same opening angle, but with the added cost of storing the moment tensors and calculating them when the tree is built. I haven't really done much optimizing, so it may not be the most efficient implementation possible. I may extend this to higher-order moments some time as well.

mikegrudic commented 3 years ago

Very cool Ben! I'll take it for a spin and if it's all looking good pull it in. Upon quick inspection, I suspect based on my numba experience that the quadrupole field summation will probably run faster if we spell out the matrix operation explicitly in a double loop, but it remains to be seen if that matters since the tree force summation is dominated by overhead.

mikegrudic commented 3 years ago

Hi @bwkeller , I pulled in your contribution and added an optional quadrupole argument to the frontend functions so it's all easy to use. Everything seems to work, but I notice that the quadrupole field summation is muuuuch slower than the extra flops would justify, to the point that it is not yet useful (it doesn't beat going to a smaller opening angle with the monopole to achieve the desired accuracy).

The first thing I'd try is getting rid of numpy matrix operations like dot, outer, etc., and just writing out the double loops. Numba seems to really reward just writing the dumbest C code possible in Python syntax. The readme notebook should be a handy benchmark. We should also probably get rid of the second recursive call of ComputeMoments in the treebuild, which isn't necessary because we've already done the work computing the moments.

mikegrudic commented 3 years ago

I also just found that merely adding the if tree.HasQuads check in the tree accel walk reduced performance by about 40%, even when quadrupoles were not in use and it added no extra flops. For this reason I have created separate treewalk functions to be used if quadrupoles are enabled, see commit. Goes to show how annoyingly precarious numba optimization is.

mikegrudic commented 3 years ago

OK in the most recent push I managed to get quadrupoles working with performance that is competitive with monopole at fixed accuracy, by trying the stuff I suggested above.

bwkeller commented 3 years ago

Holy cow, awesome! I hope I didn't cause too much work for you.

On Thu, Jun 24, 2021 at 2:48 AM Mike Grudić @.***> wrote:

OK in the most recent push I managed to get quadrupoles working with performance that is competitive with monopole at fixed accuracy, by trying the stuff I suggested above.

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://github.com/mikegrudic/pytreegrav/pull/8#issuecomment-867250253, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAHEJKTIE3ZCVHXVKVKFOOLTUJ6F5ANCNFSM4662L3YA .

mikegrudic commented 3 years ago

not at all I started this project for fun and/or procrastination