ALCHEMI-Toolkit wrapper for Orb

[vsimkus]

Summary

• Adds OrbWrapper, a nvalchemi-toolkit integration that lets Orb models plug into nvalchemi dynamics engines (FIRE, MD, etc.) and pipeline composition (e.g. Orb + DFT-D3; conservative-models only)
• Adds README examples how to use the wrapper

[vsimkus]

• ALCHEMI-Toolkit wrapper for Orb #167 👈 (View in Graphite)
• Fix DirectRegressor compilation + other internal changes #166
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[socket-security]

Review the following changes in direct dependencies. Learn more about Socket for GitHub.

Diff	Package	Supply Chain Security	Vulnerability	Quality	Maintenance	License
	pypi/​nvalchemi-toolkit@​0.1.0
	pypi/​ruff@​0.16.1 ⏵ 0.15.22
	pypi/​nvalchemi-toolkit-ops@​0.3.1 ⏵ 0.4.1

View full report

[socket-security]

Warning

Review the following alerts detected in dependencies.

According to your organization's Security Policy, it is recommended to resolve "Warn" alerts. Learn more about Socket for GitHub.

Action	Severity	Alert  (click "▶" to expand/collapse)
Warn		Obfuscated code: pypi loguru is 90.0% likely obfuscated Confidence: 0.90 Location: Package overview From: ? → pypi/nvalchemi-toolkit@0.1.0 → pypi/loguru@0.7.3 ℹ Read more on: This package | This alert | What is obfuscated code? Next steps: Take a moment to review the security alert above. Review the linked package source code to understand the potential risk. Ensure the package is not malicious before proceeding. If you're unsure how to proceed, reach out to your security team or ask the Socket team for help at support@socket.dev. Suggestion: Packages should not obfuscate their code. Consider not using packages with obfuscated code. Mark the package as acceptable risk. To ignore this alert only in this pull request, reply with the comment @SocketSecurity ignore pypi/loguru@0.7.3. You can also ignore all packages with @SocketSecurity ignore-all. To ignore an alert for all future pull requests, use Socket's Dashboard to change the triage state of this alert.

View full report

[benrhodes26]

LGTM 🌿

[timduignan]

I think I've found an issue, opening a PR internally for review/discussion so can we hold of on merging until thats resolved thanks.

[timduignan]

Ok nevermind I think the problem was elsewhere sorry.

[timduignan]

Heads-up on a performance issue: OrbWrapper.forward calls self.model.forward(...) directly (both branches). That bypasses nn.Module.__call__ — and since pretrained.py compiles the model via model.compile(...) (which installs the optimized graph on __call__), the model ends up running eager inside the wrapper even when loaded with compile=True. Calling it directly also skips any forward/pre-forward module hooks.

Switching both sites to self.model(...) fixes it. We hit this in our internal copy and measured ~1.63× on a water NPT run (H100, conservative orbmol model, compiled): ~7.95 → ~13.0 steps/s, back-to-back on one node — which brings the high-level NPT.run path level with a hand-rolled run_npt_step loop. The compute_forces/compute_stress/fp64_energy kwargs still forward through __call__, so it's a drop-in change.