I was playing with the fast multipole method this summer, mainly by translating the fantastic pybbfmm library from torch to jax. showcase: simulation of an infection spread over the British Isles (10M agents, 3s to simulate with jax, 3x over pytorch). - ThreadSky

About ThreadSky

maxxxzdn.bsky.social • 100 days ago

I was playing with the fast multipole method this summer, mainly by translating the fantastic pybbfmm library from torch to jax.

showcase: simulation of an infection spread over the British Isles (10M agents, 3s to simulate with jax, 3x over pytorch).

Comments

mniepert.bsky.social•99 days ago

I also played with the same bb multiple pole implementation a few months ago, hoping to build an end-to-end learnable multipole MPNN. So learn a kernel function for 3d point clouds to model long-range interactions. Unfortunately I had to give up due to lack of time and the complex code. :)

maxxxzdn.bsky.social•99 days ago

Ah nice to know, our motivation pretty much aligned! I came in the end to a prototype, even benchmarked it on LRGB, but unfortunately, it paled to transformers in any way so I had to move on :(

maxxxzdn.bsky.social•99 days ago

Although I admit the way I approached it was not optimal (it required kernel evaluation at every point, which blew up memory), I will likely do it differently if I have time again.

maxxxzdn.bsky.social•100 days ago

btw uploading a gif here was quite a challenge but this tool seems to work: https://mozzius-gif_upload.web.val.run/

maxxxzdn.bsky.social•100 days ago

learned a ton by working on it, mainly on the ML systems side, but result-wise the GIF is the only outcome 😁

maxxxzdn.bsky.social•100 days ago

the code turned out to be a mishmash of torch and jax, mainly due to jax's not supporting dynamic arrays which doesn't stack well with trees. one day, perhaps?

link to the original repo: https://andyljones.com/pybbfmm/index.html

Posting Rules

Be respectful to others
No spam or self-promotion
Stay on topic
Follow Bluesky's terms of service

Comments

Posting Rules

Reply