From 200edb60081c75a8f019ce37485cb542256a1554 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 03:10:47 -0500 Subject: [PATCH 01/19] use tf.variable to store bias_atom_polar --- deepmd/tf/fit/polar.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 3091fd56ca..326f0ffbd3 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -27,8 +27,14 @@ from deepmd.tf.loss.tensor import ( TensorLoss, ) + +from deepmd.tf.utils.errors import ( + GraphWithoutTensorError, +) + from deepmd.tf.utils.graph import ( get_fitting_net_variables_from_graph_def, + get_tensor_by_name_from_graph, ) from deepmd.tf.utils.network import ( one_layer, @@ -429,6 +435,17 @@ def build( atype = input_dict.get("atype", None) nframes = input_dict.get("nframes") start_index = 0 + + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): + # self.t_bias_atom_polar = tf.constant(self.constant_matrix, name="t_bias_atom_polar", dtype=GLOBAL_TF_FLOAT_PRECISION) + self.t_bias_atom_polar = tf.get_variable( + "t_bias_atom_polar", + self.constant_matrix.shape, + dtype=GLOBAL_TF_FLOAT_PRECISION, + trainable=False, + initializer=tf.constant_initializer(self.constant_matrix), + ) + inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) if nframes is None: @@ -504,6 +521,7 @@ def build( # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] + # final_layer = final_layer + tf.slice(self.t_bias_atom_polar, [sel_type_idx], [1]) * tf.eye( final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], @@ -551,6 +569,15 @@ def init_variables( self.fitting_net_variables = get_fitting_net_variables_from_graph_def( graph_def, suffix=suffix ) + try: + self.bias_atom_polar = get_tensor_by_name_from_graph( + graph, f"fitting_attr{suffix}/t_bias_atom_polar" + ) + except GraphWithoutTensorError: + raise RuntimeError("No bias_atom_polar in the graph.") + # print("No bias_atom_polar in the graph.") + # for compatibility, old models has no t_bias_atom_e + # pass def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """Receive the mixed precision setting. From ac4a4d5278c80a47a1be50a3a557adea0bfca438 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 03:12:00 -0500 Subject: [PATCH 02/19] read variable t_bias_atom_polar when shift_diag=True --- deepmd/tf/fit/polar.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 326f0ffbd3..abe3b8767f 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -521,8 +521,8 @@ def build( # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] - # final_layer = final_layer + tf.slice(self.t_bias_atom_polar, [sel_type_idx], [1]) * tf.eye( - final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye( + final_layer = final_layer + tf.slice(self.t_bias_atom_polar, [sel_type_idx], [1]) * tf.eye( + # final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], dtype=GLOBAL_TF_FLOAT_PRECISION, @@ -569,15 +569,13 @@ def init_variables( self.fitting_net_variables = get_fitting_net_variables_from_graph_def( graph_def, suffix=suffix ) - try: - self.bias_atom_polar = get_tensor_by_name_from_graph( - graph, f"fitting_attr{suffix}/t_bias_atom_polar" - ) - except GraphWithoutTensorError: - raise RuntimeError("No bias_atom_polar in the graph.") - # print("No bias_atom_polar in the graph.") - # for compatibility, old models has no t_bias_atom_e - # pass + if self.shift_diag + try: + self.bias_atom_polar = get_tensor_by_name_from_graph( + graph, f"fitting_attr{suffix}/t_bias_atom_polar" + ) + except GraphWithoutTensorError: + raise RuntimeError("No bias_atom_polar in the graph.") def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """Receive the mixed precision setting. From bb4edf07585d61934f25e5f783cc219a7652e695 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Feb 2025 08:17:51 +0000 Subject: [PATCH 03/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/polar.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index abe3b8767f..5e25a10c6e 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -27,11 +27,9 @@ from deepmd.tf.loss.tensor import ( TensorLoss, ) - from deepmd.tf.utils.errors import ( GraphWithoutTensorError, ) - from deepmd.tf.utils.graph import ( get_fitting_net_variables_from_graph_def, get_tensor_by_name_from_graph, From 489ac47a07aa835b3b887860088034b50de3e19c Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 03:19:05 -0500 Subject: [PATCH 04/19] clean commented lines --- deepmd/tf/fit/polar.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index abe3b8767f..ccce0d7548 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -437,7 +437,6 @@ def build( start_index = 0 with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): - # self.t_bias_atom_polar = tf.constant(self.constant_matrix, name="t_bias_atom_polar", dtype=GLOBAL_TF_FLOAT_PRECISION) self.t_bias_atom_polar = tf.get_variable( "t_bias_atom_polar", self.constant_matrix.shape, @@ -522,7 +521,6 @@ def build( sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] final_layer = final_layer + tf.slice(self.t_bias_atom_polar, [sel_type_idx], [1]) * tf.eye( - # final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], dtype=GLOBAL_TF_FLOAT_PRECISION, From 4296757aacda22c3b5a488b1ed28518b71b0edeb Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 03:34:17 -0500 Subject: [PATCH 05/19] use self.t_bias_atom_polar for type_embedding --- deepmd/tf/fit/polar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 48b3acd7e1..2fe85dbc1c 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -466,7 +466,7 @@ def build( # nframes x nloc_masked constant_matrix = tf.reshape( tf.reshape( - tf.tile(tf.repeat(self.constant_matrix, natoms[2:]), [nframes]), + tf.tile(tf.repeat(self.t_bias_atom_polar, natoms[2:]), [nframes]), [nframes, -1], )[nloc_mask], [nframes, -1], From fc0f88f45bf47f2c8916c359c09860f1f651b1a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yifan=20Li=E6=9D=8E=E4=B8=80=E5=B8=86?= Date: Thu, 6 Feb 2025 03:37:31 -0500 Subject: [PATCH 06/19] fix syntax error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Yifan Li李一帆 --- deepmd/tf/fit/polar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 2fe85dbc1c..ee33feb0af 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -565,7 +565,7 @@ def init_variables( self.fitting_net_variables = get_fitting_net_variables_from_graph_def( graph_def, suffix=suffix ) - if self.shift_diag + if self.shift_diag: try: self.bias_atom_polar = get_tensor_by_name_from_graph( graph, f"fitting_attr{suffix}/t_bias_atom_polar" From 4390144186ddc012a059fcbdcc2b68ce6b3f98c8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Feb 2025 08:40:23 +0000 Subject: [PATCH 07/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/polar.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index ee33feb0af..8b4eebc9da 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -466,7 +466,9 @@ def build( # nframes x nloc_masked constant_matrix = tf.reshape( tf.reshape( - tf.tile(tf.repeat(self.t_bias_atom_polar, natoms[2:]), [nframes]), + tf.tile( + tf.repeat(self.t_bias_atom_polar, natoms[2:]), [nframes] + ), [nframes, -1], )[nloc_mask], [nframes, -1], @@ -518,7 +520,9 @@ def build( # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] - final_layer = final_layer + tf.slice(self.t_bias_atom_polar, [sel_type_idx], [1]) * tf.eye( + final_layer = final_layer + tf.slice( + self.t_bias_atom_polar, [sel_type_idx], [1] + ) * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], dtype=GLOBAL_TF_FLOAT_PRECISION, From b794731109f644cbd3777bbdc134a0a6c3bde90a Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 14:32:46 -0500 Subject: [PATCH 08/19] use warning when t_bias_atom_polar is not found --- deepmd/tf/fit/polar.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 2fe85dbc1c..b7fad2f519 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -571,7 +571,8 @@ def init_variables( graph, f"fitting_attr{suffix}/t_bias_atom_polar" ) except GraphWithoutTensorError: - raise RuntimeError("No bias_atom_polar in the graph.") + warnings.warn("You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue.") + pass def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """Receive the mixed precision setting. From 22e1fd50bb6ff90adf9d533d0bfb8bda4b1414d9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Feb 2025 19:36:10 +0000 Subject: [PATCH 09/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/polar.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 020f52bf05..e7dcc57f71 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -575,7 +575,9 @@ def init_variables( graph, f"fitting_attr{suffix}/t_bias_atom_polar" ) except GraphWithoutTensorError: - warnings.warn("You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue.") + warnings.warn( + "You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue." + ) pass def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: From e54c56c7be97d3d108df9d0fc34a134fd9f07f0d Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 14:46:29 -0500 Subject: [PATCH 10/19] remove unnecessary pass --- deepmd/tf/fit/polar.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index e7dcc57f71..326e2c0eaa 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -578,7 +578,6 @@ def init_variables( warnings.warn( "You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue." ) - pass def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """Receive the mixed precision setting. From e8bdc392b6627533c9b7295edb69b084ff3c9d14 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 15:13:21 -0500 Subject: [PATCH 11/19] use stacklevel=2 for warning --- deepmd/tf/fit/polar.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 326e2c0eaa..37c7438f33 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -576,7 +576,8 @@ def init_variables( ) except GraphWithoutTensorError: warnings.warn( - "You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue." + "You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue.", + stacklevel=2 ) def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: From 36d46778f7d6c44c07fee5cd225e44c31c0f7887 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Feb 2025 20:15:11 +0000 Subject: [PATCH 12/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/polar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 37c7438f33..4a23a3f0e5 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -577,7 +577,7 @@ def init_variables( except GraphWithoutTensorError: warnings.warn( "You are trying to read a model trained with shift_diag=True, but the mean of the diagonal terms of the polarizability is not stored in the graph. This will lead to wrong inference results. You may train your model with the latest DeePMD-kit to avoid this issue.", - stacklevel=2 + stacklevel=2, ) def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: From 72c50f6297928460af8558cdabb0ed95f3d65f5c Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 17:10:33 -0500 Subject: [PATCH 13/19] add bias_atom_polar in serialize and deserialize --- deepmd/tf/fit/polar.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 37c7438f33..86c1ef8c2e 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -640,6 +640,9 @@ def serialize(self, suffix: str) -> dict: variables=self.fitting_net_variables, suffix=suffix, ), + "@variables": { + "bias_atom_polar": self.bias_atom_polar.reshape(-1), + }, "type_map": self.type_map, } return data @@ -667,6 +670,7 @@ def deserialize(cls, data: dict, suffix: str): data["nets"], suffix=suffix, ) + fitting.bias_atom_polar = data["@variables"]["bias_atom_polar"].ravel() return fitting From 4bcca9933755bac580e882e274a8a5ff8b9d26c9 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 18:29:54 -0500 Subject: [PATCH 14/19] rename constant_matrix to bias_atom_polar --- deepmd/tf/fit/polar.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index dd5ee6c127..194ee556e2 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -155,7 +155,7 @@ def __init__( if not isinstance(self.sel_type, list): self.sel_type = [self.sel_type] self.sel_type = sorted(self.sel_type) - self.constant_matrix = np.zeros( # pylint: disable=no-explicit-dtype + self.bias_atom_polar = np.zeros( # pylint: disable=no-explicit-dtype self.ntypes ) # self.ntypes x 1, store the average diagonal value # if type(self.diag_shift) is not list: @@ -280,7 +280,7 @@ def compute_output_stats(self, all_stat) -> None: ) atom_polar, _, _, _ = np.linalg.lstsq(matrix, bias, rcond=None) for itype in range(len(self.sel_type)): - self.constant_matrix[self.sel_type[itype]] = np.mean( + self.bias_atom_polar[self.sel_type[itype]] = np.mean( np.diagonal(atom_polar[itype].reshape((3, 3))) ) @@ -437,10 +437,10 @@ def build( with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): self.t_bias_atom_polar = tf.get_variable( "t_bias_atom_polar", - self.constant_matrix.shape, + self.bias_atom_polar.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, - initializer=tf.constant_initializer(self.constant_matrix), + initializer=tf.constant_initializer(self.bias_atom_polar), ) inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) @@ -464,7 +464,7 @@ def build( ) if self.shift_diag: # nframes x nloc_masked - constant_matrix = tf.reshape( + bias_atom_polar = tf.reshape( tf.reshape( tf.tile( tf.repeat(self.t_bias_atom_polar, natoms[2:]), [nframes] @@ -542,7 +542,7 @@ def build( final_layer *= tf.expand_dims(tf.expand_dims(scale, -1), -1) if self.shift_diag: final_layer += tf.expand_dims( - tf.expand_dims(constant_matrix, -1), -1 + tf.expand_dims(bias_atom_polar, -1), -1 ) * tf.eye(3, batch_shape=[1, 1], dtype=GLOBAL_TF_FLOAT_PRECISION) outs = final_layer From 5ca01b4825375d10e6ede195d6f81f516faa0af0 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 20:05:44 -0500 Subject: [PATCH 15/19] rename bias_atom_polar to constant_matrix to keep consistent with other backends --- deepmd/tf/fit/polar.py | 35 +++++++++++++++++++++-------------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 194ee556e2..289ee8dd15 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -155,7 +155,7 @@ def __init__( if not isinstance(self.sel_type, list): self.sel_type = [self.sel_type] self.sel_type = sorted(self.sel_type) - self.bias_atom_polar = np.zeros( # pylint: disable=no-explicit-dtype + self.constant_matrix = np.zeros( # pylint: disable=no-explicit-dtype self.ntypes ) # self.ntypes x 1, store the average diagonal value # if type(self.diag_shift) is not list: @@ -280,7 +280,7 @@ def compute_output_stats(self, all_stat) -> None: ) atom_polar, _, _, _ = np.linalg.lstsq(matrix, bias, rcond=None) for itype in range(len(self.sel_type)): - self.bias_atom_polar[self.sel_type[itype]] = np.mean( + self.constant_matrix[self.sel_type[itype]] = np.mean( np.diagonal(atom_polar[itype].reshape((3, 3))) ) @@ -435,12 +435,12 @@ def build( start_index = 0 with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): - self.t_bias_atom_polar = tf.get_variable( - "t_bias_atom_polar", - self.bias_atom_polar.shape, + self.t_constant_matrix = tf.get_variable( + "t_constant_matrix", + self.constant_matrix.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, - initializer=tf.constant_initializer(self.bias_atom_polar), + initializer=tf.constant_initializer(self.constant_matrix), ) inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) @@ -464,10 +464,10 @@ def build( ) if self.shift_diag: # nframes x nloc_masked - bias_atom_polar = tf.reshape( + constant_matrix = tf.reshape( tf.reshape( tf.tile( - tf.repeat(self.t_bias_atom_polar, natoms[2:]), [nframes] + tf.repeat(self.t_constant_matrix, natoms[2:]), [nframes] ), [nframes, -1], )[nloc_mask], @@ -521,7 +521,7 @@ def build( sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] final_layer = final_layer + tf.slice( - self.t_bias_atom_polar, [sel_type_idx], [1] + self.t_constant_matrix, [sel_type_idx], [1] ) * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], @@ -542,7 +542,7 @@ def build( final_layer *= tf.expand_dims(tf.expand_dims(scale, -1), -1) if self.shift_diag: final_layer += tf.expand_dims( - tf.expand_dims(bias_atom_polar, -1), -1 + tf.expand_dims(constant_matrix, -1), -1 ) * tf.eye(3, batch_shape=[1, 1], dtype=GLOBAL_TF_FLOAT_PRECISION) outs = final_layer @@ -571,8 +571,8 @@ def init_variables( ) if self.shift_diag: try: - self.bias_atom_polar = get_tensor_by_name_from_graph( - graph, f"fitting_attr{suffix}/t_bias_atom_polar" + self.constant_matrix = get_tensor_by_name_from_graph( + graph, f"fitting_attr{suffix}/t_constant_matrix" ) except GraphWithoutTensorError: warnings.warn( @@ -641,7 +641,13 @@ def serialize(self, suffix: str) -> dict: suffix=suffix, ), "@variables": { - "bias_atom_polar": self.bias_atom_polar.reshape(-1), + "fparam_avg": None, + "fparam_inv_std": None, + "aparam_avg": None, + "aparam_inv_std": None, + "case_embd": None, + "scale": self.scale.reshape(-1, 1), + "constant_matrix": self.constant_matrix.reshape(-1), }, "type_map": self.type_map, } @@ -661,6 +667,7 @@ def deserialize(cls, data: dict, suffix: str): Model The deserialized model """ + print("executing deserialize") data = data.copy() check_version_compatibility( data.pop("@version", 1), 4, 1 @@ -670,7 +677,7 @@ def deserialize(cls, data: dict, suffix: str): data["nets"], suffix=suffix, ) - fitting.bias_atom_polar = data["@variables"]["bias_atom_polar"].ravel() + fitting.constant_matrix = data["@variables"]["constant_matrix"].ravel() return fitting From 7312a43d2ba27d5c850252cd5fe0ca888fd614e0 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 20:07:24 -0500 Subject: [PATCH 16/19] pop bias_atom_e from PolarFitting of the dp backend --- deepmd/dpmodel/fitting/polarizability_fitting.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 8acb818a46..2734f60c73 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -195,6 +195,7 @@ def serialize(self) -> dict: data["shift_diag"] = self.shift_diag data["@variables"]["scale"] = to_numpy_array(self.scale) data["@variables"]["constant_matrix"] = to_numpy_array(self.constant_matrix) + data["@variables"].pop("bias_atom_e") return data @classmethod From 4dfb83acfaadaf3a2dabf0e0c2d1de1c08cd6217 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 20:10:22 -0500 Subject: [PATCH 17/19] remove remaining print --- deepmd/tf/fit/polar.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 289ee8dd15..c8fd4e86e8 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -667,7 +667,6 @@ def deserialize(cls, data: dict, suffix: str): Model The deserialized model """ - print("executing deserialize") data = data.copy() check_version_compatibility( data.pop("@version", 1), 4, 1 From 9b7ef69d3ca576c4bb783c8898bbf9d2661f5109 Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 20:51:55 -0500 Subject: [PATCH 18/19] pop bias_atom_e from variables for dpmodel in test_tf_consistent_with_ref --- source/tests/consistent/common.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 7ee8477a3b..35230e2140 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -354,6 +354,9 @@ def test_tf_consistent_with_ref(self) -> None: data1.pop("@version") data2.pop("@version") + if tf_obj.__class__.__name__.startswith("Polar"): + data1["@variables"].pop("bias_atom_e") + np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( From c9baf668dd603ede6b388694eed2ef4ef7e954ba Mon Sep 17 00:00:00 2001 From: Yifan Li Date: Thu, 6 Feb 2025 20:52:27 -0500 Subject: [PATCH 19/19] do not pop bias_atom_e from dpmodel --- deepmd/dpmodel/fitting/polarizability_fitting.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 2734f60c73..8acb818a46 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -195,7 +195,6 @@ def serialize(self) -> dict: data["shift_diag"] = self.shift_diag data["@variables"]["scale"] = to_numpy_array(self.scale) data["@variables"]["constant_matrix"] = to_numpy_array(self.constant_matrix) - data["@variables"].pop("bias_atom_e") return data @classmethod