From 0133c8672215c03945eaef5c92a265ff7ae61350 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 20 Jul 2023 17:40:29 -0400 Subject: [PATCH 1/4] support generalized force loss Signed-off-by: Jinzhe Zeng --- deepmd/loss/ener.py | 97 +++++++++++++++++++++++++++++++++++++++- deepmd/utils/argcheck.py | 50 ++++++++++++++++----- doc/data/system.md | 1 + pyproject.toml | 2 +- 4 files changed, 137 insertions(+), 13 deletions(-) diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index 9c97638931..e42886f29e 100644 --- a/deepmd/loss/ener.py +++ b/deepmd/loss/ener.py @@ -27,8 +27,42 @@ class EnerStdLoss(Loss): Parameters ---------- + starter_learning_rate : float + The learning rate at the start of the training. + start_pref_e : float + The prefactor of energy loss at the start of the training. + limit_pref_e : float + The prefactor of energy loss at the end of the training. + start_pref_f : float + The prefactor of force loss at the start of the training. + limit_pref_f : float + The prefactor of force loss at the end of the training. + start_pref_v : float + The prefactor of virial loss at the start of the training. + limit_pref_v : float + The prefactor of virial loss at the end of the training. + start_pref_ae : float + The prefactor of atomic energy loss at the start of the training. + limit_pref_ae : float + The prefactor of atomic energy loss at the end of the training. + start_pref_pf : float + The prefactor of atomic prefactor force loss at the start of the training. + limit_pref_pf : float + The prefactor of atomic prefactor force loss at the end of the training. + relative_f : float + If provided, relative force error will be used in the loss. The difference + of force will be normalized by the magnitude of the force in the label with + a shift given by relative_f enable_atom_ener_coeff : bool if true, the energy will be computed as \sum_i c_i E_i + start_pref_gf : float + The prefactor of generalized force loss at the start of the training. + limit_pref_gf : float + The prefactor of generalized force loss at the end of the training. + numb_generalized_coord : int + The dimension of generalized coordinates. + **kwargs + Other keyword arguments. """ def __init__( @@ -46,6 +80,9 @@ def __init__( limit_pref_pf: float = 0.0, relative_f: Optional[float] = None, enable_atom_ener_coeff: bool = False, + start_pref_gf: float = 0.0, + limit_pref_gf: float = 0.0, + numb_generalized_coord=0, **kwargs, ) -> None: self.starter_learning_rate = starter_learning_rate @@ -61,11 +98,19 @@ def __init__( self.limit_pref_pf = limit_pref_pf self.relative_f = relative_f self.enable_atom_ener_coeff = enable_atom_ener_coeff + self.start_pref_gf = start_pref_gf + self.limit_pref_gf = limit_pref_gf + self.numb_generalized_coord = numb_generalized_coord self.has_e = self.start_pref_e != 0.0 or self.limit_pref_e != 0.0 self.has_f = self.start_pref_f != 0.0 or self.limit_pref_f != 0.0 self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0 self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0 self.has_pf = self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0 + self.has_gf = self.start_pref_gf != 0.0 or self.limit_pref_gf != 0.0 + if self.has_gf and self.numb_generalized_coord < 1: + raise RuntimeError( + "When generalized force loss is used, the dimension of generalized coordinates should be larger than 0" + ) # data required add_data_requirement("energy", 1, atomic=False, must=False, high_prec=True) add_data_requirement("force", 3, atomic=True, must=False, high_prec=False) @@ -74,6 +119,15 @@ def __init__( add_data_requirement( "atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3 ) + # drdq: the partial derivative of atomic coordinates w.r.t. generalized coordinates + # TODO: could numb_generalized_coord decided from the training data? + add_data_requirement( + "drdq", + self.numb_generalized_coord * 3, + atomic=True, + must=False, + high_prec=False, + ) if self.enable_atom_ener_coeff: add_data_requirement( "atom_ener_coeff", @@ -94,11 +148,13 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): virial_hat = label_dict["virial"] atom_ener_hat = label_dict["atom_ener"] atom_pref = label_dict["atom_pref"] + drdq = label_dict["drdq"] find_energy = label_dict["find_energy"] find_force = label_dict["find_force"] find_virial = label_dict["find_virial"] find_atom_ener = label_dict["find_atom_ener"] find_atom_pref = label_dict["find_atom_pref"] + find_drdq = label_dict["find_drdq"] if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as @@ -117,7 +173,7 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): tf.square(energy - energy_hat), name="l2_" + suffix ) - if self.has_f or self.has_pf or self.relative_f: + if self.has_f or self.has_pf or self.relative_f or self.has_gf: force_reshape = tf.reshape(force, [-1]) force_hat_reshape = tf.reshape(force_hat, [-1]) diff_f = force_hat_reshape - force_reshape @@ -139,6 +195,16 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): name="l2_pref_force_" + suffix, ) + if self.has_gf: + drdq = label_dict["drdq"] + drdq_reshape = tf.reshape(drdq, [-1]) + gen_force_hat = tf.einsum("bij,bi->bj", drdq_reshape, force_hat_reshape) + gen_force = tf.einsum("bij,bi->bj", drdq_reshape, force_reshape) + diff_gen_force = gen_force_hat - gen_force + l2_gen_force_loss = tf.reduce_mean( + tf.square(diff_gen_force), name="l2_gen_force_" + suffix + ) + if self.has_v: virial_reshape = tf.reshape(virial, [-1]) virial_hat_reshape = tf.reshape(virial_hat, [-1]) @@ -202,6 +268,15 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): / self.starter_learning_rate ) ) + pref_gf = global_cvt_2_tf_float( + find_drdq + * ( + self.limit_pref_gf + + (self.start_pref_gf - self.limit_pref_gf) + * learning_rate + / self.starter_learning_rate + ) + ) l2_loss = 0 more_loss = {} @@ -220,6 +295,9 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): if self.has_pf: l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss) more_loss["l2_pref_force_loss"] = l2_pref_force_loss + if self.has_gf: + l2_loss += global_cvt_2_ener_float(pref_gf * l2_gen_force_loss) + more_loss["l2_gen_force_loss"] = l2_gen_force_loss # only used when tensorboard was set as true self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss)) @@ -238,6 +316,18 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): "l2_virial_loss_" + suffix, tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0]), ) + if self.has_ae: + self.l2_loss_atom_ener_summary = tf.summary.scalar( + "l2_atom_ener_loss_" + suffix, tf.sqrt(l2_atom_ener_loss) + ) + if self.has_pf: + self.l2_loss_pref_force_summary = tf.summary.scalar( + "l2_pref_force_loss_" + suffix, tf.sqrt(l2_pref_force_loss) + ) + if self.has_gf: + self.l2_loss_gf_summary = tf.summary.scalar( + "l2_gen_force_loss_" + suffix, tf.sqrt(l2_gen_force_loss) + ) self.l2_l = l2_loss self.l2_more = more_loss @@ -252,8 +342,9 @@ def eval(self, sess, feed_dict, natoms): self.l2_more["l2_virial_loss"] if self.has_v else placeholder, self.l2_more["l2_atom_ener_loss"] if self.has_ae else placeholder, self.l2_more["l2_pref_force_loss"] if self.has_pf else placeholder, + self.l2_more["l2_gf_loss"] if self.has_gf else placeholder, ] - error, error_e, error_f, error_v, error_ae, error_pf = run_sess( + error, error_e, error_f, error_v, error_ae, error_pf, error_gf = run_sess( sess, run_data, feed_dict=feed_dict ) results = {"natoms": natoms[0], "rmse": np.sqrt(error)} @@ -267,6 +358,8 @@ def eval(self, sess, feed_dict, natoms): results["rmse_v"] = np.sqrt(error_v) / natoms[0] if self.has_pf: results["rmse_pf"] = np.sqrt(error_pf) + if self.has_gf: + results["rmse_gf"] = np.sqrt(error_gf) return results diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 5a246ba5aa..017fcdf360 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -918,7 +918,7 @@ def pairwise_dprc() -> Argument: # --- Learning rate configurations: --- # def learning_rate_exp(): - doc_start_lr = "The learning rate the start of the training." + doc_start_lr = "The learning rate at the start of the training." doc_stop_lr = "The desired learning rate at the end of the training." doc_decay_steps = ( "The learning rate is decaying every this number of training steps." @@ -977,8 +977,12 @@ def learning_rate_dict_args(): # --- Loss configurations: --- # -def start_pref(item): - return f"The prefactor of {item} loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the {item} label should be provided by file {item}.npy in each data system. If both start_pref_{item} and limit_pref_{item} are set to 0, then the {item} will be ignored." +def start_pref(item, label=None, abbr=None): + if label is None: + label = item + if abbr is None: + abbr = item + return f"The prefactor of {item} loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the {label} label should be provided by file {label}.npy in each data system. If both start_pref_{abbr} and limit_pref_{abbr} are set to 0, then the {item} will be ignored." def limit_pref(item): @@ -986,16 +990,21 @@ def limit_pref(item): def loss_ener(): - doc_start_pref_e = start_pref("energy") + doc_start_pref_e = start_pref("energy", abbr="e") doc_limit_pref_e = limit_pref("energy") - doc_start_pref_f = start_pref("force") + doc_start_pref_f = start_pref("force", abbr="f") doc_limit_pref_f = limit_pref("force") - doc_start_pref_v = start_pref("virial") + doc_start_pref_v = start_pref("virial", abbr="v") doc_limit_pref_v = limit_pref("virial") - doc_start_pref_ae = start_pref("atom_ener") - doc_limit_pref_ae = limit_pref("atom_ener") - doc_start_pref_pf = start_pref("atom_pref") - doc_limit_pref_pf = limit_pref("atom_pref") + doc_start_pref_ae = start_pref("atomic energy", label="atom_ener", abbr="ae") + doc_limit_pref_ae = limit_pref("atomic energy") + doc_start_pref_pf = start_pref( + "atomic prefactor force", label="atom_pref", abbr="pf" + ) + doc_limit_pref_pf = limit_pref("atomic prefactor force") + doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf") + doc_limit_pref_gf = limit_pref("generalized force") + doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used." doc_relative_f = "If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label." doc_enable_atom_ener_coeff = "If true, the energy will be computed as \\sum_i c_i E_i. c_i should be provided by file atom_ener_coeff.npy in each data system, otherwise it's 1." return [ @@ -1077,6 +1086,27 @@ def loss_ener(): default=False, doc=doc_enable_atom_ener_coeff, ), + Argument( + "start_pref_gf", + float, + optional=True, + default=0.0, + doc=doc_start_pref_gf, + ), + Argument( + "limit_pref_gf", + float, + optional=True, + default=0.0, + doc=doc_start_pref_gf, + ), + Argument( + "numb_generalized_coord", + int, + optional=True, + default=0, + doc=doc_numb_generalized_coord, + ), ] diff --git a/doc/data/system.md b/doc/data/system.md index a81016f4cb..0ecd0e9119 100644 --- a/doc/data/system.md +++ b/doc/data/system.md @@ -33,6 +33,7 @@ dipole | Frame dipole | dipole.raw | A atomic_dipole | Atomic dipole | atomic_dipole.raw | Any | Nframes \* Natoms \* 3 | polarizability | Frame polarizability | polarizability.raw | Any | Nframes \* 9 | in the order `XX XY XZ YX YY YZ ZX ZY ZZ` atomic_polarizability | Atomic polarizability | atomic_polarizability.raw| Any | Nframes \* Natoms \* 9 | in the order `XX XY XZ YX YY YZ ZX ZY ZZ` +drdq | Partial derivative of atomic coordinates with respect to generalized coordinates | drdq.raw | 1 | Nframes \* Natoms \* 3 \* Ngen_coords | In general, we always use the following convention of units: diff --git a/pyproject.toml b/pyproject.toml index 36cc8e2f43..8bb8224eac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ dependencies = [ 'numpy', 'scipy', 'pyyaml', - 'dargs >= 0.2.9', + 'dargs >= 0.3.5', 'python-hostlist >= 1.21', 'typing_extensions; python_version < "3.7"', 'importlib_metadata>=1.4; python_version < "3.8"', From 3007652b057daf4eff576ba90a40f895c2c61f84 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 20 Jul 2023 18:06:14 -0400 Subject: [PATCH 2/4] fix errors Signed-off-by: Jinzhe Zeng --- deepmd/loss/ener.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index e42886f29e..ffe7b84091 100644 --- a/deepmd/loss/ener.py +++ b/deepmd/loss/ener.py @@ -148,13 +148,14 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): virial_hat = label_dict["virial"] atom_ener_hat = label_dict["atom_ener"] atom_pref = label_dict["atom_pref"] - drdq = label_dict["drdq"] find_energy = label_dict["find_energy"] find_force = label_dict["find_force"] find_virial = label_dict["find_virial"] find_atom_ener = label_dict["find_atom_ener"] find_atom_pref = label_dict["find_atom_pref"] - find_drdq = label_dict["find_drdq"] + if self.has_gf: + drdq = label_dict["drdq"] + find_drdq = label_dict["find_drdq"] if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as @@ -268,15 +269,16 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): / self.starter_learning_rate ) ) - pref_gf = global_cvt_2_tf_float( - find_drdq - * ( - self.limit_pref_gf - + (self.start_pref_gf - self.limit_pref_gf) - * learning_rate - / self.starter_learning_rate + if self.has_gf: + pref_gf = global_cvt_2_tf_float( + find_drdq + * ( + self.limit_pref_gf + + (self.start_pref_gf - self.limit_pref_gf) + * learning_rate + / self.starter_learning_rate + ) ) - ) l2_loss = 0 more_loss = {} From f6852319728fc2d1aa2d6bf9a9ce7f8b52c09693 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 21 Jul 2023 17:44:45 -0400 Subject: [PATCH 3/4] fix errors; add examples Signed-off-by: Jinzhe Zeng --- deepmd/loss/ener.py | 16 +- deepmd/utils/argcheck.py | 2 +- examples/dprc/data/set.000/drdq.npy | Bin 0 -> 6728 bytes examples/dprc/generalized_force/input.json | 222 +++++++++++++++++++++ source/tests/test_examples.py | 2 + 5 files changed, 236 insertions(+), 6 deletions(-) create mode 100644 examples/dprc/data/set.000/drdq.npy create mode 100644 examples/dprc/generalized_force/input.json diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index ffe7b84091..4cc7619d50 100644 --- a/deepmd/loss/ener.py +++ b/deepmd/loss/ener.py @@ -82,7 +82,7 @@ def __init__( enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, - numb_generalized_coord=0, + numb_generalized_coord: int = 0, **kwargs, ) -> None: self.starter_learning_rate = starter_learning_rate @@ -198,9 +198,15 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): if self.has_gf: drdq = label_dict["drdq"] - drdq_reshape = tf.reshape(drdq, [-1]) - gen_force_hat = tf.einsum("bij,bi->bj", drdq_reshape, force_hat_reshape) - gen_force = tf.einsum("bij,bi->bj", drdq_reshape, force_reshape) + force_reshape_nframes = tf.reshape(force, [-1, natoms[0] * 3]) + force_hat_reshape_nframes = tf.reshape(force_hat, [-1, natoms[0] * 3]) + drdq_reshape = tf.reshape( + drdq, [-1, natoms[0] * 3, self.numb_generalized_coord] + ) + gen_force_hat = tf.einsum( + "bij,bi->bj", drdq_reshape, force_hat_reshape_nframes + ) + gen_force = tf.einsum("bij,bi->bj", drdq_reshape, force_reshape_nframes) diff_gen_force = gen_force_hat - gen_force l2_gen_force_loss = tf.reduce_mean( tf.square(diff_gen_force), name="l2_gen_force_" + suffix @@ -344,7 +350,7 @@ def eval(self, sess, feed_dict, natoms): self.l2_more["l2_virial_loss"] if self.has_v else placeholder, self.l2_more["l2_atom_ener_loss"] if self.has_ae else placeholder, self.l2_more["l2_pref_force_loss"] if self.has_pf else placeholder, - self.l2_more["l2_gf_loss"] if self.has_gf else placeholder, + self.l2_more["l2_gen_force_loss"] if self.has_gf else placeholder, ] error, error_e, error_f, error_v, error_ae, error_pf, error_gf = run_sess( sess, run_data, feed_dict=feed_dict diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 017fcdf360..8398cb5420 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1098,7 +1098,7 @@ def loss_ener(): float, optional=True, default=0.0, - doc=doc_start_pref_gf, + doc=doc_limit_pref_gf, ), Argument( "numb_generalized_coord", diff --git a/examples/dprc/data/set.000/drdq.npy b/examples/dprc/data/set.000/drdq.npy new file mode 100644 index 0000000000000000000000000000000000000000..1569cc37e6f79959bfdc5ea854c55257e0bd6e96 GIT binary patch literal 6728 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$-Itms>rkXkmwF+bcE{0J64gqK7ZIy3aC+&|kXmPl8>bU)rUuVw#dO2Z# z^yZA|62=?s=lA5#%jsIW-@CYH- Date: Fri, 21 Jul 2023 18:28:30 -0400 Subject: [PATCH 4/4] add tests Signed-off-by: Jinzhe Zeng --- source/tests/test_loss_gf.py | 80 ++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 source/tests/test_loss_gf.py diff --git a/source/tests/test_loss_gf.py b/source/tests/test_loss_gf.py new file mode 100644 index 0000000000..04f40d943b --- /dev/null +++ b/source/tests/test_loss_gf.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import tensorflow as tf + +from deepmd.loss import ( + EnerStdLoss, +) + + +class TestLossGf(tf.test.TestCase): + def setUp(self): + self.loss = EnerStdLoss( + 1e-3, + start_pref_e=0, + limit_pref_e=0, + start_pref_f=0, + limit_pref_f=0, + start_pref_v=0, + limit_pref_v=0, + start_pref_ae=0, + limit_pref_ae=0, + start_pref_av=0, + limit_pref_av=0, + start_pref_gf=1, + limit_pref_gf=1, + numb_generalized_coord=2, + ) + + def test_build_loss(self): + natoms = tf.constant([6, 6]) + model_dict = { + "energy": tf.zeros((1, 1), dtype=tf.float64), + "force": tf.random.uniform((1, 6 * 3), dtype=tf.float64), + "virial": tf.zeros((1, 9), dtype=tf.float64), + "atom_ener": tf.zeros((1, 6), dtype=tf.float64), + "atom_virial": tf.zeros((1, 6 * 9), dtype=tf.float64), + "atom_pref": tf.zeros((1, 6), dtype=tf.float64), + } + label_dict = { + "energy": tf.zeros((1, 1), dtype=tf.float64), + "force": tf.random.uniform((1, 6 * 3), dtype=tf.float64), + "virial": tf.zeros((1, 9), dtype=tf.float64), + "atom_ener": tf.zeros((1, 6), dtype=tf.float64), + "atom_virial": tf.zeros((1, 6 * 9), dtype=tf.float64), + "atom_pref": tf.zeros((1, 6), dtype=tf.float64), + "find_energy": 0.0, + "find_force": 1.0, + "find_virial": 0.0, + "find_atom_ener": 0.0, + "find_atom_virial": 0.0, + "find_atom_pref": 0.0, + "drdq": tf.random.uniform((1, 6 * 3 * 2), dtype=tf.float64), + "find_drdq": 1.0, + } + t_total_loss, more_loss = self.loss.build( + tf.constant(1e-3), + natoms, + model_dict, + label_dict, + "_test_gf_loss", + ) + with self.cached_session() as sess: + total_loss, gen_force_loss, force, force_hat, drdq = sess.run( + [ + t_total_loss, + more_loss["l2_gen_force_loss"], + model_dict["force"], + label_dict["force"], + label_dict["drdq"], + ] + ) + + force = np.reshape(force, (1, 6 * 3)) + force_hat = np.reshape(force_hat, (1, 6 * 3)) + drdq = np.reshape(drdq, (1, 6 * 3, 2)) + gen_force = np.dot(force, drdq) + gen_force_hat = np.dot(force_hat, drdq) + np_loss = np.mean(np.square(gen_force - gen_force_hat)) + assert np.allclose(total_loss, np_loss, 6) + assert np.allclose(gen_force_loss, np_loss, 6)