MLIAP Unified Generalized Interface #4673

codeformore · 2025-08-04T16:43:15Z

Summary

This implementation introduces the ability to extract named per-atom "extra properties" from machine-learned interatomic potentials (MLIAPs) implementing the unified interface, such as uncertainty metrics from ensemble models. It includes both infrastructure support in the MLIAP data handling and a new compute command for use in LAMMPS input scripts.

Key features include:

New compute style: compute extraProperty/atom
Kokkos support: Added as compute extraProperty/atom/kk
Supports hybrid pair styles via the hybridIndex keyword on extraProperty/atom
Automatic property retrieval through MLIAP infrastructure

Author(s)

Caden Fischer (Los Alamos National Laboratory) - caden.fischer@gmail.com
Deepsana Shahi (Los Alamos National Laboratory) - deepsanashahi@gmail.com

Licensing

By submitting this pull request, I agree, that my contribution will be included in LAMMPS and redistributed under either the GNU General Public License version 2 (GPL v2) or the GNU Lesser General Public License version 2.1 (LGPL v2.1).

Implementation Notes

Key Components Implemented

ExtraProperties Class
- Added a new ExtraProperties class instantiated within MLIAPData.
- Internally manages a std::unordered_map<std::string, double**>, where:
  - std::string is the name of the property (e.g., "energy_std")
  - double** points to the associated per-atom data (dimensioned by number of atoms × property dimension)
- The class is responsible for lifetime management of the memory associated with these properties.
compute extraProperty/atom Command
- A new compute style (extraProperty/atom and extraProperty/atom/kk) registers a request for a specific property name and its dimensionality.
- The compute interacts with the ExtraProperties instance in MLIAPData to access the requested per-atom vectors.
Integration with MLIAP Unified Interface
- Introduced a new virtual method compute_extra_property() in the MLIAPUnifiedInterface.
- LAMMPS uses this to request extra property data by name during force evaluation.
- Implementers of this interface populate the data using update_extra_property(), ensuring the requested property is made available to LAMMPS via ExtraProperties.

Testing Strategy

Test Case: Used ensemble models trained with lanl/HIPPYNN to compute and expose standard deviations for properties such as potential energy.
Input Decks: Custom input scripts provided for testing this feature. These include:
- Requests for energy_std in.mliap.unified.hippynn.ensemble.Al.txt
- Kokkos variant of energy_std in.mliap.unified.hippynn.ensemble.kk.Al.txt
- Use of hybrid pair styles in.mliap.unified.hippynn.ensemble.kk.Al.txt
Test Matrix:
- Backends: Tested on both non-Kokkos and Kokkos LAMMPS builds
- Parallel Configurations:
  - MPI ranks: 1, 2, 4, 8
  - Nodes: 1 or 2
  - Kokkos: 1 GPU per MPI task in Kokkos runs

All runs completed successfully and produced the expected per-atom extra property output.

Post Submission Checklist

The feature or features in this pull request is complete
Licensing information is complete
Corresponding author information is complete
The source code follows the LAMMPS formatting guidelines
Suitable new documentation files and/or updates to the existing docs are included
The added/updated documentation is integrated and tested with the documentation build system
The feature has been verified to work with the conventional build system
The feature has been verified to work with the CMake based build system
One or more example input decks are included

akohlmey

@codeformore Thanks for your contribution. Before even looking at any specific details of the implementation, I have some rather formal and fundamental suggestions for changes concerning the naming conventions of styles and keywords, and a few other LAMMPS programming style issues (especially for header file "hygiene" which is extremely important for a project as large as LAMMPS).

There will likely be suggestions from me on more details later on and also from other LAMMPS developers. However, I thought it is best to do this in stages and not overwhelm you with a huge pile of comments.

src/pair_hybrid.h

src/ML-IAP/pair_mliap.h

src/ML-IAP/pair_mliap.cpp

src/ML-IAP/compute_extra_property_atom.h

src/ML-IAP/mliap_data.cpp

src/ML-IAP/compute_extra_property_atom.cpp

akohlmey · 2025-08-04T17:17:21Z

Caden Fischer - caden.fischer@gmail.com Deepsana Shahi

Please add your affiliation.

codeformore · 2025-08-04T17:34:10Z

Tagging Ben Nebgen (Los Alamos National Laboratory) @bnebgen-LANL and Nick Lubbers (Los Alamos National Laboratory) @lubbersnick who worked with us on this feature.

akohlmey · 2025-08-04T21:17:30Z

@codeformore here is one more thing that occurred to me: your description of the ExtraProperties class has a lot of similarities with what fix property/atom does. You could delegate the entire data management to that fix (just create a suitable fix instance internally) since it supports per-atom vectors and arrays with integer and double data. As a bonus, you have built-in support for accessing the defined custom properties. Of course, you will still have to fill it with the desired data, but it could lower the amount of maintenance and debugging work to stick with some existing and well tested functionality.

… definitions

codeformore · 2025-08-05T16:33:40Z

Rebased onto develop branch.

…ange

codeformore · 2025-08-05T16:58:18Z

@codeformore Thanks for your contribution. Before even looking at any specific details of the implementation, I have some rather formal and fundamental suggestions for changes concerning the naming conventions of styles and keywords, and a few other LAMMPS programming style issues (especially for header file "hygiene" which is extremely important for a project as large as LAMMPS).

There will likely be suggestions from me on more details later on and also from other LAMMPS developers. However, I thought it is best to do this in stages and not overwhelm you with a huge pile of comments.

I've pushed updates that address the feedback provided in the previous review. Please let me know if there's anything I missed. Thanks again for your time and thoughtful feedback.

bnebgen-LANL · 2025-08-05T17:04:19Z

@codeformore here is one more thing that occurred to me: your description of the ExtraProperties class has a lot of similarities with what fix property/atom does. You could delegate the entire data management to that fix (just create a suitable fix instance internally) since it supports per-atom vectors and arrays with integer and double data. As a bonus, you have built-in support for accessing the defined custom properties. Of course, you will still have to fill it with the desired data, but it could lower the amount of maintenance and debugging work to stick with some existing and well tested functionality.

Axel, thanks for the idea. You are probably correct in that doing it via a fix would possibly make things easier. We'd like to hold off on making this shift until another time since we've got this feature working, if that is OK. We aren't seeing any performance degradation doing it as a compute.

akohlmey · 2025-08-05T18:24:46Z

Axel, thanks for the idea. You are probably correct in that doing it via a fix would possibly make things easier. We'd like to hold off on making this shift until another time since we've got this feature working

@bnebgen-LANL Sorry, but you are missing my point. The issue is not compute versus fix. A compute is the appropriate class to provide the information you want to make available. My point is that your compute is duplicating functionality that already exists, i.e. mapping per-atom arrays to keywords. So, to reduce code duplication, you would just create a fix property/atom instance internally from within the compute and can then store data in it and refer to it.

I personally don't worry so much about performance, but I do worry about the effort required to maintain LAMMPS. The more duplicated functionality there is, the more work it is to keep things working and up-to-date. What is worse is that the interfaces are usually not exactly the same, which results in somewhat inconsistent behavior and that is something we want to reduce, if not avoid, for the sake of the users.

rbberger · 2025-08-05T19:37:36Z

src/ML-IAP/compute_mliap_property_atom.cpp

+  data = castedPair->data;
+
+  //Register the extra_property compute with the data class
+  data->register_extra_property(property_name, size_peratom_cols);


As @akohlmey mentioned this can likely be replaced by something that internally creates a fix property/atom. See https://github.com/lammps/lammps/blob/develop/src/RHEO/compute_rheo_interface.cpp#L66 for an example on how to create a d2_some_name X array with X columns and retrieve its index via atom->find_custom.

Thanks Richard. Sorry to be a pain, but is there an example in kokkos? We would want to see the correct way to do this on both sides.

@stanmoore1 It seems to me that fix property/atom/kk is missing full support for IARRAY and DARRAY and isn't creating any dual views for those. So I would expect the following in atom_kokkos.h:

DAT::tdual_int_3d k_iarray; DAT::tdual_float_3d k_darray;

Also

DAT::tdual_int_2d k_ivector;

seems to be missing? Same story?

@akohlmey ~~think we are also missing something equivalent to find_custom in atom_kokkos.h that gives us the kokkos views so we can utilize device data.~~ Nevermind, it's the same index as on the host and the k_* members are public, so all you need is to cast atom to AtomKokkos* and access the public field with the same index.

@bnebgen-LANL I'm not seeing a straight forward example that you can just copy. It looks more like the Kokkos code could be extended to a more general solution and you would be the first client of it.

Thanks for the detailed suggestions and discussion.

I looked into replacing the ExtraProperties infrastructure with a fix-based approach using fix property/atom as suggested. While that pattern works well on the CPU implementation, it appears that the Kokkos version (fix property/atom/kk) currently lacks full support for DARRAY and IARRAY, including the necessary tdual_* views (as mentioned above). Without that, I don't see a straightforward way to access custom per-atom double arrays on device (please tell me if my assumption is incorrect), which we rely on for the modifications in the MLIAP unified interface.

Given this, I don't believe the fix-based pattern is currently viable on the Kokkos side without additional infrastructure. I'm happy to adapt to a unified fix property/atom interface in the future if support is added—but for now, the ExtraProperties approach provides the functionality we need for both non-Kokkos and Kokkos runs and keeps the CPU and Kokkos implementations similar.

Please let me know if you'd recommend a different approach given these constraints. I'm open to alternatives if there's a better-supported path. Thanks again for the feedback.

It should be fairly trivial to add the missing dual views in fix property/atom/kk. We already have something similar in SPARTA.

bnebgen-LANL · 2025-08-06T19:54:10Z

I'm happy to go with whatever the lammps maintainers want. However, Caden is only at LANL until Friday. Is there a way we could move forward with the merge and if/when Stan or somebody else implements the appropriate dual views, I can revisit this to implement it through the fix methodology?

akohlmey · 2025-08-06T20:00:04Z

Is there a way we could move forward with the merge

@bnebgen-LANL What is the rush? We generally prefer to iterate pull requests until everybody is satisfied. At the moment, the PR doesn't even pass the automated testing. So far, we only looked at some formal details. There hasn't even been a proper review, manual tests, checks of the documentation and so on. There also are some sweeping changes coming (KOKKOS with mixed and single precision, build system updates and more. Plus we have the LAMMPS workshop next week.
So there are likely weeks until this is going to be ready, even without the elimination of duplicated functionality.

akohlmey · 2025-11-06T01:50:13Z

@bnebgen-LANL this has been idle for three months. Any updates?

If there are no plans to bring this pull request into a shape that it has merge conflicts resolved and passes the automated tests (or corrects false positives in the tests), then I suggest we close this for now and you resubmit when it is closer to being complete and reviewable.

bnebgen-LANL · 2025-11-06T13:35:07Z

Sorry. The person working on this was out sick for a while. We should make more progress soon.

athomps · 2025-11-19T20:43:35Z

Hi @bnebgen-LANL, sorry for not looking at this sooner. Based on this discussion, in order for this to be accepted in to LAMMPS, I see two alternatives:

Leverage existing fix properties/atom rather than creating new compute extraProperty/atom command.
Leverage existing fix pair command. This was created to extract per-atom uncertainty metric from pair_style pace/extrapolation. It would be relatively easy to add an interface to extract pointer to uncertainties from MLIAPData::extra_properties, see PairPACEExtrapolation::extract_peratom()

athomps

Given the difficulties identified with the new compute extraProperty/atom command, I think it best to withdraw this PR until a clear path forward has been identified. [I have changed my mind on this, see below]

athomps

The best path forward for this is to leverage the existing fix pair command, which was created by @sjplimp for exactly this kind of use case. It so far just has one use instance: to extract per-atom uncertainty metric from pair_style pace/extrapolation. It would be relatively easy to add a function PairMLIAP::extract_peratom() to extract pointers to uncertainties and other per-atom quantities from MLIAPData::extra_properties, see PairPACEExtrapolation::extract_peratom(). The atom uncertainties currently accessed from the compute class ComputeMLIAPPropertyAtom can just as well or better be accessed from the FixPair class. Since this can be done fairly quickly, I recommend keeping this PR open.

codeformore requested review from athomps and stanmoore1 as code owners August 4, 2025 16:43

akohlmey assigned stanmoore1 Aug 4, 2025

akohlmey added enhancement mliap_package kokkos_package labels Aug 4, 2025

akohlmey requested a review from sjplimp August 4, 2025 16:52

akohlmey requested changes Aug 4, 2025

View reviewed changes

akohlmey added this to LAMMPS Pull Requests Aug 4, 2025

akohlmey moved this to In Progress in LAMMPS Pull Requests Aug 4, 2025

codeformore added 14 commits August 5, 2025 10:32

Add compute_extra_property to MLIAPInterface

38eb82f

Change extra_properties data structure to std::UnorderedMap

c78f9f8

Add compute_extra_property to unified interface

93dd1e7

Add Kokkos compute for extra properties

874b0f6

Add ExtraProperty class to mliap_data

12cf322

Complete extraProperties on CPU

8bb7b31

Turn ExtraPropertiesKokkos into a class and seperate declarations and…

7288c03

… definitions

Fix uninitalized nlocal in compute_extra_property

a19b257

Fix reference to CPU data instead of Kokkos data

780360e

Add documentation for compute extraProperty/atom

402be9f

Remake references each iteration in compute_extra_properties

401164c

Add compute extra property compatiable with hybrid pair style mliaps

f8c4307

Update documentation of compute extraProperty/atom

b48c8d1

Add include for unordered_map to mliap_data.h

98d2569

codeformore force-pushed the generalized-interface branch from bb295a7 to 98d2569 Compare August 5, 2025 16:33

Fix std::string passing

a34bed8

codeformore added 4 commits August 5, 2025 10:41

Remove compute includes in pair styles

1736fc1

Add include for compute_extra_properties in pair_mliap

d53b293

Rename extraProperty/atom to mliap/property/atom

4317158

Update documentation for extraProperty/atom to mliap/property/atom ch…

c1fc52c

…ange

rbberger reviewed Aug 5, 2025

View reviewed changes

rbberger mentioned this pull request Aug 13, 2025

Extend fix property/atom/kk with support for ivector, iarray and darray #4690

Draft

11 tasks

athomps requested changes Nov 19, 2025

View reviewed changes

athomps moved this from In Progress to Has Problems in LAMMPS Pull Requests Nov 19, 2025

athomps self-requested a review November 21, 2025 17:58

athomps requested changes Nov 21, 2025

View reviewed changes

MLIAP Unified Generalized Interface #4673

Are you sure you want to change the base?

MLIAP Unified Generalized Interface #4673

Uh oh!

Conversation

codeformore commented Aug 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

akohlmey left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

akohlmey commented Aug 4, 2025

Uh oh!

codeformore commented Aug 4, 2025

Uh oh!

akohlmey commented Aug 4, 2025

Uh oh!

codeformore commented Aug 5, 2025

Uh oh!

codeformore commented Aug 5, 2025

Uh oh!

bnebgen-LANL commented Aug 5, 2025

Uh oh!

akohlmey commented Aug 5, 2025

Uh oh!

rbberger Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

bnebgen-LANL Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

rbberger Aug 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rbberger Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

codeformore Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

stanmoore1 Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

bnebgen-LANL commented Aug 6, 2025

Uh oh!

akohlmey commented Aug 6, 2025

Uh oh!

akohlmey commented Nov 6, 2025

Uh oh!

bnebgen-LANL commented Nov 6, 2025

Uh oh!

athomps commented Nov 19, 2025

Uh oh!

athomps left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

athomps left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

codeformore commented Aug 4, 2025 •

edited

Loading

rbberger Aug 6, 2025 •

edited

Loading

athomps left a comment •

edited

Loading

athomps left a comment •

edited

Loading