From 2597b1c26d057c577885bc271edd3c2577c529b8 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 16 Jan 2024 20:08:58 +0800 Subject: [PATCH 1/7] fix output def issues (some restrictions from torch.jit) --- deepmd_utils/model_format/__init__.py | 4 +++ deepmd_utils/model_format/output_def.py | 45 +++++++++++++++---------- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py index 72dd7b59ee..961b483730 100644 --- a/deepmd_utils/model_format/__init__.py +++ b/deepmd_utils/model_format/__init__.py @@ -26,6 +26,8 @@ OutputVariableDef, VariableDef, fitting_check_output, + get_deriv_name, + get_reduce_name, model_check_output, ) from .se_e2_a import ( @@ -55,4 +57,6 @@ "VariableDef", "model_check_output", "fitting_check_output", + "get_reduce_name", + "get_deriv_name", ] diff --git a/deepmd_utils/model_format/output_def.py b/deepmd_utils/model_format/output_def.py index 7feb24a145..12bce53b82 100644 --- a/deepmd_utils/model_format/output_def.py +++ b/deepmd_utils/model_format/output_def.py @@ -3,7 +3,6 @@ Dict, List, Tuple, - Union, ) @@ -38,7 +37,7 @@ def __init__( **kwargs, ): super().__init__(*args, **kwargs) - self.md = cls.output_def(self) + self.md = self.output_def() def __call__( self, @@ -77,7 +76,7 @@ def __init__( **kwargs, ): super().__init__(*args, **kwargs) - self.md = cls.output_def(self) + self.md = self.output_def() def __call__( self, @@ -113,7 +112,7 @@ class VariableDef: def __init__( self, name: str, - shape: Union[List[int], Tuple[int]], + shape: list[int], atomic: bool = True, ): self.name = name @@ -149,12 +148,18 @@ class OutputVariableDef(VariableDef): def __init__( self, name: str, - shape: Union[List[int], Tuple[int]], + shape: List[int], reduciable: bool = False, differentiable: bool = False, ): - # fitting output must be atomic - super().__init__(name, shape, atomic=True) + ## fitting output must be atomic + ## Here we cannot use super because it does not pass jit + # super().__init__(name, shape, atomic=True) + ## the work around is the following + self.name = name + self.shape = list(shape) + self.atomic = True + # self.reduciable = reduciable self.differentiable = differentiable if not self.reduciable and self.differentiable: @@ -176,13 +181,13 @@ class FittingOutputDef: def __init__( self, - var_defs: List[OutputVariableDef] = [], + var_defs: List[OutputVariableDef], ): self.var_defs = {vv.name: vv for vv in var_defs} def __getitem__( self, - key, + key: str, ) -> OutputVariableDef: return self.var_defs[key] @@ -224,10 +229,16 @@ def __init__( ]: self.var_defs.update(ii) - def __getitem__(self, key) -> VariableDef: + def __getitem__( + self, + key: str, + ) -> VariableDef: return self.var_defs[key] - def get_data(self, key) -> Dict[str, VariableDef]: + def get_data( + self, + key: str, + ) -> Dict[str, VariableDef]: return self.var_defs def keys(self): @@ -246,17 +257,17 @@ def keys_derv_c(self): return self.def_derv_c.keys() -def get_reduce_name(name): +def get_reduce_name(name: str) -> str: return name + "_redu" -def get_deriv_name(name): +def get_deriv_name(name: str) -> Tuple[str, str]: return name + "_derv_r", name + "_derv_c" def do_reduce( - def_outp, -): + def_outp: FittingOutputDef, +) -> Dict[str, VariableDef]: def_redu = {} for kk, vv in def_outp.get_data().items(): if vv.reduciable: @@ -266,8 +277,8 @@ def do_reduce( def do_derivative( - def_outp, -): + def_outp: FittingOutputDef, +) -> Dict[str, VariableDef]: def_derv_r = {} def_derv_c = {} for kk, vv in def_outp.get_data().items(): From 95cc1dac3d9f56b5a9a40d8970ae1398431249e6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 16 Jan 2024 20:48:27 +0800 Subject: [PATCH 2/7] allow the final dim of shape def to be -1. add UTs for check_var --- source/tests/test_output_def.py | 41 +++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/source/tests/test_output_def.py b/source/tests/test_output_def.py index e0c56784da..bb3aab18a1 100644 --- a/source/tests/test_output_def.py +++ b/source/tests/test_output_def.py @@ -11,6 +11,10 @@ fitting_check_output, model_check_output, ) +from deepmd_utils.model_format.output_def import ( + VariableDef, + check_var, +) class TestDef(unittest.TestCase): @@ -243,3 +247,40 @@ def call(self): ff = Foo(shape=[nf, nloc, 2]) ff() self.assertIn("not matching", context.exception) + + def test_check_var(self): + var_def = VariableDef("foo", [2, 3], atomic=True) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5, 6]), var_def) + self.assertIn("length not matching", context.exception) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5]), var_def) + self.assertIn("shape not matching", context.exception) + check_var(np.zeros([2, 3, 2, 3]), var_def) + + var_def = VariableDef("foo", [2, 3], atomic=False) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5]), var_def) + self.assertIn("length not matching", context.exception) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4]), var_def) + self.assertIn("shape not matching", context.exception) + check_var(np.zeros([2, 2, 3]), var_def) + + var_def = VariableDef("foo", [2, -1], atomic=True) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5, 6]), var_def) + self.assertIn("length not matching", context.exception) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5]), var_def) + self.assertIn("shape not matching", context.exception) + check_var(np.zeros([2, 3, 2, 8]), var_def) + + var_def = VariableDef("foo", [2, -1], atomic=False) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4, 5]), var_def) + self.assertIn("length not matching", context.exception) + with self.assertRaises(ValueError) as context: + check_var(np.zeros([2, 3, 4]), var_def) + self.assertIn("shape not matching", context.exception) + check_var(np.zeros([2, 2, 8]), var_def) From b8cb289a43afc36611b7184f940ab5f3115cebdf Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 16 Jan 2024 20:48:53 +0800 Subject: [PATCH 3/7] add missing files --- deepmd_utils/model_format/output_def.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/deepmd_utils/model_format/output_def.py b/deepmd_utils/model_format/output_def.py index 12bce53b82..5fb045f292 100644 --- a/deepmd_utils/model_format/output_def.py +++ b/deepmd_utils/model_format/output_def.py @@ -6,19 +6,31 @@ ) +def check_shape( + shape: List[int], + def_shape: List[int], +): + """Check if the shape satisfies the defined shape.""" + assert len(shape) == len(def_shape) + if def_shape[-1] == -1: + if list(shape[:-1]) != def_shape[:-1]: + raise ValueError(f"{shape[:-1]} shape not matching def {def_shape[:-1]}") + else: + if list(shape) != def_shape: + raise ValueError(f"{shape} shape not matching def {def_shape}") + + def check_var(var, var_def): if var_def.atomic: # var.shape == [nf, nloc, *var_def.shape] if len(var.shape) != len(var_def.shape) + 2: raise ValueError(f"{var.shape[2:]} length not matching def {var_def.shape}") - if list(var.shape[2:]) != var_def.shape: - raise ValueError(f"{var.shape[2:]} not matching def {var_def.shape}") + check_shape(list(var.shape[2:]), var_def.shape) else: # var.shape == [nf, *var_def.shape] if len(var.shape) != len(var_def.shape) + 1: raise ValueError(f"{var.shape[1:]} length not matching def {var_def.shape}") - if list(var.shape[1:]) != var_def.shape: - raise ValueError(f"{var.shape[1:]} not matching def {var_def.shape}") + check_shape(list(var.shape[1:]), var_def.shape) def model_check_output(cls): From 578e819055e93071f5f62153821bdeb7d6147721 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 18 Jan 2024 00:06:50 +0800 Subject: [PATCH 4/7] remove variable def, not help in jit --- deepmd_utils/model_format/__init__.py | 2 - deepmd_utils/model_format/output_def.py | 72 +++++++++---------------- source/tests/test_output_def.py | 21 ++++++-- 3 files changed, 42 insertions(+), 53 deletions(-) diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py index 961b483730..253bca3507 100644 --- a/deepmd_utils/model_format/__init__.py +++ b/deepmd_utils/model_format/__init__.py @@ -24,7 +24,6 @@ FittingOutputDef, ModelOutputDef, OutputVariableDef, - VariableDef, fitting_check_output, get_deriv_name, get_reduce_name, @@ -54,7 +53,6 @@ "ModelOutputDef", "FittingOutputDef", "OutputVariableDef", - "VariableDef", "model_check_output", "fitting_check_output", "get_reduce_name", diff --git a/deepmd_utils/model_format/output_def.py b/deepmd_utils/model_format/output_def.py index 5fb045f292..268dc21ea6 100644 --- a/deepmd_utils/model_format/output_def.py +++ b/deepmd_utils/model_format/output_def.py @@ -104,35 +104,7 @@ def __call__( return wrapper -class VariableDef: - """Defines the shape and other properties of a variable. - - Parameters - ---------- - name - Name of the output variable. Notice that the xxxx_redu, - xxxx_derv_c, xxxx_derv_r are reserved names that should - not be used to define variables. - shape - The shape of the variable. e.g. energy should be [1], - dipole should be [3], polarizabilty should be [3,3]. - atomic - If the variable is defined for each atom. - - """ - - def __init__( - self, - name: str, - shape: list[int], - atomic: bool = True, - ): - self.name = name - self.shape = list(shape) - self.atomic = atomic - - -class OutputVariableDef(VariableDef): +class OutputVariableDef: """Defines the shape and other properties of the one output variable. It is assume that the fitting network output variables for each @@ -163,15 +135,11 @@ def __init__( shape: List[int], reduciable: bool = False, differentiable: bool = False, + atomic: bool = True, ): - ## fitting output must be atomic - ## Here we cannot use super because it does not pass jit - # super().__init__(name, shape, atomic=True) - ## the work around is the following self.name = name self.shape = list(shape) - self.atomic = True - # + self.atomic = atomic self.reduciable = reduciable self.differentiable = differentiable if not self.reduciable and self.differentiable: @@ -232,7 +200,7 @@ def __init__( self.def_outp = fit_defs self.def_redu = do_reduce(self.def_outp) self.def_derv_r, self.def_derv_c = do_derivative(self.def_outp) - self.var_defs = {} + self.var_defs: Dict[str, OutputVariableDef] = {} for ii in [ self.def_outp.get_data(), self.def_redu, @@ -244,13 +212,13 @@ def __init__( def __getitem__( self, key: str, - ) -> VariableDef: + ) -> OutputVariableDef: return self.var_defs[key] def get_data( self, key: str, - ) -> Dict[str, VariableDef]: + ) -> Dict[str, OutputVariableDef]: return self.var_defs def keys(self): @@ -279,23 +247,35 @@ def get_deriv_name(name: str) -> Tuple[str, str]: def do_reduce( def_outp: FittingOutputDef, -) -> Dict[str, VariableDef]: - def_redu = {} +) -> Dict[str, OutputVariableDef]: + def_redu: Dict[str, OutputVariableDef] = {} for kk, vv in def_outp.get_data().items(): if vv.reduciable: rk = get_reduce_name(kk) - def_redu[rk] = VariableDef(rk, vv.shape, atomic=False) + def_redu[rk] = OutputVariableDef( + rk, vv.shape, reduciable=False, differentiable=False, atomic=False + ) return def_redu def do_derivative( def_outp: FittingOutputDef, -) -> Dict[str, VariableDef]: - def_derv_r = {} - def_derv_c = {} +) -> Tuple[Dict[str, OutputVariableDef], Dict[str, OutputVariableDef]]: + def_derv_r: Dict[str, OutputVariableDef] = {} + def_derv_c: Dict[str, OutputVariableDef] = {} for kk, vv in def_outp.get_data().items(): if vv.differentiable: rkr, rkc = get_deriv_name(kk) - def_derv_r[rkr] = VariableDef(rkr, [*vv.shape, 3], atomic=True) - def_derv_c[rkc] = VariableDef(rkc, [*vv.shape, 3, 3], atomic=False) + def_derv_r[rkr] = OutputVariableDef( + rkr, + vv.shape + [3], # noqa: RUF005 + reduciable=False, + differentiable=False, + ) + def_derv_c[rkc] = OutputVariableDef( + rkc, + vv.shape + [3, 3], # noqa: RUF005 + reduciable=True, + differentiable=False, + ) return def_derv_r, def_derv_c diff --git a/source/tests/test_output_def.py b/source/tests/test_output_def.py index bb3aab18a1..cfe1198185 100644 --- a/source/tests/test_output_def.py +++ b/source/tests/test_output_def.py @@ -12,11 +12,22 @@ model_check_output, ) from deepmd_utils.model_format.output_def import ( - VariableDef, check_var, ) +class VariableDef: + def __init__( + self, + name: str, + shape: list[int], + atomic: bool = True, + ): + self.name = name + self.shape = list(shape) + self.atomic = atomic + + class TestDef(unittest.TestCase): def test_model_output_def(self): defs = [ @@ -85,7 +96,7 @@ def test_model_output_def(self): self.assertEqual(md["foo"].atomic, True) self.assertEqual(md["energy_redu"].atomic, False) self.assertEqual(md["energy_derv_r"].atomic, True) - self.assertEqual(md["energy_derv_c"].atomic, False) + self.assertEqual(md["energy_derv_c"].atomic, True) def test_raise_no_redu_deriv(self): with self.assertRaises(ValueError) as context: @@ -108,7 +119,7 @@ def call(self): "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros([nf, 1]), "energy_derv_r": np.zeros([nf, nloc, 1, 3]), - "energy_derv_c": np.zeros([nf, 1, 3, 3]), + "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), } ff = Foo() @@ -133,7 +144,7 @@ def call(self): return { "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros([nf, 1]), - "energy_derv_c": np.zeros([nf, 1, 3, 3]), + "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), } ff = Foo() @@ -165,7 +176,7 @@ def call(self): "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros(self.shape_rd), "energy_derv_r": np.zeros(self.shape_dr), - "energy_derv_c": np.zeros([nf, 1, 3, 3]), + "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), } ff = Foo() From 81b00488b6a2dcfe6bb2589ef5122e83bdf9916a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 18 Jan 2024 00:31:43 +0800 Subject: [PATCH 5/7] change nloc to nall in dervs. they are defined for extended atoms --- source/tests/test_output_def.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/source/tests/test_output_def.py b/source/tests/test_output_def.py index cfe1198185..21d5c3543b 100644 --- a/source/tests/test_output_def.py +++ b/source/tests/test_output_def.py @@ -105,6 +105,7 @@ def test_raise_no_redu_deriv(self): def test_model_decorator(self): nf = 2 nloc = 3 + nall = 4 @model_check_output class Foo(NativeOP): @@ -118,8 +119,8 @@ def call(self): return { "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros([nf, 1]), - "energy_derv_r": np.zeros([nf, nloc, 1, 3]), - "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), + "energy_derv_r": np.zeros([nf, nall, 1, 3]), + "energy_derv_c": np.zeros([nf, nall, 1, 3, 3]), } ff = Foo() @@ -128,6 +129,7 @@ def call(self): def test_model_decorator_keyerror(self): nf = 2 nloc = 3 + nall = 4 @model_check_output class Foo(NativeOP): @@ -144,7 +146,7 @@ def call(self): return { "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros([nf, 1]), - "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), + "energy_derv_c": np.zeros([nf, nall, 1, 3, 3]), } ff = Foo() @@ -155,13 +157,14 @@ def call(self): def test_model_decorator_shapeerror(self): nf = 2 nloc = 3 + nall = 4 @model_check_output class Foo(NativeOP): def __init__( self, shape_rd=[nf, 1], - shape_dr=[nf, nloc, 1, 3], + shape_dr=[nf, nall, 1, 3], ): self.shape_rd, self.shape_dr = shape_rd, shape_dr @@ -176,7 +179,7 @@ def call(self): "energy": np.zeros([nf, nloc, 1]), "energy_redu": np.zeros(self.shape_rd), "energy_derv_r": np.zeros(self.shape_dr), - "energy_derv_c": np.zeros([nf, nloc, 1, 3, 3]), + "energy_derv_c": np.zeros([nf, nall, 1, 3, 3]), } ff = Foo() @@ -207,6 +210,7 @@ def call(self): def test_fitting_decorator(self): nf = 2 nloc = 3 + nall = 4 @fitting_check_output class Foo(NativeOP): From 20af845644166df97e63abd0df4cb96cabfcc03e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 17 Jan 2024 18:01:35 -0500 Subject: [PATCH 6/7] fix list[int] --- source/tests/test_output_def.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/tests/test_output_def.py b/source/tests/test_output_def.py index 21d5c3543b..f0c194e46d 100644 --- a/source/tests/test_output_def.py +++ b/source/tests/test_output_def.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from typing import List import numpy as np @@ -20,7 +21,7 @@ class VariableDef: def __init__( self, name: str, - shape: list[int], + shape: List[int], atomic: bool = True, ): self.name = name From 5b31c00f7dab76d81e3851f744d7c86af319fdf2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 17 Jan 2024 23:01:56 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/test_output_def.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/test_output_def.py b/source/tests/test_output_def.py index f0c194e46d..82d1b13a80 100644 --- a/source/tests/test_output_def.py +++ b/source/tests/test_output_def.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest -from typing import List +from typing import ( + List, +) import numpy as np