From 64b784f05368313c0b5629a143e6fcf4265d25b9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 10 Sep 2024 12:06:58 +0800 Subject: [PATCH 1/7] support preset bias of atomic model output --- deepmd/pt/model/model/__init__.py | 10 ++++ deepmd/utils/argcheck.py | 10 ++++ source/tests/pt/model/test_get_model.py | 70 +++++++++++++++++++++++++ 3 files changed, 90 insertions(+) create mode 100644 source/tests/pt/model/test_get_model.py diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index ed9cfcd7c2..851ac6b59f 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -176,6 +176,15 @@ def get_standard_model(model_params): fitting = BaseFitting(**fitting_net) atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) + default_preset_out_bias = [None] * len(model_params["type_map"]) + preset_out_bias = model_params.get("preset_out_bias") + preset_out_bias = ( + preset_out_bias if preset_out_bias is not None else default_preset_out_bias + ) + if len(preset_out_bias) != len(model_params["type_map"]): + raise ValueError( + "length of the preset_out_bias should be the same as the type_map" + ) if fitting_net["type"] == "dipole": modelcls = DipoleModel @@ -196,6 +205,7 @@ def get_standard_model(model_params): type_map=model_params["type_map"], atom_exclude_types=atom_exclude_types, pair_exclude_types=pair_exclude_types, + preset_out_bias=preset_out_bias, ) model.model_def_script = json.dumps(model_params_old) return model diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index c2f483e715..5b9c8eb95a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1771,6 +1771,9 @@ def model_args(exclude_hybrid=False): doc_spin = "The settings for systems with spin." doc_atom_exclude_types = "Exclude the atomic contribution of the listed atom types" doc_pair_exclude_types = "The atom pairs of the listed types are not treated to be neighbors, i.e. they do not see each other." + doc_preset_out_bias = ( + "The preset bias of the atomic output. The set_davg_zero should be set to true." + ) doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " "If not set or set to 'RANDOM', the fitting net will be randomly initialized." @@ -1833,6 +1836,13 @@ def model_args(exclude_hybrid=False): default=[], doc=doc_only_pt_supported + doc_atom_exclude_types, ), + Argument( + "preset_out_bias", + list, + optional=True, + default=None, + doc=doc_only_pt_supported + doc_preset_out_bias, + ), Argument( "srtab_add_bias", bool, diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py new file mode 100644 index 0000000000..555ea7dbe3 --- /dev/null +++ b/source/tests/pt/model/test_get_model.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import unittest + +import torch + +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils import ( + env, +) + +dtype = torch.float64 + +model_se_e2_a = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "se_e2_a", + "sel": [46, 92, 4], + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 16, + "seed": 1, + }, + "fitting_net": { + "neuron": [24, 24, 24], + "resnet_dt": True, + "seed": 1, + }, + "data_stat_nbatch": 20, + "atom_exclude_types": [1], + "pair_exclude_types": [[1, 2]], + "preset_out_bias": [ + None, + [1.0], + [3.0], + ], +} + + +class TestGetModel(unittest.TestCase): + def test_model_attr(self): + model_params = copy.deepcopy(model_se_e2_a) + self.model = get_model(model_params).to(env.DEVICE) + atomic_model = self.model.atomic_model + self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) + self.assertEqual(atomic_model.preset_out_bias, [None, [1.0], [3.0]]) + self.assertEqual(atomic_model.atom_exclude_types, [1]) + self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) + + def test_notset_model_attr(self): + model_params = copy.deepcopy(model_se_e2_a) + model_params.pop("atom_exclude_types") + model_params.pop("pair_exclude_types") + model_params.pop("preset_out_bias") + self.model = get_model(model_params).to(env.DEVICE) + atomic_model = self.model.atomic_model + self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) + self.assertEqual(atomic_model.preset_out_bias, [None, None, None]) + self.assertEqual(atomic_model.atom_exclude_types, []) + self.assertEqual(atomic_model.pair_exclude_types, []) + + def test_preset_wrong_len(self): + model_params = copy.deepcopy(model_se_e2_a) + model_params["preset_out_bias"] = [None] + with self.assertRaises(ValueError): + self.model = get_model(model_params).to(env.DEVICE) From 0ec6fe1f12e385fa12e778612ad67bd1644a817d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 10 Sep 2024 13:49:58 +0800 Subject: [PATCH 2/7] fix bugs --- deepmd/pt/model/model/__init__.py | 26 ++++++++++++++++++------ deepmd/utils/argcheck.py | 6 ++---- source/tests/pt/model/test_get_model.py | 27 +++++++++++++++++-------- 3 files changed, 41 insertions(+), 18 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 851ac6b59f..905fe413c4 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -26,6 +26,9 @@ from deepmd.pt.model.task import ( BaseFitting, ) +from deepmd.pt.utils.utils import ( + to_torch_tensor, +) from deepmd.utils.spin import ( Spin, ) @@ -151,6 +154,21 @@ def get_zbl_model(model_params): ) +def _convert_preset_out_bias_to_torch_tensor(preset_out_bias, type_map): + if preset_out_bias is not None: + for kk in preset_out_bias.keys(): + if len(preset_out_bias[kk]) != len(type_map): + raise ValueError( + "length of the preset_out_bias should be the same as the type_map" + ) + for jj in range(len(preset_out_bias[kk])): + if preset_out_bias[kk][jj] is not None: + preset_out_bias[kk][jj] = to_torch_tensor( + np.array(preset_out_bias[kk][jj]) + ) + return preset_out_bias + + def get_standard_model(model_params): model_params_old = model_params model_params = copy.deepcopy(model_params) @@ -178,13 +196,9 @@ def get_standard_model(model_params): pair_exclude_types = model_params.get("pair_exclude_types", []) default_preset_out_bias = [None] * len(model_params["type_map"]) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = ( - preset_out_bias if preset_out_bias is not None else default_preset_out_bias + preset_out_bias = _convert_preset_out_bias_to_torch_tensor( + preset_out_bias, model_params["type_map"] ) - if len(preset_out_bias) != len(model_params["type_map"]): - raise ValueError( - "length of the preset_out_bias should be the same as the type_map" - ) if fitting_net["type"] == "dipole": modelcls = DipoleModel diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 5b9c8eb95a..2ed29fca20 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1771,9 +1771,7 @@ def model_args(exclude_hybrid=False): doc_spin = "The settings for systems with spin." doc_atom_exclude_types = "Exclude the atomic contribution of the listed atom types" doc_pair_exclude_types = "The atom pairs of the listed types are not treated to be neighbors, i.e. they do not see each other." - doc_preset_out_bias = ( - "The preset bias of the atomic output. The set_davg_zero should be set to true." - ) + doc_preset_out_bias = "The preset bias of the atomic output. Is provided as a dict. Taking the energy model that has three atom types for example, the preset_out_bias may be given as `{ 'energy': [None, 0., 1.] }`. In this case the bias of type 1 and 2 are set to 0. and 1., respectively.The set_davg_zero should be set to true." doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " "If not set or set to 'RANDOM', the fitting net will be randomly initialized." @@ -1838,7 +1836,7 @@ def model_args(exclude_hybrid=False): ), Argument( "preset_out_bias", - list, + dict, optional=True, default=None, doc=doc_only_pt_supported + doc_preset_out_bias, diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py index 555ea7dbe3..051cabc4d9 100644 --- a/source/tests/pt/model/test_get_model.py +++ b/source/tests/pt/model/test_get_model.py @@ -33,11 +33,13 @@ "data_stat_nbatch": 20, "atom_exclude_types": [1], "pair_exclude_types": [[1, 2]], - "preset_out_bias": [ - None, - [1.0], - [3.0], - ], + "preset_out_bias": { + "energy": [ + None, + [1.0], + [3.0], + ] + }, } @@ -47,7 +49,16 @@ def test_model_attr(self): self.model = get_model(model_params).to(env.DEVICE) atomic_model = self.model.atomic_model self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) - self.assertEqual(atomic_model.preset_out_bias, [None, [1.0], [3.0]]) + self.assertEqual( + atomic_model.preset_out_bias, + { + "energy": [ + None, + torch.tensor([1.0], dtype=dtype, device=env.DEVICE), + torch.tensor([3.0], dtype=dtype, device=env.DEVICE), + ] + }, + ) self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) @@ -59,12 +70,12 @@ def test_notset_model_attr(self): self.model = get_model(model_params).to(env.DEVICE) atomic_model = self.model.atomic_model self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) - self.assertEqual(atomic_model.preset_out_bias, [None, None, None]) + self.assertEqual(atomic_model.preset_out_bias, None) self.assertEqual(atomic_model.atom_exclude_types, []) self.assertEqual(atomic_model.pair_exclude_types, []) def test_preset_wrong_len(self): model_params = copy.deepcopy(model_se_e2_a) - model_params["preset_out_bias"] = [None] + model_params["preset_out_bias"] = {"energy": [None]} with self.assertRaises(ValueError): self.model = get_model(model_params).to(env.DEVICE) From fad0596c8a07a9b3ffb4fca0dc257813f8b0bb64 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 10 Sep 2024 14:07:59 +0800 Subject: [PATCH 3/7] rm unused variable --- deepmd/pt/model/model/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 905fe413c4..3609ecdd8b 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -194,7 +194,6 @@ def get_standard_model(model_params): fitting = BaseFitting(**fitting_net) atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) - default_preset_out_bias = [None] * len(model_params["type_map"]) preset_out_bias = model_params.get("preset_out_bias") preset_out_bias = _convert_preset_out_bias_to_torch_tensor( preset_out_bias, model_params["type_map"] From bcc3b75bd8de73cda921625b28f3ad0e550e9194 Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Tue, 10 Sep 2024 14:44:16 +0800 Subject: [PATCH 4/7] Update deepmd/pt/model/model/__init__.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> --- deepmd/pt/model/model/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 3609ecdd8b..a316dc001a 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -156,7 +156,8 @@ def get_zbl_model(model_params): def _convert_preset_out_bias_to_torch_tensor(preset_out_bias, type_map): if preset_out_bias is not None: - for kk in preset_out_bias.keys(): + if preset_out_bias is not None: + for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): raise ValueError( "length of the preset_out_bias should be the same as the type_map" From f2885d537a29b2823a4e41223d1d874c9c876de2 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 10 Sep 2024 13:45:11 -0400 Subject: [PATCH 5/7] Update deepmd/pt/model/model/__init__.py Signed-off-by: Jinzhe Zeng --- deepmd/pt/model/model/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index a316dc001a..cb6a022794 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -155,7 +155,6 @@ def get_zbl_model(model_params): def _convert_preset_out_bias_to_torch_tensor(preset_out_bias, type_map): - if preset_out_bias is not None: if preset_out_bias is not None: for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): From 32970202915fb5595233d38a0ab370f09d3c613b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 11 Sep 2024 09:02:20 +0800 Subject: [PATCH 6/7] fix type hint of preset_out_bias. adopt other suggestions --- deepmd/pt/model/atomic_model/base_atomic_model.py | 5 +++-- deepmd/pt/model/model/__init__.py | 11 +++-------- deepmd/pt/utils/stat.py | 6 +++--- deepmd/utils/argcheck.py | 5 +++-- source/tests/pt/model/test_get_model.py | 5 +++-- 5 files changed, 15 insertions(+), 17 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index d73c794c73..3cc66397b4 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -11,6 +11,7 @@ Union, ) +import numpy as np import torch from deepmd.dpmodel.atomic_model import ( @@ -66,7 +67,7 @@ class BaseAtomicModel(torch.nn.Module, BaseAtomicModel_): of the atomic model. Implemented by removing the pairs from the nlist. rcond : float, optional The condition number for the regression of atomic energy. - preset_out_bias : Dict[str, List[Optional[torch.Tensor]]], optional + preset_out_bias : Dict[str, List[Optional[np.array]]], optional Specifying atomic energy contribution in vacuum. Given by key:value pairs. The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] @@ -80,7 +81,7 @@ def __init__( atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [], rcond: Optional[float] = None, - preset_out_bias: Optional[Dict[str, torch.Tensor]] = None, + preset_out_bias: Optional[Dict[str, np.array]] = None, ): torch.nn.Module.__init__(self) BaseAtomicModel_.__init__(self) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 3609ecdd8b..9c131b4e8f 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -26,9 +26,6 @@ from deepmd.pt.model.task import ( BaseFitting, ) -from deepmd.pt.utils.utils import ( - to_torch_tensor, -) from deepmd.utils.spin import ( Spin, ) @@ -154,7 +151,7 @@ def get_zbl_model(model_params): ) -def _convert_preset_out_bias_to_torch_tensor(preset_out_bias, type_map): +def _convert_preset_out_bias_to_array(preset_out_bias, type_map): if preset_out_bias is not None: for kk in preset_out_bias.keys(): if len(preset_out_bias[kk]) != len(type_map): @@ -163,9 +160,7 @@ def _convert_preset_out_bias_to_torch_tensor(preset_out_bias, type_map): ) for jj in range(len(preset_out_bias[kk])): if preset_out_bias[kk][jj] is not None: - preset_out_bias[kk][jj] = to_torch_tensor( - np.array(preset_out_bias[kk][jj]) - ) + preset_out_bias[kk][jj] = np.array(preset_out_bias[kk][jj]) return preset_out_bias @@ -195,7 +190,7 @@ def get_standard_model(model_params): atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_torch_tensor( + preset_out_bias = _convert_preset_out_bias_to_array( preset_out_bias, model_params["type_map"] ) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 8adf21c127..7ddfacb541 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -242,7 +242,7 @@ def compute_output_stats( keys: Union[str, List[str]] = ["energy"], stat_file_path: Optional[DPPath] = None, rcond: Optional[float] = None, - preset_bias: Optional[Dict[str, List[Optional[torch.Tensor]]]] = None, + preset_bias: Optional[Dict[str, List[Optional[np.array]]]] = None, model_forward: Optional[Callable[..., torch.Tensor]] = None, atomic_output: Optional[FittingOutputDef] = None, ): @@ -264,7 +264,7 @@ def compute_output_stats( The path to the stat file. rcond : float, optional The condition number for the regression of atomic energy. - preset_bias : Dict[str, List[Optional[torch.Tensor]]], optional + preset_bias : Dict[str, List[Optional[np.array]]], optional Specifying atomic energy contribution in vacuum. Given by key:value pairs. The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] @@ -405,7 +405,7 @@ def compute_output_stats_global( ntypes: int, keys: List[str], rcond: Optional[float] = None, - preset_bias: Optional[Dict[str, List[Optional[torch.Tensor]]]] = None, + preset_bias: Optional[Dict[str, List[Optional[np.array]]]] = None, model_pred: Optional[Dict[str, np.ndarray]] = None, atomic_output: Optional[FittingOutputDef] = None, ): diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 2ed29fca20..4eab9d87df 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4,6 +4,7 @@ import warnings from typing import ( Callable, + Dict, List, Optional, Union, @@ -1771,7 +1772,7 @@ def model_args(exclude_hybrid=False): doc_spin = "The settings for systems with spin." doc_atom_exclude_types = "Exclude the atomic contribution of the listed atom types" doc_pair_exclude_types = "The atom pairs of the listed types are not treated to be neighbors, i.e. they do not see each other." - doc_preset_out_bias = "The preset bias of the atomic output. Is provided as a dict. Taking the energy model that has three atom types for example, the preset_out_bias may be given as `{ 'energy': [None, 0., 1.] }`. In this case the bias of type 1 and 2 are set to 0. and 1., respectively.The set_davg_zero should be set to true." + doc_preset_out_bias = "The preset bias of the atomic output. Is provided as a dict. Taking the energy model that has three atom types for example, the preset_out_bias may be given as `{ 'energy': [null, 0., 1.] }`. In this case the bias of type 1 and 2 are set to 0. and 1., respectively.The set_davg_zero should be set to true." doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " "If not set or set to 'RANDOM', the fitting net will be randomly initialized." @@ -1836,7 +1837,7 @@ def model_args(exclude_hybrid=False): ), Argument( "preset_out_bias", - dict, + Dict[str, Optional[float]], optional=True, default=None, doc=doc_only_pt_supported + doc_preset_out_bias, diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py index 051cabc4d9..c433597d5a 100644 --- a/source/tests/pt/model/test_get_model.py +++ b/source/tests/pt/model/test_get_model.py @@ -2,6 +2,7 @@ import copy import unittest +import numpy as np import torch from deepmd.pt.model.model import ( @@ -54,8 +55,8 @@ def test_model_attr(self): { "energy": [ None, - torch.tensor([1.0], dtype=dtype, device=env.DEVICE), - torch.tensor([3.0], dtype=dtype, device=env.DEVICE), + np.array([1.0]), + np.array([3.0]), ] }, ) From 1afca1c8c411e6b7b93e8b7a58e52c7f3c6d52c1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 11 Sep 2024 11:33:24 +0800 Subject: [PATCH 7/7] array->ndarray --- deepmd/pt/model/atomic_model/base_atomic_model.py | 6 +++--- deepmd/pt/utils/stat.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 3cc66397b4..4742fe66a3 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -67,9 +67,9 @@ class BaseAtomicModel(torch.nn.Module, BaseAtomicModel_): of the atomic model. Implemented by removing the pairs from the nlist. rcond : float, optional The condition number for the regression of atomic energy. - preset_out_bias : Dict[str, List[Optional[np.array]]], optional + preset_out_bias : Dict[str, List[Optional[np.ndarray]]], optional Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.array of output shape. + The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descrptor should be set. @@ -81,7 +81,7 @@ def __init__( atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [], rcond: Optional[float] = None, - preset_out_bias: Optional[Dict[str, np.array]] = None, + preset_out_bias: Optional[Dict[str, np.ndarray]] = None, ): torch.nn.Module.__init__(self) BaseAtomicModel_.__init__(self) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 7ddfacb541..6de70eb175 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -187,8 +187,8 @@ def model_forward_auto_batch_size(*args, **kwargs): def _make_preset_out_bias( ntypes: int, - ibias: List[Optional[np.array]], -) -> Optional[np.array]: + ibias: List[Optional[np.ndarray]], +) -> Optional[np.ndarray]: """Make preset out bias. output: @@ -242,7 +242,7 @@ def compute_output_stats( keys: Union[str, List[str]] = ["energy"], stat_file_path: Optional[DPPath] = None, rcond: Optional[float] = None, - preset_bias: Optional[Dict[str, List[Optional[np.array]]]] = None, + preset_bias: Optional[Dict[str, List[Optional[np.ndarray]]]] = None, model_forward: Optional[Callable[..., torch.Tensor]] = None, atomic_output: Optional[FittingOutputDef] = None, ): @@ -264,9 +264,9 @@ def compute_output_stats( The path to the stat file. rcond : float, optional The condition number for the regression of atomic energy. - preset_bias : Dict[str, List[Optional[np.array]]], optional + preset_bias : Dict[str, List[Optional[np.ndarray]]], optional Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.array of output shape. + The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descrptor should be set. model_forward : Callable[..., torch.Tensor], optional @@ -405,7 +405,7 @@ def compute_output_stats_global( ntypes: int, keys: List[str], rcond: Optional[float] = None, - preset_bias: Optional[Dict[str, List[Optional[np.array]]]] = None, + preset_bias: Optional[Dict[str, List[Optional[np.ndarray]]]] = None, model_pred: Optional[Dict[str, np.ndarray]] = None, atomic_output: Optional[FittingOutputDef] = None, ):