[Paper] NeuralDMD

[jacobbieker]

Arxiv/Blog/Paper Link

https://arxiv.org/pdf/2507.03094

Code: https://github.com/as2c/NeuralDMD/tree/main

Detailed Description

Interesting paper and code on a model that can do data assimilation as well as black hole reconstruction! Would be interesting to port to PyTorch and include here as another way of data assimilation.

[citronella3alain]

Hi, @jacobbieker, I'd be interested in tackling this issue. As it seems to me, the first step is translating the jax code into pytorch and then adding it as another model in the repo, correct? Thanks

[citronella3alain]

I've succeeded running the jax inference code for the weather data assimilation example in NeuralDMD; the output looks as expected. However, when I run the existing training code, be it on top of the provided weights / train from scratch, the results are abysmal; let me dig deeper.

My plan is to first make sure the jax side of things work out correctly before porting to torch; this way I'll have a point of reference to compare against.

[jacobbieker]

Great! Yeah, would be good to know why the results are not great before adding it. I think having that reference is a really good idea!

[citronella3alain]

Hm, I'm having trouble figuring out what's wrong training-side, so I created an issue in the NeuralDMD repo for now: as2c/NeuralDMD#1

In the meantime, I'll look at other good first issues here if that's fine with you @jacobbieker ?

[jacobbieker]

Yeah, no worries! A lot of the [Paper] tagged issues would also probably qualify as good first issues, they are mostly standalone if any other paper seems interesting!

[glitch401]

Hi @jacobbieker, I'm interested in contributing and was looking at this issue. I had a quick question about the current status, if you don't mind.

I saw @citronella3alain's helpful comments from July about the difficulties with training the original JAX model. To follow up on their investigation, I checked the upstream as2c/NeuralDMD repository. I noticed that the issue they opened (as2c/NeuralDMD#1) is still open, and it looks like the last code commit was from before their comment.

I was just wondering if this means the original training bug is likely still present? Am I understanding the situation correctly, or am I perhaps missing some context? Just wanted to check before diving in

[jacobbieker]

Hi,

I haven't really looked into it, so I'm not sure. Possibly asking on that issue might help? If you want to dive in here, there are quite a few other model implementation issues that might not have that problem.

Thanks,
Jacob

[glitch401]

Hi @jacobbieker,

Thanks so much for the quick response and the encouragement! I did leave a comment on the issue there.
I am reviewing other paper implementation issues that might be good to start off with