-
Notifications
You must be signed in to change notification settings - Fork 484
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
Hybrid descriptor #3365
Merged
Merged
Hybrid descriptor #3365
Changes from all commits
Commits
Show all changes
15 commits
Select commit
Hold shift + click to select a range
503bd89
hybrid descriptor
njzjz 74a1cbb
Merge remote-tracking branch 'origin/devel' into hybrid-descriptor
njzjz 14abf14
update sel
njzjz cdaa76f
update sel
njzjz 4c9b921
add jit test; fix jit
njzjz e736b38
bugfix
njzjz 63e4add
handle nlist
njzjz ccd1501
docs
njzjz 5f3f8bb
fix codeql
njzjz 36b099c
make jit happy
njzjz 18a0938
Make the error message more informative
njzjz 1a255f2
add the type hint of list
njzjz 317cd8a
support mix mixed_types and no mixed_types
njzjz 05a0d6c
output other output
njzjz 35f68e2
pass jit
njzjz File filter
Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,242 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
from typing import ( | ||
Any, | ||
Dict, | ||
List, | ||
Optional, | ||
Union, | ||
) | ||
|
||
import numpy as np | ||
|
||
from deepmd.dpmodel.common import ( | ||
NativeOP, | ||
) | ||
from deepmd.dpmodel.descriptor.base_descriptor import ( | ||
BaseDescriptor, | ||
) | ||
from deepmd.dpmodel.utils.nlist import ( | ||
nlist_distinguish_types, | ||
) | ||
from deepmd.utils.path import ( | ||
DPPath, | ||
) | ||
from deepmd.utils.version import ( | ||
check_version_compatibility, | ||
) | ||
|
||
|
||
@BaseDescriptor.register("hybrid") | ||
class DescrptHybrid(BaseDescriptor, NativeOP): | ||
"""Concate a list of descriptors to form a new descriptor. | ||
|
||
Parameters | ||
---------- | ||
list : list : List[Union[BaseDescriptor, Dict[str, Any]]] | ||
Build a descriptor from the concatenation of the list of descriptors. | ||
The descriptor can be either an object or a dictionary. | ||
""" | ||
|
||
def __init__( | ||
self, | ||
list: List[Union[BaseDescriptor, Dict[str, Any]]], | ||
) -> None: | ||
super().__init__() | ||
# warning: list is conflict with built-in list | ||
descrpt_list = list | ||
if descrpt_list == [] or descrpt_list is None: | ||
raise RuntimeError( | ||
"cannot build descriptor from an empty list of descriptors." | ||
) | ||
formatted_descript_list = [] | ||
for ii in descrpt_list: | ||
if isinstance(ii, BaseDescriptor): | ||
formatted_descript_list.append(ii) | ||
elif isinstance(ii, dict): | ||
formatted_descript_list.append(BaseDescriptor(**ii)) | ||
else: | ||
raise NotImplementedError | ||
self.descrpt_list = formatted_descript_list | ||
self.numb_descrpt = len(self.descrpt_list) | ||
for ii in range(1, self.numb_descrpt): | ||
assert ( | ||
self.descrpt_list[ii].get_ntypes() == self.descrpt_list[0].get_ntypes() | ||
), f"number of atom types in {ii}th descrptor {self.descrpt_list[0].__class__.__name__} does not match others" | ||
# if hybrid sel is larger than sub sel, the nlist needs to be cut for each type | ||
hybrid_sel = self.get_sel() | ||
self.nlist_cut_idx: List[np.ndarray] = [] | ||
if self.mixed_types() and not all( | ||
descrpt.mixed_types() for descrpt in self.descrpt_list | ||
): | ||
self.sel_no_mixed_types = np.max( | ||
[ | ||
descrpt.get_sel() | ||
for descrpt in self.descrpt_list | ||
if not descrpt.mixed_types() | ||
], | ||
axis=0, | ||
).tolist() | ||
else: | ||
self.sel_no_mixed_types = None | ||
for ii in range(self.numb_descrpt): | ||
if self.mixed_types() == self.descrpt_list[ii].mixed_types(): | ||
hybrid_sel = self.get_sel() | ||
else: | ||
assert self.sel_no_mixed_types is not None | ||
hybrid_sel = self.sel_no_mixed_types | ||
sub_sel = self.descrpt_list[ii].get_sel() | ||
start_idx = np.cumsum(np.pad(hybrid_sel, (1, 0), "constant"))[:-1] | ||
end_idx = start_idx + np.array(sub_sel) | ||
cut_idx = np.concatenate( | ||
[range(ss, ee) for ss, ee in zip(start_idx, end_idx)] | ||
) | ||
self.nlist_cut_idx.append(cut_idx) | ||
|
||
def get_rcut(self) -> float: | ||
"""Returns the cut-off radius.""" | ||
return np.max([descrpt.get_rcut() for descrpt in self.descrpt_list]).item() | ||
|
||
def get_sel(self) -> List[int]: | ||
"""Returns the number of selected atoms for each type.""" | ||
if self.mixed_types(): | ||
return [ | ||
np.max( | ||
[descrpt.get_nsel() for descrpt in self.descrpt_list], axis=0 | ||
).item() | ||
] | ||
else: | ||
return np.max( | ||
[descrpt.get_sel() for descrpt in self.descrpt_list], axis=0 | ||
).tolist() | ||
|
||
def get_ntypes(self) -> int: | ||
"""Returns the number of element types.""" | ||
return self.descrpt_list[0].get_ntypes() | ||
|
||
def get_dim_out(self) -> int: | ||
"""Returns the output dimension.""" | ||
return np.sum([descrpt.get_dim_out() for descrpt in self.descrpt_list]).item() | ||
|
||
def get_dim_emb(self) -> int: | ||
"""Returns the output dimension.""" | ||
return np.sum([descrpt.get_dim_emb() for descrpt in self.descrpt_list]).item() | ||
|
||
def mixed_types(self): | ||
"""Returns if the descriptor requires a neighbor list that distinguish different | ||
atomic types or not. | ||
""" | ||
return any(descrpt.mixed_types() for descrpt in self.descrpt_list) | ||
|
||
def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None): | ||
"""Update mean and stddev for descriptor elements.""" | ||
for descrpt in self.descrpt_list: | ||
descrpt.compute_input_stats(merged, path) | ||
|
||
def call( | ||
self, | ||
coord_ext, | ||
atype_ext, | ||
nlist, | ||
mapping: Optional[np.ndarray] = None, | ||
): | ||
"""Compute the descriptor. | ||
|
||
Parameters | ||
---------- | ||
coord_ext | ||
The extended coordinates of atoms. shape: nf x (nallx3) | ||
atype_ext | ||
The extended aotm types. shape: nf x nall | ||
nlist | ||
The neighbor list. shape: nf x nloc x nnei | ||
mapping | ||
The index mapping, not required by this descriptor. | ||
|
||
Returns | ||
------- | ||
descriptor | ||
The descriptor. shape: nf x nloc x (ng x axis_neuron) | ||
gr | ||
The rotationally equivariant and permutationally invariant single particle | ||
representation. shape: nf x nloc x ng x 3. | ||
g2 | ||
The rotationally invariant pair-partical representation. | ||
h2 | ||
The rotationally equivariant pair-partical representation. | ||
sw | ||
The smooth switch function. | ||
""" | ||
out_descriptor = [] | ||
out_gr = [] | ||
out_g2 = [] | ||
out_h2 = None | ||
out_sw = None | ||
if self.sel_no_mixed_types is not None: | ||
nl_distinguish_types = nlist_distinguish_types( | ||
nlist, | ||
atype_ext, | ||
self.sel_no_mixed_types, | ||
) | ||
else: | ||
nl_distinguish_types = None | ||
for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx): | ||
# cut the nlist to the correct length | ||
if self.mixed_types() == descrpt.mixed_types(): | ||
nl = nlist[:, :, nci] | ||
else: | ||
# mixed_types is True, but descrpt.mixed_types is False | ||
assert nl_distinguish_types is not None | ||
nl = nl_distinguish_types[:, :, nci] | ||
odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping) | ||
out_descriptor.append(odescriptor) | ||
if gr is not None: | ||
out_gr.append(gr) | ||
if g2 is not None: | ||
out_g2.append(g2) | ||
if self.get_rcut() == descrpt.get_rcut(): | ||
out_h2 = h2 | ||
out_sw = sw | ||
|
||
out_descriptor = np.concatenate(out_descriptor, axis=-1) | ||
out_gr = np.concatenate(out_gr, axis=-2) if out_gr else None | ||
out_g2 = np.concatenate(out_g2, axis=-1) if out_g2 else None | ||
return out_descriptor, out_gr, out_g2, out_h2, out_sw | ||
|
||
@classmethod | ||
def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: | ||
"""Update the selection and perform neighbor statistics. | ||
|
||
Parameters | ||
---------- | ||
global_jdata : dict | ||
The global data, containing the training section | ||
local_jdata : dict | ||
The local data refer to the current class | ||
""" | ||
local_jdata_cpy = local_jdata.copy() | ||
local_jdata_cpy["list"] = [ | ||
BaseDescriptor.update_sel(global_jdata, sub_jdata) | ||
for sub_jdata in local_jdata["list"] | ||
] | ||
return local_jdata_cpy | ||
|
||
def serialize(self) -> dict: | ||
return { | ||
"@class": "Descriptor", | ||
"type": "hybrid", | ||
"@version": 1, | ||
"list": [descrpt.serialize() for descrpt in self.descrpt_list], | ||
} | ||
|
||
@classmethod | ||
def deserialize(cls, data: dict) -> "DescrptHybrid": | ||
data = data.copy() | ||
class_name = data.pop("@class") | ||
assert class_name == "Descriptor" | ||
class_type = data.pop("type") | ||
assert class_type == "hybrid" | ||
check_version_compatibility(data.pop("@version"), 1, 1) | ||
obj = cls( | ||
list=[BaseDescriptor.deserialize(ii) for ii in data["list"]], | ||
) | ||
return obj |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Check warning
Code scanning / CodeQL
Variable defined multiple times Warning