Skip to content

feat(pt): add atomic_weight to dipole model inference #4628

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

Closed
ChiahsinChu wants to merge 13 commits into from
Closed

feat(pt): add atomic_weight to dipole model inference #4628

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
275fb52
add atomic_weight for model inference
ChiahsinChu Mar 3, 2025
b5356d9
remove test tensor
ChiahsinChu Mar 3, 2025
205e506
fix bug in model frozen
ChiahsinChu Mar 3, 2025
50a8a21
deal with case without fitting net
ChiahsinChu Mar 3, 2025
a47b1e0
Merge branch 'devel' into devel-dipole_with_atomic_weight
ChiahsinChu Mar 4, 2025
3aaccc7
Merge branch 'devel' into devel-dipole_with_atomic_weight
ChiahsinChu Mar 5, 2025
3cbaa61
Merge branch 'devel' into devel-dipole_with_atomic_weight
ChiahsinChu Mar 8, 2025
3763d8e
Merge branch 'devel' into devel-dipole_with_atomic_weight
ChiahsinChu Mar 11, 2025
32a64e1
move `atomic_weight` to atomic model level
ChiahsinChu Mar 11, 2025
ad063bd
(np) add atomic_weight to atomic model and dipole model; add UT for p…
ChiahsinChu Mar 11, 2025
778af5b
add doc for atomic_weight
ChiahsinChu Mar 11, 2025
088b252
fix bug in jax2tf model convert
ChiahsinChu Mar 11, 2025
66d3e0f
Merge branch 'devel' into devel-dipole_with_atomic_weight
ChiahsinChu Mar 14, 2025
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
4 changes: 4 additions & 0 deletions deepmd/pt/model/model/dipole_model.py
View file
Open in desktop
Copy link
Member

@njzjz njzjz Mar 5, 2025

Choose a reason for hiding this comment

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

It's a bit confusing since it breaks the consistency of APIs between different models and backends.

Copy link
Contributor Author

@ChiahsinChu ChiahsinChu Mar 5, 2025

Choose a reason for hiding this comment

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

Would it be better to add this keyword for all models and backends?

Copy link
Member

@njzjz njzjz Mar 5, 2025

Choose a reason for hiding this comment

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

Not sure what is the best way

Copy link
Collaborator

@wanghan-iapcm wanghan-iapcm Mar 6, 2025

Choose a reason for hiding this comment

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

I would suggest incorporating the atomic weight into the base atomic model, following a similar approach to how atom masks are managed. This would ensure that atomic weight is consistently supported across all atomic modes.

Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def forward(
fparam: Optional[torch.Tensor] = None,
aparam: Optional[torch.Tensor] = None,
do_atomic_virial: bool = False,
atomic_weight: Optional[torch.Tensor] = None,
) -> dict[str, torch.Tensor]:
model_ret = self.forward_common(
coord,
Expand All @@ -68,6 +69,7 @@ def forward(
fparam=fparam,
aparam=aparam,
do_atomic_virial=do_atomic_virial,
atomic_weight=atomic_weight,
)
if self.get_fitting_net() is not None:
model_predict = {}
Expand Down Expand Up @@ -98,6 +100,7 @@ def forward_lower(
fparam: Optional[torch.Tensor] = None,
aparam: Optional[torch.Tensor] = None,
do_atomic_virial: bool = False,
atomic_weight: Optional[torch.Tensor] = None,
):
model_ret = self.forward_common_lower(
extended_coord,
Expand All @@ -108,6 +111,7 @@ def forward_lower(
aparam=aparam,
do_atomic_virial=do_atomic_virial,
extra_nlist_sort=self.need_sorted_nlist_for_lower(),
atomic_weight=atomic_weight,
)
if self.get_fitting_net() is not None:
model_predict = {}
Expand Down
12 changes: 12 additions & 0 deletions deepmd/pt/model/model/make_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ def forward_common(
fparam: Optional[torch.Tensor] = None,
aparam: Optional[torch.Tensor] = None,
do_atomic_virial: bool = False,
atomic_weight: Optional[torch.Tensor] = None,
) -> dict[str, torch.Tensor]:
"""Return model prediction.

Expand Down Expand Up @@ -188,6 +189,7 @@ def forward_common(
do_atomic_virial=do_atomic_virial,
fparam=fp,
aparam=ap,
atomic_weight=atomic_weight,
)
model_predict = communicate_extended_output(
model_predict_lower,
Expand Down Expand Up @@ -242,6 +244,7 @@ def forward_common_lower(
do_atomic_virial: bool = False,
comm_dict: Optional[dict[str, torch.Tensor]] = None,
extra_nlist_sort: bool = False,
atomic_weight: Optional[torch.Tensor] = None,
):
"""Return model prediction. Lower interface that takes
extended atomic coordinates and types, nlist, and mapping
Expand Down Expand Up @@ -293,6 +296,15 @@ def forward_common_lower(
aparam=ap,
comm_dict=comm_dict,
)
# add weight to atomic_output
if hasattr(self.atomic_model, "fitting_net"):
if hasattr(self.atomic_model.fitting_net, "var_name"):
kw = self.atomic_model.fitting_net.var_name
Copy link
Member

@njzjz njzjz Mar 5, 2025

Choose a reason for hiding this comment

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

I feel that self.atomic_model.fitting_output_def could be used here

if atomic_weight is not None:
# atomic_weight: nf x nloc x dim
atomic_ret[kw] = atomic_ret[kw] * atomic_weight.view(
[atomic_ret[kw].shape[0], atomic_ret[kw].shape[1], -1]
)
model_predict = fit_output_to_model_output(
atomic_ret,
self.atomic_output_def(),
Expand Down