diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index 3d90c738ae..7475b0b07f 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -42,6 +42,9 @@ from .polar_atomic_model import ( DPPolarAtomicModel, ) +from .property_atomic_model import ( + DPPropertyAtomicModel, +) __all__ = [ "BaseAtomicModel", @@ -54,4 +57,5 @@ "LinearEnergyAtomicModel", "PairTabAtomicModel", "make_base_atomic_model", + "DPPropertyAtomicModel", ] diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index 6f69f8dfb6..88c42d7d07 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np from deepmd.dpmodel.fitting.property_fitting import ( PropertyFittingNet, ) @@ -15,3 +16,23 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" ) super().__init__(descriptor, fitting, type_map, **kwargs) + + def apply_out_stat( + self, + ret: dict[str, np.ndarray], + atype: np.ndarray, + ): + """Apply the stat to each atomic output. + + Parameters + ---------- + ret + The returned dict by the forward_atomic method + atype + The atom types. nf x nloc + + """ + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + for kk in self.bias_keys: + ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0] + return ret diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index 16fdedd36e..9bd07bd349 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.atomic_model.dp_atomic_model import ( - DPAtomicModel, +from deepmd.dpmodel.atomic_model import ( + DPPropertyAtomicModel, ) from deepmd.dpmodel.model.base_model import ( BaseModel, @@ -13,7 +13,7 @@ make_model, ) -DPPropertyModel_ = make_model(DPAtomicModel) +DPPropertyModel_ = make_model(DPPropertyAtomicModel) @BaseModel.register("property") diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d9744246d7..8d1f9783ab 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -779,9 +779,15 @@ def test_property( tuple[list[np.ndarray], list[int]] arrays with results and their shapes """ - data.add("property", dp.task_dim, atomic=False, must=True, high_prec=True) - if has_atom_property: - data.add("atom_property", dp.task_dim, atomic=True, must=False, high_prec=True) + property_name = dp.get_property_name() + property_dim = dp.get_property_dim() + assert isinstance(property_name, list) + assert isinstance(property_dim, list) + assert sum(property_dim) == dp.task_dim + for (name, dim) in zip(property_name, property_dim): + data.add(name, dim, atomic=False, must=True, high_prec=True) + if has_atom_property: + data.add(f"atom_{name}", dim, atomic=True, must=False, high_prec=True) if dp.get_dim_fparam() > 0: data.add( @@ -832,6 +838,18 @@ def test_property( aproperty = ret[1] aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim]) + concat_property = [] + concat_aproperty = [] + for (name, dim) in zip(property_name, property_dim): + test_data[name] = test_data[name].reshape([numb_test, dim]) + test_data[f"atom_{name}"] = test_data[f"atom_{name}"].reshape( + [numb_test, natoms * dim] + ) + concat_property.append(test_data[name]) + concat_aproperty.append(test_data[f"atom_{name}"]) + test_data["property"] = np.concatenate(concat_property, axis=1) + test_data["atom_property"] = np.concatenate(concat_aproperty, axis=1) + diff_property = property - test_data["property"][:numb_test] mae_property = mae(diff_property) rmse_property = rmse(diff_property) diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 389a0e8512..f4dd6d8b3f 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -141,5 +141,13 @@ def get_intensive(self) -> bool: """Get whether the property is intensive.""" return self.deep_eval.get_intensive() + def get_property_name(self) -> Union[list[str], str]: + """Get the name of the property.""" + return self.deep_eval.get_property_name() + + def get_property_dim(self) -> Union[list[int], int]: + """Get the dimension of the property.""" + return self.deep_eval.get_property_dim() + __all__ = ["DeepProperty"] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 59b833d34c..8fd496d5d8 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -184,6 +184,14 @@ def get_dim_aparam(self) -> int: def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() + def get_property_name(self) -> Union[list[str], str]: + """Get the name of the property.""" + return self.dp.model["Default"].get_property_name() + + def get_property_dim(self) -> Union[list[int], int]: + """Get the dimension of the property.""" + return self.dp.model["Default"].get_property_dim() + @property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 4581a2bc3e..bc4cebffaa 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Optional, + Union ) import torch @@ -78,6 +79,16 @@ def get_intensive(self) -> bool: """Get whether the property is intensive.""" return self.model_output_def()["property"].intensive + @torch.jit.export + def get_property_name(self) -> Union[list[str], str]: + """Get the name of the property.""" + return self.get_fitting_net().property_name + + @torch.jit.export + def get_property_dim(self) -> Union[list[int], int]: + """Get the dimension of the property.""" + return self.get_fitting_net().property_dim + @torch.jit.export def forward_lower( self, diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 7f0796d08f..37b3e7b0a8 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -84,6 +84,7 @@ def __init__( intensive: bool = False, bias_method: str = "normal", property_name: Union[str, list] = "property", + property_dim: Union[int, list] = 1, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -97,7 +98,12 @@ def __init__( self.task_dim = task_dim self.intensive = intensive self.bias_method = bias_method + if isinstance(property_name, str): + property_name = [property_name] self.property_name = property_name + if isinstance(property_dim, int): + property_dim = [property_dim] + self.property_dim = property_dim super().__init__( var_name="property", ntypes=ntypes, diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 54b44d618f..07e6cecbaa 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -296,7 +296,8 @@ def compute_output_stats( for key in keys: if atomic_output.var_defs[key].sub_var_name is not None: sub_keys.extend(atomic_output.var_defs[key].sub_var_name) - keys.extend(sub_keys) + del keys + keys = sub_keys new_keys = [ ii for ii in keys diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index f4b318cd4d..725cba5974 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -332,6 +332,7 @@ "type": "property", "task_dim": 3, "property_name": ["band_property"], + "property_dim": 3, "neuron": [24, 24, 24], "resnet_dt": True, "bias_method": "normal", diff --git a/source/tests/pt/property/input.json b/source/tests/pt/property/input.json index 947eb32626..3efe6b775d 100644 --- a/source/tests/pt/property/input.json +++ b/source/tests/pt/property/input.json @@ -28,6 +28,7 @@ "type": "property", "intensive": true, "property_name": "band_property", + "property_dim": 3, "task_dim": 3, "neuron": [ 100,