Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat: pairtab model pytorch #174

Closed
wants to merge 18 commits into from
Closed

Conversation

anyangml
Copy link
Collaborator

@anyangml anyangml commented Jan 23, 2024

The goal of this PR is to migrate the PairTabModel from TF to PT:

  • Build a concrete AtomicModel class of PairTabModel based on this PR.
  • Reimplement the C++ code (pair_tab.cc) in pytorch.
  • Define an example input.json

Notes:

  1. The fscale is not implemented in this method.

@anyangml anyangml requested a review from iProzd January 23, 2024 16:11
@anyangml anyangml requested review from wanghan-iapcm and iProzd and removed request for iProzd and wanghan-iapcm January 24, 2024 06:53
Copy link
Collaborator

@wanghan-iapcm wanghan-iapcm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please provide UT for the atomic model?

  1. check the output energy is as expected. e.g. in the case you have a few atoms in the system you can manually calculate the energy and check if it is consistent with your code.
  2. check if you code is jitable. see for example
    def test_jit(

deepmd_pt/model/model/pair_tab.py Outdated Show resolved Hide resolved
@anyangml
Copy link
Collaborator Author

Could you please provide UT for the atomic model?

  1. check the output energy is as expected. e.g. in the case you have a few atoms in the system you can manually calculate the energy and check if it is consistent with your code.
  2. check if you code is jitable. see for example
    def test_jit(

optimized the for loop implementation, checked the calculation locally, going to figure out jit and unittests tomorrow.

@anyangml
Copy link
Collaborator Author

Could you please provide UT for the atomic model?

  1. check the output energy is as expected. e.g. in the case you have a few atoms in the system you can manually calculate the energy and check if it is consistent with your code.
  2. check if you code is jitable. see for example
    def test_jit(

I tested jit with the following code, no error msg

file_path = "test.txt"

model = PairTabModel(
            tab_file = file_path,
            rcut =  0.1,
            sel = 2
        )
scrp_md = torch.jit.script(model)
torch.jit.save(scrp_md,"test.pth",{})

Copy link
Collaborator

@wanghan-iapcm wanghan-iapcm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What would be the behavior of your model

  • if the r-range of the input table is smaller than the cut-off distance? is the interaction between the pairs beyond upper bound of the table and smaller than rcut zero?
  • if the r-range of the input table is larger than the cut-off distance?


atomic_energy = 0.5 * torch.sum(torch.where(nlist != -1, raw_atomic_energy, torch.zeros_like(raw_atomic_energy)) ,dim=-1)

return {"atomic_energy": atomic_energy}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the key should be "energy", please check your output def.
you may want to use this decorator to ensure the correctness of your atomic model output

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the key. Not sure about the decorator, this atomic model has no Fitting

@anyangml
Copy link
Collaborator Author

  • range of the input table is smaller than the cut-off distance? is the interaction between the pairs beyond upper bound of the table and smaller than rcut zero?
  • if the r-range of the input table is larger than the cut-of

@wanghan-iapcm

  1. if r-range of the input table is smaller than the cut-off distance, there is not enough information to calculate the energy unless we want to do extrapolation. raise error
  2. if the interaction between the pairs beyond upper bound of the table and smaller than rcut , the energy is not zero but we don't have information from the table to do the calculation. If we implement a check on 1, then this won't be a problem.
  3. if the r-range of the input table is larger than the cut-off distance, I don't see a problem here. if the pairwise distance is beyond cutoff, it will be labeled as -1 in the nlist, and the energy for that pair is masked to 0.

wanghan-iapcm pushed a commit to deepmodeling/deepmd-kit that referenced this pull request Jan 31, 2024
Migrated from this
[PR](dptech-corp/deepmd-pytorch#174). This is to
reimplement the PairTab Model in Pytorch.

Notes:

1. Different from the tensorflow version, the pytorch version abstracts
away all the post energy conversion operations (force, virial).
2. Added extrapolation when `rcut` > `rmax`. The pytorch version
overwrite energy beyond extrapolation endpoint to `0`. These features
are not available in the tensorflow version. The extrapolation uses a
cubic spline form, the 1st order derivation for the starting point is
estimated using the last two rows in the user defined table. See example
below:


![img_v3_027k_b50c690d-dc2d-4803-bd2c-2e73aa3c73fg](https://github.com/deepmodeling/deepmd-kit/assets/137014849/f3efa4d3-795e-4ff8-acdc-642227f0e19c)


![img_v3_027k_8de38597-ef4e-4e5b-989e-dbd13cc93fag](https://github.com/deepmodeling/deepmd-kit/assets/137014849/493da26d-f01d-4dd0-8520-ea2d84e7b548)


![img_v3_027k_f8268564-3f5d-49e6-91d6-169a61d9347g](https://github.com/deepmodeling/deepmd-kit/assets/137014849/b8ad4d4d-a4a4-40f0-94d1-810006e7175b)


![img_v3_027k_3966ef67-dd5e-4f48-992e-c2763311451g](https://github.com/deepmodeling/deepmd-kit/assets/137014849/27f31e79-13c8-4ce8-9911-b4cc0ac8188c)

---------

Co-authored-by: Anyang Peng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@anyangml anyangml closed this Feb 20, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants