ZBL short-range potential #335

frostedoyster · 2024-09-11T16:14:29Z

Contributor (creator of pull-request) checklist

Tests updated (for new features and bugfixes)?
Documentation updated (for new features)?

📚 Documentation preview 📚: https://metatrain--335.org.readthedocs.build/en/335/

PicoCentauri

Wonderful and very useful addition. I think one can tackle the unit conversion.

PicoCentauri · 2024-09-12T17:26:23Z

src/metatrain/experimental/gap/trainer.py

@@ -71,7 +71,7 @@ def train(

        logger.info("Subtracting composition energies")
        # this acts in-place on train_y
-        remove_composition(
+        remove_additive(


Why do you change the name? Isn't it really to remove the composition?

Now it can remove any type of additive contribution (e.g. ZBL, which is handled in the same way as the composition)

Okay, I thought about a better name, but makes sense now.

PicoCentauri · 2024-09-12T17:27:18Z

src/metatrain/utils/additive/zbl.py

+        if dataset_info.length_unit != "angstrom":
+            raise ValueError(
+                "ZBL only supports angstrom units, but a "
+                f"{dataset_info.length_unit} unit was provided."


I think we can do unit conversion in the end no? Should be fairly trivial. The values are there in metatensor atomistic.

I don't know... I think it's a bit tricky and I thought the best way for now would be to enforce Angstroms and eV. Where would you do the unit conversion (both at training time and at prediction time)? And what if the user doesn't specify any units?

I wonder if @Luthaf has any ideas for this

Yeah, I agree that it might be a bit painful, but since we have units in our config I think we should support them. I would pass the units when you actually compute the potential using get_pairwise_zbl. We could save the length unit and the energy unit as property.

Could the ZBL class take the units as __init__ parameters, do the conversion when instantiating itself and then store values in the right units?

So, metatrain gives you the opportunity to train without any units. Do we error out in that case? And do we want to have a units module in metatrain, or can we perhaps recycle the one from ASE?

So, metatrain gives you the opportunity to train without any units. Do we error out in that case?

yes, I would error out in this case

And do we want to have a units module in metatrain, or can we perhaps recycle the one from ASE?

I don't really mind, we can also use the code from metatensor: https://docs.metatensor.org/latest/atomistic/reference/models/index.html#metatensor.torch.atomistic.unit_conversion_factor

PicoCentauri · 2024-09-12T17:27:36Z

src/metatrain/utils/additive/zbl.py

+            if target.unit != "eV":
+                raise ValueError(
+                    "ZBL only supports eV units, but a "
+                    f"{target.unit} output was provided."
+                )


same as above

src/metatrain/utils/additive/zbl.py

abmazitov · 2024-09-24T14:45:50Z

Looks good to me, I don't see any potential issues with PET at first glance
As the tests are fixed, we can merge this I think
But please wait until #344, it is once step from being merged

frostedoyster · 2024-09-25T14:43:01Z

IMO, since multiple people need this, we should merge as is, without worrying about handling different units or neighbor lists. We can open issues however to keep track of those

abmazitov · 2024-09-26T21:31:46Z

Okay, so I tried to run PET + ZBL, and there are a few things to note.

ZBL only supports the cutoff which is 2x larger than the largest covalent radius in the dataset. In the case of MAD dataset, it results in the minimal cutoff of ~4.9 A, which is already large than our current 4.5 A cutoff. Maybe we should change the factor to at least 1.5? In this case the minimal cutoff will be 3.66, at it seems more reasonable.
This check for ZBL minimal cutoff happens only during the evaluations, i.e. after the model was trained. Perhaps it should happen in the beginning of the training.
There is a large block of converting the NL from half to full format in the ZBL code:

# convert to full NL
half_nl = system.get_neighbor_list(nl_option)
half_nl_samples = half_nl.samples.values
half_nl_values = half_nl.values
nl = TensorBlock(
    samples=Labels(
        names=half_nl.samples.names,
        values=torch.concatenate(
            [
                half_nl_samples,
                torch.concatenate(
                    [
                        half_nl_samples[:, 1].unsqueeze(-1),
                        half_nl_samples[:, 0].unsqueeze(-1),
                        -half_nl_samples[:, 2:5],
                    ],
                    dim=1,
                ),
            ]
        ),
    ),
    components=half_nl.components,
    properties=half_nl.properties,
    values=torch.concatenate(
        [
            half_nl_values,
            -half_nl_values,
        ],
    ),
)

Maybe it makes sense to transfer this code to utilities? It might be useful to have something like this as a utility for other applications, and clean-up the code a bit.

frostedoyster · 2024-09-27T05:37:04Z

Thanks a lot @abmazitov for testing this out. Unfortunately the factor of 2 is not something arbitrary but it's given by the necessity of making the potential continuous while re-using the neighbor list of the ML potential. I thought the values were pretty safe in a general case, but you have large atoms and a small cutoff so it goes over by a bit... I think the best solution is to make ZBL request its own NL, as painful as it might be, which would also spare us the NL conversion step that you highlighted.

Finally, the fact that the error is only raised after training is finished is a huge bug and I will be investigating it

src/metatrain/utils/additive/zbl.py

frostedoyster · 2024-09-27T11:37:13Z

Things to be done:

ZBL needs to be able to request its own NL in case it's larger than the one from the model or a different kind (ZBL always wants full, right now we're using a slow converter)
Improve the covalent radii design that we're inheriting from ASE, i.e., the elements for which ASE devs didn't find a covalent radius are listed as "0.2 A". And find a solution in case the user wants to use ZBL with these elements
Carefully check the units. The lengthscales are reasonable but the height of the repulsive wall is suspicious

src/metatrain/utils/neighbor_lists.py

Luthaf

This looks good to me. Regarding unit conversion, since the code will currently error if the units are wrong there is no way to silently introduce mistakes, so I'm happy to merge as is for now and open an issue for unit conversion in ZBL.

Luthaf · 2024-10-09T14:23:12Z

docs/src/dev-docs/utils/additive/index.rst

+Data
+====


Title seems wrong for this section

Luthaf · 2024-10-09T14:23:44Z

examples/ase/run_ase.py

-#    We have to import ``rascaline.torch`` even though it is not used explicitly in this
-#    tutorial. The SOAP-BPNN model contains compiled extensions and therefore the import
-#    is required.
-#


Well spotted, thanks!

ZBL potential (not integrated into any architecture yet)

498e566

frostedoyster requested a review from PicoCentauri September 11, 2024 16:14

frostedoyster force-pushed the zbl branch 2 times, most recently from 396d35d to a850270 Compare September 12, 2024 03:44

Fix docs?

55847c7

frostedoyster force-pushed the zbl branch from a850270 to 55847c7 Compare September 12, 2024 11:02

PicoCentauri reviewed Sep 12, 2024

View reviewed changes

Integrate ZBL into models

396a67c

frostedoyster requested a review from abmazitov as a code owner September 24, 2024 13:50

abmazitov approved these changes Sep 24, 2024

View reviewed changes

frostedoyster added 2 commits September 24, 2024 17:56

Incorporate into SOAP-BPNN trainer

bfe2959

Merge branch 'main' into zbl

f193378

frostedoyster changed the title ~~ZBL potential (not integrated into any architecture yet)~~ ZBL short-range potential Sep 24, 2024

frostedoyster added 3 commits September 24, 2024 18:09

Integrate into GAP and Alchemical trainers

b15ab35

Fix gradient training

78f476c

Add ZBL tutorial

ee028dc

frostedoyster requested review from abmazitov, PicoCentauri and Luthaf September 25, 2024 09:24

frostedoyster added 4 commits September 25, 2024 12:09

Add explanation for the H-H curve

52ffa9a

Fix PET hyperparameter

35b7b41

Fix GAP tests

d62fbc2

Fix bugg

3ae490d

frostedoyster force-pushed the zbl branch from 478eb8f to 3ae490d Compare September 25, 2024 14:32

Improve tutorial

5b297c5

PicoCentauri reviewed Sep 27, 2024

View reviewed changes

src/metatrain/utils/additive/zbl.py Show resolved Hide resolved

New function to find all requested neighbor lists in a model

43dbc23

frostedoyster force-pushed the zbl branch from 54bcca5 to 43dbc23 Compare October 1, 2024 11:58

Luthaf reviewed Oct 1, 2024

View reviewed changes

src/metatrain/utils/neighbor_lists.py Show resolved Hide resolved

frostedoyster and others added 5 commits October 1, 2024 14:15

Add NL request for ZBL

18c87ad

Fix PET

603daa0

Warn about covalent radii not known by ASE

1b6726b

Fix ZBL test

b9cf4f9

Merge branch 'main' into zbl

a596b1b

frostedoyster requested review from Luthaf and PicoCentauri October 1, 2024 15:22

frostedoyster and others added 3 commits October 7, 2024 17:56

Update interaction ranges + docs + more testing

b3072e6

Merge branch 'main' into zbl

82f55f0

Merge branch 'main' into zbl

51d0108

Luthaf approved these changes Oct 9, 2024

View reviewed changes

Change docs section title

02dd91c

frostedoyster merged commit cc84568 into main Oct 9, 2024
12 checks passed

frostedoyster deleted the zbl branch October 9, 2024 16:12

frostedoyster mentioned this pull request Oct 9, 2024

ZBL follow-up #355

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ZBL short-range potential #335

ZBL short-range potential #335

frostedoyster commented Sep 11, 2024 •

edited

Loading

PicoCentauri left a comment

PicoCentauri Sep 12, 2024

frostedoyster Sep 16, 2024

PicoCentauri Sep 16, 2024

PicoCentauri Sep 12, 2024

frostedoyster Sep 16, 2024 •

edited

Loading

PicoCentauri Sep 16, 2024

Luthaf Sep 17, 2024

frostedoyster Sep 20, 2024

Luthaf Oct 9, 2024

PicoCentauri Sep 12, 2024

abmazitov commented Sep 24, 2024

frostedoyster commented Sep 25, 2024

abmazitov commented Sep 26, 2024

frostedoyster commented Sep 27, 2024

frostedoyster commented Sep 27, 2024 •

edited

Loading

Luthaf left a comment

Luthaf Oct 9, 2024

Luthaf Oct 9, 2024

frostedoyster Oct 9, 2024

		Data
		====

ZBL short-range potential #335

ZBL short-range potential #335

Conversation

frostedoyster commented Sep 11, 2024 • edited Loading

Contributor (creator of pull-request) checklist

PicoCentauri left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frostedoyster Sep 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

abmazitov commented Sep 24, 2024

frostedoyster commented Sep 25, 2024

abmazitov commented Sep 26, 2024

frostedoyster commented Sep 27, 2024

frostedoyster commented Sep 27, 2024 • edited Loading

Luthaf left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frostedoyster commented Sep 11, 2024 •

edited

Loading

frostedoyster Sep 16, 2024 •

edited

Loading

frostedoyster commented Sep 27, 2024 •

edited

Loading