From ae64a72b00d7b687cb5880648a64005ef5cd0693 Mon Sep 17 00:00:00 2001 From: wangyan Date: Fri, 1 Nov 2024 22:13:27 +0800 Subject: [PATCH 01/12] Add compression API --- .../dpmodel/atomic_model/base_atomic_model.py | 25 +++++++++++++++ .../dpmodel/atomic_model/dp_atomic_model.py | 31 +++++++++++++++++++ deepmd/dpmodel/model/make_model.py | 30 ++++++++++++++++++ 3 files changed, 86 insertions(+) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 4e7620bdda..102a0d3d4c 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -141,6 +141,31 @@ def change_type_map( self.out_bias = self.out_bias[:, remap_index, :] self.out_std = self.out_std[:, remap_index, :] + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call descriptor enable_compression() + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + raise NotImplementedError("This atomi model doesn't support compression!") + def forward_common_atomic( self, extended_coord: np.ndarray, diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index a621ece27e..18f0499f2c 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -85,6 +85,37 @@ def has_message_passing(self) -> bool: def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" return self.descriptor.need_sorted_nlist_for_lower() + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call descriptor enable_compression() + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.descriptor.enable_compression( + min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) def forward_atomic( self, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b6379573e1..32cd0ddce6 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -185,6 +185,36 @@ def model_output_type(self) -> list[str]: if vv.category == OutputVariableCategory.OUT ] return vars + + def enable_compression( + self, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call atomic_model enable_compression() + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.atomic_model.enable_compression( + self.get_min_nbor_dist(), + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) def call( self, From d1d49c394eeab7d962821053e6e5e1b4f7908dcf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 1 Nov 2024 14:16:00 +0000 Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../dpmodel/atomic_model/base_atomic_model.py | 24 ++++++++--------- .../dpmodel/atomic_model/dp_atomic_model.py | 26 +++++++++---------- deepmd/dpmodel/model/make_model.py | 24 ++++++++--------- 3 files changed, 36 insertions(+), 38 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 102a0d3d4c..c6596f0405 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -150,19 +150,19 @@ def enable_compression( check_frequency: int = -1, ) -> None: """Call descriptor enable_compression() - + Parameters - ---------- - min_nbor_dist - The nearest distance between atoms - table_extrapolate - The scale of model extrapolation - table_stride_1 - The uniform stride of the first table - table_stride_2 - The uniform stride of the second table - check_frequency - The overflow check frequency + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency """ raise NotImplementedError("This atomi model doesn't support compression!") diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 18f0499f2c..13b08a8677 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -85,7 +85,7 @@ def has_message_passing(self) -> bool: def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the atomic model needs sorted nlist when using `forward_lower`.""" return self.descriptor.need_sorted_nlist_for_lower() - + def enable_compression( self, min_nbor_dist: float, @@ -95,19 +95,19 @@ def enable_compression( check_frequency: int = -1, ) -> None: """Call descriptor enable_compression() - + Parameters - ---------- - min_nbor_dist - The nearest distance between atoms - table_extrapolate - The scale of model extrapolation - table_stride_1 - The uniform stride of the first table - table_stride_2 - The uniform stride of the second table - check_frequency - The overflow check frequency + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency """ self.descriptor.enable_compression( min_nbor_dist, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 32cd0ddce6..135d4db3bb 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -185,7 +185,7 @@ def model_output_type(self) -> list[str]: if vv.category == OutputVariableCategory.OUT ] return vars - + def enable_compression( self, table_extrapolate: float = 5, @@ -194,19 +194,17 @@ def enable_compression( check_frequency: int = -1, ) -> None: """Call atomic_model enable_compression() - + Parameters - ---------- - min_nbor_dist - The nearest distance between atoms - table_extrapolate - The scale of model extrapolation - table_stride_1 - The uniform stride of the first table - table_stride_2 - The uniform stride of the second table - check_frequency - The overflow check frequency + ---------- + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency """ self.atomic_model.enable_compression( self.get_min_nbor_dist(), From f2c292729c1743924a6bf0654c2e1b48115eb304 Mon Sep 17 00:00:00 2001 From: wangyan Date: Fri, 1 Nov 2024 22:26:43 +0800 Subject: [PATCH 03/12] eFirst line nd with a period --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 2 +- deepmd/dpmodel/atomic_model/dp_atomic_model.py | 2 +- deepmd/dpmodel/model/make_model.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 102a0d3d4c..4fabb00894 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -149,7 +149,7 @@ def enable_compression( table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call descriptor enable_compression() + """Call descriptor enable_compression(). Parameters ---------- diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 18f0499f2c..5227360e93 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -94,7 +94,7 @@ def enable_compression( table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call descriptor enable_compression() + """Call descriptor enable_compression(). Parameters ---------- diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 32cd0ddce6..64ef3b8724 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -193,7 +193,7 @@ def enable_compression( table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call atomic_model enable_compression() + """Call atomic_model enable_compression(). Parameters ---------- From 8699c668a9421ffb6e00b7b56f236c3c356812b2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 1 Nov 2024 14:33:29 +0000 Subject: [PATCH 04/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 2 +- deepmd/dpmodel/model/make_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 4d7ecbd391..e48f4b19d8 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -150,7 +150,7 @@ def enable_compression( check_frequency: int = -1, ) -> None: """Call descriptor enable_compression(). - + Parameters ---------- min_nbor_dist diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b4ed65fb0e..95d97262df 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -194,7 +194,7 @@ def enable_compression( check_frequency: int = -1, ) -> None: """Call atomic_model enable_compression(). - + Parameters ---------- table_extrapolate From 6402f892bd71416e6294e345d646a8ef9d602dcb Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 10:17:41 +0800 Subject: [PATCH 05/12] fix atomic model compression --- .../dpmodel/atomic_model/base_atomic_model.py | 25 ------------------- .../atomic_model/make_base_atomic_model.py | 25 +++++++++++++++++++ 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index e48f4b19d8..4e7620bdda 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -141,31 +141,6 @@ def change_type_map( self.out_bias = self.out_bias[:, remap_index, :] self.out_std = self.out_std[:, remap_index, :] - def enable_compression( - self, - min_nbor_dist: float, - table_extrapolate: float = 5, - table_stride_1: float = 0.01, - table_stride_2: float = 0.1, - check_frequency: int = -1, - ) -> None: - """Call descriptor enable_compression(). - - Parameters - ---------- - min_nbor_dist - The nearest distance between atoms - table_extrapolate - The scale of model extrapolation - table_stride_1 - The uniform stride of the first table - table_stride_2 - The uniform stride of the second table - check_frequency - The overflow check frequency - """ - raise NotImplementedError("This atomi model doesn't support compression!") - def forward_common_atomic( self, extended_coord: np.ndarray, diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 99a92c23a4..a4c38518a3 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -148,6 +148,31 @@ def change_type_map( ) -> None: pass + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call descriptor enable_compression(). + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + raise NotImplementedError("This atomi model doesn't support compression!") + def make_atom_mask( self, atype: t_tensor, From 9a1ce86e66de9dd49e4900f1d2dc47bf1e649c77 Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 10:18:05 +0800 Subject: [PATCH 06/12] fix base model compression --- deepmd/dpmodel/model/base_model.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 777697b4b7..8403d179ed 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -190,6 +190,30 @@ def update_sel( model_type = local_jdata.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) return cls.update_sel(train_data, type_map, local_jdata) + + def enable_compression( + self, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call descriptor enable_compression(). + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + raise NotImplementedError("This atomi model doesn't support compression!") @classmethod def get_model(cls, model_params: dict) -> "BaseBaseModel": From 74e15f711e875f76674cac1992f789cfccbc8b55 Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 10:20:57 +0800 Subject: [PATCH 07/12] add pt atomic model enable_compression() --- .../pt/model/atomic_model/dp_atomic_model.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 48c8d0d859..b75b2c9d35 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -159,6 +159,37 @@ def deserialize(cls, data) -> "DPAtomicModel": data["fitting"] = fitting_obj obj = super().deserialize(data) return obj + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call descriptor enable_compression(). + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.descriptor.enable_compression( + min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) def forward_atomic( self, From 9ac52049ef9f3f559d39f8573f2562bf14864207 Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 10:21:30 +0800 Subject: [PATCH 08/12] add pt base model enable_compression() --- deepmd/pt/model/model/make_model.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index a9d5e26060..44d8d63ad5 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -97,6 +97,34 @@ def model_output_type(self) -> list[str]: if vv.category == OutputVariableCategory.OUT.value: vars.append(kk) return vars + + def enable_compression( + self, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Call atomic_model enable_compression(). + + Parameters + ---------- + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.atomic_model.enable_compression( + self.get_min_nbor_dist(), + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) # cannot use the name forward. torch script does not work def forward_common( From 6e935863f0cb03c8c057cefe917a20f48cfbd904 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 2 Nov 2024 02:23:03 +0000 Subject: [PATCH 09/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/base_model.py | 4 +--- deepmd/pt/model/atomic_model/dp_atomic_model.py | 2 +- deepmd/pt/model/model/make_model.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 8403d179ed..3ee700ec16 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -190,7 +190,7 @@ def update_sel( model_type = local_jdata.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) return cls.update_sel(train_data, type_map, local_jdata) - + def enable_compression( self, table_extrapolate: float = 5, @@ -202,8 +202,6 @@ def enable_compression( Parameters ---------- - min_nbor_dist - The nearest distance between atoms table_extrapolate The scale of model extrapolation table_stride_1 diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index b75b2c9d35..4f53fb69f9 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -159,7 +159,7 @@ def deserialize(cls, data) -> "DPAtomicModel": data["fitting"] = fitting_obj obj = super().deserialize(data) return obj - + def enable_compression( self, min_nbor_dist: float, diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 44d8d63ad5..709fce8150 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -97,7 +97,7 @@ def model_output_type(self) -> list[str]: if vv.category == OutputVariableCategory.OUT.value: vars.append(kk) return vars - + def enable_compression( self, table_extrapolate: float = 5, From 4d106a101a5ebf795175d225fe8c678bb0aa466c Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 15:14:37 +0800 Subject: [PATCH 10/12] fix syntax errors --- deepmd/dpmodel/model/base_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 8403d179ed..e800d5cd36 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -198,7 +198,7 @@ def enable_compression( table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call descriptor enable_compression(). + """Enable model compression by tabulation. Parameters ---------- @@ -213,7 +213,7 @@ def enable_compression( check_frequency The overflow check frequency """ - raise NotImplementedError("This atomi model doesn't support compression!") + raise NotImplementedError("This atomic model doesn't support compression!") @classmethod def get_model(cls, model_params: dict) -> "BaseBaseModel": From a9232494c9259a0445cd6c446bd8f5ac39f2c7fb Mon Sep 17 00:00:00 2001 From: wangyan Date: Sat, 2 Nov 2024 15:45:13 +0800 Subject: [PATCH 11/12] add compression API for LinearEnergyAtomicModel --- .../atomic_model/linear_atomic_model.py | 32 +++++++++++++++++++ .../model/atomic_model/linear_atomic_model.py | 32 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 224fdd145c..541a5ae69c 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -148,6 +148,38 @@ def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]: key=lambda x: (x[1], x[0]), ) return [p[0] for p in zipped], [p[1] for p in zipped] + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Compress model. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + for model in self.models: + model.enable_compression( + min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) def forward_atomic( self, diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 570fcdcc43..0ba284e458 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -183,6 +183,38 @@ def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]: sorted_rcuts: list[float] = outer_sorted[:, 0].tolist() sorted_sels: list[int] = outer_sorted[:, 1].to(torch.int64).tolist() return sorted_rcuts, sorted_sels + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Compress model. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + for model in self.models: + model.enable_compression( + min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ) def forward_atomic( self, From 999400a0b581bb3319a2706f93fa2120ede14b34 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 2 Nov 2024 07:46:49 +0000 Subject: [PATCH 12/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 2 +- deepmd/pt/model/atomic_model/linear_atomic_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 541a5ae69c..485f82cb72 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -148,7 +148,7 @@ def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]: key=lambda x: (x[1], x[0]), ) return [p[0] for p in zipped], [p[1] for p in zipped] - + def enable_compression( self, min_nbor_dist: float, diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 0ba284e458..33c2eb6c59 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -183,7 +183,7 @@ def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]: sorted_rcuts: list[float] = outer_sorted[:, 0].tolist() sorted_sels: list[int] = outer_sorted[:, 1].to(torch.int64).tolist() return sorted_rcuts, sorted_sels - + def enable_compression( self, min_nbor_dist: float,