Add Jacobian Symmetry physicality benchmark#713
Conversation
ElliottKasoar
added
the
new benchmark
There was a problem hiding this comment.
it would be great to have some form of interactivity here, do you have any ideas what you'd like to do? one idea would be a bar chart (bar for each molecule), like in the water slab dipoles, which you can click on to see the structures
There was a problem hiding this comment.
Went with the bar chart idea: click a model's cell → bar chart of λ per structure, then click a bar → a heatmap of that structure's antisymmetric Jacobian showing where the symmetry breaks. Also added N/A markers for structures a model can't run.
| # ORB models compile their forward pass with torch.compile/dynamo. Unlike | ||
| # other benchmarks, this one cycles through structures with many different | ||
| # atom counts and immediately hammers each freshly-compiled graph with up to | ||
| # 2*3N forward calls, which most likely triggers a low-level SIGABRT inside | ||
| # torchinductor's dlopen'd kernels. Disabling compilation avoids it. | ||
| # torch._dynamo.config reads this env var at import time, so it must be set | ||
| # before torch (or anything importing it) is loaded anywhere in the process. | ||
| os.environ["TORCH_COMPILE_DISABLE"] = "1" |
There was a problem hiding this comment.
ive managed to get around this issue by changing frame.calc = copy(calc) to frame.calc = calc. you can then do frame.calc = None before appending the frame, like in S30L/plf547 benchmarks
There was a problem hiding this comment.
opinions on this @ElliottKasoar, i think this is safe to do?
|
Thanks for the PR! just left a few comments for now, mostly minor |
2d02a5e to
eb6572e
Compare
3 participants
Pre-review checklist for PR author
PR author must check the checkboxes below when creating the PR.
Summary
Adds a "Jacobian Symmetry" physicality benchmark that tests whether a model's predicted forces are conservative (derivatives of a single potential energy). The force Jacobian is built by central finite differences and the antisymmetric fraction of its Frobenius norm, lambda = ||J_anti|| / ||J||, is reported (mean and max over 10 diverse structures). lambda = 0 for perfectly conservative forces. Based on Bigi, Langer & Ceriotti, arXiv:2412.11569.
Linked issue
Resolves #712
Progress
Testing
Tested on all 21 registered models
New decorators/callbacks
None
AI usage disclosure
I developed this benchmark with some assistance from Claude Code. I have personally reviewed, tested, and understand all of the code.