From 134b36a07bec3268a60dd21191ec7f9ebf01c8db Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 27 Feb 2025 21:39:02 +0800 Subject: [PATCH 1/7] add plugin for data modifier --- deepmd/dpmodel/modifier/__init__.py | 8 + deepmd/dpmodel/modifier/base_modifier.py | 74 +++ deepmd/tf/entrypoints/train.py | 24 +- deepmd/tf/modifier/__init__.py | 12 + deepmd/tf/modifier/base_modifier.py | 13 + deepmd/tf/modifier/dipole_charge.py | 492 ++++++++++++++++++ source/tests/tf/test_data_modifier.py | 12 +- source/tests/tf/test_data_modifier_shuffle.py | 10 +- 8 files changed, 620 insertions(+), 25 deletions(-) create mode 100644 deepmd/dpmodel/modifier/__init__.py create mode 100644 deepmd/dpmodel/modifier/base_modifier.py create mode 100644 deepmd/tf/modifier/__init__.py create mode 100644 deepmd/tf/modifier/base_modifier.py create mode 100644 deepmd/tf/modifier/dipole_charge.py diff --git a/deepmd/dpmodel/modifier/__init__.py b/deepmd/dpmodel/modifier/__init__.py new file mode 100644 index 0000000000..d4e8ab56e3 --- /dev/null +++ b/deepmd/dpmodel/modifier/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_modifier import ( + make_base_modifier, +) + +__all__ = [ + "make_base_modifier", +] diff --git a/deepmd/dpmodel/modifier/base_modifier.py b/deepmd/dpmodel/modifier/base_modifier.py new file mode 100644 index 0000000000..9edc4722e1 --- /dev/null +++ b/deepmd/dpmodel/modifier/base_modifier.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import inspect +from abc import ( + ABC, + abstractmethod, +) + +from deepmd.utils.plugin import ( + PluginVariant, + make_plugin_registry, +) + + +def make_base_modifier() -> type[object]: + class BaseModifier(ABC, PluginVariant, make_plugin_registry("modifier")): + """Base class for data modifier.""" + + def __new__(cls, *args, **kwargs): + if cls is BaseModifier: + cls = cls.get_class_by_type(kwargs["type"]) + return super().__new__(cls) + + @abstractmethod + def serialize(self) -> dict: + """Serialize the modifier. + + Returns + ------- + dict + The serialized data + """ + pass + + @classmethod + def deserialize(cls, data: dict) -> "BaseModifier": + """Deserialize the modifier. + + Parameters + ---------- + data : dict + The serialized data + + Returns + ------- + BaseModel + The deserialized modifier + """ + if inspect.isabstract(cls): + return cls.get_class_by_type(data["type"]).deserialize(data) + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + + @classmethod + def get_modifier(cls, modifier_params: dict) -> "BaseModifier": + """Get the modifier by the parameters. + + By default, all the parameters are directly passed to the constructor. + If not, override this method. + + Parameters + ---------- + modifier_params : dict + The modifier parameters + + Returns + ------- + BaseModifier + The modifier + """ + modifier_params = modifier_params.copy() + modifier_params.pop("type", None) + modifier = cls(**modifier_params) + return modifier + + return BaseModifier diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 1762f1049a..b12e4fe1af 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -4,6 +4,7 @@ Can handle local or distributed training. """ +import copy import json import logging import time @@ -20,12 +21,12 @@ reset_default_tf_session_config, tf, ) -from deepmd.tf.infer.data_modifier import ( - DipoleChargeModifier, -) from deepmd.tf.model.model import ( Model, ) +from deepmd.tf.modifier import ( + BaseModifier, +) from deepmd.tf.train.run_options import ( RunOptions, ) @@ -275,18 +276,13 @@ def _do_work( def get_modifier(modi_data=None): - modifier: Optional[DipoleChargeModifier] + modifier: Optional[BaseModifier] if modi_data is not None: - if modi_data["type"] == "dipole_charge": - modifier = DipoleChargeModifier( - modi_data["model_name"], - modi_data["model_charge_map"], - modi_data["sys_charge_map"], - modi_data["ewald_h"], - modi_data["ewald_beta"], - ) - else: - raise RuntimeError("unknown modifier type " + str(modi_data["type"])) + modifier_params = copy.deepcopy(modi_data) + modifier_type = modifier_params.pop("type") + modifier = BaseModifier.get_class_by_type(modifier_type).get_modifier( + modifier_params + ) else: modifier = None return modifier diff --git a/deepmd/tf/modifier/__init__.py b/deepmd/tf/modifier/__init__.py new file mode 100644 index 0000000000..2441b693bc --- /dev/null +++ b/deepmd/tf/modifier/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_modifier import ( + BaseModifier, +) +from .dipole_charge import ( + DipoleChargeModifier, +) + +__all__ = [ + "BaseModifier", + "DipoleChargeModifier", +] diff --git a/deepmd/tf/modifier/base_modifier.py b/deepmd/tf/modifier/base_modifier.py new file mode 100644 index 0000000000..4e214e0835 --- /dev/null +++ b/deepmd/tf/modifier/base_modifier.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.modifier.base_modifier import ( + make_base_modifier, +) +from deepmd.tf.infer import ( + DeepPot, +) + + +class BaseModifier(DeepPot, make_base_modifier()): + def __init__(self, *args, **kwargs) -> None: + """Construct a basic model for different tasks.""" + DeepPot.__init__(self, *args, **kwargs) diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py new file mode 100644 index 0000000000..43ca4143eb --- /dev/null +++ b/deepmd/tf/modifier/dipole_charge.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os + +import numpy as np + +import deepmd.tf.op # noqa: F401 +from deepmd.tf.common import ( + make_default_mesh, + select_idx_map, +) +from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, + op_module, + tf, +) +from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole +from deepmd.tf.infer.ewald_recp import ( + EwaldRecp, +) +from deepmd.tf.modifier.base_modifier import ( + BaseModifier, +) +from deepmd.tf.utils.data import ( + DeepmdData, +) +from deepmd.tf.utils.sess import ( + run_sess, +) + + +@BaseModifier.register("dipole_charge") +class DipoleChargeModifier(DeepDipole, BaseModifier): + """Parameters + ---------- + model_name + The model file for the DeepDipole model + model_charge_map + Gives the amount of charge for the wfcc + sys_charge_map + Gives the amount of charge for the real atoms + ewald_h + Grid spacing of the reciprocal part of Ewald sum. Unit: A + ewald_beta + Splitting parameter of the Ewald sum. Unit: A^{-1} + """ + + def __new__(cls, *args, **kwargs): + model_file = kwargs.get("model_name", None) + if model_file is None: + raise TypeError("Missing required argument: 'model_name'") + return super().__new__(cls, model_file) + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1, + ewald_beta: float = 1, + ) -> None: + """Constructor.""" + # the dipole model is loaded with prefix 'dipole_charge' + self.modifier_prefix = "dipole_charge" + # init dipole model + DeepDipole.__init__( + self, model_name, load_prefix=self.modifier_prefix, default_tf_graph=True + ) + self.model_name = model_name + self.model_charge_map = model_charge_map + self.sys_charge_map = sys_charge_map + self.sel_type = list(self.get_sel_type()) + # init ewald recp + self.ewald_h = ewald_h + self.ewald_beta = ewald_beta + self.er = EwaldRecp(self.ewald_h, self.ewald_beta) + # dimension of dipole + self.ext_dim = 3 + self.t_ndesc = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "descrpt_attr/ndescrpt:0") + ) + self.t_sela = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "descrpt_attr/sel:0") + ) + [self.ndescrpt, self.sel_a] = run_sess(self.sess, [self.t_ndesc, self.t_sela]) + self.sel_r = [0 for ii in range(len(self.sel_a))] + self.nnei_a = np.cumsum(self.sel_a)[-1] + self.nnei_r = np.cumsum(self.sel_r)[-1] + self.nnei = self.nnei_a + self.nnei_r + self.ndescrpt_a = self.nnei_a * 4 + self.ndescrpt_r = self.nnei_r * 1 + assert self.ndescrpt == self.ndescrpt_a + self.ndescrpt_r + self.force = None + self.ntypes = len(self.sel_a) + + def serialize(self) -> dict: + """Serialize the modifier. + + Returns + ------- + dict + The serialized data + """ + data = { + "@class": "Modifier", + "type": self.modifier_prefix, + "@version": 3, + "model_name": self.model_name, + "model_charge_map": self.model_charge_map, + "sys_charge_map": self.sys_charge_map, + "ewald_h": self.ewald_h, + "ewald_beta": self.ewald_beta, + } + return data + + @classmethod + def deserialize(cls, data: dict) -> "BaseModifier": + """Deserialize the modifier. + + Parameters + ---------- + data : dict + The serialized data + + Returns + ------- + BaseModel + The deserialized modifier + """ + data = data.copy() + modifier = cls(**data) + return modifier + + def build_fv_graph(self) -> tf.Tensor: + """Build the computational graph for the force and virial inference.""" + with tf.variable_scope("modifier_attr"): + t_mdl_name = tf.constant(self.model_name, name="mdl_name", dtype=tf.string) + t_modi_type = tf.constant( + self.modifier_prefix, name="type", dtype=tf.string + ) + t_mdl_charge_map = tf.constant( + " ".join([str(ii) for ii in self.model_charge_map]), + name="mdl_charge_map", + dtype=tf.string, + ) + t_sys_charge_map = tf.constant( + " ".join([str(ii) for ii in self.sys_charge_map]), + name="sys_charge_map", + dtype=tf.string, + ) + t_ewald_h = tf.constant(self.ewald_h, name="ewald_h", dtype=tf.float64) + t_ewald_b = tf.constant( + self.ewald_beta, name="ewald_beta", dtype=tf.float64 + ) + with self.graph.as_default(): + return self._build_fv_graph_inner() + + def _build_fv_graph_inner(self): + self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_ef") + nf = 10 + nfxnas = 64 * nf + nfxna = 192 * nf + nf = -1 + nfxnas = -1 + nfxna = -1 + self.t_box_reshape = tf.reshape(self.t_box, [-1, 9]) + t_nframes = tf.shape(self.t_box_reshape)[0] + + # (nframes x natoms) x ndescrpt + self.descrpt = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "o_rmat:0") + ) + self.descrpt_deriv = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "o_rmat_deriv:0") + ) + self.nlist = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "o_nlist:0") + ) + self.rij = self.graph.get_tensor_by_name( + os.path.join(self.modifier_prefix, "o_rij:0") + ) + # self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt]) + # self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3]) + + # nframes x (natoms_sel x 3) + self.t_ef_reshape = tf.reshape(self.t_ef, [t_nframes, -1]) + # nframes x (natoms x 3) + self.t_ef_reshape = self._enrich(self.t_ef_reshape, dof=3) + # (nframes x natoms) x 3 + self.t_ef_reshape = tf.reshape(self.t_ef_reshape, [nfxna, 3]) + # nframes x (natoms_sel x 3) + self.t_tensor_reshape = tf.reshape(self.t_tensor, [t_nframes, -1]) + # nframes x (natoms x 3) + self.t_tensor_reshape = self._enrich(self.t_tensor_reshape, dof=3) + # (nframes x natoms) x 3 + self.t_tensor_reshape = tf.reshape(self.t_tensor_reshape, [nfxna, 3]) + # (nframes x natoms) x ndescrpt + [self.t_ef_d] = tf.gradients( + self.t_tensor_reshape, self.descrpt, self.t_ef_reshape + ) + # nframes x (natoms x ndescrpt) + self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt]) + # t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite + self.t_ef_d_oppo = -self.t_ef_d + + force = op_module.prod_force_se_a( + self.t_ef_d_oppo, + self.descrpt_deriv, + self.nlist, + self.t_natoms, + n_a_sel=self.nnei_a, + n_r_sel=self.nnei_r, + ) + virial, atom_virial = op_module.prod_virial_se_a( + self.t_ef_d_oppo, + self.descrpt_deriv, + self.rij, + self.nlist, + self.t_natoms, + n_a_sel=self.nnei_a, + n_r_sel=self.nnei_r, + ) + force = tf.identity(force, name="o_dm_force") + virial = tf.identity(virial, name="o_dm_virial") + atom_virial = tf.identity(atom_virial, name="o_dm_av") + return force, virial, atom_virial + + def _enrich(self, dipole, dof=3): + coll = [] + sel_start_idx = 0 + for type_i in range(self.ntypes): + if type_i in self.sel_type: + di = tf.slice( + dipole, + [0, sel_start_idx * dof], + [-1, self.t_natoms[2 + type_i] * dof], + ) + sel_start_idx += self.t_natoms[2 + type_i] + else: + di = tf.zeros( + [tf.shape(dipole)[0], self.t_natoms[2 + type_i] * dof], + dtype=GLOBAL_TF_FLOAT_PRECISION, + ) + coll.append(di) + return tf.concat(coll, axis=1) + + def _slice_descrpt_deriv(self, deriv): + coll = [] + start_idx = 0 + for type_i in range(self.ntypes): + if type_i in self.sel_type: + di = tf.slice( + deriv, + [0, start_idx * self.ndescrpt], + [-1, self.t_natoms[2 + type_i] * self.ndescrpt], + ) + coll.append(di) + start_idx += self.t_natoms[2 + type_i] + return tf.concat(coll, axis=1) + + def eval( + self, + coord: np.ndarray, + box: np.ndarray, + atype: np.ndarray, + eval_fv: bool = True, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Evaluate the modification. + + Parameters + ---------- + coord + The coordinates of atoms + box + The simulation region. PBC is assumed + atype + The atom types + eval_fv + Evaluate force and virial + + Returns + ------- + tot_e + The energy modification + tot_f + The force modification + tot_v + The virial modification + """ + atype = np.array(atype, dtype=int) + coord, atype, imap = self.sort_input(coord, atype) + # natoms = coord.shape[1] // 3 + natoms = atype.size + nframes = coord.shape[0] + box = np.reshape(box, [nframes, 9]) + atype = np.reshape(atype, [natoms]) + sel_idx_map = select_idx_map(atype, self.sel_type) + nsel = len(sel_idx_map) + # setup charge + charge = np.zeros([natoms]) # pylint: disable=no-explicit-dtype + for ii in range(natoms): + charge[ii] = self.sys_charge_map[atype[ii]] + charge = np.tile(charge, [nframes, 1]) + + # add wfcc + all_coord, all_charge, dipole = self._extend_system(coord, box, atype, charge) + + # print('compute er') + batch_size = 5 + tot_e = [] + all_f = [] + all_v = [] + for ii in range(0, nframes, batch_size): + e, f, v = self.er.eval( + all_coord[ii : ii + batch_size], + all_charge[ii : ii + batch_size], + box[ii : ii + batch_size], + ) + tot_e.append(e) + all_f.append(f) + all_v.append(v) + tot_e = np.concatenate(tot_e, axis=0) + all_f = np.concatenate(all_f, axis=0) + all_v = np.concatenate(all_v, axis=0) + # print('finish er') + # reshape + tot_e.reshape([nframes, 1]) + + tot_f = None + tot_v = None + if self.force is None: + self.force, self.virial, self.av = self.build_fv_graph() + if eval_fv: + # compute f + ext_f = all_f[:, natoms * 3 :] + corr_f = [] + corr_v = [] + corr_av = [] + for ii in range(0, nframes, batch_size): + f, v, av = self._eval_fv( + coord[ii : ii + batch_size], + box[ii : ii + batch_size], + atype, + ext_f[ii : ii + batch_size], + ) + corr_f.append(f) + corr_v.append(v) + corr_av.append(av) + corr_f = np.concatenate(corr_f, axis=0) + corr_v = np.concatenate(corr_v, axis=0) + corr_av = np.concatenate(corr_av, axis=0) + tot_f = all_f[:, : natoms * 3] + corr_f + for ii in range(nsel): + orig_idx = sel_idx_map[ii] + tot_f[:, orig_idx * 3 : orig_idx * 3 + 3] += ext_f[ + :, ii * 3 : ii * 3 + 3 + ] + tot_f = self.reverse_map(np.reshape(tot_f, [nframes, -1, 3]), imap) + # reshape + tot_f = tot_f.reshape([nframes, natoms, 3]) + # compute v + dipole3 = np.reshape(dipole, [nframes, nsel, 3]) + ext_f3 = np.reshape(ext_f, [nframes, nsel, 3]) + ext_f3 = np.transpose(ext_f3, [0, 2, 1]) + # fd_corr_v = -np.matmul(ext_f3, dipole3).T.reshape([nframes, 9]) + # fd_corr_v = -np.matmul(ext_f3, dipole3) + # fd_corr_v = np.transpose(fd_corr_v, [0, 2, 1]).reshape([nframes, 9]) + fd_corr_v = -np.matmul(ext_f3, dipole3).reshape([nframes, 9]) + # print(all_v, '\n', corr_v, '\n', fd_corr_v) + tot_v = all_v + corr_v + fd_corr_v + # reshape + tot_v = tot_v.reshape([nframes, 9]) + + return tot_e, tot_f, tot_v + + def _eval_fv(self, coords, cells, atom_types, ext_f): + # reshape the inputs + cells = np.reshape(cells, [-1, 9]) + nframes = cells.shape[0] + coords = np.reshape(coords, [nframes, -1]) + natoms = coords.shape[1] // 3 + + # sort inputs + coords, atom_types, imap, sel_at, sel_imap = self.sort_input( + coords, atom_types, sel_atoms=self.get_sel_type() + ) + + # make natoms_vec and default_mesh + natoms_vec = self.make_natoms_vec(atom_types) + assert natoms_vec[0] == natoms + default_mesh = make_default_mesh(True, False) + + # evaluate + tensor = [] + feed_dict_test = {} + feed_dict_test[self.t_natoms] = natoms_vec + feed_dict_test[self.t_type] = np.tile(atom_types, [nframes, 1]).reshape([-1]) + feed_dict_test[self.t_coord] = coords.reshape([-1]) + feed_dict_test[self.t_box] = cells.reshape([-1]) + feed_dict_test[self.t_mesh] = default_mesh.reshape([-1]) + feed_dict_test[self.t_ef] = ext_f.reshape([-1]) + # print(run_sess(self.sess, tf.shape(self.t_tensor), feed_dict = feed_dict_test)) + fout, vout, avout = run_sess( + self.sess, [self.force, self.virial, self.av], feed_dict=feed_dict_test + ) + # print('fout: ', fout.shape, fout) + fout = self.reverse_map(np.reshape(fout, [nframes, -1, 3]), imap) + fout = np.reshape(fout, [nframes, -1]) + return fout, vout, avout + + def _extend_system(self, coord, box, atype, charge): + natoms = coord.shape[1] // 3 + nframes = coord.shape[0] + # sel atoms and setup ref coord + sel_idx_map = select_idx_map(atype, self.sel_type) + nsel = len(sel_idx_map) + coord3 = coord.reshape([nframes, natoms, 3]) + ref_coord = coord3[:, sel_idx_map, :] + ref_coord = np.reshape(ref_coord, [nframes, nsel * 3]) + + batch_size = 8 + all_dipole = [] + for ii in range(0, nframes, batch_size): + dipole = DeepDipole.eval( + self, coord[ii : ii + batch_size], box[ii : ii + batch_size], atype + ) + all_dipole.append(dipole) + dipole = np.concatenate(all_dipole, axis=0) + assert dipole.shape[0] == nframes + dipole = np.reshape(dipole, [nframes, nsel * 3]) + + wfcc_coord = ref_coord + dipole + # wfcc_coord = dipole + wfcc_charge = np.zeros([nsel]) # pylint: disable=no-explicit-dtype + for ii in range(nsel): + orig_idx = self.sel_type.index(atype[sel_idx_map[ii]]) + wfcc_charge[ii] = self.model_charge_map[orig_idx] + wfcc_charge = np.tile(wfcc_charge, [nframes, 1]) + + wfcc_coord = np.reshape(wfcc_coord, [nframes, nsel * 3]) + wfcc_charge = np.reshape(wfcc_charge, [nframes, nsel]) + + all_coord = np.concatenate((coord, wfcc_coord), axis=1) + all_charge = np.concatenate((charge, wfcc_charge), axis=1) + + return all_coord, all_charge, dipole + + def modify_data(self, data: dict, data_sys: DeepmdData) -> None: + """Modify data. + + Parameters + ---------- + data + Internal data of DeepmdData. + Be a dict, has the following keys + - coord coordinates + - box simulation box + - type atom types + - find_energy tells if data has energy + - find_force tells if data has force + - find_virial tells if data has virial + - energy energy + - force force + - virial virial + data_sys : DeepmdData + The data system. + """ + if ( + "find_energy" not in data + and "find_force" not in data + and "find_virial" not in data + ): + return + + get_nframes = None + coord = data["coord"][:get_nframes, :] + if not data_sys.pbc: + raise RuntimeError("Open systems (nopbc) are not supported") + box = data["box"][:get_nframes, :] + atype = data["type"][:get_nframes, :] + atype = atype[0] + nframes = coord.shape[0] + + tot_e, tot_f, tot_v = self.eval(coord, box, atype) + + # print(tot_f[:,0]) + + if "find_energy" in data and data["find_energy"] == 1.0: + data["energy"] -= tot_e.reshape(data["energy"].shape) + if "find_force" in data and data["find_force"] == 1.0: + data["force"] -= tot_f.reshape(data["force"].shape) + if "find_virial" in data and data["find_virial"] == 1.0: + data["virial"] -= tot_v.reshape(data["virial"].shape) diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index db11fa5c2d..7a7558793e 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -7,7 +7,7 @@ GLOBAL_NP_FLOAT_PRECISION, tf, ) -from deepmd.tf.infer.data_modifier import ( +from deepmd.tf.modifier import ( DipoleChargeModifier, ) from deepmd.tf.train.run_options import ( @@ -97,11 +97,11 @@ def test_fv(self) -> None: def _test_fv(self) -> None: dcm = DipoleChargeModifier( - str(tests_path / os.path.join(modifier_datapath, "dipole.pb")), - [-8], - [6, 1], - 1, - 0.25, + model_name=str(tests_path / os.path.join(modifier_datapath, "dipole.pb")), + model_charge_map=[-8], + sys_charge_map=[6, 1], + ewald_h=1, + ewald_beta=0.25, ) data = Data() coord, box, atype = data.get_data() diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index 002b4f5746..d0d4972919 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -197,11 +197,11 @@ def test_z_dipole(self) -> None: def test_modify(self) -> None: dcm = DipoleChargeModifier( - os.path.join(modifier_datapath, "dipole.pb"), - [-1, -3], - [1, 1, 1, 1, 1], - 1, - 0.25, + model_name=os.path.join(modifier_datapath, "dipole.pb"), + model_charge_map=[-1, -3], + sys_charge_map=[1, 1, 1, 1, 1], + ewald_h=1, + ewald_beta=0.25, ) ve0, vf0, vv0 = dcm.eval(self.coords0, self.box0, self.atom_types0) ve1, vf1, vv1 = dcm.eval(self.coords1, self.box1, self.atom_types1) From b60f3a86cce1703683be4b6e79f4a211d927c029 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Mon, 3 Mar 2025 08:33:31 +0800 Subject: [PATCH 2/7] fix bug in argcheck --- deepmd/utils/argcheck.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a00cfb047a..899d4d4fda 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -55,6 +55,8 @@ doc_dos = "Fit a density of states model. The total density of states / site-projected density of states labels should be provided by `dos.npy` or `atom_dos.npy` in each data system. The file has number of frames lines and number of energy grid columns (times number of atoms in `atom_dos.npy`). See `loss` parameter." doc_dipole = "Fit an atomic dipole model. Global dipole labels or atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file either has number of frames lines and 3 times of number of selected atoms columns, or has number of frames lines and 3 columns. See `loss` parameter." doc_polar = "Fit an atomic polarizability model. Global polarizazbility labels or atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file with has number of frames lines and 9 times of number of selected atoms columns, or has number of frames lines and 9 columns. See `loss` parameter." +# modifier +doc_dipole_charge = "Use WFCC to model the electronic structure of the system. Correct the long-range interaction." def list_to_doc(xx): @@ -1782,6 +1784,10 @@ def fitting_variant_type_args(): # --- Modifier configurations: --- # +modifier_args_plugin = ArgsPlugin() + + +@modifier_args_plugin.register("dipole_charge", doc=doc_dipole_charge) def modifier_dipole_charge(): doc_model_name = "The name of the frozen dipole model file." doc_model_charge_map = f"The charge of the WFCC. The list length should be the same as the {make_link('sel_type', 'model[standard]/fitting_net[dipole]/sel_type')}. " @@ -1802,14 +1808,9 @@ def modifier_dipole_charge(): def modifier_variant_type_args(): doc_modifier_type = "The type of modifier." - doc_dipole_charge = "Use WFCC to model the electronic structure of the system. Correct the long-range interaction." return Variant( "type", - [ - Argument( - "dipole_charge", dict, modifier_dipole_charge(), doc=doc_dipole_charge - ), - ], + modifier_args_plugin.get_all_argument(), optional=False, doc=doc_modifier_type, ) From 41606aea96aed0ce9b726c0d23aa970f4b9c1d37 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Tue, 4 Mar 2025 08:31:43 +0800 Subject: [PATCH 3/7] remove old data modifier module --- deepmd/tf/__init__.py | 2 +- deepmd/tf/infer/__init__.py | 4 - deepmd/tf/infer/data_modifier.py | 444 ------------------------------- deepmd/tf/infer/deep_eval.py | 2 +- 4 files changed, 2 insertions(+), 450 deletions(-) delete mode 100644 deepmd/tf/infer/data_modifier.py diff --git a/deepmd/tf/__init__.py b/deepmd/tf/__init__.py index 933729fde2..cca5e54e7a 100644 --- a/deepmd/tf/__init__.py +++ b/deepmd/tf/__init__.py @@ -21,7 +21,7 @@ DeepEval, DeepPotential, ) -from .infer.data_modifier import ( +from .modifier import ( DipoleChargeModifier, ) diff --git a/deepmd/tf/infer/__init__.py b/deepmd/tf/infer/__init__.py index de8a77976e..ca9464ec43 100644 --- a/deepmd/tf/infer/__init__.py +++ b/deepmd/tf/infer/__init__.py @@ -8,9 +8,6 @@ DeepEval, ) -from .data_modifier import ( - DipoleChargeModifier, -) from .deep_dipole import ( DeepDipole, ) @@ -43,7 +40,6 @@ "DeepPot", "DeepPotential", "DeepWFC", - "DipoleChargeModifier", "EwaldRecp", "calc_model_devi", ] diff --git a/deepmd/tf/infer/data_modifier.py b/deepmd/tf/infer/data_modifier.py deleted file mode 100644 index ddb1af68d7..0000000000 --- a/deepmd/tf/infer/data_modifier.py +++ /dev/null @@ -1,444 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import os - -import numpy as np - -import deepmd.tf.op # noqa: F401 -from deepmd.tf.common import ( - make_default_mesh, - select_idx_map, -) -from deepmd.tf.env import ( - GLOBAL_TF_FLOAT_PRECISION, - op_module, - tf, -) -from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole -from deepmd.tf.infer.ewald_recp import ( - EwaldRecp, -) -from deepmd.tf.utils.data import ( - DeepmdData, -) -from deepmd.tf.utils.sess import ( - run_sess, -) - - -class DipoleChargeModifier(DeepDipole): - """Parameters - ---------- - model_name - The model file for the DeepDipole model - model_charge_map - Gives the amount of charge for the wfcc - sys_charge_map - Gives the amount of charge for the real atoms - ewald_h - Grid spacing of the reciprocal part of Ewald sum. Unit: A - ewald_beta - Splitting parameter of the Ewald sum. Unit: A^{-1} - """ - - def __init__( - self, - model_name: str, - model_charge_map: list[float], - sys_charge_map: list[float], - ewald_h: float = 1, - ewald_beta: float = 1, - ) -> None: - """Constructor.""" - # the dipole model is loaded with prefix 'dipole_charge' - self.modifier_prefix = "dipole_charge" - # init dipole model - DeepDipole.__init__( - self, model_name, load_prefix=self.modifier_prefix, default_tf_graph=True - ) - self.model_name = model_name - self.model_charge_map = model_charge_map - self.sys_charge_map = sys_charge_map - self.sel_type = list(self.get_sel_type()) - # init ewald recp - self.ewald_h = ewald_h - self.ewald_beta = ewald_beta - self.er = EwaldRecp(self.ewald_h, self.ewald_beta) - # dimension of dipole - self.ext_dim = 3 - self.t_ndesc = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "descrpt_attr/ndescrpt:0") - ) - self.t_sela = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "descrpt_attr/sel:0") - ) - [self.ndescrpt, self.sel_a] = run_sess(self.sess, [self.t_ndesc, self.t_sela]) - self.sel_r = [0 for ii in range(len(self.sel_a))] - self.nnei_a = np.cumsum(self.sel_a)[-1] - self.nnei_r = np.cumsum(self.sel_r)[-1] - self.nnei = self.nnei_a + self.nnei_r - self.ndescrpt_a = self.nnei_a * 4 - self.ndescrpt_r = self.nnei_r * 1 - assert self.ndescrpt == self.ndescrpt_a + self.ndescrpt_r - self.force = None - self.ntypes = len(self.sel_a) - - def build_fv_graph(self) -> tf.Tensor: - """Build the computational graph for the force and virial inference.""" - with tf.variable_scope("modifier_attr"): - t_mdl_name = tf.constant(self.model_name, name="mdl_name", dtype=tf.string) - t_modi_type = tf.constant( - self.modifier_prefix, name="type", dtype=tf.string - ) - t_mdl_charge_map = tf.constant( - " ".join([str(ii) for ii in self.model_charge_map]), - name="mdl_charge_map", - dtype=tf.string, - ) - t_sys_charge_map = tf.constant( - " ".join([str(ii) for ii in self.sys_charge_map]), - name="sys_charge_map", - dtype=tf.string, - ) - t_ewald_h = tf.constant(self.ewald_h, name="ewald_h", dtype=tf.float64) - t_ewald_b = tf.constant( - self.ewald_beta, name="ewald_beta", dtype=tf.float64 - ) - with self.graph.as_default(): - return self._build_fv_graph_inner() - - def _build_fv_graph_inner(self): - self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_ef") - nf = 10 - nfxnas = 64 * nf - nfxna = 192 * nf - nf = -1 - nfxnas = -1 - nfxna = -1 - self.t_box_reshape = tf.reshape(self.t_box, [-1, 9]) - t_nframes = tf.shape(self.t_box_reshape)[0] - - # (nframes x natoms) x ndescrpt - self.descrpt = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "o_rmat:0") - ) - self.descrpt_deriv = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "o_rmat_deriv:0") - ) - self.nlist = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "o_nlist:0") - ) - self.rij = self.graph.get_tensor_by_name( - os.path.join(self.modifier_prefix, "o_rij:0") - ) - # self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt]) - # self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3]) - - # nframes x (natoms_sel x 3) - self.t_ef_reshape = tf.reshape(self.t_ef, [t_nframes, -1]) - # nframes x (natoms x 3) - self.t_ef_reshape = self._enrich(self.t_ef_reshape, dof=3) - # (nframes x natoms) x 3 - self.t_ef_reshape = tf.reshape(self.t_ef_reshape, [nfxna, 3]) - # nframes x (natoms_sel x 3) - self.t_tensor_reshape = tf.reshape(self.t_tensor, [t_nframes, -1]) - # nframes x (natoms x 3) - self.t_tensor_reshape = self._enrich(self.t_tensor_reshape, dof=3) - # (nframes x natoms) x 3 - self.t_tensor_reshape = tf.reshape(self.t_tensor_reshape, [nfxna, 3]) - # (nframes x natoms) x ndescrpt - [self.t_ef_d] = tf.gradients( - self.t_tensor_reshape, self.descrpt, self.t_ef_reshape - ) - # nframes x (natoms x ndescrpt) - self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt]) - # t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite - self.t_ef_d_oppo = -self.t_ef_d - - force = op_module.prod_force_se_a( - self.t_ef_d_oppo, - self.descrpt_deriv, - self.nlist, - self.t_natoms, - n_a_sel=self.nnei_a, - n_r_sel=self.nnei_r, - ) - virial, atom_virial = op_module.prod_virial_se_a( - self.t_ef_d_oppo, - self.descrpt_deriv, - self.rij, - self.nlist, - self.t_natoms, - n_a_sel=self.nnei_a, - n_r_sel=self.nnei_r, - ) - force = tf.identity(force, name="o_dm_force") - virial = tf.identity(virial, name="o_dm_virial") - atom_virial = tf.identity(atom_virial, name="o_dm_av") - return force, virial, atom_virial - - def _enrich(self, dipole, dof=3): - coll = [] - sel_start_idx = 0 - for type_i in range(self.ntypes): - if type_i in self.sel_type: - di = tf.slice( - dipole, - [0, sel_start_idx * dof], - [-1, self.t_natoms[2 + type_i] * dof], - ) - sel_start_idx += self.t_natoms[2 + type_i] - else: - di = tf.zeros( - [tf.shape(dipole)[0], self.t_natoms[2 + type_i] * dof], - dtype=GLOBAL_TF_FLOAT_PRECISION, - ) - coll.append(di) - return tf.concat(coll, axis=1) - - def _slice_descrpt_deriv(self, deriv): - coll = [] - start_idx = 0 - for type_i in range(self.ntypes): - if type_i in self.sel_type: - di = tf.slice( - deriv, - [0, start_idx * self.ndescrpt], - [-1, self.t_natoms[2 + type_i] * self.ndescrpt], - ) - coll.append(di) - start_idx += self.t_natoms[2 + type_i] - return tf.concat(coll, axis=1) - - def eval( - self, - coord: np.ndarray, - box: np.ndarray, - atype: np.ndarray, - eval_fv: bool = True, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Evaluate the modification. - - Parameters - ---------- - coord - The coordinates of atoms - box - The simulation region. PBC is assumed - atype - The atom types - eval_fv - Evaluate force and virial - - Returns - ------- - tot_e - The energy modification - tot_f - The force modification - tot_v - The virial modification - """ - atype = np.array(atype, dtype=int) - coord, atype, imap = self.sort_input(coord, atype) - # natoms = coord.shape[1] // 3 - natoms = atype.size - nframes = coord.shape[0] - box = np.reshape(box, [nframes, 9]) - atype = np.reshape(atype, [natoms]) - sel_idx_map = select_idx_map(atype, self.sel_type) - nsel = len(sel_idx_map) - # setup charge - charge = np.zeros([natoms]) # pylint: disable=no-explicit-dtype - for ii in range(natoms): - charge[ii] = self.sys_charge_map[atype[ii]] - charge = np.tile(charge, [nframes, 1]) - - # add wfcc - all_coord, all_charge, dipole = self._extend_system(coord, box, atype, charge) - - # print('compute er') - batch_size = 5 - tot_e = [] - all_f = [] - all_v = [] - for ii in range(0, nframes, batch_size): - e, f, v = self.er.eval( - all_coord[ii : ii + batch_size], - all_charge[ii : ii + batch_size], - box[ii : ii + batch_size], - ) - tot_e.append(e) - all_f.append(f) - all_v.append(v) - tot_e = np.concatenate(tot_e, axis=0) - all_f = np.concatenate(all_f, axis=0) - all_v = np.concatenate(all_v, axis=0) - # print('finish er') - # reshape - tot_e.reshape([nframes, 1]) - - tot_f = None - tot_v = None - if self.force is None: - self.force, self.virial, self.av = self.build_fv_graph() - if eval_fv: - # compute f - ext_f = all_f[:, natoms * 3 :] - corr_f = [] - corr_v = [] - corr_av = [] - for ii in range(0, nframes, batch_size): - f, v, av = self._eval_fv( - coord[ii : ii + batch_size], - box[ii : ii + batch_size], - atype, - ext_f[ii : ii + batch_size], - ) - corr_f.append(f) - corr_v.append(v) - corr_av.append(av) - corr_f = np.concatenate(corr_f, axis=0) - corr_v = np.concatenate(corr_v, axis=0) - corr_av = np.concatenate(corr_av, axis=0) - tot_f = all_f[:, : natoms * 3] + corr_f - for ii in range(nsel): - orig_idx = sel_idx_map[ii] - tot_f[:, orig_idx * 3 : orig_idx * 3 + 3] += ext_f[ - :, ii * 3 : ii * 3 + 3 - ] - tot_f = self.reverse_map(np.reshape(tot_f, [nframes, -1, 3]), imap) - # reshape - tot_f = tot_f.reshape([nframes, natoms, 3]) - # compute v - dipole3 = np.reshape(dipole, [nframes, nsel, 3]) - ext_f3 = np.reshape(ext_f, [nframes, nsel, 3]) - ext_f3 = np.transpose(ext_f3, [0, 2, 1]) - # fd_corr_v = -np.matmul(ext_f3, dipole3).T.reshape([nframes, 9]) - # fd_corr_v = -np.matmul(ext_f3, dipole3) - # fd_corr_v = np.transpose(fd_corr_v, [0, 2, 1]).reshape([nframes, 9]) - fd_corr_v = -np.matmul(ext_f3, dipole3).reshape([nframes, 9]) - # print(all_v, '\n', corr_v, '\n', fd_corr_v) - tot_v = all_v + corr_v + fd_corr_v - # reshape - tot_v = tot_v.reshape([nframes, 9]) - - return tot_e, tot_f, tot_v - - def _eval_fv(self, coords, cells, atom_types, ext_f): - # reshape the inputs - cells = np.reshape(cells, [-1, 9]) - nframes = cells.shape[0] - coords = np.reshape(coords, [nframes, -1]) - natoms = coords.shape[1] // 3 - - # sort inputs - coords, atom_types, imap, sel_at, sel_imap = self.sort_input( - coords, atom_types, sel_atoms=self.get_sel_type() - ) - - # make natoms_vec and default_mesh - natoms_vec = self.make_natoms_vec(atom_types) - assert natoms_vec[0] == natoms - default_mesh = make_default_mesh(True, False) - - # evaluate - tensor = [] - feed_dict_test = {} - feed_dict_test[self.t_natoms] = natoms_vec - feed_dict_test[self.t_type] = np.tile(atom_types, [nframes, 1]).reshape([-1]) - feed_dict_test[self.t_coord] = coords.reshape([-1]) - feed_dict_test[self.t_box] = cells.reshape([-1]) - feed_dict_test[self.t_mesh] = default_mesh.reshape([-1]) - feed_dict_test[self.t_ef] = ext_f.reshape([-1]) - # print(run_sess(self.sess, tf.shape(self.t_tensor), feed_dict = feed_dict_test)) - fout, vout, avout = run_sess( - self.sess, [self.force, self.virial, self.av], feed_dict=feed_dict_test - ) - # print('fout: ', fout.shape, fout) - fout = self.reverse_map(np.reshape(fout, [nframes, -1, 3]), imap) - fout = np.reshape(fout, [nframes, -1]) - return fout, vout, avout - - def _extend_system(self, coord, box, atype, charge): - natoms = coord.shape[1] // 3 - nframes = coord.shape[0] - # sel atoms and setup ref coord - sel_idx_map = select_idx_map(atype, self.sel_type) - nsel = len(sel_idx_map) - coord3 = coord.reshape([nframes, natoms, 3]) - ref_coord = coord3[:, sel_idx_map, :] - ref_coord = np.reshape(ref_coord, [nframes, nsel * 3]) - - batch_size = 8 - all_dipole = [] - for ii in range(0, nframes, batch_size): - dipole = DeepDipole.eval( - self, coord[ii : ii + batch_size], box[ii : ii + batch_size], atype - ) - all_dipole.append(dipole) - dipole = np.concatenate(all_dipole, axis=0) - assert dipole.shape[0] == nframes - dipole = np.reshape(dipole, [nframes, nsel * 3]) - - wfcc_coord = ref_coord + dipole - # wfcc_coord = dipole - wfcc_charge = np.zeros([nsel]) # pylint: disable=no-explicit-dtype - for ii in range(nsel): - orig_idx = self.sel_type.index(atype[sel_idx_map[ii]]) - wfcc_charge[ii] = self.model_charge_map[orig_idx] - wfcc_charge = np.tile(wfcc_charge, [nframes, 1]) - - wfcc_coord = np.reshape(wfcc_coord, [nframes, nsel * 3]) - wfcc_charge = np.reshape(wfcc_charge, [nframes, nsel]) - - all_coord = np.concatenate((coord, wfcc_coord), axis=1) - all_charge = np.concatenate((charge, wfcc_charge), axis=1) - - return all_coord, all_charge, dipole - - def modify_data(self, data: dict, data_sys: DeepmdData) -> None: - """Modify data. - - Parameters - ---------- - data - Internal data of DeepmdData. - Be a dict, has the following keys - - coord coordinates - - box simulation box - - type atom types - - find_energy tells if data has energy - - find_force tells if data has force - - find_virial tells if data has virial - - energy energy - - force force - - virial virial - data_sys : DeepmdData - The data system. - """ - if ( - "find_energy" not in data - and "find_force" not in data - and "find_virial" not in data - ): - return - - get_nframes = None - coord = data["coord"][:get_nframes, :] - if not data_sys.pbc: - raise RuntimeError("Open systems (nopbc) are not supported") - box = data["box"][:get_nframes, :] - atype = data["type"][:get_nframes, :] - atype = atype[0] - nframes = coord.shape[0] - - tot_e, tot_f, tot_v = self.eval(coord, box, atype) - - # print(tot_f[:,0]) - - if "find_energy" in data and data["find_energy"] == 1.0: - data["energy"] -= tot_e.reshape(data["energy"].shape) - if "find_force" in data and data["find_force"] == 1.0: - data["force"] -= tot_f.reshape(data["force"].shape) - if "find_virial" in data and data["find_virial"] == 1.0: - data["virial"] -= tot_v.reshape(data["virial"].shape) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 8e5e3deea5..d594a47115 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -139,7 +139,7 @@ def __init__( # looks ugly... if self.modifier_type == "dipole_charge": - from deepmd.tf.infer.data_modifier import ( + from deepmd.tf.modifier import ( DipoleChargeModifier, ) From 8c8629019ad42880c07680283d69166d8a68630d Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Tue, 4 Mar 2025 09:33:02 +0800 Subject: [PATCH 4/7] update data modifier in unit tests --- source/tests/tf/test_data_modifier_shuffle.py | 6 +++--- source/tests/tf/test_dipolecharge.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index d0d4972919..49b46ead6a 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -8,12 +8,12 @@ GLOBAL_NP_FLOAT_PRECISION, tf, ) -from deepmd.tf.infer.data_modifier import ( - DipoleChargeModifier, -) from deepmd.tf.infer.deep_dipole import ( DeepDipole, ) +from deepmd.tf.modifier import ( + DipoleChargeModifier, +) from deepmd.tf.train.run_options import ( RunOptions, ) diff --git a/source/tests/tf/test_dipolecharge.py b/source/tests/tf/test_dipolecharge.py index d4ad3254bc..1253ff631f 100644 --- a/source/tests/tf/test_dipolecharge.py +++ b/source/tests/tf/test_dipolecharge.py @@ -7,7 +7,7 @@ from deepmd.tf.env import ( GLOBAL_NP_FLOAT_PRECISION, ) -from deepmd.tf.infer import ( +from deepmd.tf.modifier import ( DipoleChargeModifier, ) from deepmd.tf.utils.convert import ( From f9c5d53039e6e87b44097edea24af0f3ac37175a Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Tue, 4 Mar 2025 11:17:21 +0800 Subject: [PATCH 5/7] update data modifier in unit tests --- source/tests/tf/test_data_modifier_shuffle.py | 6 +++--- source/tests/tf/test_dipolecharge.py | 8 ++++++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index d0d4972919..49b46ead6a 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -8,12 +8,12 @@ GLOBAL_NP_FLOAT_PRECISION, tf, ) -from deepmd.tf.infer.data_modifier import ( - DipoleChargeModifier, -) from deepmd.tf.infer.deep_dipole import ( DeepDipole, ) +from deepmd.tf.modifier import ( + DipoleChargeModifier, +) from deepmd.tf.train.run_options import ( RunOptions, ) diff --git a/source/tests/tf/test_dipolecharge.py b/source/tests/tf/test_dipolecharge.py index d4ad3254bc..71c46446f6 100644 --- a/source/tests/tf/test_dipolecharge.py +++ b/source/tests/tf/test_dipolecharge.py @@ -7,7 +7,7 @@ from deepmd.tf.env import ( GLOBAL_NP_FLOAT_PRECISION, ) -from deepmd.tf.infer import ( +from deepmd.tf.modifier import ( DipoleChargeModifier, ) from deepmd.tf.utils.convert import ( @@ -32,7 +32,11 @@ def setUpClass(cls) -> None: "dipolecharge_d.pb", ) cls.dp = DipoleChargeModifier( - "dipolecharge_d.pb", [-1.0, -3.0], [1.0, 1.0, 1.0, 1.0, 1.0], 4.0, 0.2 + model_name="dipolecharge_d.pb", + model_charge_map=[-1.0, -3.0], + sys_charge_map=[1.0, 1.0, 1.0, 1.0, 1.0], + ewald_h=4.0, + ewald_beta=0.2, ) def setUp(self) -> None: From 502b2046bd33efdeb25f40da344504c5ef1ca382 Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Tue, 11 Mar 2025 21:40:18 +0800 Subject: [PATCH 6/7] (tf) update DipoleChargeModifier.__new__ --- deepmd/tf/modifier/dipole_charge.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/deepmd/tf/modifier/dipole_charge.py b/deepmd/tf/modifier/dipole_charge.py index 43ca4143eb..d40c9ccd2f 100644 --- a/deepmd/tf/modifier/dipole_charge.py +++ b/deepmd/tf/modifier/dipole_charge.py @@ -44,11 +44,8 @@ class DipoleChargeModifier(DeepDipole, BaseModifier): Splitting parameter of the Ewald sum. Unit: A^{-1} """ - def __new__(cls, *args, **kwargs): - model_file = kwargs.get("model_name", None) - if model_file is None: - raise TypeError("Missing required argument: 'model_name'") - return super().__new__(cls, model_file) + def __new__(cls, *args, model_name=None, **kwargs): + return super().__new__(cls, model_name) def __init__( self, From e6bca315b48310614268d190844d3eb0311e98ce Mon Sep 17 00:00:00 2001 From: Jia-Xin Zhu Date: Thu, 20 Mar 2025 14:04:58 +0800 Subject: [PATCH 7/7] feat(pt): add plugin for data modifier --- deepmd/pt/modifier/__init__.py | 8 +++++ deepmd/pt/modifier/base_modifier.py | 56 +++++++++++++++++++++++++++++ deepmd/pt/train/training.py | 53 ++++++++++++++++++++++++--- deepmd/pt/utils/stat.py | 8 +++-- 4 files changed, 119 insertions(+), 6 deletions(-) create mode 100644 deepmd/pt/modifier/__init__.py create mode 100644 deepmd/pt/modifier/base_modifier.py diff --git a/deepmd/pt/modifier/__init__.py b/deepmd/pt/modifier/__init__.py new file mode 100644 index 0000000000..bfa1540ce9 --- /dev/null +++ b/deepmd/pt/modifier/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_modifier import ( + BaseModifier, +) + +__all__ = [ + "BaseModifier", +] diff --git a/deepmd/pt/modifier/base_modifier.py b/deepmd/pt/modifier/base_modifier.py new file mode 100644 index 0000000000..be9246a2ad --- /dev/null +++ b/deepmd/pt/modifier/base_modifier.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.dpmodel.modifier.base_modifier import ( + make_base_modifier, +) + + +class BaseModifier(torch.nn.Module, make_base_modifier()): + def __init__(self, *args, **kwargs) -> None: + """Construct a basic model for different tasks.""" + torch.nn.Module.__init__(self) + + def modify_data(self, data: dict) -> None: + """Modify data. + + Parameters + ---------- + data + Internal data of DeepmdData. + Be a dict, has the following keys + - coord coordinates + - box simulation box + - atype atom types + - find_energy tells if data has energy + - find_force tells if data has force + - find_virial tells if data has virial + - energy energy + - force force + - virial virial + """ + if ( + "find_energy" not in data + and "find_force" not in data + and "find_virial" not in data + ): + return + + get_nframes = None + coord = data["coord"][:get_nframes, :] + if data["box"] is None: + box = None + else: + box = data["box"][:get_nframes, :] + atype = data["atype"][:get_nframes, :] + atype = atype[0] + nframes = coord.shape[0] + + tot_e, tot_f, tot_v = self.forward(coord, atype, box, False, None, None) + + if "find_energy" in data and data["find_energy"] == 1.0: + data["energy"] -= tot_e.reshape(data["energy"].shape) + if "find_force" in data and data["find_force"] == 1.0: + data["force"] -= tot_f.reshape(data["force"].shape) + if "find_virial" in data and data["find_virial"] == 1.0: + data["virial"] -= tot_v.reshape(data["virial"].shape) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 3fd0f9a41f..66903b08c9 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -36,6 +36,9 @@ get_model, get_zbl_model, ) +from deepmd.pt.modifier import ( + BaseModifier, +) from deepmd.pt.optimizer import ( KFOptimizerWrapper, LKFOptimizer, @@ -134,6 +137,16 @@ def __init__( ) self.num_model = len(self.model_keys) + # modifier for the training data + modifier_params = model_params.get("modifier", None) + if modifier_params is not None: + assert self.multi_task is False, ( + "Modifier is not supported for multi-task training!" + ) + self.modifier = get_data_modifier(modifier_params) + else: + self.modifier = None + # Iteration config self.num_steps = training_params["numb_steps"] self.disp_file = training_params.get("disp_file", "lcurve.out") @@ -215,12 +228,26 @@ def single_model_stat( _stat_file_path, _data_requirement, finetune_has_new_type=False, + modifier=None, ): _data_requirement += get_additional_data_requirement(_model) _training_data.add_data_requirement(_data_requirement) if _validation_data is not None: _validation_data.add_data_requirement(_data_requirement) + # modify data + if modifier is not None: + log.info(f"Using {modifier.modifier_type} as data modifier") + for _data in [_training_data, _validation_data]: + if _data is not None: + all_sampled = make_stat_input( + _data.systems, + _data.dataloaders, + -1, + ) + for sampled in all_sampled: + modifier.modify_data(sampled) + @functools.lru_cache def get_sample(): sampled = make_stat_input( @@ -315,6 +342,7 @@ def get_lr(lr_params): finetune_has_new_type=self.finetune_links["Default"].get_has_new_type() if self.finetune_links is not None else False, + modifier=self.modifier, ) ( self.training_dataloader, @@ -353,6 +381,7 @@ def get_lr(lr_params): ].get_has_new_type() if self.finetune_links is not None else False, + modifier=self.modifier, ) ( self.training_dataloader[model_key], @@ -1043,10 +1072,13 @@ def save_model(self, save_path, lr=0.0, step=0) -> None: optim_state_dict = deepcopy(self.optimizer.state_dict()) for item in optim_state_dict["param_groups"]: item["lr"] = float(item["lr"]) - torch.save( - {"model": module.state_dict(), "optimizer": optim_state_dict}, - save_path, - ) + save_dict = { + "model": module.state_dict(), + "optimizer": optim_state_dict, + } + if self.modifier is not None: + save_dict["data_modifier"] = self.modifier.state_dict() + torch.save(save_dict, save_path) checkpoint_dir = save_path.parent checkpoint_files = [ f @@ -1352,3 +1384,16 @@ def model_change_out_bias( f"to {to_numpy_array(new_bias).reshape(-1)!s}." ) return _model + + +def get_data_modifier(_modifier_params: dict[str, Any]): + modifier_params = deepcopy(_modifier_params) + try: + modifier_type = modifier_params.pop("type") + except KeyError: + raise ValueError("Data modifier type not specified!") from None + return ( + BaseModifier.get_class_by_type(modifier_type) + .get_modifier(modifier_params) + .to(DEVICE) + ) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index cf6892b49d..bb265a6b0d 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -47,12 +47,16 @@ def make_stat_input(datasets, dataloaders, nbatches): - a list of dicts, each of which contains data from a system """ lst = [] - log.info(f"Packing data for statistics from {len(datasets)} systems") + if nbatches > 0: + log.info(f"Packing data for statistics from {len(datasets)} systems") for i in range(len(datasets)): sys_stat = {} with torch.device("cpu"): iterator = iter(dataloaders[i]) - numb_batches = min(nbatches, len(dataloaders[i])) + if nbatches == -1: + numb_batches = len(dataloaders[i]) + else: + numb_batches = min(nbatches, len(dataloaders[i])) for _ in range(numb_batches): try: stat_data = next(iterator)