diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index d88c4c3af5..8d27fbcac4 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -48,12 +48,15 @@ class LinearEnergyAtomicModel(BaseAtomicModel): type_map : list[str] Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + weights : Optional[Union[str,list[float]]] + Weights of the models. If str, must be `sum` or `mean`. If list, must be a list of float. """ def __init__( self, models: list[BaseAtomicModel], type_map: list[str], + weights: Optional[Union[str, list[float]]] = "mean", **kwargs, ): super().__init__(type_map, **kwargs) @@ -89,6 +92,16 @@ def __init__( ) self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE) # pylint: disable=no-explicit-dtype + if isinstance(weights, str): + assert weights in ["sum", "mean"] + elif isinstance(weights, list): + assert len(weights) == len(models) + else: + raise ValueError( + f"'weights' must be a string ('sum' or 'mean') or a list of float of length {len(models)}." + ) + self.weights = weights + def mixed_types(self) -> bool: """If true, the model 1. assumes total number of atoms aligned across frames; @@ -320,7 +333,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel": data = copy.deepcopy(data) - check_version_compatibility(data.get("@version", 2), 2, 1) + check_version_compatibility(data.pop("@version", 2), 2, 1) data.pop("@class", None) data.pop("type", None) models = [ @@ -331,16 +344,42 @@ def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel": return super().deserialize(data) def _compute_weight( - self, extended_coord, extended_atype, nlists_ + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlists_: list[torch.Tensor], ) -> list[torch.Tensor]: """This should be a list of user defined weights that matches the number of models to be combined.""" nmodels = len(self.models) nframes, nloc, _ = nlists_[0].shape - return [ - torch.ones((nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE) - / nmodels - for _ in range(nmodels) - ] + if isinstance(self.weights, str): + if self.weights == "sum": + return [ + torch.ones( + (nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE + ) + for _ in range(nmodels) + ] + elif self.weights == "mean": + return [ + torch.ones( + (nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE + ) + / nmodels + for _ in range(nmodels) + ] + else: + raise ValueError( + "`weights` must be 'sum' or 'mean' when provided as a string." + ) + elif isinstance(self.weights, list): + return [ + torch.ones((nframes, nloc, 1), dtype=torch.float64, device=env.DEVICE) + * w + for w in self.weights + ] + else: + raise NotImplementedError def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" @@ -365,7 +404,9 @@ def get_sel_type(self) -> list[int]: return torch.unique( torch.cat( [ - torch.as_tensor(model.get_sel_type(), dtype=torch.int32) + torch.as_tensor( + model.get_sel_type(), dtype=torch.int64, device=env.DEVICE + ) for model in self.models ] ) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 1c81d42013..26aefa6201 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -36,6 +36,9 @@ from .dos_model import ( DOSModel, ) +from .dp_linear_model import ( + LinearEnergyModel, +) from .dp_model import ( DPModelCommon, ) @@ -105,6 +108,62 @@ def get_spin_model(model_params): return SpinEnergyModel(backbone_model=backbone_model, spin=spin) +def get_linear_model(model_params): + model_params = copy.deepcopy(model_params) + weights = model_params.get("weights", "mean") + list_of_models = [] + ntypes = len(model_params["type_map"]) + for sub_model_params in model_params["models"]: + if "descriptor" in sub_model_params: + # descriptor + sub_model_params["descriptor"]["ntypes"] = ntypes + sub_model_params["descriptor"]["type_map"] = copy.deepcopy( + model_params["type_map"] + ) + descriptor = BaseDescriptor(**sub_model_params["descriptor"]) + # fitting + fitting_net = sub_model_params.get("fitting_net", {}) + fitting_net["type"] = fitting_net.get("type", "ener") + fitting_net["ntypes"] = descriptor.get_ntypes() + fitting_net["type_map"] = copy.deepcopy(model_params["type_map"]) + fitting_net["mixed_types"] = descriptor.mixed_types() + if fitting_net["type"] in ["dipole", "polar"]: + fitting_net["embedding_width"] = descriptor.get_dim_emb() + fitting_net["dim_descrpt"] = descriptor.get_dim_out() + grad_force = "direct" not in fitting_net["type"] + if not grad_force: + fitting_net["out_dim"] = descriptor.get_dim_emb() + if "ener" in fitting_net["type"]: + fitting_net["return_energy"] = True + fitting = BaseFitting(**fitting_net) + list_of_models.append( + DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"]) + ) + + else: # must be pairtab + assert ( + "type" in sub_model_params and sub_model_params["type"] == "pairtab" + ), "Sub-models in LinearEnergyModel must be a DPModel or a PairTable Model" + list_of_models.append( + PairTabAtomicModel( + sub_model_params["tab_file"], + sub_model_params["rcut"], + sub_model_params["sel"], + type_map=model_params["type_map"], + ) + ) + + atom_exclude_types = model_params.get("atom_exclude_types", []) + pair_exclude_types = model_params.get("pair_exclude_types", []) + return LinearEnergyModel( + models=list_of_models, + type_map=model_params["type_map"], + weights=weights, + atom_exclude_types=atom_exclude_types, + pair_exclude_types=pair_exclude_types, + ) + + def get_zbl_model(model_params): model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) @@ -247,6 +306,8 @@ def get_model(model_params): return get_zbl_model(model_params) else: return get_standard_model(model_params) + elif model_type == "linear_ener": + return get_linear_model(model_params) else: return BaseModel.get_class_by_type(model_type).get_model(model_params) @@ -265,4 +326,5 @@ def get_model(model_params): "DPZBLModel", "make_model", "make_hessian_model", + "LinearEnergyModel", ] diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py new file mode 100644 index 0000000000..ef2e84bd19 --- /dev/null +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) +from typing import ( + Optional, +) + +import torch + +from deepmd.pt.model.atomic_model import ( + LinearEnergyAtomicModel, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + +DPLinearModel_ = make_model(LinearEnergyAtomicModel) + + +@BaseModel.register("linear_ener") +class LinearEnergyModel(DPLinearModel_): + model_type = "ener" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def translated_output_def(self): + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_energy": deepcopy(out_def_data["energy"]), + "energy": deepcopy(out_def_data["energy_redu"]), + } + if self.do_grad_r("energy"): + output_def["force"] = deepcopy(out_def_data["energy_derv_r"]) + output_def["force"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"]) + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"]) + output_def["atom_virial"].squeeze(-3) + if "mask" in out_def_data: + output_def["mask"] = deepcopy(out_def_data["mask"]) + return output_def + + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy"): + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3) + else: + model_predict["force"] = model_ret["dforce"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ): + model_ret = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), + ) + + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy"): + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -3 + ) + else: + assert model_ret["dforce"] is not None + model_predict["dforce"] = model_ret["dforce"] + return model_predict + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[list[str]], + local_jdata: dict, + ) -> tuple[dict, Optional[float]]: + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms + local_jdata : dict + The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms + """ + local_jdata_cpy = local_jdata.copy() + type_map = local_jdata_cpy["type_map"] + min_nbor_dist = None + for idx, sub_model in enumerate(local_jdata_cpy["models"]): + if "tab_file" not in sub_model: + sub_model, temp_min = DPModelCommon.update_sel( + train_data, type_map, local_jdata["models"][idx] + ) + if min_nbor_dist is None or temp_min <= min_nbor_dist: + min_nbor_dist = temp_min + return local_jdata_cpy, min_nbor_dist diff --git a/doc/model/linear.md b/doc/model/linear.md index 3891559d90..47fdd1750b 100644 --- a/doc/model/linear.md +++ b/doc/model/linear.md @@ -1,7 +1,7 @@ -## Linear model {{ tensorflow_icon }} +## Linear model {{ tensorflow_icon }} {{ pytorch_icon }} :::{note} -**Supported backends**: TensorFlow {{ tensorflow_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} ::: One can linearly combine existing models with arbitrary coefficients: diff --git a/examples/water/d3/dftd3.txt b/examples/water/d3/dftd3.txt index bbc9726134..09e5fb697a 100644 --- a/examples/water/d3/dftd3.txt +++ b/examples/water/d3/dftd3.txt @@ -97,4 +97,4 @@ 9.700000000000001066e+00 -1.186747936398473687e-05 -7.637113677130612127e-06 -5.528293849956352819e-06 9.800000000000000711e+00 -1.114523618469756001e-05 -7.174288601187318493e-06 -5.194401230658985063e-06 9.900000000000000355e+00 -1.047381249252528874e-05 -6.743886368019750717e-06 -4.883815978498405921e-06 -1.000000000000000000e+01 0.000000000000000e00e+00 0.000000000000000e00e+00 0.000000000000000e00e+00 +1.000000000000000000e+01 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 diff --git a/examples/water/d3/input_pt.json b/examples/water/d3/input_pt.json new file mode 100644 index 0000000000..c2d9304a7e --- /dev/null +++ b/examples/water/d3/input_pt.json @@ -0,0 +1,96 @@ +{ + "_comment1": " model parameters", + "model": { + "type": "linear_ener", + "weights": "sum", + "type_map": [ + "O", + "H" + ], + "models": [ + { + "descriptor": { + "type": "se_atten", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "precision": "float64", + "seed": 1, + "_comment2": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float64", + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + }, + { + "type": "pairtab", + "tab_file": "dftd3.txt", + "rcut": 10.0, + "sel": 534 + } + ] + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": [ + "../data/data_0/", + "../data/data_1/", + "../data/data_2/" + ], + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} diff --git a/examples/water/linear/input_pt.json b/examples/water/linear/input_pt.json new file mode 100644 index 0000000000..e8d8e07136 --- /dev/null +++ b/examples/water/linear/input_pt.json @@ -0,0 +1,124 @@ +{ + "_comment1": " model parameters", + "model": { + "type": "linear_ener", + "weights": "sum", + "type_map": [ + "O", + "H" + ], + "models": [ + { + "descriptor": { + "type": "se_atten", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "precision": "float64", + "seed": 1, + "_comment2": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float64", + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + }, + { + "descriptor": { + "type": "se_atten", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "precision": "float64", + "seed": 1, + "_comment2": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float64", + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + } + ] + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": [ + "../data/data_0/", + "../data/data_1/", + "../data/data_2/" + ], + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} diff --git a/examples/water/zbl/input.json b/examples/water/zbl/input.json index cb5602d92d..54586ca0cf 100644 --- a/examples/water/zbl/input.json +++ b/examples/water/zbl/input.json @@ -10,7 +10,7 @@ "H" ], "descriptor": { - "type": "se_e2_a", + "type": "se_atten_v2", "sel": [ 46, 92 diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 6abb482824..246e767f01 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -34,7 +34,9 @@ p_examples / "water" / "hybrid" / "input.json", p_examples / "water" / "dplr" / "train" / "dw.json", p_examples / "water" / "dplr" / "train" / "ener.json", + p_examples / "water" / "d3" / "input_pt.json", p_examples / "water" / "linear" / "input.json", + p_examples / "water" / "linear" / "input_pt.json", p_examples / "nopbc" / "train" / "input.json", p_examples / "water_tensor" / "dipole" / "dipole_input.json", p_examples / "water_tensor" / "polar" / "polar_input.json", diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index 6aec895041..2d391c7115 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -98,6 +98,7 @@ "data_stat_nbatch": 20, } + model_spin = { "type_map": ["O", "H", "B"], "descriptor": { diff --git a/source/tests/universal/common/cases/model/model.py b/source/tests/universal/common/cases/model/model.py index c31f5cd889..cee69d9d6c 100644 --- a/source/tests/universal/common/cases/model/model.py +++ b/source/tests/universal/common/cases/model/model.py @@ -28,6 +28,25 @@ def setUpClass(cls) -> None: cls.epsilon_dict = {} +class LinearEnerModelTest(ModelTestCase): + @classmethod + def setUpClass(cls) -> None: + cls.expected_rcut = 5.0 + cls.expected_type_map = ["O", "H"] + cls.expected_dim_fparam = 0 + cls.expected_dim_aparam = 0 + cls.expected_sel_type = [0, 1] + cls.expected_aparam_nall = False + cls.expected_model_output_type = ["energy", "mask"] + cls.model_output_equivariant = [] + cls.expected_sel = [46, 92] + cls.expected_sel_mix = sum(cls.expected_sel) + cls.expected_has_message_passing = False + cls.aprec_dict = {} + cls.rprec_dict = {} + cls.epsilon_dict = {} + + class DipoleModelTest(ModelTestCase): @classmethod def setUpClass(cls) -> None: diff --git a/source/tests/universal/pt/model/test_model.py b/source/tests/universal/pt/model/test_model.py index 41df0cf762..81c32eb94c 100644 --- a/source/tests/universal/pt/model/test_model.py +++ b/source/tests/universal/pt/model/test_model.py @@ -21,6 +21,7 @@ DOSModel, DPZBLModel, EnergyModel, + LinearEnergyModel, PolarModel, PropertyModel, SpinEnergyModel, @@ -43,6 +44,7 @@ DipoleModelTest, DosModelTest, EnerModelTest, + LinearEnerModelTest, PolarModelTest, PropertyModelTest, SpinEnerModelTest, @@ -803,3 +805,100 @@ def setUpClass(cls): cls.expected_sel_type = ft.get_sel_type() cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + + +@parameterized( + des_parameterized=( + ( + *[(param_func, DescrptDPA1) for param_func in DescriptorParamDPA1List], + *[(param_func, DescrptDPA2) for param_func in DescriptorParamDPA2List], + (DescriptorParamHybridMixed, DescrptHybrid), + (DescriptorParamHybridMixedTTebd, DescrptHybrid), + ), # descrpt_class_param & class + ((FittingParamEnergy, EnergyFittingNet),), # fitting_class_param & class + ), + fit_parameterized=( + ( + (DescriptorParamDPA1, DescrptDPA1), + (DescriptorParamDPA2, DescrptDPA2), + ), # descrpt_class_param & class + ( + *[(param_func, EnergyFittingNet) for param_func in FittingParamEnergyList], + ), # fitting_class_param & class + ), +) +class TestLinearEnergyModelPT(unittest.TestCase, LinearEnerModelTest, PTTestCase): + @property + def modules_to_test(self): + skip_test_jit = getattr(self, "skip_test_jit", False) + modules = PTTestCase.modules_to_test.fget(self) + if not skip_test_jit: + # for Model, we can test script module API + modules += [ + self._script_module + if hasattr(self, "_script_module") + else self.script_module + ] + return modules + + @classmethod + def setUpClass(cls): + LinearEnerModelTest.setUpClass() + (DescriptorParam, Descrpt) = cls.param[0] + (FittingParam, Fitting) = cls.param[1] + # set special precision + cls.aprec_dict["test_smooth"] = 1e-5 + cls.input_dict_ds = DescriptorParam( + len(cls.expected_type_map), + cls.expected_rcut, + cls.expected_rcut / 2, + cls.expected_sel, + cls.expected_type_map, + ) + + # set skip tests + skiptest, skip_reason = skip_model_tests(cls) + if skiptest: + raise cls.skipTest(cls, skip_reason) + + ds1, ds2 = Descrpt(**cls.input_dict_ds), Descrpt(**cls.input_dict_ds) + cls.input_dict_ft = FittingParam( + ntypes=len(cls.expected_type_map), + dim_descrpt=ds1.get_dim_out(), + mixed_types=ds1.mixed_types(), + type_map=cls.expected_type_map, + ) + ft1 = Fitting( + **cls.input_dict_ft, + ) + ft2 = Fitting( + **cls.input_dict_ft, + ) + dp_model1 = DPAtomicModel( + ds1, + ft1, + type_map=cls.expected_type_map, + ) + dp_model2 = DPAtomicModel( + ds2, + ft2, + type_map=cls.expected_type_map, + ) + cls.module = LinearEnergyModel( + [dp_model1, dp_model2], + type_map=cls.expected_type_map, + ) + # only test jit API once for different models + if ( + DescriptorParam not in defalut_des_param + or FittingParam not in defalut_fit_param + ): + cls.skip_test_jit = True + else: + with torch.jit.optimized_execution(False): + cls._script_module = torch.jit.script(cls.module) + cls.output_def = cls.module.translated_output_def() + cls.expected_has_message_passing = ds1.has_message_passing() + cls.expected_dim_fparam = ft1.get_dim_fparam() + cls.expected_dim_aparam = ft1.get_dim_aparam() + cls.expected_sel_type = ft1.get_sel_type()