Skip to content

support combining frozen models into a pairwise DPRc model #2902

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.

Sign up for GitHub

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

Merged
wanghan-iapcm merged 3 commits into from
Oct 7, 2023
+11 −7
Merged

support combining frozen models into a pairwise DPRc model #2902

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3a3f38d
pairwise_dprc + frozen
njzjz Oct 4, 2023
42d8410
allow QM models have fewer types
njzjz Oct 6, 2023
ed23f1a
revert must
njzjz Oct 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 11 additions & 7 deletions deepmd/model/pairwise_dprc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,6 @@
TypeEmbedNet,
)

from .ener import (
EnerModel,
)


class PairwiseDPRc(Model):
"""Pairwise Deep Potential - Range Correction."""
Expand Down Expand Up @@ -87,13 +83,13 @@ def __init__(
padding=True,
)

self.qm_model = EnerModel(
self.qm_model = Model(
**qm_model,
type_map=type_map,
type_embedding=self.typeebd,
compress=compress,
)
self.qmmm_model = EnerModel(
self.qmmm_model = Model(
**qmmm_model,
type_map=type_map,
type_embedding=self.typeebd,
Expand Down Expand Up @@ -187,6 +183,14 @@ def build(

mesh_mixed_type = make_default_mesh(False, True)

# allow loading a frozen QM model that has only QM types
# Note: here we don't map the type between models, so
# the type of the frozen model must be the same as
# the first Ntypes of the current model
if self.get_ntypes() > self.qm_model.get_ntypes():
natoms_qm = tf.slice(natoms_qm, [0], [self.qm_model.get_ntypes() + 2])
assert self.get_ntypes() == self.qmmm_model.get_ntypes()

qm_dict = self.qm_model.build(
coord_qm,
atype_qm,
Expand Down Expand Up @@ -301,7 +305,7 @@ def get_rcut(self):
return max(self.qm_model.get_rcut(), self.qmmm_model.get_rcut())

def get_ntypes(self) -> int:
return self.qm_model.get_ntypes()
return self.ntypes

def data_stat(self, data):
self.qm_model.data_stat(data)
Expand Down