enh(metatomic): Vesin first pass #1

HaoZeke · 2025-09-15T12:33:59Z

Reopened.

Very rough WIP. To be updated with https://gitlab.com/gromacs/gromacs/-/merge_requests/4323
(after tests run). Modeled (currently) after the nnpot and plumed/eon integrations via vesin.

For instructions on optimal usage / reproduction consider details here.

HaoZeke · 2025-09-15T12:38:36Z

needs CI imported from https://github.com/gromacs/gromacs/blob/main/.github/workflows/build_cmake.yml

PicoCentauri

Very cool. I took a quick look on one file. I will continue later.

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

src/gromacs/applied_forces/metatomic/metatomic_mdmodule.h

PicoCentauri

I think we also need the TopologyPreprocessor to remove the LJ, charges etc as the NNpot and the QMMM module is doing it.

Otherwise it looks already almost functional!

src/gromacs/applied_forces/metatomic/metatomic_options.cpp

as is done in nnpot

Co-authored-by: PicoCentauri <[email protected]>

Co-authored-by: luthaf <[email protected]>

NVIDIA/cuda-samples#333

cmake/gmxManageMetatomic.cmake

HaoZeke · 2025-11-06T02:38:53Z

First compiling build... (links too which is nice)

Co-authored-by: PicoCentauri <[email protected]>

... and cleanup

HaoZeke · 2025-11-07T11:32:04Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

+    // "Broadcast" the static atom numbers to all ranks: all ranks must call this.
+    if (mpiComm_.isParallel())
+    {
+        // resize/allocate on non-main ranks and broadcast raw bytes of vector
+        nblock_abc(mpiComm_.isMainRank(), mpiComm_.comm(), static_cast<std::size_t>(n_atoms), &atomNumbers_);
+    }


This is a bit unclear to me, it works under the assumption that main knows all the atom numbers, something I can't seem to figure out.

https://manual.gromacs.org/current/reference-manual/algorithms/parallelization-domain-decomp.html#domain-decomposition-flow-chart

Looks like main will have the right set of atoms?

Co-authored-by: PicoCentauri <[email protected]>

HaoZeke · 2025-11-07T16:08:28Z

@Luthaf this is ready for a drive-by, @PicoCentauri and I found a bug in the LJ forces for NNPot and this, which we're opening upstream.

PicoCentauri · 2025-11-08T16:09:45Z

This is the issue for reference.

https://gitlab.com/gromacs/gromacs/-/issues/5477#note_2876613421

cmake/gmxManageMetatomic.cmake

src/external/tng_io/src/tests/CMakeLists.txt

src/gromacs/applied_forces/metatomic/CMakeLists.txt

Luthaf · 2025-11-12T09:48:34Z

...pplied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputDefaultValuesWhenActive.xml

+metatomic-input_group        = System
+metatomic-model              = 
+metatomic-extensions         = 
+metatomic-device             = 
+metatomic-check_consistency  = true


Naming question, should these use metatomic_check_consistency or metatomic-check-consistency instead? It is a bit weird to have a mix of dash and underscore.

I think all - should be gromacs style.

Luthaf · 2025-11-12T10:07:21Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

+    // Its ONLY job is to update the `idxLookup_` map, which maps
+    // the contiguous ML atom index (0..n_atoms-1) to the


I don't understand what you mean here. If all ranks have access to all positions, then we can give all atoms to the model in each rank, and use selected_atoms to pass the sparse index/what the model in a given rank should be doing.

Luthaf · 2025-11-12T10:08:36Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

+        gromacs_scalar_type = torch::kFloat64;
+    }
+    // This blob_options is for reading GROMACS memory
+    auto blob_options = torch::TensorOptions().dtype(gromacs_scalar_type).device(torch::kCPU);


do we need to check something to know if the data is on CPU? Or will the GPU data live in separate variables?

@PicoCentauri and I shall investigate / ping lukas

Luthaf · 2025-11-12T10:09:20Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

+        auto torch_cell =
+                torch::from_blob(&box_, { 3, 3 }, blob_options).to(this->dtype_).to(this->device_);
+
+        auto torch_pbc = torch::tensor(std::vector<uint8_t>{ 1, 1, 1 },


are we sure that gromacs never have mixed pbc?

No I should do something like in the nnpot, i.e.

void TorchModel::preparePbcType(PbcType* pbcType) { torch::Tensor pbcTensor = torch::tensor({ true, true, true }, torch::TensorOptions().dtype(torch::kBool)); if (*pbcType == PbcType::XY) { pbcTensor[2] = false; } else if (*pbcType != PbcType::Xyz) { GMX_THROW(InconsistentInputError( "Option use_pbc was set to true, but PBC type is not supported.")); } pbcTensor = pbcTensor.to(device_); inputs_.push_back(pbcTensor); }

Luthaf · 2025-11-12T10:10:08Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

+        auto torch_pbc = torch::tensor(std::vector<uint8_t>{ 1, 1, 1 },
+                                       torch::TensorOptions().dtype(torch::kBool).device(this->device_));
+        auto torch_types =
+                torch::tensor(atomNumbers_, torch::TensorOptions().dtype(torch::kInt32)).to(this->device_);


So this is using the atomic number as atomic type, right? How would this handle CG simulations?

I think those need separate code, and have to register as supporting coarse-grained force fields

xref:
[1]: https://manual.gromacs.org/documentation/2019-rc1/reference-manual/functions/force-field.html#coarse-grained-force-fields
[2]: tpxio.cpp --> tpxv_RestrictedBendingAndCombinedAngleTorsionPotentials, /**< potentials for supporting coarse-grained force fields */

And I think then also the atomNumbers_ will be different. I will check how we get them actually from the defined topology.

Okay they come from the atomtypes directive in the topolgy file. I checked a couple of predefined topologies that live here:

https://github.com/metatensor/gromacs/tree/main/share/top

and they seem to define the atomnumbers in the canonical way

gromacs/share/top/amber99sb-ildn.ff/ffnonbonded.itp

Lines 1 to 6 in eaa82eb

[ atomtypes ]

; name at.num mass charge ptype sigma epsilon

Br 35 79.90 0.0000 A 3.95559e-01 1.33888e+00 ; Converted from parm99.dat

C 6 12.01 0.0000 A 3.39967e-01 3.59824e-01

CA 6 12.01 0.0000 A 3.39967e-01 3.59824e-01

CB 6 12.01 0.0000 A 3.39967e-01 3.59824e-01

The only thing we might want to filter out with a warning is atom number 0 as they use as these are dummy particles for some forcefields

gromacs/share/top/amber99sb-ildn.ff/ffnonbonded.itp

Lines 70 to 74 in eaa82eb

; MW=Dummy mass for tip4p/EW/5p water extra point charge

MW 0 0.0000 0.0000 D 0.00000e+00 0.00000e+00

; Dummy masses for rigid CH3 and NH3 groups

MCH3 0 0.0000 0.0000 A 0.00000e+00 0.00000e+00

MNH3 0 0.0000 0.0000 A 0.00000e+00 0.00000e+00

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

Co-authored-by: Guillaume Fraux <[email protected]>

Co-authored-by: Guillaume Fraux <[email protected]> Update src/gromacs/applied_forces/metatomic/metatomic_options.cpp Co-authored-by: Guillaume Fraux <[email protected]> Update src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.h Co-authored-by: Guillaume Fraux <[email protected]> Update src/gromacs/applied_forces/metatomic/metatomic_topologypreprocessor.cpp Co-authored-by: Guillaume Fraux <[email protected]>

PicoCentauri · 2025-11-12T14:54:32Z

...pplied_forces/metatomic/tests/refdata/MetatomicOptionsTest_OutputDefaultValuesWhenActive.xml

+metatomic-input_group        = System
+metatomic-model              = 
+metatomic-extensions         = 
+metatomic-device             = 
+metatomic-check_consistency  = true


I think all - should be gromacs style.

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.h

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

Co-authored-by: Guillaume Fraux <[email protected]>

PicoCentauri · 2025-11-13T10:51:29Z

src/gromacs/applied_forces/metatomic/metatomic_forceprovider_stub.cpp

+#include "gromacs/utility/basedefinitions.h"
+#include "gromacs/utility/exceptions.h"
+
+#include "metatomic_forceprovider.h"


This will cause problems when building without metatomic because we will not find this header file.

Luthaf and others added 8 commits September 15, 2025 14:31

Intitial scaffolding of metatensor interface

43e0fe7

[MAINT] Fix merge cruft

a3fc5c9

chore(metatensor): orchestrate vesin design

9fbf93f

chore(build): try to find metatensor

8932089

chore(metatomic): trace nnpot through the code

11f4ff2

chore(metatomic): register as an energy output

2f9b51d

chore(metatomic): rename for the new ecosystem

1ac6b17

chore(metatomic): regex rename for metatomic

e0d7995

PicoCentauri reviewed Sep 15, 2025

View reviewed changes

HaoZeke commented Sep 15, 2025

View reviewed changes

src/gromacs/applied_forces/metatomic/metatomic_mdmodule.h Outdated Show resolved Hide resolved

PicoCentauri reviewed Sep 15, 2025

View reviewed changes

src/gromacs/applied_forces/metatomic/metatomic_options.cpp Show resolved Hide resolved

HaoZeke and others added 8 commits September 15, 2025 21:24

chore(metatomic): add topology processing

8d58263

as is done in nnpot

maint(metatomic): docstring update and formatting

f6837a5

maint(params): use the right variables

799454d

doc(metatomic): add a note wrt tpr

18a2a95

Co-authored-by: PicoCentauri <[email protected]>

chore(metatomic): rework cmake dependencies

0138afc

chore(metatomic): try to use more git and fallback

3d964d5

chore(metatomic): kang from LAMMPS

bbbff69

Co-authored-by: luthaf <[email protected]>

chore(openmp): fix for cuda compilers

6a69669

NVIDIA/cuda-samples#333

HaoZeke force-pushed the useVesin branch from 2ddf513 to 6a69669 Compare September 15, 2025 21:19

chore(metatomic): resolve more compilation errors

c23f94a

HaoZeke force-pushed the useVesin branch from 7877e92 to b894fc3 Compare September 15, 2025 22:06

chore(metatomic): try to fix compilers

58bea16

HaoZeke force-pushed the useVesin branch from b894fc3 to 58bea16 Compare September 15, 2025 22:11

PicoCentauri reviewed Sep 16, 2025

View reviewed changes

cmake/gmxManageMetatomic.cmake Show resolved Hide resolved

chore(metatomic): grab vesin through cmake

9212e26

HaoZeke force-pushed the useVesin branch from 8756722 to 9212e26 Compare September 16, 2025 15:18

HaoZeke added 2 commits September 16, 2025 18:00

chore(metatomic): remove caffe2 mkl workaround

b6e3d9e

chore(metatomic): easier imports

e24d9a0

HaoZeke requested a review from PicoCentauri November 6, 2025 02:38

HaoZeke and others added 10 commits November 6, 2025 16:23

chore(mtaopt): rework in the style of nnpot

4ee2d02

Co-authored-by: PicoCentauri <[email protected]>

chore(mdmta): register everywhere

899126d

tst(mtaopt): add reference tests for options

1744fa6

bug(mtaforce): ensure compatible extensions

15438cd

bug(mtaforce): fix vesin calling conventions

ecb1b31

chore(mtaforce): use the model dtype

51c945e

chore(mtaforce): cleanup debugging

4c71a0d

chore(mtapot): update for mpi usage and efficiency

5d78bf3

chore(mtaforce): defensively handle dtypes

71a85e6

... and cleanup

feat(mtapot): use more internal helpers

26583ee

HaoZeke commented Nov 7, 2025

View reviewed changes

HaoZeke and others added 2 commits November 7, 2025 14:51

chore(mtapot): cleanup a bit

dbe5379

chore(typo): fixup

82c786d

Co-authored-by: PicoCentauri <[email protected]>

Luthaf reviewed Nov 12, 2025

View reviewed changes

HaoZeke force-pushed the useVesin branch 2 times, most recently from 0cc95c2 to 15ab6cc Compare November 12, 2025 13:53

chore(cmake): rework for older versions and clean

fcd9e5f

Co-authored-by: Guillaume Fraux <[email protected]>

HaoZeke force-pushed the useVesin branch 2 times, most recently from 26a7592 to 4168411 Compare November 12, 2025 14:09

HaoZeke force-pushed the useVesin branch from 4168411 to f4b0bb7 Compare November 12, 2025 14:16

PicoCentauri reviewed Nov 12, 2025

View reviewed changes

HaoZeke mentioned this pull request Nov 13, 2025

enh(metatomic): use more ergonomic torch TheochemUI/eOn#269

Open

Update src/gromacs/applied_forces/metatomic/metatomic_forceprovider.h

c619614

Co-authored-by: Guillaume Fraux <[email protected]>

PicoCentauri reviewed Nov 13, 2025

View reviewed changes

HaoZeke mentioned this pull request Nov 14, 2025

enh(metatomic): Add uncertainty and use metatomic_torch::Module TheochemUI/eOn#270

Open

		// Its ONLY job is to update the `idxLookup_` map, which maps
		// the contiguous ML atom index (0..n_atoms-1) to the

	[ atomtypes ]
	; name at.num mass charge ptype sigma epsilon
	Br 35 79.90 0.0000 A 3.95559e-01 1.33888e+00 ; Converted from parm99.dat
	C 6 12.01 0.0000 A 3.39967e-01 3.59824e-01
	CA 6 12.01 0.0000 A 3.39967e-01 3.59824e-01
	CB 6 12.01 0.0000 A 3.39967e-01 3.59824e-01

	; MW=Dummy mass for tip4p/EW/5p water extra point charge
	MW 0 0.0000 0.0000 D 0.00000e+00 0.00000e+00
	; Dummy masses for rigid CH3 and NH3 groups
	MCH3 0 0.0000 0.0000 A 0.00000e+00 0.00000e+00
	MNH3 0 0.0000 0.0000 A 0.00000e+00 0.00000e+00

enh(metatomic): Vesin first pass #1

Are you sure you want to change the base?

enh(metatomic): Vesin first pass #1

Uh oh!

Conversation

HaoZeke commented Sep 15, 2025

Uh oh!

HaoZeke commented Sep 15, 2025

Uh oh!

PicoCentauri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

PicoCentauri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

HaoZeke commented Nov 6, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HaoZeke commented Nov 7, 2025

Uh oh!

PicoCentauri commented Nov 8, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants