diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index 9c97638931..4cc7619d50 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: int = 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", @@ -99,6 +153,9 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): find_virial = label_dict["find_virial"] find_atom_ener = label_dict["find_atom_ener"] find_atom_pref = label_dict["find_atom_pref"] + 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 @@ -117,7 +174,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 +196,22 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): name="l2_pref_force_" + suffix, ) + if self.has_gf: + drdq = label_dict["drdq"] + 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 + ) + if self.has_v: virial_reshape = tf.reshape(virial, [-1]) virial_hat_reshape = tf.reshape(virial_hat, [-1]) @@ -202,6 +275,16 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): / 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 = {} @@ -220,6 +303,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 +324,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 +350,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_gen_force_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 +366,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..8398cb5420 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_limit_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/examples/dprc/data/set.000/drdq.npy b/examples/dprc/data/set.000/drdq.npy new file mode 100644 index 0000000000..1569cc37e6 Binary files /dev/null and b/examples/dprc/data/set.000/drdq.npy differ diff --git a/examples/dprc/generalized_force/input.json b/examples/dprc/generalized_force/input.json new file mode 100644 index 0000000000..01f2a52a6d --- /dev/null +++ b/examples/dprc/generalized_force/input.json @@ -0,0 +1,222 @@ +{ + "_comment": " model parameters", + "model": { + "type_map": [ + "C", + "P", + "O", + "H", + "OW", + "HW" + ], + "type_embedding": { + "neuron": [ + 8 + ], + "precision": "float32" + }, + "descriptor": { + "type": "hybrid", + "list": [ + { + "type": "se_e2_a", + "sel": [ + 6, + 1, + 6, + 11, + 0, + 0 + ], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 16, + "precision": "float32", + "exclude_types": [ + [ + 0, + 4 + ], + [ + 0, + 5 + ], + [ + 1, + 4 + ], + [ + 1, + 5 + ], + [ + 2, + 4 + ], + [ + 2, + 5 + ], + [ + 3, + 4 + ], + [ + 3, + 5 + ], + [ + 4, + 4 + ], + [ + 4, + 5 + ], + [ + 5, + 5 + ] + ], + "seed": 1 + }, + { + "type": "se_atten", + "stripped_type_embedding": true, + "sel": 150, + "rcut_smth": 5.80, + "rcut": 6.00, + "attn_layer": 0, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": false, + "axis_neuron": 12, + "set_davg_zero": true, + "exclude_types": [ + [ + 0, + 0 + ], + [ + 0, + 1 + ], + [ + 0, + 2 + ], + [ + 0, + 3 + ], + [ + 1, + 1 + ], + [ + 1, + 2 + ], + [ + 1, + 3 + ], + [ + 2, + 2 + ], + [ + 2, + 3 + ], + [ + 3, + 3 + ], + [ + 4, + 4 + ], + [ + 4, + 5 + ], + [ + 5, + 5 + ] + ], + "precision": "float32", + "seed": 1 + } + ] + }, + "fitting_net": { + "type": "ener", + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float32", + "atom_ener": [ + null, + null, + null, + null, + 0.0, + 0.0 + ], + "seed": 1 + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment": "that's all" + }, + + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "start_pref_gf": 10, + "limit_pref_gf": 100, + "numb_generalized_coord": 1, + "_comment": " that's all" + }, + + "training": { + "training_data": { + "systems": [ + "../data" + ], + "batch_size": "auto", + "_comment": "that's all" + }, + "numb_steps": 10000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment": "that's all" + }, + + "_comment": "that's all" +} 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"', diff --git a/source/tests/test_examples.py b/source/tests/test_examples.py index 0b8082e71a..c9f71e4a81 100644 --- a/source/tests/test_examples.py +++ b/source/tests/test_examples.py @@ -38,7 +38,9 @@ p_examples / "zinc_protein" / "zinc_se_a_mask.json", p_examples / "dos" / "train" / "input.json", p_examples / "spin" / "se_e2_a" / "input.json", + p_examples / "dprc" / "normal" / "input.json", p_examples / "dprc" / "pairwise" / "input.json", + p_examples / "dprc" / "generalized_force" / "input.json", ) 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)