From 464326822337688ebd01540a5a42abb594766fa3 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Wed, 6 May 2026 00:06:39 +0800 Subject: [PATCH 01/24] feat(dpa3): decouple charge_spin from fparam --- deepmd/calculator.py | 8 +- .../dpmodel/atomic_model/base_atomic_model.py | 20 + .../dpmodel/atomic_model/dp_atomic_model.py | 66 ++-- .../atomic_model/linear_atomic_model.py | 2 + .../atomic_model/make_base_atomic_model.py | 1 + .../atomic_model/pairtab_atomic_model.py | 1 + deepmd/dpmodel/descriptor/dpa1.py | 1 + deepmd/dpmodel/descriptor/dpa2.py | 1 + deepmd/dpmodel/descriptor/dpa3.py | 26 +- deepmd/dpmodel/descriptor/hybrid.py | 5 +- .../descriptor/make_base_descriptor.py | 1 + deepmd/dpmodel/descriptor/se_e2_a.py | 1 + deepmd/dpmodel/descriptor/se_r.py | 1 + deepmd/dpmodel/descriptor/se_t.py | 1 + deepmd/dpmodel/descriptor/se_t_tebd.py | 1 + deepmd/dpmodel/model/ener_model.py | 4 + deepmd/dpmodel/model/make_model.py | 20 + deepmd/dpmodel/utils/batch.py | 2 +- deepmd/dpmodel/utils/lmdb_data.py | 4 +- deepmd/dpmodel/utils/stat.py | 5 +- deepmd/entrypoints/test.py | 13 + deepmd/infer/deep_eval.py | 16 + deepmd/pt/infer/deep_eval.py | 30 +- .../model/atomic_model/base_atomic_model.py | 22 ++ .../pt/model/atomic_model/dp_atomic_model.py | 48 ++- .../model/atomic_model/linear_atomic_model.py | 2 + .../atomic_model/pairtab_atomic_model.py | 1 + deepmd/pt/model/descriptor/dpa1.py | 1 + deepmd/pt/model/descriptor/dpa2.py | 1 + deepmd/pt/model/descriptor/dpa3.py | 35 +- deepmd/pt/model/descriptor/hybrid.py | 5 +- deepmd/pt/model/descriptor/se_a.py | 1 + deepmd/pt/model/descriptor/se_r.py | 1 + deepmd/pt/model/descriptor/se_t.py | 1 + deepmd/pt/model/descriptor/se_t_tebd.py | 1 + deepmd/pt/model/model/ener_model.py | 4 + deepmd/pt/model/model/make_model.py | 19 + deepmd/pt/train/training.py | 28 +- deepmd/pt/train/wrapper.py | 2 + deepmd/pt/utils/stat.py | 3 +- deepmd/pt_expt/infer/deep_eval.py | 150 +++++++- deepmd/pt_expt/model/dipole_model.py | 337 +++++++++-------- deepmd/pt_expt/model/dos_model.py | 297 ++++++++------- deepmd/pt_expt/model/dp_linear_model.py | 15 +- deepmd/pt_expt/model/dp_zbl_model.py | 345 +++++++++--------- deepmd/pt_expt/model/ener_model.py | 17 +- deepmd/pt_expt/model/make_model.py | 14 + deepmd/pt_expt/model/polar_model.py | 297 ++++++++------- deepmd/pt_expt/model/property_model.py | 311 ++++++++-------- deepmd/pt_expt/model/spin_ener_model.py | 9 + deepmd/pt_expt/train/training.py | 45 ++- deepmd/pt_expt/train/wrapper.py | 2 + deepmd/utils/argcheck.py | 18 +- source/tests/consistent/descriptor/common.py | 49 ++- .../tests/consistent/descriptor/test_dpa3.py | 16 +- source/tests/consistent/model/test_ener.py | 85 +++-- source/tests/pt/model/test_dpa3.py | 70 +++- source/tests/pt_expt/descriptor/test_dpa3.py | 83 +++++ .../dpmodel/descriptor/test_descriptor.py | 3 + .../universal/dpmodel/model/test_model.py | 11 +- 60 files changed, 1662 insertions(+), 917 deletions(-) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 1d8e955de7..356bfeb9ce 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -137,8 +137,14 @@ def calculate( fparam = self.atoms.info.get("fparam", None) aparam = self.atoms.info.get("aparam", None) + charge_spin = self.atoms.info.get("charge_spin", None) e, f, v = self.dp.eval( - coords=coord, cells=cell, atom_types=atype, fparam=fparam, aparam=aparam + coords=coord, + cells=cell, + atom_types=atype, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, )[:3] self.results["energy"] = e[0][0] # see https://gitlab.com/ase/ase/-/merge_requests/2485 diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 1120078bb2..3987d1264b 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -156,6 +156,18 @@ def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return None + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return False + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return False + + def get_default_chg_spin(self) -> list[float] | None: + """Get the default charge_spin values.""" + return None + def reinit_atom_exclude( self, exclude_types: list[int] = [], @@ -231,6 +243,7 @@ def forward_common_atomic( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Common interface for atomic inference. @@ -279,6 +292,7 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) ret_dict = self.apply_out_stat(ret_dict, atype) @@ -307,6 +321,7 @@ def call( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: return self.forward_common_atomic( extended_coord, @@ -315,6 +330,7 @@ def call( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) def get_intensive(self) -> bool: @@ -519,6 +535,7 @@ def model_forward( box: np.ndarray | None, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> dict[str, np.ndarray]: # Get reference array to determine the target array type and device # Use out_bias as reference since it's always present @@ -538,6 +555,8 @@ def model_forward( fparam = xp.asarray(fparam, device=device) if aparam is not None: aparam = xp.asarray(aparam, device=device) + if charge_spin is not None: + charge_spin = xp.asarray(charge_spin, device=device) ( extended_coord, @@ -559,6 +578,7 @@ def model_forward( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) # Convert outputs back to numpy arrays return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()} diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 466e3ddd95..131c6a5ed6 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -79,6 +79,28 @@ def __init__( ) super().init_out_stat() + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return self.add_chg_spin_ebd + + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + if self.add_chg_spin_ebd: + return self.descriptor.get_dim_chg_spin() + return 0 + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + if self.add_chg_spin_ebd: + return self.descriptor.has_default_chg_spin() + return False + + def get_default_chg_spin(self) -> list[float] | None: + """Get the default charge_spin values.""" + if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): + return self.descriptor.get_default_chg_spin() + return None + def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" return self.fitting_net.output_def() @@ -157,6 +179,7 @@ def forward_atomic( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Models' atomic predictions. @@ -174,6 +197,8 @@ def forward_atomic( frame parameter. nf x ndf aparam atomic parameter. nf x nloc x nda + charge_spin + charge and spin parameter for descriptor. nf x 2 Returns ------- @@ -184,37 +209,28 @@ def forward_atomic( nframes, nloc, nnei = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) - # Handle default fparam if fitting net supports it - if ( - hasattr(self.fitting_net, "get_dim_fparam") - and self.fitting_net.get_dim_fparam() > 0 - and fparam is None - ): - # use default fparam - from deepmd.dpmodel.array_api import ( - array_api_compat, - ) - - default_fparam = self.fitting_net.get_default_fparam() - assert default_fparam is not None - xp = array_api_compat.array_namespace(extended_coord) - default_fparam_array = xp.asarray( - default_fparam, - dtype=extended_coord.dtype, - device=array_api_compat.device(extended_coord), - ) - fparam_input_for_des = xp.tile( - xp.reshape(default_fparam_array, (1, -1)), (nframes, 1) - ) - else: - fparam_input_for_des = fparam + # Handle default charge_spin if descriptor supports it + if self.add_chg_spin_ebd and charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is not None: + from deepmd.dpmodel.array_api import ( + array_api_compat, + ) + + xp = array_api_compat.array_namespace(extended_coord) + cs_array = xp.asarray( + default_cs, + dtype=extended_coord.dtype, + device=array_api_compat.device(extended_coord), + ) + charge_spin = xp.tile(xp.reshape(cs_array, (1, -1)), (nframes, 1)) descriptor, rot_mat, g2, h2, sw = self.descriptor( extended_coord, extended_atype, nlist, mapping=mapping, - fparam=fparam_input_for_des if self.add_chg_spin_ebd else None, + charge_spin=charge_spin if self.add_chg_spin_ebd else None, ) ret = self.fitting_net( descriptor, diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 3ed9077df7..8ee05b1575 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -224,6 +224,7 @@ def forward_atomic( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return atomic prediction. @@ -280,6 +281,7 @@ def forward_atomic( mapping, fparam, aparam, + charge_spin=charge_spin, )["energy"] ) weights = self._compute_weight(extended_coord, extended_atype, nlists_) diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 3e48e88c87..7118aa5d7a 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -138,6 +138,7 @@ def fwd( mapping: t_tensor | None = None, fparam: t_tensor | None = None, aparam: t_tensor | None = None, + charge_spin: t_tensor | None = None, ) -> dict[str, t_tensor]: pass diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 51c370eca0..2daae42fbd 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -253,6 +253,7 @@ def forward_atomic( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) nframes, nloc, nnei = nlist.shape diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index bc2a04a836..1e7d29b004 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -500,6 +500,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> Array: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 2fa765f04b..2338abd382 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -831,6 +831,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, Array, Array, Array, Array]: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 5f5aea50e5..8e4a4ba5ad 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -359,6 +359,7 @@ def __init__( use_loc_mapping: bool = True, type_map: list[str] | None = None, add_chg_spin_ebd: bool = False, + default_chg_spin: list[float] | None = None, ) -> None: super().__init__() @@ -414,6 +415,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.use_econf_tebd = use_econf_tebd self.add_chg_spin_ebd = add_chg_spin_ebd + self.default_chg_spin = default_chg_spin + if self.add_chg_spin_ebd and self.default_chg_spin is not None: + assert len(self.default_chg_spin) == 2, ( + "default_chg_spin must have exactly 2 values [charge, spin]" + ) self.use_tebd_bias = use_tebd_bias self.use_loc_mapping = use_loc_mapping self.type_map = type_map @@ -480,6 +486,18 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input.""" + return 2 if self.add_chg_spin_ebd else 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether default charge_spin values are set.""" + return self.default_chg_spin is not None + + def get_default_chg_spin(self) -> list[float] | None: + """Returns the default charge_spin values.""" + return self.default_chg_spin + def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth @@ -616,6 +634,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, Array, Array, Array, Array]: """Compute the descriptor. @@ -668,13 +687,13 @@ def call( ) if self.add_chg_spin_ebd: - assert fparam is not None + assert charge_spin is not None assert self.chg_embedding is not None assert self.spin_embedding is not None chg_tebd = self.chg_embedding.call() spin_tebd = self.spin_embedding.call() - charge = xp.astype(fparam[:, 0], xp.int64) + 100 - spin = xp.astype(fparam[:, 1], xp.int64) + charge = xp.astype(charge_spin[:, 0], xp.int64) + 100 + spin = xp.astype(charge_spin[:, 1], xp.int64) chg_ebd = xp.reshape( xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0), (nframes, self.tebd_dim), @@ -718,6 +737,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "use_loc_mapping": self.use_loc_mapping, "add_chg_spin_ebd": self.add_chg_spin_ebd, + "default_chg_spin": self.default_chg_spin, "type_map": self.type_map, "type_embedding": self.type_embedding.serialize(), } diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index b15fbc15d2..01a8f31bde 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -276,6 +276,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[ Array, Array | None, @@ -332,7 +333,9 @@ def call( # mixed_types is True, but descrpt.mixed_types is False assert nl_distinguish_types is not None nl = nl_distinguish_types[:, :, nci] - odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping) + odescriptor, gr, g2, h2, sw = descrpt( + coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin + ) out_descriptor.append(odescriptor) if gr is not None: out_gr.append(gr) diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 47245898ce..c03948be54 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -187,6 +187,7 @@ def fwd( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> Array: """Calculate descriptor.""" pass diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 8997412325..81ee2cc3e2 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -399,6 +399,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> Array: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index b5ba7a282f..9548b1de9f 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -371,6 +371,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> Array: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index e599669068..510be5e2cf 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -346,6 +346,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index 2d36994d61..6060f208cb 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -354,6 +354,7 @@ def call( nlist: Array, mapping: Array | None = None, fparam: Array | None = None, + charge_spin: Array | None = None, ) -> tuple[Array, Array]: """Compute the descriptor. diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 57b518d75d..63a2344c5c 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -86,6 +86,7 @@ def call( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: model_ret = self.call_common( @@ -94,6 +95,7 @@ def call( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -119,6 +121,7 @@ def call_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: model_ret = self.call_common_lower( @@ -128,6 +131,7 @@ def call_lower( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index fb77838b4c..9b8a07a0b3 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -77,6 +77,7 @@ def model_call_from_call_lower( aparam: Array | None = None, do_atomic_virial: bool = False, coord_corr_for_virial: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -146,6 +147,7 @@ def model_call_from_call_lower( "fparam": fp, "aparam": ap, "do_atomic_virial": do_atomic_virial, + "charge_spin": charge_spin, } if extended_coord_corr is not None: call_lower_kwargs["extended_coord_corr"] = extended_coord_corr @@ -266,6 +268,7 @@ def call_common( aparam: Array | None = None, do_atomic_virial: bool = False, coord_corr_for_virial: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction. @@ -312,6 +315,7 @@ def call_common( aparam=ap, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, + charge_spin=charge_spin, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -326,6 +330,7 @@ def call_common_lower( aparam: Array | None = None, do_atomic_virial: bool = False, extended_coord_corr: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping @@ -379,6 +384,7 @@ def call_common_lower( aparam=ap, do_atomic_virial=do_atomic_virial, extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -393,6 +399,7 @@ def forward_common_atomic( aparam: Array | None = None, do_atomic_virial: bool = False, extended_coord_corr: Array | None = None, + charge_spin: Array | None = None, ) -> dict[str, Array]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, @@ -401,6 +408,7 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) return fit_output_to_model_output( atomic_ret, @@ -700,6 +708,18 @@ def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return self.atomic_model.get_default_fparam() + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return self.atomic_model.has_chg_spin_ebd() + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return self.atomic_model.has_default_chg_spin() + + def get_default_chg_spin(self) -> list[float] | None: + """Get the default charge_spin values.""" + return self.atomic_model.get_default_chg_spin() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/utils/batch.py b/deepmd/dpmodel/utils/batch.py index 204ae9771f..2cbf8a72ff 100644 --- a/deepmd/dpmodel/utils/batch.py +++ b/deepmd/dpmodel/utils/batch.py @@ -11,7 +11,7 @@ _DROP_KEYS = {"default_mesh", "sid", "fid"} # Keys that belong to model input (everything else is label). -_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam"} +_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam", "charge_spin"} def normalize_batch(batch: dict[str, Any]) -> dict[str, Any]: diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 29253263a6..dc207f4aa1 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -667,8 +667,8 @@ def __getitem__(self, index: int) -> dict[str, Any]: np.repeat(frame[req_key], repeat).reshape(-1).astype(req_dtype) ) - # Add find_* for fparam/aparam/spin if not already set - for extra_key in ["fparam", "aparam", "spin"]: + # Add find_* for fparam/aparam/spin/charge_spin if not already set + for extra_key in ["fparam", "aparam", "spin", "charge_spin"]: if f"find_{extra_key}" not in frame: frame[f"find_{extra_key}"] = ( np.float32(1.0) if extra_key in frame else np.float32(0.0) diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py index 2c170da705..ca00f6c064 100644 --- a/deepmd/dpmodel/utils/stat.py +++ b/deepmd/dpmodel/utils/stat.py @@ -215,8 +215,11 @@ def _compute_model_predict( box = to_numpy_array(system["box"]) fparam = to_numpy_array(system.get("fparam", None)) aparam = to_numpy_array(system.get("aparam", None)) + charge_spin = to_numpy_array(system.get("charge_spin", None)) - sample_predict = model_forward(coord, atype, box, fparam=fparam, aparam=aparam) + sample_predict = model_forward( + coord, atype, box, fparam=fparam, aparam=aparam, charge_spin=charge_spin + ) for kk in keys: model_predict[kk].append( sample_predict[kk] # already numpy from model_forward diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 604857c837..f589f25989 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -590,6 +590,14 @@ def test_ener( ) if dp.get_dim_aparam() > 0: data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) + if dp.has_chg_spin_ebd(): + data.add( + "charge_spin", + 2, + atomic=False, + must=not dp.has_default_chg_spin(), + high_prec=False, + ) if dp.has_spin: data.add("spin", 3, atomic=True, must=True, high_prec=False) data.add("force_mag", 3, atomic=True, must=False, high_prec=False) @@ -631,6 +639,10 @@ def test_ener( aparam = test_data["aparam"][:numb_test] else: aparam = None + if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 1.0) != 0.0: + charge_spin = test_data["charge_spin"][:numb_test] + else: + charge_spin = None ret = dp.eval( coord, @@ -642,6 +654,7 @@ def test_ener( efield=efield, mixed_type=mixed_type, spin=spin, + charge_spin=charge_spin, ) energy = ret[0] force = ret[1] diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 807414fa5d..557f3ddd23 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -166,6 +166,14 @@ def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return False + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return False + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return False + @abstractmethod def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" @@ -451,6 +459,14 @@ def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return self.deep_eval.has_default_fparam() + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return self.deep_eval.has_chg_spin_ebd() + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return self.deep_eval.has_default_chg_spin() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.deep_eval.get_dim_aparam() diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 2e30b8574a..6128e2633b 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -252,6 +252,20 @@ def has_default_fparam(self) -> bool: # for compatibility with old models return False + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + try: + return self.dp.model["Default"].has_chg_spin_ebd() + except AttributeError: + return False + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + try: + return self.dp.model["Default"].has_default_chg_spin() + except AttributeError: + return False + def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() @@ -344,6 +358,7 @@ def eval( atomic: bool = False, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -393,7 +408,7 @@ def eval( request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, fparam, aparam, request_defs + coords, cells, atom_types, fparam, aparam, charge_spin, request_defs ) else: out = self._eval_func(self._eval_model_spin, numb_test, natoms)( @@ -403,6 +418,7 @@ def eval( np.array(kwargs["spin"]), fparam, aparam, + charge_spin, request_defs, ) return dict( @@ -504,6 +520,7 @@ def _eval_model( atom_types: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, + charge_spin: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) @@ -546,6 +563,10 @@ def _eval_model( ) else: aparam_input = None + if charge_spin is not None: + charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2)) + else: + charge_spin_input = None do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C for x in request_defs ) @@ -556,6 +577,7 @@ def _eval_model( do_atomic_virial=do_atomic_virial, fparam=fparam_input, aparam=aparam_input, + charge_spin=charge_spin_input, ) if isinstance(batch_output, tuple): batch_output = batch_output[0] @@ -582,6 +604,7 @@ def _eval_model_spin( spins: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, + charge_spin: np.ndarray | None, request_defs: list[OutputVariableDef], ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) @@ -624,6 +647,10 @@ def _eval_model_spin( ) else: aparam_input = None + if charge_spin is not None: + charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2)) + else: + charge_spin_input = None do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs @@ -636,6 +663,7 @@ def _eval_model_spin( do_atomic_virial=do_atomic_virial, fparam=fparam_input, aparam=aparam_input, + charge_spin=charge_spin_input, ) if isinstance(batch_output, tuple): batch_output = batch_output[0] diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index bfc67cf82b..dfd4ce60fc 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -189,6 +189,18 @@ def get_default_fparam(self) -> torch.Tensor | None: """Get the default frame parameters.""" return None + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return False + + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return False + + def get_default_chg_spin(self) -> torch.Tensor | None: + """Get the default charge_spin values.""" + return None + def _make_wrapped_sampler( self, sampled_func: Callable[[], list[dict]], @@ -305,6 +317,7 @@ def forward_common_atomic( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Common interface for atomic inference. @@ -356,6 +369,7 @@ def forward_common_atomic( fparam=fparam, aparam=aparam, comm_dict=comm_dict, + charge_spin=charge_spin, ) ret_dict = self.apply_out_stat(ret_dict, atype) @@ -386,6 +400,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: return self.forward_common_atomic( extended_coord, @@ -395,6 +410,7 @@ def forward( fparam=fparam, aparam=aparam, comm_dict=comm_dict, + charge_spin=charge_spin, ) def change_type_map( @@ -621,6 +637,7 @@ def model_forward( box: torch.Tensor | None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: with ( torch.no_grad() @@ -638,6 +655,10 @@ def model_forward( mixed_types=self.mixed_types(), box=box, ) + if charge_spin is not None and not isinstance( + charge_spin, torch.Tensor + ): + charge_spin = to_torch_tensor(charge_spin) atomic_ret = self.forward_common_atomic( extended_coord, extended_atype, @@ -645,6 +666,7 @@ def model_forward( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) return {kk: vv.detach() for kk, vv in atomic_ret.items()} diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index efb2a532e5..9b932aae72 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -244,6 +244,7 @@ def forward_atomic( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -273,20 +274,11 @@ def forward_atomic( if self.do_grad_r() or self.do_grad_c(): extended_coord.requires_grad_(True) - # Handle default fparam if fitting net supports it - if ( - hasattr(self.fitting_net, "get_dim_fparam") - and self.fitting_net.get_dim_fparam() > 0 - and fparam is None - ): - # use default fparam - default_fparam_tensor = self.fitting_net.get_default_fparam() - assert default_fparam_tensor is not None - fparam_input_for_des = torch.tile( - default_fparam_tensor.unsqueeze(0), [nframes, 1] - ) - else: - fparam_input_for_des = fparam + # Handle default chg_spin if descriptor supports it + if self.add_chg_spin_ebd and charge_spin is None: + default_cs_tensor = self.descriptor.get_default_chg_spin() + if default_cs_tensor is not None: + charge_spin = torch.tile(default_cs_tensor.unsqueeze(0), [nframes, 1]) descriptor, rot_mat, g2, h2, sw = self.descriptor( extended_coord, @@ -294,7 +286,7 @@ def forward_atomic( nlist, mapping=mapping, comm_dict=comm_dict, - fparam=fparam_input_for_des if self.add_chg_spin_ebd else None, + charge_spin=charge_spin if self.add_chg_spin_ebd else None, ) assert descriptor is not None if self.enable_eval_descriptor_hook: @@ -394,6 +386,32 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> torch.Tensor | None: return self.fitting_net.get_default_fparam() + @torch.jit.export + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return self.add_chg_spin_ebd + + @torch.jit.export + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + if self.add_chg_spin_ebd: + return self.descriptor.get_dim_chg_spin() + return 0 + + @torch.jit.export + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + if self.add_chg_spin_ebd: + return self.descriptor.has_default_chg_spin() + return False + + @torch.jit.export + def get_default_chg_spin(self) -> torch.Tensor | None: + """Get the default charge_spin values as a tensor.""" + if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): + return self.descriptor.get_default_chg_spin() + return None + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting_net.get_dim_aparam() diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 4c415658e2..5c0f616634 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -233,6 +233,7 @@ def forward_atomic( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -292,6 +293,7 @@ def forward_atomic( fparam, aparam, comm_dict=comm_dict, + charge_spin=charge_spin, )["energy"] ) weights = self._compute_weight(extended_coord, extended_atype, nlists_) diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index e838a7a24d..5750f7cfd1 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -271,6 +271,7 @@ def forward_atomic( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc, nnei = nlist.shape extended_coord = extended_coord.view(nframes, -1, 3) diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index b4ff2dd46f..ba96b05d42 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -672,6 +672,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 89aff38168..cf88ec7de2 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -717,6 +717,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index 0c6982afe5..1befcbdea0 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -122,8 +122,13 @@ def __init__( use_loc_mapping: bool = True, type_map: list[str] | None = None, add_chg_spin_ebd: bool = False, + default_chg_spin: list[float] | None = None, ) -> None: super().__init__() + if default_chg_spin is not None: + assert len(default_chg_spin) == 2, ( + "default_chg_spin must be a list of length 2 [charge, spin]." + ) def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: if isinstance(sub_data, dict): @@ -177,6 +182,7 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: self.use_econf_tebd = use_econf_tebd self.add_chg_spin_ebd = add_chg_spin_ebd + self.default_chg_spin = default_chg_spin self.use_loc_mapping = use_loc_mapping self.use_tebd_bias = use_tebd_bias self.type_map = type_map @@ -249,6 +255,27 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut + @torch.jit.export + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + return 2 if self.add_chg_spin_ebd else 0 + + @torch.jit.export + def has_default_chg_spin(self) -> bool: + """Check if the descriptor has default charge_spin values.""" + return self.default_chg_spin is not None + + @torch.jit.export + def get_default_chg_spin(self) -> torch.Tensor | None: + """Get the default charge_spin values as a tensor.""" + if self.default_chg_spin is None: + return None + return torch.tensor( + self.default_chg_spin, + dtype=self.prec, + device=env.DEVICE, + ) + def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth @@ -427,6 +454,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "use_loc_mapping": self.use_loc_mapping, "add_chg_spin_ebd": self.add_chg_spin_ebd, + "default_chg_spin": self.default_chg_spin, "type_map": self.type_map, "type_embedding": self.type_embedding.embedding.serialize(), } @@ -504,6 +532,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -555,11 +584,11 @@ def forward( node_ebd_ext = self.type_embedding(extended_atype) if self.add_chg_spin_ebd: - assert fparam is not None + assert charge_spin is not None assert self.chg_embedding is not None assert self.spin_embedding is not None - charge = fparam[:, 0].to(dtype=torch.int64) + 100 - spin = fparam[:, 1].to(dtype=torch.int64) + charge = charge_spin[:, 0].to(dtype=torch.int64) + 100 + spin = charge_spin[:, 1].to(dtype=torch.int64) chg_ebd = self.chg_embedding(charge) spin_ebd = self.spin_embedding(spin) sys_cs_embd = self.act( diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 55c1f9d2e3..e9a3258ba5 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -269,6 +269,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -332,7 +333,9 @@ def forward( nl = nl_distinguish_types[ :, :, self.nlist_cut_idx[ii].to(atype_ext.device) ] - odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping) + odescriptor, gr, g2, h2, sw = descrpt( + coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin + ) out_descriptor.append(odescriptor) if gr is not None: out_gr.append(gr) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index d4ee032d49..8bf1308e30 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -309,6 +309,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 92ef4c800d..f824a0c44c 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -428,6 +428,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 9e0de85f49..c518740de2 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -344,6 +344,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index f3c7544549..9c38cc2fff 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -442,6 +442,7 @@ def forward( mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 1680d1e258..28387553fb 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -72,6 +72,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -80,6 +81,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} @@ -115,6 +117,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -126,6 +129,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 83e0209ad8..0eacf85037 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -139,6 +139,7 @@ def forward_common( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, coord_corr_for_virial: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -204,6 +205,7 @@ def forward_common( fparam=fp, aparam=ap, extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, ) model_predict = communicate_extended_output( model_predict_lower, @@ -259,6 +261,7 @@ def forward_common_lower( comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, extended_coord_corr: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping @@ -311,6 +314,7 @@ def forward_common_lower( fparam=fp, aparam=ap, comm_dict=comm_dict, + charge_spin=charge_spin, ) model_predict = fit_output_to_model_output( atomic_ret, @@ -551,6 +555,21 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> torch.Tensor | None: return self.atomic_model.get_default_fparam() + @torch.jit.export + def has_chg_spin_ebd(self) -> bool: + """Check if the model has charge spin embedding.""" + return self.atomic_model.has_chg_spin_ebd() + + @torch.jit.export + def has_default_chg_spin(self) -> bool: + """Check if the model has default charge_spin values.""" + return self.atomic_model.has_default_chg_spin() + + @torch.jit.export + def get_default_chg_spin(self) -> torch.Tensor | None: + """Get the default charge_spin values.""" + return self.atomic_model.get_default_chg_spin() + @torch.jit.export def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index e674eb4b33..c9325b1030 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1779,13 +1779,20 @@ def get_data( "box", "fparam", "aparam", + "charge_spin", ] input_dict = dict.fromkeys(input_keys) label_dict = {} for item_key in batch_data: if item_key in input_keys: - if item_key != "fparam" or batch_data["find_fparam"] != 0.0: - input_dict[item_key] = batch_data[item_key] + if item_key == "fparam" and batch_data.get("find_fparam", 1.0) == 0.0: + continue + if ( + item_key == "charge_spin" + and batch_data.get("find_charge_spin", 1.0) == 0.0 + ): + continue + input_dict[item_key] = batch_data[item_key] else: if item_key not in ["sid", "fid"]: label_dict[item_key] = batch_data[item_key] @@ -1899,6 +1906,23 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: DataRequirementItem("spin", ndof=3, atomic=True, must=True) ] additional_data_requirement += spin_requirement_items + has_chg_spin_ebd = getattr(_model, "has_chg_spin_ebd", False) + if callable(has_chg_spin_ebd): + has_chg_spin_ebd = has_chg_spin_ebd() + if has_chg_spin_ebd: + has_default_cs = _model.has_default_chg_spin() + cs_default = ( + _model.get_default_chg_spin().cpu().numpy() if has_default_cs else 0.0 + ) + additional_data_requirement.append( + DataRequirementItem( + "charge_spin", + ndof=2, + atomic=False, + must=not has_default_cs, + default=cs_default, + ) + ) return additional_data_requirement diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index ddb4a4323d..1d741dd534 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -165,6 +165,7 @@ def forward( do_atomic_virial: bool = False, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> tuple[Any, Any, Any]: if not self.multi_task: task_key = "Default" @@ -179,6 +180,7 @@ def forward( "do_atomic_virial": do_atomic_virial, "fparam": fparam, "aparam": aparam, + "charge_spin": charge_spin, } has_spin = getattr(self.model[task_key], "has_spin", False) if callable(has_spin): diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index dc4d43f508..f8c3685b78 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -189,6 +189,7 @@ def _compute_model_predict( ) fparam = system.get("fparam", None) aparam = system.get("aparam", None) + charge_spin = system.get("charge_spin", None) def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: return auto_batch_size.execute_all( @@ -200,7 +201,7 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: ) sample_predict = model_forward_auto_batch_size( - coord, atype, box, fparam=fparam, aparam=aparam + coord, atype, box, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) for kk in keys: model_predict[kk].append( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index f2fe908297..bd586e19b4 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -413,6 +413,21 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "mixed_types": model.mixed_types(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), + "has_chg_spin_ebd": ( + model.has_chg_spin_ebd() + if hasattr(model, "has_chg_spin_ebd") + else False + ), + "has_default_chg_spin": ( + model.has_default_chg_spin() + if hasattr(model, "has_default_chg_spin") + else False + ), + "default_chg_spin": ( + model.get_default_chg_spin() + if hasattr(model, "get_default_chg_spin") + else None + ), "is_spin": self._is_spin, } if self._is_spin: @@ -435,6 +450,7 @@ def _eager_runner_spin( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: ext_coord = ext_coord.detach().requires_grad_(True) return model.forward_common_lower( @@ -445,6 +461,7 @@ def _eager_runner_spin( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=True, ) @@ -458,6 +475,7 @@ def _eager_runner( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: ext_coord = ext_coord.detach().requires_grad_(True) return model.forward_common_lower( @@ -467,6 +485,7 @@ def _eager_runner( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=True, ) @@ -496,6 +515,18 @@ def get_dim_aparam(self) -> int: return self._dpmodel.get_dim_aparam() return int(self.metadata["dim_aparam"]) + def has_chg_spin_ebd(self) -> bool: + """Check whether the model uses a dedicated charge_spin input.""" + if self._dpmodel is not None and hasattr(self._dpmodel, "has_chg_spin_ebd"): + return bool(self._dpmodel.has_chg_spin_ebd()) + return bool(self.metadata.get("has_chg_spin_ebd", False)) + + def has_default_chg_spin(self) -> bool: + """Check whether the model has a default charge_spin fallback.""" + if self._dpmodel is not None and hasattr(self._dpmodel, "has_default_chg_spin"): + return bool(self._dpmodel.has_default_chg_spin()) + return bool(self.metadata.get("has_default_chg_spin", False)) + @property def model_type(self) -> type["DeepEvalWrapper"]: """The evaluator of the model type.""" @@ -568,6 +599,7 @@ def eval( atomic: bool = False, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -623,11 +655,24 @@ def eval( if spins is not None: spins = np.array(spins) out = self._eval_func(self._eval_model_spin, numb_test, natoms)( - coords, cells, atom_types, spins, fparam, aparam, request_defs + coords, + cells, + atom_types, + spins, + fparam, + aparam, + request_defs, + charge_spin, ) else: out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, fparam, aparam, request_defs + coords, + cells, + atom_types, + fparam, + aparam, + request_defs, + charge_spin, ) return dict( zip( @@ -917,6 +962,7 @@ def _prepare_inputs( atom_types: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, + charge_spin: np.ndarray | None = None, ) -> tuple: """Prepare tensor inputs for model evaluation. @@ -924,7 +970,7 @@ def _prepare_inputs( ------- tuple (ext_coord_t, ext_atype_t, nlist_t, mapping_t, - fparam_t, aparam_t, nframes, natoms) + fparam_t, aparam_t, charge_spin_t, nframes, natoms) """ nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -1008,6 +1054,32 @@ def _prepare_inputs( else: aparam_t = None + # charge_spin handling: dedicated input, separate from fparam. + if charge_spin is not None: + charge_spin_t = torch.tensor( + np.asarray(charge_spin).reshape(nframes, 2), + dtype=torch.float64, + device=DEVICE, + ) + elif self.metadata.get("has_chg_spin_ebd", False): + default_cs = self.metadata.get("default_chg_spin") + if default_cs is not None: + if hasattr(default_cs, "cpu"): + default_cs = default_cs.cpu().numpy() + charge_spin_t = ( + torch.tensor(default_cs, dtype=torch.float64, device=DEVICE) + .unsqueeze(0) + .expand(nframes, -1) + .contiguous() + ) + else: + raise ValueError( + "charge_spin is required for this model (add_chg_spin_ebd=True) " + "but was not provided, and no default_chg_spin is set." + ) + else: + charge_spin_t = None + return ( ext_coord_t, ext_atype_t, @@ -1015,6 +1087,7 @@ def _prepare_inputs( mapping_t, fparam_t, aparam_t, + charge_spin_t, nframes, natoms, ) @@ -1027,6 +1100,7 @@ def _eval_model( fparam: np.ndarray | None, aparam: np.ndarray | None, request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: ( ext_coord_t, @@ -1035,9 +1109,10 @@ def _eval_model( mapping_t, fparam_t, aparam_t, + charge_spin_t, nframes, natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam) + ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) # Call the model (forward_common_lower interface, internal keys) if self._is_pt2: @@ -1046,11 +1121,23 @@ def _eval_model( # It also filters non-tensor args automatically, matching the # export-time signature where None args were excluded. model_ret = self._pt2_runner( - ext_coord_t, ext_atype_t, nlist_t, mapping_t, fparam_t, aparam_t + ext_coord_t, + ext_atype_t, + nlist_t, + mapping_t, + fparam_t, + aparam_t, + charge_spin_t, ) else: model_ret = self.exported_module( - ext_coord_t, ext_atype_t, nlist_t, mapping_t, fparam_t, aparam_t + ext_coord_t, + ext_atype_t, + nlist_t, + mapping_t, + fparam_t, + aparam_t, + charge_spin=charge_spin_t, ) # Apply communicate_extended_output to map extended atoms → local atoms @@ -1093,6 +1180,7 @@ def _eval_model_spin( fparam: np.ndarray | None, aparam: np.ndarray | None, request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]: nframes = coords.shape[0] if len(atom_types.shape) == 1: @@ -1182,6 +1270,32 @@ def _eval_model_spin( else: aparam_t = None + # charge_spin handling: dedicated input, separate from fparam. + if charge_spin is not None: + charge_spin_t = torch.tensor( + np.asarray(charge_spin).reshape(nframes, 2), + dtype=torch.float64, + device=DEVICE, + ) + elif self.metadata.get("has_chg_spin_ebd", False): + default_cs = self.metadata.get("default_chg_spin") + if default_cs is not None: + if hasattr(default_cs, "cpu"): + default_cs = default_cs.cpu().numpy() + charge_spin_t = ( + torch.tensor(default_cs, dtype=torch.float64, device=DEVICE) + .unsqueeze(0) + .expand(nframes, -1) + .contiguous() + ) + else: + raise ValueError( + "charge_spin is required for this model (add_chg_spin_ebd=True) " + "but was not provided, and no default_chg_spin is set." + ) + else: + charge_spin_t = None + # Call the model with spin (7 args) if self._is_pt2: model_ret = self._pt2_runner( @@ -1192,6 +1306,7 @@ def _eval_model_spin( mapping_t, fparam_t, aparam_t, + charge_spin_t, ) else: model_ret = self.exported_module( @@ -1202,6 +1317,7 @@ def _eval_model_spin( mapping_t, fparam_t, aparam_t, + charge_spin=charge_spin_t, ) # Apply communicate_extended_output to map extended atoms → local atoms @@ -1357,6 +1473,7 @@ def eval_descriptor( atom_types: np.ndarray, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate descriptor. @@ -1402,19 +1519,19 @@ def eval_descriptor( mapping_t, fparam_t, _aparam_t, + charge_spin_t, _nframes, _natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam) + ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) with torch.no_grad(): - fparam_for_des = ( - fparam_t if getattr(dp_am, "add_chg_spin_ebd", False) else None - ) descriptor, *_ = dp_am.descriptor( ext_coord_t, ext_atype_t, nlist_t, mapping=mapping_t, - fparam=fparam_for_des, + charge_spin=charge_spin_t + if getattr(dp_am, "add_chg_spin_ebd", False) + else None, ) return descriptor.detach().cpu().numpy() @@ -1425,6 +1542,7 @@ def eval_fitting_last_layer( atom_types: np.ndarray, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate the last hidden layer of the fitting network. @@ -1470,19 +1588,19 @@ def eval_fitting_last_layer( mapping_t, fparam_t, aparam_t, + charge_spin_t, _nframes, natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam) + ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) with torch.no_grad(): - fparam_for_des = ( - fparam_t if getattr(dp_am, "add_chg_spin_ebd", False) else None - ) descriptor, rot_mat, g2, h2, _sw = dp_am.descriptor( ext_coord_t, ext_atype_t, nlist_t, mapping=mapping_t, - fparam=fparam_for_des, + charge_spin=charge_spin_t + if getattr(dp_am, "add_chg_spin_ebd", False) + else None, ) atype = ext_atype_t[:, :natoms] fitting_net = dp_am.fitting_net diff --git a/deepmd/pt_expt/model/dipole_model.py b/deepmd/pt_expt/model/dipole_model.py index 4b0e570ecb..679966d3d4 100644 --- a/deepmd/pt_expt/model/dipole_model.py +++ b/deepmd/pt_expt/model/dipole_model.py @@ -1,162 +1,175 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPDipoleAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("dipole") -class DipoleModel(DPModelCommon, DPDipoleModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPDipoleModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["dipole"] = model_ret["dipole"] - model_predict["global_dipole"] = model_ret["dipole_redu"] - if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: - model_predict["force"] = model_ret["dipole_derv_r"] - if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: - model_predict["virial"] = model_ret["dipole_derv_c_redu"] - if do_atomic_virial and model_ret["dipole_derv_c"] is not None: - model_predict["atom_virial"] = model_ret["dipole_derv_c"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["dipole"] = model_ret["dipole"] - model_predict["global_dipole"] = model_ret["dipole_redu"] - if self.do_grad_r("dipole") and model_ret.get("dipole_derv_r") is not None: - model_predict["extended_force"] = model_ret["dipole_derv_r"] - if self.do_grad_c("dipole") and model_ret.get("dipole_derv_c_redu") is not None: - model_predict["virial"] = model_ret["dipole_derv_c_redu"] - if do_atomic_virial and model_ret.get("dipole_derv_c") is not None: - model_predict["extended_virial"] = model_ret["dipole_derv_c"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "dipole": out_def_data["dipole"], - "global_dipole": out_def_data["dipole_redu"], - } - if self.do_grad_r("dipole"): - output_def["force"] = out_def_data["dipole_derv_r"] - output_def["force"].squeeze(-2) - if self.do_grad_c("dipole"): - output_def["virial"] = out_def_data["dipole_derv_c_redu"] - output_def["virial"].squeeze(-2) - output_def["atom_virial"] = out_def_data["dipole_derv_c"] - output_def["atom_virial"].squeeze(-2) - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDipoleAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dipole") +class DipoleModel(DPModelCommon, DPDipoleModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDipoleModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: + model_predict["force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret["dipole_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret.get("dipole_derv_r") is not None: + model_predict["extended_force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret.get("dipole_derv_c_redu") is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret.get("dipole_derv_c") is not None: + model_predict["extended_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "dipole": out_def_data["dipole"], + "global_dipole": out_def_data["dipole_redu"], + } + if self.do_grad_r("dipole"): + output_def["force"] = out_def_data["dipole_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("dipole"): + output_def["virial"] = out_def_data["dipole_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["dipole_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/dos_model.py b/deepmd/pt_expt/model/dos_model.py index 219c22e753..d74676b6ea 100644 --- a/deepmd/pt_expt/model/dos_model.py +++ b/deepmd/pt_expt/model/dos_model.py @@ -1,142 +1,155 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPDOSAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("dos") -class DOSModel(DPModelCommon, DPDOSModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPDOSModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_dos"] = model_ret["dos"] - model_predict["dos"] = model_ret["dos_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_dos"] = model_ret["dos"] - model_predict["dos"] = model_ret["dos_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "atom_dos": out_def_data["dos"], - "dos": out_def_data["dos_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDOSAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dos") +class DOSModel(DPModelCommon, DPDOSModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_dos": out_def_data["dos"], + "dos": out_def_data["dos_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 0ac75659b0..08515960a5 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -47,6 +47,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( @@ -55,6 +56,7 @@ def forward( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -78,6 +80,7 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( @@ -87,6 +90,7 @@ def forward_lower( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -130,6 +134,7 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -142,6 +147,7 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: extended_coord = extended_coord.detach().requires_grad_(True) nlist = _pad_nlist_for_export(nlist) @@ -152,6 +158,7 @@ def fn( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) @@ -160,7 +167,13 @@ def fn( model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) try: traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, ) finally: model.need_sorted_nlist_for_lower = _orig_need_sort diff --git a/deepmd/pt_expt/model/dp_zbl_model.py b/deepmd/pt_expt/model/dp_zbl_model.py index baa30c4ce0..be7ffcf70e 100644 --- a/deepmd/pt_expt/model/dp_zbl_model.py +++ b/deepmd/pt_expt/model/dp_zbl_model.py @@ -1,166 +1,179 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - DPZBLLinearEnergyAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("zbl") -class DPZBLModel(DPModelCommon, DPZBLModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPZBLModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial and model_ret["energy_derv_c"] is not None: - model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy") and model_ret.get("energy_derv_c_redu") is not None: - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial and model_ret.get("energy_derv_c") is not None: - model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( - -2 - ) - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "atom_energy": out_def_data["energy"], - "energy": out_def_data["energy_redu"], - } - if self.do_grad_r("energy"): - output_def["force"] = out_def_data["energy_derv_r"] - output_def["force"].squeeze(-2) - if self.do_grad_c("energy"): - output_def["virial"] = out_def_data["energy_derv_c_redu"] - output_def["virial"].squeeze(-2) - output_def["atom_virial"] = out_def_data["energy_derv_c"] - output_def["atom_virial"].squeeze(-2) - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - # Force `_format_nlist`'s sort branch into the compiled graph so the - # exported model tolerates oversized nlists at runtime — see - # make_model.py for the full rationale. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + DPZBLLinearEnergyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("zbl") +class DPZBLModel(DPModelCommon, DPZBLModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPZBLModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret["energy_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret.get("energy_derv_c_redu") is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret.get("energy_derv_c") is not None: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -2 + ) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_energy": out_def_data["energy"], + "energy": out_def_data["energy_redu"], + } + if self.do_grad_r("energy"): + output_def["force"] = out_def_data["energy_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = out_def_data["energy_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["energy_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # Force `_format_nlist`'s sort branch into the compiled graph so the + # exported model tolerates oversized nlists at runtime — see + # make_model.py for the full rationale. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index beb91c4ec4..684bc5b07e 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -57,6 +57,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( @@ -65,6 +66,7 @@ def forward( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -90,6 +92,7 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( @@ -99,6 +102,7 @@ def forward_lower( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -144,6 +148,7 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -175,6 +180,7 @@ def forward_lower_exportable( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, **make_fx_kwargs, ) @@ -191,9 +197,16 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: model_ret = traced( - extended_coord, extended_atype, nlist, mapping, fparam, aparam + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, ) model_predict: dict[str, torch.Tensor] = {} model_predict["atom_energy"] = model_ret["energy"] @@ -211,5 +224,5 @@ def fn( return model_predict return make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam + extended_coord, extended_atype, nlist, mapping, fparam, aparam, charge_spin ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index b28b81ffb1..aa195631c1 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -61,6 +61,7 @@ def _cal_hessian_ext( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, create_graph: bool = False, ) -> torch.Tensor: """Compute hessian of reduced output w.r.t. extended coordinates. @@ -112,6 +113,7 @@ def _cal_hessian_ext( mapping[ii] if mapping is not None else None, fparam[ii] if fparam is not None else None, aparam[ii] if aparam is not None else None, + charge_spin[ii] if charge_spin is not None else None, ) hess = torch.autograd.functional.hessian( wrapper, @@ -142,6 +144,7 @@ def __init__( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, ) -> None: self.model = model self.kk = kk @@ -152,6 +155,7 @@ def __init__( self.mapping = mapping self.fparam = fparam self.aparam = aparam + self.charge_spin = charge_spin def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: """Compute scalar reduced energy for one frame, one component. @@ -174,6 +178,9 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: mapping=self.mapping.unsqueeze(0) if self.mapping is not None else None, fparam=self.fparam.unsqueeze(0) if self.fparam is not None else None, aparam=self.aparam.unsqueeze(0) if self.aparam is not None else None, + charge_spin=self.charge_spin.unsqueeze(0) + if self.charge_spin is not None + else None, ) # atomic_ret[kk]: [1, nloc, *def] atom_energy = atomic_ret[self.kk][0] # [nloc, *def] @@ -278,6 +285,7 @@ def forward_common_atomic( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, extended_coord_corr: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: @@ -288,6 +296,7 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) model_ret = fit_output_to_model_output( atomic_ret, @@ -316,6 +325,7 @@ def forward_common_atomic( mapping, fparam, aparam, + charge_spin=charge_spin, create_graph=self.training, ) return model_ret @@ -328,6 +338,7 @@ def forward_common_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -367,6 +378,7 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: extended_coord = extended_coord.detach().requires_grad_(True) nlist = _pad_nlist_for_export(nlist) @@ -377,6 +389,7 @@ def fn( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) @@ -397,6 +410,7 @@ def fn( mapping, fparam, aparam, + charge_spin, ) finally: model.need_sorted_nlist_for_lower = _orig_need_sort diff --git a/deepmd/pt_expt/model/polar_model.py b/deepmd/pt_expt/model/polar_model.py index dd6b1c5d0f..cafddf80b0 100644 --- a/deepmd/pt_expt/model/polar_model.py +++ b/deepmd/pt_expt/model/polar_model.py @@ -1,142 +1,155 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPPolarAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("polar") -class PolarModel(DPModelCommon, DPPolarModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPPolarModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["polar"] = model_ret["polarizability"] - model_predict["global_polar"] = model_ret["polarizability_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["polar"] = model_ret["polarizability"] - model_predict["global_polar"] = model_ret["polarizability_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "polar": out_def_data["polarizability"], - "global_polar": out_def_data["polarizability_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPolarAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("polar") +class PolarModel(DPModelCommon, DPPolarModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPolarModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "polar": out_def_data["polarizability"], + "global_polar": out_def_data["polarizability_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/property_model.py b/deepmd/pt_expt/model/property_model.py index 223f8e5d78..db809da131 100644 --- a/deepmd/pt_expt/model/property_model.py +++ b/deepmd/pt_expt/model/property_model.py @@ -1,149 +1,162 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPPropertyAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("property") -class PropertyModel(DPModelCommon, DPPropertyModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPPropertyModel_.__init__(self, *args, **kwargs) - - def get_var_name(self) -> str: - """Get the name of the property.""" - return self.get_fitting_net().var_name - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - var_name = self.get_var_name() - model_predict = {} - model_predict[f"atom_{var_name}"] = model_ret[var_name] - model_predict[var_name] = model_ret[f"{var_name}_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - var_name = self.get_var_name() - model_predict = {} - model_predict[f"atom_{var_name}"] = model_ret[var_name] - model_predict[var_name] = model_ret[f"{var_name}_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - var_name = self.get_var_name() - output_def = { - f"atom_{var_name}": out_def_data[var_name], - var_name: out_def_data[f"{var_name}_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, extended_atype, nlist, mapping, fparam, aparam - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPropertyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("property") +class PropertyModel(DPModelCommon, DPPropertyModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPropertyModel_.__init__(self, *args, **kwargs) + + def get_var_name(self) -> str: + """Get the name of the property.""" + return self.get_fitting_net().var_name + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + var_name = self.get_var_name() + output_def = { + f"atom_{var_name}": out_def_data[var_name], + var_name: out_def_data[f"{var_name}_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/spin_ener_model.py b/deepmd/pt_expt/model/spin_ener_model.py index e96d0fbaf1..f24770f19e 100644 --- a/deepmd/pt_expt/model/spin_ener_model.py +++ b/deepmd/pt_expt/model/spin_ener_model.py @@ -48,6 +48,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( @@ -57,6 +58,7 @@ def forward( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -81,6 +83,7 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( @@ -91,6 +94,7 @@ def forward_lower( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -119,6 +123,7 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -151,6 +156,7 @@ def forward_lower_exportable( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, **make_fx_kwargs, ) @@ -168,6 +174,7 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: model_ret = traced( extended_coord, @@ -177,6 +184,7 @@ def fn( mapping, fparam, aparam, + charge_spin, ) model_predict: dict[str, torch.Tensor] = {} model_predict["atom_energy"] = model_ret["energy"] @@ -203,4 +211,5 @@ def fn( mapping, fparam, aparam, + charge_spin, ) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 5692b019cd..039981c4eb 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -138,6 +138,28 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: "aparam", _model.get_dim_aparam(), atomic=True, must=True ) ) + has_chg_spin_ebd = getattr(_model, "has_chg_spin_ebd", False) + if callable(has_chg_spin_ebd): + has_chg_spin_ebd = has_chg_spin_ebd() + if has_chg_spin_ebd: + has_default_cs = _model.has_default_chg_spin() + if has_default_cs: + default_cs = _model.get_default_chg_spin() + if hasattr(default_cs, "cpu"): + default_cs = default_cs.cpu().numpy() + else: + default_cs = np.asarray(default_cs) + else: + default_cs = 0.0 + additional_data_requirement.append( + DataRequirementItem( + "charge_spin", + ndof=2, + atomic=False, + must=not has_default_cs, + default=default_cs, + ) + ) return additional_data_requirement @@ -193,6 +215,7 @@ def _trace_and_compile( mapping: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None = None, compile_opts: dict[str, Any] | None = None, ) -> torch.nn.Module: """Symbolic-trace ``forward_lower`` and compile with inductor + dynamic=True. @@ -231,6 +254,7 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: extended_coord = extended_coord.detach().requires_grad_(True) return model.forward_lower( @@ -240,6 +264,7 @@ def fn( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) # Pick a trace-time nframes that's unlikely to collide with any other @@ -270,6 +295,7 @@ def _expand(t: torch.Tensor | None) -> torch.Tensor | None: mapping = _expand(mapping) fparam = _expand(fparam) aparam = _expand(aparam) + charge_spin = _expand(charge_spin) # Decompose silu_backward into primitive ops (sigmoid + mul + ...) # so that inductor can compile the graph without requiring a @@ -286,7 +312,7 @@ def _expand(t: torch.Tensor | None) -> torch.Tensor | None: tracing_mode="symbolic", _allow_non_fake_inputs=True, decomposition_table=decomp_table, - )(ext_coord, ext_atype, nlist, mapping, fparam, aparam) + )(ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin) # make_fx inserts aten.detach.default for saved tensors used in the # decomposed autograd.grad backward ops. These detach nodes break @@ -343,6 +369,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: from deepmd.dpmodel.utils.nlist import ( @@ -381,7 +408,7 @@ def forward( ext_coord = ext_coord.detach().requires_grad_(True) result = self.compiled_forward_lower( - ext_coord, ext_atype, nlist, mapping, fparam, aparam + ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin ) # Translate forward_lower keys -> forward keys. @@ -981,6 +1008,7 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: fparam = inp.get("fparam") aparam = inp.get("aparam") + charge_spin = inp.get("charge_spin") compiled_lower = _trace_and_compile( model, @@ -990,7 +1018,8 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: mapping, fparam, aparam, - compile_opts, + charge_spin=charge_spin, + compile_opts=compile_opts, ) wrapper_mod.model[task_key] = _CompiledModel(model, compiled_lower) @@ -1036,6 +1065,16 @@ def get_data( batch = normalize_batch(data_sys.get_batch()) input_dict, label_dict = split_batch(batch) + # Drop optional inputs whose find_* flag is False so the model sees None. + for opt_key in ("fparam", "charge_spin"): + find_key = f"find_{opt_key}" + if ( + opt_key in input_dict + and find_key in label_dict + and not bool(label_dict[find_key]) + ): + input_dict.pop(opt_key) + # Convert numpy values to torch tensors. for dd in (input_dict, label_dict): for key, val in dd.items(): diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index f67efe8a8e..2a301a92e3 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -170,6 +170,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, cur_lr: float | torch.Tensor | None = None, label: dict[str, torch.Tensor] | None = None, task_key: str | None = None, @@ -189,6 +190,7 @@ def forward( "do_atomic_virial": do_atomic_virial, "fparam": fparam, "aparam": aparam, + "charge_spin": charge_spin, } model_pred = self.model[task_key](**input_dict) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 0213c01f9c..1232026a76 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1375,8 +1375,15 @@ def descrpt_dpa3_args() -> list[Argument]: ) doc_add_chg_spin_ebd = ( "Whether to add charge and spin embedding to the descriptor. " - "When enabled, fparam is expected to have 2 values (charge, spin) " - "which are embedded and added to the type embedding." + "When enabled, the dedicated `charge_spin` input (shape [nframes, 2], " + "[charge, spin]) is embedded and added to the type embedding. " + "When `charge_spin` is missing in the input data, `default_chg_spin` " + "is used as a fallback if provided." + ) + doc_default_chg_spin = ( + "Default charge and spin values used as fallback when `charge_spin` " + "is not provided in the input data. Must be a list of length 2 " + "[charge, spin]. Only used when `add_chg_spin_ebd` is True." ) doc_activation_function = f"The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}." doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." @@ -1410,6 +1417,13 @@ def descrpt_dpa3_args() -> list[Argument]: default=False, doc=doc_add_chg_spin_ebd, ), + Argument( + "default_chg_spin", + list[float], + optional=True, + default=None, + doc=doc_default_chg_spin, + ), Argument( "activation_function", str, diff --git a/source/tests/consistent/descriptor/common.py b/source/tests/consistent/descriptor/common.py index 33bf7312de..16f6e0e1ca 100644 --- a/source/tests/consistent/descriptor/common.py +++ b/source/tests/consistent/descriptor/common.py @@ -103,6 +103,7 @@ def eval_dp_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( coords.reshape(1, -1, 3), @@ -119,7 +120,12 @@ def eval_dp_descriptor( distinguish_types=(not mixed_types), ) return dp_obj( - ext_coords, ext_atype, nlist=nlist, mapping=mapping, fparam=fparam + ext_coords, + ext_atype, + nlist=nlist, + mapping=mapping, + fparam=fparam, + charge_spin=charge_spin, ) def eval_pt_descriptor( @@ -131,6 +137,7 @@ def eval_pt_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pt( torch.from_numpy(coords).to(PT_DEVICE).reshape(1, -1, 3), @@ -149,10 +156,20 @@ def eval_pt_descriptor( fparam_pt = ( torch.from_numpy(fparam).to(PT_DEVICE) if fparam is not None else None ) + charge_spin_pt = ( + torch.from_numpy(charge_spin).to(PT_DEVICE) + if charge_spin is not None + else None + ) return [ x.detach().cpu().numpy() if torch.is_tensor(x) else x for x in pt_obj( - ext_coords, ext_atype, nlist=nlist, mapping=mapping, fparam=fparam_pt + ext_coords, + ext_atype, + nlist=nlist, + mapping=mapping, + fparam=fparam_pt, + charge_spin=charge_spin_pt, ) ] @@ -165,6 +182,7 @@ def eval_pt_expt_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( torch.from_numpy(coords).to(PT_DEVICE).reshape(1, -1, 3), @@ -183,10 +201,20 @@ def eval_pt_expt_descriptor( fparam_pt = ( torch.from_numpy(fparam).to(PT_DEVICE) if fparam is not None else None ) + charge_spin_pt = ( + torch.from_numpy(charge_spin).to(PT_DEVICE) + if charge_spin is not None + else None + ) return [ x.detach().cpu().numpy() if torch.is_tensor(x) else x for x in pt_expt_obj( - ext_coords, ext_atype, nlist=nlist, mapping=mapping, fparam=fparam_pt + ext_coords, + ext_atype, + nlist=nlist, + mapping=mapping, + fparam=fparam_pt, + charge_spin=charge_spin_pt, ) ] @@ -199,6 +227,7 @@ def eval_jax_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( jnp.array(coords).reshape(1, -1, 3), @@ -215,10 +244,16 @@ def eval_jax_descriptor( distinguish_types=(not mixed_types), ) fparam_jax = jnp.array(fparam) if fparam is not None else None + charge_spin_jax = jnp.array(charge_spin) if charge_spin is not None else None return [ np.asarray(x) if isinstance(x, jnp.ndarray) else x for x in jax_obj( - ext_coords, ext_atype, nlist=nlist, mapping=mapping, fparam=fparam_jax + ext_coords, + ext_atype, + nlist=nlist, + mapping=mapping, + fparam=fparam_jax, + charge_spin=charge_spin_jax, ) ] @@ -231,6 +266,7 @@ def eval_pd_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pd( paddle.to_tensor(coords).to(PD_DEVICE).reshape([1, -1, 3]), @@ -265,6 +301,7 @@ def eval_array_api_strict_descriptor( box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( array_api_strict.asarray(coords.reshape(1, -1, 3)), @@ -283,6 +320,9 @@ def eval_array_api_strict_descriptor( fparam_array_api = ( array_api_strict.asarray(fparam) if fparam is not None else None ) + charge_spin_array_api = ( + array_api_strict.asarray(charge_spin) if charge_spin is not None else None + ) return [ to_numpy_array(x) if hasattr(x, "__array_namespace__") else x for x in array_api_strict_obj( @@ -291,6 +331,7 @@ def eval_array_api_strict_descriptor( nlist=nlist, mapping=mapping, fparam=fparam_array_api, + charge_spin=charge_spin_array_api, ) ] diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index b067ca94dc..2c9ac9428d 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -395,8 +395,8 @@ def setUp(self) -> None: _precision, add_chg_spin_ebd, ) = self.param - # fparam for charge=5, spin=1 when add_chg_spin_ebd is True - self.fparam = ( + # charge_spin for charge=5, spin=1 when add_chg_spin_ebd is True + self.charge_spin = ( np.array([[5, 1]], dtype=GLOBAL_NP_FLOAT_PRECISION) if add_chg_spin_ebd else None @@ -420,7 +420,7 @@ def eval_dp(self, dp_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def eval_pt(self, pt_obj: Any) -> Any: @@ -431,7 +431,7 @@ def eval_pt(self, pt_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def eval_pd(self, pd_obj: Any) -> Any: @@ -442,7 +442,7 @@ def eval_pd(self, pd_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def eval_jax(self, jax_obj: Any) -> Any: @@ -453,7 +453,7 @@ def eval_jax(self, jax_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def eval_pt_expt(self, pt_expt_obj: Any) -> Any: @@ -464,7 +464,7 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: @@ -475,7 +475,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: self.atype, self.box, mixed_types=True, - fparam=self.fparam, + charge_spin=self.charge_spin, ) def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index def9f67f32..d62f84bea8 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -2016,57 +2016,56 @@ def raise_error(): @parameterized( - ("no_fparam", "explicit_fparam", "default_fparam"), # fparam_mode + ("no_chg_spin", "explicit_chg_spin", "default_chg_spin"), # cs_mode ) @unittest.skipUnless(INSTALLED_PT and INSTALLED_PT_EXPT, "PT and PT_EXPT are required") class TestEnerChgSpinEbdFparam(unittest.TestCase): - """Test dp/pt/pt_expt model forward consistency for add_chg_spin_ebd with three fparam modes. + """Test dp/pt/pt_expt model forward consistency for add_chg_spin_ebd with three modes. - - no_fparam: numb_fparam=0, add_chg_spin_ebd=False (baseline) - - explicit_fparam: numb_fparam=2, add_chg_spin_ebd=True, fparam provided - - default_fparam: numb_fparam=2, default_fparam set, add_chg_spin_ebd=True, fparam=None + - no_chg_spin: add_chg_spin_ebd=False (baseline) + - explicit_chg_spin: add_chg_spin_ebd=True, charge_spin provided + - default_chg_spin: add_chg_spin_ebd=True, default_chg_spin=[5,1], charge_spin=None """ def setUp(self) -> None: - (self.fparam_mode,) = self.param + (self.cs_mode,) = self.param - add_chg_spin_ebd = self.fparam_mode != "no_fparam" + add_chg_spin_ebd = self.cs_mode != "no_chg_spin" fitting_cfg: dict[str, Any] = { "neuron": [10, 10], "precision": "float64", "seed": 1, } - if self.fparam_mode != "no_fparam": - fitting_cfg["numb_fparam"] = 2 - if self.fparam_mode == "default_fparam": - fitting_cfg["default_fparam"] = [5, 1] + descriptor_cfg: dict[str, Any] = { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 3, + "e_rcut": 6.0, + "e_rcut_smth": 5.0, + "e_sel": 10, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "precision": "float64", + "seed": 1, + "add_chg_spin_ebd": add_chg_spin_ebd, + } + if self.cs_mode == "default_chg_spin": + descriptor_cfg["default_chg_spin"] = [5.0, 1.0] data = model_args().normalize_value( { "type_map": ["O", "H"], - "descriptor": { - "type": "dpa3", - "repflow": { - "n_dim": 20, - "e_dim": 10, - "a_dim": 8, - "nlayers": 3, - "e_rcut": 6.0, - "e_rcut_smth": 5.0, - "e_sel": 10, - "a_rcut": 4.0, - "a_rcut_smth": 3.5, - "a_sel": 8, - "axis_neuron": 4, - "update_angle": True, - "update_style": "res_residual", - "update_residual": 0.1, - "update_residual_init": "const", - }, - "precision": "float64", - "seed": 1, - "add_chg_spin_ebd": add_chg_spin_ebd, - }, + "descriptor": descriptor_cfg, "fitting_net": fitting_cfg, }, trim_pattern="_*", @@ -2106,15 +2105,15 @@ def setUp(self) -> None: dtype=GLOBAL_NP_FLOAT_PRECISION, ).reshape(1, 9) - # fparam: charge=5, spin=1 - if self.fparam_mode == "explicit_fparam": - self.fparam_np = np.array([[5, 1]], dtype=GLOBAL_NP_FLOAT_PRECISION) + # charge_spin: charge=5, spin=1; only set in explicit mode. + if self.cs_mode == "explicit_chg_spin": + self.charge_spin_np = np.array([[5, 1]], dtype=GLOBAL_NP_FLOAT_PRECISION) else: - self.fparam_np = None + self.charge_spin_np = None def test_forward_consistency(self) -> None: dp_ret = self.dp_model( - self.coords, self.atype, box=self.box, fparam=self.fparam_np + self.coords, self.atype, box=self.box, charge_spin=self.charge_spin_np ) pt_ret = { kk: torch_to_numpy(vv) @@ -2122,7 +2121,7 @@ def test_forward_consistency(self) -> None: numpy_to_torch(self.coords), numpy_to_torch(self.atype), box=numpy_to_torch(self.box), - fparam=numpy_to_torch(self.fparam_np), + charge_spin=numpy_to_torch(self.charge_spin_np), do_atomic_virial=True, ).items() } @@ -2134,7 +2133,7 @@ def test_forward_consistency(self) -> None: coord_t, pt_expt_numpy_to_torch(self.atype), box=pt_expt_numpy_to_torch(self.box), - fparam=pt_expt_numpy_to_torch(self.fparam_np), + charge_spin=pt_expt_numpy_to_torch(self.charge_spin_np), do_atomic_virial=True, ).items() } @@ -2144,12 +2143,12 @@ def test_forward_consistency(self) -> None: pt_ret[key], rtol=1e-10, atol=1e-10, - err_msg=f"dp vs pt mismatch in {key} (mode={self.fparam_mode})", + err_msg=f"dp vs pt mismatch in {key} (mode={self.cs_mode})", ) np.testing.assert_allclose( dp_ret[key], pe_ret[key], rtol=1e-10, atol=1e-10, - err_msg=f"dp vs pt_expt mismatch in {key} (mode={self.fparam_mode})", + err_msg=f"dp vs pt_expt mismatch in {key} (mode={self.cs_mode})", ) diff --git a/source/tests/pt/model/test_dpa3.py b/source/tests/pt/model/test_dpa3.py index 12b0be4532..cf00561379 100644 --- a/source/tests/pt/model/test_dpa3.py +++ b/source/tests/pt/model/test_dpa3.py @@ -55,7 +55,7 @@ def test_consistency( nme, prec, ect, - add_chg_spin, + cs_mode, ) in itertools.product( [True, False], # update_angle ["res_residual"], # update_style @@ -66,13 +66,20 @@ def test_consistency( [1, 2], # n_multi_edge_message ["float64"], # precision [False], # use_econf_tebd - [False, True], # add_chg_spin_ebd + ["no_chg_spin", "explicit_chg_spin", "default_chg_spin"], ): dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) if prec == "float64": atol = 1e-8 # marginal GPU test cases... + add_chg_spin = cs_mode != "no_chg_spin" + default_chg_spin = [5.0, 1.0] if cs_mode == "default_chg_spin" else None + # Descriptor.forward does not apply default_chg_spin fallback + # (that lives in dp_atomic_model). When add_chg_spin_ebd is on, + # tests must always pass an explicit charge_spin tensor. + need_cs_input = add_chg_spin + repflow = RepFlowArgs( n_dim=20, e_dim=10, @@ -105,27 +112,28 @@ def test_consistency( use_econf_tebd=ect, type_map=["O", "H"] if ect else None, add_chg_spin_ebd=add_chg_spin, + default_chg_spin=default_chg_spin, seed=GLOBAL_SEED, ).to(env.DEVICE) dd0.repflows.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) dd0.repflows.stddev = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) - # Prepare fparam if needed - fparam = None - fparam_np = None - if add_chg_spin: - fparam = torch.tensor([[5, 1]], dtype=dtype, device=env.DEVICE).expand( - nf, -1 - ) - fparam_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) + # Prepare charge_spin per mode. + charge_spin = None + charge_spin_np = None + if need_cs_input: + charge_spin = torch.tensor( + [[5, 1]], dtype=dtype, device=env.DEVICE + ).expand(nf, -1) + charge_spin_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) rd0, _, _, _, _ = dd0( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), torch.tensor(self.nlist, dtype=int, device=env.DEVICE), torch.tensor(self.mapping, dtype=int, device=env.DEVICE), - fparam=fparam, + charge_spin=charge_spin, ) # serialization dd1 = DescrptDPA3.deserialize(dd0.serialize()) @@ -134,7 +142,7 @@ def test_consistency( torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), torch.tensor(self.nlist, dtype=int, device=env.DEVICE), torch.tensor(self.mapping, dtype=int, device=env.DEVICE), - fparam=fparam, + charge_spin=charge_spin, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(), @@ -149,7 +157,7 @@ def test_consistency( self.atype_ext, self.nlist, self.mapping, - fparam=fparam_np, + charge_spin=charge_spin_np, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(), @@ -158,6 +166,42 @@ def test_consistency( atol=atol, ) + # default_chg_spin should match explicit when value is the same. + if cs_mode == "default_chg_spin": + dd_explicit = DescrptDPA3( + self.nt, + repflow=repflow, + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + add_chg_spin_ebd=True, + default_chg_spin=None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + dd_explicit.repflows.mean = torch.tensor( + davg, dtype=dtype, device=env.DEVICE + ) + dd_explicit.repflows.stddev = torch.tensor( + dstd, dtype=dtype, device=env.DEVICE + ) + cs = torch.tensor([[5, 1]], dtype=dtype, device=env.DEVICE).expand( + nf, -1 + ) + rd_explicit, _, _, _, _ = dd_explicit( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + torch.tensor(self.mapping, dtype=int, device=env.DEVICE), + charge_spin=cs, + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd_explicit.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + def test_jit( self, ) -> None: diff --git a/source/tests/pt_expt/descriptor/test_dpa3.py b/source/tests/pt_expt/descriptor/test_dpa3.py index ef4b479724..dc9594b49b 100644 --- a/source/tests/pt_expt/descriptor/test_dpa3.py +++ b/source/tests/pt_expt/descriptor/test_dpa3.py @@ -128,6 +128,89 @@ def test_consistency(self, ua, ruri, acr, acer, acus, nme) -> None: atol=atol, ) + @pytest.mark.parametrize("cs_mode", ["explicit_chg_spin", "default_chg_spin"]) + def test_consistency_chg_spin(self, cs_mode) -> None: + rng = np.random.default_rng(GLOBAL_SEED) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(rng.normal(size=(self.nt, nnei, 4))) + + prec = "float64" + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + atol = 1e-8 + + default_chg_spin = [5.0, 1.0] if cs_mode == "default_chg_spin" else None + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=8, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei - 1, + axis_neuron=4, + update_angle=True, + update_style="res_residual", + update_residual_init="const", + smooth_edge_update=True, + ) + + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + exclude_types=[], + precision=prec, + add_chg_spin_ebd=True, + default_chg_spin=default_chg_spin, + seed=GLOBAL_SEED, + ).to(self.device) + dd0.repflows.mean = torch.tensor(davg, dtype=dtype, device=self.device) + dd0.repflows.stddev = torch.tensor(dstd, dtype=dtype, device=self.device) + + # descriptor.forward does not apply default_chg_spin fallback; + # always pass an explicit charge_spin tensor here. + charge_spin = torch.tensor([[5, 1]], dtype=dtype, device=self.device).expand( + nf, -1 + ) + charge_spin_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) + + coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) + atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) + nlist_t = torch.tensor(self.nlist, dtype=int, device=self.device) + mapping_t = torch.tensor(self.mapping, dtype=int, device=self.device) + + rd0, _, _, _, _ = dd0( + coord_ext, atype_ext, nlist_t, mapping_t, charge_spin=charge_spin + ) + # serialization round-trip preserves default_chg_spin + dd1 = DescrptDPA3.deserialize(dd0.serialize()) + rd1, _, _, _, _ = dd1( + coord_ext, atype_ext, nlist_t, mapping_t, charge_spin=charge_spin + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + ) + # vs dpmodel + dd2 = DPDescrptDPA3.deserialize(dd0.serialize()) + rd2, _, _, _, _ = dd2.call( + self.coord_ext, + self.atype_ext, + self.nlist, + self.mapping, + charge_spin=charge_spin_np, + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), rd2, rtol=rtol, atol=atol + ) + @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision def test_exportable(self, prec) -> None: rng = np.random.default_rng(GLOBAL_SEED) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 6f5a337b4d..2b778d41c3 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -489,6 +489,7 @@ def DescriptorParamDPA3( precision="float64", use_loc_mapping=True, add_chg_spin_ebd=False, + default_chg_spin=None, ): input_dict = { # kwargs for repformer @@ -537,6 +538,7 @@ def DescriptorParamDPA3( "use_tebd_bias": False, "use_loc_mapping": use_loc_mapping, "add_chg_spin_ebd": add_chg_spin_ebd, + "default_chg_spin": default_chg_spin, "type_map": type_map, "seed": GLOBAL_SEED, } @@ -564,6 +566,7 @@ def DescriptorParamDPA3( "precision": ("float64",), "use_loc_mapping": (True, False), "add_chg_spin_ebd": (False, True), + "default_chg_spin": (None, [5.0, 1.0]), } ), ) diff --git a/source/tests/universal/dpmodel/model/test_model.py b/source/tests/universal/dpmodel/model/test_model.py index b5c6bd82ee..65ac116807 100644 --- a/source/tests/universal/dpmodel/model/test_model.py +++ b/source/tests/universal/dpmodel/model/test_model.py @@ -63,13 +63,10 @@ def skip_model_tests(test_obj): if test_obj.input_dict_ds.get("add_chg_spin_ebd", False): - import inspect - - (FittingParam, _) = test_obj.param[1] - sig = inspect.signature(FittingParam) - numb_param = sig.parameters.get("numb_param") - if numb_param is None or numb_param.default != 2: - return True, "add_chg_spin_ebd requires numb_fparam=2" + # The universal model driver does not feed `charge_spin` directly; + # rely on `default_chg_spin` fallback inside dp_atomic_model. + if test_obj.input_dict_ds.get("default_chg_spin") is None: + return True, "add_chg_spin_ebd requires default_chg_spin in universal tests" if not test_obj.input_dict_ds.get( "smooth_type_embedding", True ) or not test_obj.input_dict_ds.get("smooth", True): From 783419bb43263d5676a9c2139a6cdbcdc1319922 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 21:40:51 +0800 Subject: [PATCH 02/24] fix diff --- deepmd/pt_expt/model/dipole_model.py | 350 ++++++++++++------------ deepmd/pt_expt/model/dos_model.py | 310 ++++++++++----------- deepmd/pt_expt/model/dp_zbl_model.py | 358 ++++++++++++------------- deepmd/pt_expt/model/polar_model.py | 310 ++++++++++----------- deepmd/pt_expt/model/property_model.py | 324 +++++++++++----------- 5 files changed, 826 insertions(+), 826 deletions(-) diff --git a/deepmd/pt_expt/model/dipole_model.py b/deepmd/pt_expt/model/dipole_model.py index 679966d3d4..a4598d4020 100644 --- a/deepmd/pt_expt/model/dipole_model.py +++ b/deepmd/pt_expt/model/dipole_model.py @@ -1,175 +1,175 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPDipoleAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("dipole") -class DipoleModel(DPModelCommon, DPDipoleModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPDipoleModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["dipole"] = model_ret["dipole"] - model_predict["global_dipole"] = model_ret["dipole_redu"] - if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: - model_predict["force"] = model_ret["dipole_derv_r"] - if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: - model_predict["virial"] = model_ret["dipole_derv_c_redu"] - if do_atomic_virial and model_ret["dipole_derv_c"] is not None: - model_predict["atom_virial"] = model_ret["dipole_derv_c"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["dipole"] = model_ret["dipole"] - model_predict["global_dipole"] = model_ret["dipole_redu"] - if self.do_grad_r("dipole") and model_ret.get("dipole_derv_r") is not None: - model_predict["extended_force"] = model_ret["dipole_derv_r"] - if self.do_grad_c("dipole") and model_ret.get("dipole_derv_c_redu") is not None: - model_predict["virial"] = model_ret["dipole_derv_c_redu"] - if do_atomic_virial and model_ret.get("dipole_derv_c") is not None: - model_predict["extended_virial"] = model_ret["dipole_derv_c"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "dipole": out_def_data["dipole"], - "global_dipole": out_def_data["dipole_redu"], - } - if self.do_grad_r("dipole"): - output_def["force"] = out_def_data["dipole_derv_r"] - output_def["force"].squeeze(-2) - if self.do_grad_c("dipole"): - output_def["virial"] = out_def_data["dipole_derv_c_redu"] - output_def["virial"].squeeze(-2) - output_def["atom_virial"] = out_def_data["dipole_derv_c"] - output_def["atom_virial"].squeeze(-2) - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, - extended_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDipoleAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPDipoleModel_ = make_model(DPDipoleAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dipole") +class DipoleModel(DPModelCommon, DPDipoleModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDipoleModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: + model_predict["force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret["dipole_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["dipole"] = model_ret["dipole"] + model_predict["global_dipole"] = model_ret["dipole_redu"] + if self.do_grad_r("dipole") and model_ret.get("dipole_derv_r") is not None: + model_predict["extended_force"] = model_ret["dipole_derv_r"] + if self.do_grad_c("dipole") and model_ret.get("dipole_derv_c_redu") is not None: + model_predict["virial"] = model_ret["dipole_derv_c_redu"] + if do_atomic_virial and model_ret.get("dipole_derv_c") is not None: + model_predict["extended_virial"] = model_ret["dipole_derv_c"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "dipole": out_def_data["dipole"], + "global_dipole": out_def_data["dipole_redu"], + } + if self.do_grad_r("dipole"): + output_def["force"] = out_def_data["dipole_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("dipole"): + output_def["virial"] = out_def_data["dipole_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["dipole_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/dos_model.py b/deepmd/pt_expt/model/dos_model.py index d74676b6ea..ed961efe37 100644 --- a/deepmd/pt_expt/model/dos_model.py +++ b/deepmd/pt_expt/model/dos_model.py @@ -1,155 +1,155 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPDOSAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("dos") -class DOSModel(DPModelCommon, DPDOSModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPDOSModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_dos"] = model_ret["dos"] - model_predict["dos"] = model_ret["dos_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_dos"] = model_ret["dos"] - model_predict["dos"] = model_ret["dos_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "atom_dos": out_def_data["dos"], - "dos": out_def_data["dos_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, - extended_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPDOSAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPDOSModel_ = make_model(DPDOSAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("dos") +class DOSModel(DPModelCommon, DPDOSModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_dos"] = model_ret["dos"] + model_predict["dos"] = model_ret["dos_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_dos": out_def_data["dos"], + "dos": out_def_data["dos_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/dp_zbl_model.py b/deepmd/pt_expt/model/dp_zbl_model.py index be7ffcf70e..9a65346782 100644 --- a/deepmd/pt_expt/model/dp_zbl_model.py +++ b/deepmd/pt_expt/model/dp_zbl_model.py @@ -1,179 +1,179 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - DPZBLLinearEnergyAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("zbl") -class DPZBLModel(DPModelCommon, DPZBLModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPZBLModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial and model_ret["energy_derv_c"] is not None: - model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy") and model_ret.get("energy_derv_c_redu") is not None: - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial and model_ret.get("energy_derv_c") is not None: - model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( - -2 - ) - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "atom_energy": out_def_data["energy"], - "energy": out_def_data["energy_redu"], - } - if self.do_grad_r("energy"): - output_def["force"] = out_def_data["energy_derv_r"] - output_def["force"].squeeze(-2) - if self.do_grad_c("energy"): - output_def["virial"] = out_def_data["energy_derv_c_redu"] - output_def["virial"].squeeze(-2) - output_def["atom_virial"] = out_def_data["energy_derv_c"] - output_def["atom_virial"].squeeze(-2) - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - - # Force `_format_nlist`'s sort branch into the compiled graph so the - # exported model tolerates oversized nlists at runtime — see - # make_model.py for the full rationale. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, - extended_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + DPZBLLinearEnergyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("zbl") +class DPZBLModel(DPModelCommon, DPZBLModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPZBLModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret["energy_derv_c"] is not None: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy") and model_ret.get("energy_derv_c_redu") is not None: + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial and model_ret.get("energy_derv_c") is not None: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -2 + ) + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "atom_energy": out_def_data["energy"], + "energy": out_def_data["energy_redu"], + } + if self.do_grad_r("energy"): + output_def["force"] = out_def_data["energy_derv_r"] + output_def["force"].squeeze(-2) + if self.do_grad_c("energy"): + output_def["virial"] = out_def_data["energy_derv_c_redu"] + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = out_def_data["energy_derv_c"] + output_def["atom_virial"].squeeze(-2) + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # Force `_format_nlist`'s sort branch into the compiled graph so the + # exported model tolerates oversized nlists at runtime — see + # make_model.py for the full rationale. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/polar_model.py b/deepmd/pt_expt/model/polar_model.py index cafddf80b0..75beb31afb 100644 --- a/deepmd/pt_expt/model/polar_model.py +++ b/deepmd/pt_expt/model/polar_model.py @@ -1,155 +1,155 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPPolarAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("polar") -class PolarModel(DPModelCommon, DPPolarModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPPolarModel_.__init__(self, *args, **kwargs) - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["polar"] = model_ret["polarizability"] - model_predict["global_polar"] = model_ret["polarizability_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - model_predict = {} - model_predict["polar"] = model_ret["polarizability"] - model_predict["global_polar"] = model_ret["polarizability_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - output_def = { - "polar": out_def_data["polarizability"], - "global_polar": out_def_data["polarizability_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, - extended_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPolarAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPPolarModel_ = make_model(DPPolarAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("polar") +class PolarModel(DPModelCommon, DPPolarModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPolarModel_.__init__(self, *args, **kwargs) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["polar"] = model_ret["polarizability"] + model_predict["global_polar"] = model_ret["polarizability_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "polar": out_def_data["polarizability"], + "global_polar": out_def_data["polarizability_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced diff --git a/deepmd/pt_expt/model/property_model.py b/deepmd/pt_expt/model/property_model.py index db809da131..629f6042ad 100644 --- a/deepmd/pt_expt/model/property_model.py +++ b/deepmd/pt_expt/model/property_model.py @@ -1,162 +1,162 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import types -from typing import ( - Any, -) - -import torch -from torch.fx.experimental.proxy_tensor import ( - make_fx, -) - -from deepmd.dpmodel.atomic_model import ( - DPPropertyAtomicModel, -) -from deepmd.dpmodel.model.dp_model import ( - DPModelCommon, -) - -from .make_model import ( - _pad_nlist_for_export, - make_model, -) -from .model import ( - BaseModel, -) - -DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(BaseModel,)) - - -@BaseModel.register("property") -class PropertyModel(DPModelCommon, DPPropertyModel_): - def __init__( - self, - *args: Any, - **kwargs: Any, - ) -> None: - DPModelCommon.__init__(self) - DPPropertyModel_.__init__(self, *args, **kwargs) - - def get_var_name(self) -> str: - """Get the name of the property.""" - return self.get_fitting_net().var_name - - def forward( - self, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - var_name = self.get_var_name() - model_predict = {} - model_predict[f"atom_{var_name}"] = model_ret[var_name] - model_predict[var_name] = model_ret[f"{var_name}_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def forward_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - ) -> dict[str, torch.Tensor]: - model_ret = self.call_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - var_name = self.get_var_name() - model_predict = {} - model_predict[f"atom_{var_name}"] = model_ret[var_name] - model_predict[var_name] = model_ret[f"{var_name}_redu"] - if "mask" in model_ret: - model_predict["mask"] = model_ret["mask"] - return model_predict - - def translated_output_def(self) -> dict[str, Any]: - out_def_data = self.model_output_def().get_data() - var_name = self.get_var_name() - output_def = { - f"atom_{var_name}": out_def_data[var_name], - var_name: out_def_data[f"{var_name}_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - - def forward_lower_exportable( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - do_atomic_virial: bool = False, - **make_fx_kwargs: Any, - ) -> torch.nn.Module: - model = self - - def fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None, - fparam: torch.Tensor | None, - aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None, - ) -> dict[str, torch.Tensor]: - extended_coord = extended_coord.detach().requires_grad_(True) - nlist = _pad_nlist_for_export(nlist) - return model.forward_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - do_atomic_virial=do_atomic_virial, - ) - - # See make_model.py for the rationale of the pad + monkeypatch. - _orig_need_sort = model.need_sorted_nlist_for_lower - model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) - try: - traced = make_fx(fn, **make_fx_kwargs)( - extended_coord, - extended_atype, - nlist, - mapping, - fparam, - aparam, - charge_spin, - ) - finally: - model.need_sorted_nlist_for_lower = _orig_need_sort - return traced +# SPDX-License-Identifier: LGPL-3.0-or-later +import types +from typing import ( + Any, +) + +import torch +from torch.fx.experimental.proxy_tensor import ( + make_fx, +) + +from deepmd.dpmodel.atomic_model import ( + DPPropertyAtomicModel, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) + +from .make_model import ( + _pad_nlist_for_export, + make_model, +) +from .model import ( + BaseModel, +) + +DPPropertyModel_ = make_model(DPPropertyAtomicModel, T_Bases=(BaseModel,)) + + +@BaseModel.register("property") +class PropertyModel(DPModelCommon, DPPropertyModel_): + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPPropertyModel_.__init__(self, *args, **kwargs) + + def get_var_name(self) -> str: + """Get the name of the property.""" + return self.get_fitting_net().var_name + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, torch.Tensor]: + model_ret = self.call_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + var_name = self.get_var_name() + model_predict = {} + model_predict[f"atom_{var_name}"] = model_ret[var_name] + model_predict[var_name] = model_ret[f"{var_name}_redu"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + var_name = self.get_var_name() + output_def = { + f"atom_{var_name}": out_def_data[var_name], + var_name: out_def_data[f"{var_name}_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_def + + def forward_lower_exportable( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + model = self + + def fn( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + extended_coord = extended_coord.detach().requires_grad_(True) + nlist = _pad_nlist_for_export(nlist) + return model.forward_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + do_atomic_virial=do_atomic_virial, + ) + + # See make_model.py for the rationale of the pad + monkeypatch. + _orig_need_sort = model.need_sorted_nlist_for_lower + model.need_sorted_nlist_for_lower = types.MethodType(lambda self: True, model) + try: + traced = make_fx(fn, **make_fx_kwargs)( + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + charge_spin, + ) + finally: + model.need_sorted_nlist_for_lower = _orig_need_sort + return traced From df501d19f67e33bb2835acec44418c54566df967 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 22:22:28 +0800 Subject: [PATCH 03/24] fix comment --- deepmd/dpmodel/descriptor/dpa3.py | 4 +- .../descriptor/make_base_descriptor.py | 12 ++++ deepmd/entrypoints/test.py | 2 +- deepmd/pt/model/descriptor/dpa1.py | 12 ++++ deepmd/pt/model/descriptor/dpa2.py | 12 ++++ deepmd/pt/model/descriptor/dpa3.py | 7 +- deepmd/pt/model/descriptor/hybrid.py | 12 ++++ deepmd/pt/model/descriptor/se_a.py | 12 ++++ deepmd/pt/model/descriptor/se_r.py | 12 ++++ deepmd/pt/model/descriptor/se_t.py | 12 ++++ deepmd/pt/model/descriptor/se_t_tebd.py | 12 ++++ deepmd/pt/train/training.py | 5 +- deepmd/pt_expt/descriptor/dpa1.py | 1 + deepmd/pt_expt/descriptor/dpa2.py | 1 + deepmd/pt_expt/descriptor/se_e2_a.py | 1 + deepmd/pt_expt/descriptor/se_r.py | 1 + deepmd/pt_expt/descriptor/se_t.py | 1 + deepmd/pt_expt/descriptor/se_t_tebd.py | 1 + deepmd/pt_expt/train/training.py | 5 +- deepmd/utils/argcheck.py | 35 ++++++++++ source/tests/consistent/descriptor/common.py | 66 ++++++++----------- 21 files changed, 172 insertions(+), 54 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 8e4a4ba5ad..1cce82baf4 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -416,8 +416,8 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.use_econf_tebd = use_econf_tebd self.add_chg_spin_ebd = add_chg_spin_ebd self.default_chg_spin = default_chg_spin - if self.add_chg_spin_ebd and self.default_chg_spin is not None: - assert len(self.default_chg_spin) == 2, ( + if self.default_chg_spin is not None and len(self.default_chg_spin) != 2: + raise ValueError( "default_chg_spin must have exactly 2 values [charge, spin]" ) self.use_tebd_bias = use_tebd_bias diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index c03948be54..5ca4686619 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -96,6 +96,18 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension of g2.""" pass + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> Any: + """Returns the default charge_spin value, or None.""" + return None + @abstractmethod def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index f589f25989..8f30d0c30c 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -639,7 +639,7 @@ def test_ener( aparam = test_data["aparam"][:numb_test] else: aparam = None - if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 1.0) != 0.0: + if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0: charge_spin = test_data["charge_spin"][:numb_test] else: charge_spin = None diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index ba96b05d42..9751db621c 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -325,6 +325,18 @@ def __init__( for param in self.parameters(): param.requires_grad = trainable + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_atten.get_rcut() diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index cf88ec7de2..d6e38999b7 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -330,6 +330,18 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: param.requires_grad = trainable self.compress = False + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index 1befcbdea0..8a78ea4860 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -4,6 +4,7 @@ ) from typing import ( Any, + Optional, ) import torch @@ -125,8 +126,8 @@ def __init__( default_chg_spin: list[float] | None = None, ) -> None: super().__init__() - if default_chg_spin is not None: - assert len(default_chg_spin) == 2, ( + if default_chg_spin is not None and len(default_chg_spin) != 2: + raise ValueError( "default_chg_spin must be a list of length 2 [charge, spin]." ) @@ -266,7 +267,7 @@ def has_default_chg_spin(self) -> bool: return self.default_chg_spin is not None @torch.jit.export - def get_default_chg_spin(self) -> torch.Tensor | None: + def get_default_chg_spin(self) -> Optional[torch.Tensor]: # noqa: UP045 """Get the default charge_spin values as a tensor.""" if self.default_chg_spin is None: return None diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index e9a3258ba5..d79678e668 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -99,6 +99,18 @@ def __init__( ).astype(np.int64) self.nlist_cut_idx.append(to_torch_tensor(cut_idx)) + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" # do not use numpy here - jit is not happy diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 8bf1308e30..d840c8c001 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -137,6 +137,18 @@ def __init__( seed=seed, ) + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.sea.get_rcut() diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index f824a0c44c..654e2e16bc 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -167,6 +167,18 @@ def __init__( ] ) + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index c518740de2..bd2a9a60cd 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -171,6 +171,18 @@ def __init__( seed=seed, ) + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.seat.get_rcut() diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 9c38cc2fff..6937bb99e8 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -210,6 +210,18 @@ def __init__( for param in self.parameters(): param.requires_grad = trainable + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_ttebd.get_rcut() diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index c9325b1030..164261be2a 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1906,10 +1906,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: DataRequirementItem("spin", ndof=3, atomic=True, must=True) ] additional_data_requirement += spin_requirement_items - has_chg_spin_ebd = getattr(_model, "has_chg_spin_ebd", False) - if callable(has_chg_spin_ebd): - has_chg_spin_ebd = has_chg_spin_ebd() - if has_chg_spin_ebd: + if _model.has_chg_spin_ebd(): has_default_cs = _model.has_default_chg_spin() cs_default = ( _model.get_default_chg_spin().cpu().numpy() if has_default_cs else 0.0 diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 01df91abd6..792b6ad06b 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -183,6 +183,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptDPA1DP.call.__wrapped__( diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index 1723df5a30..c0193b1eb6 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -233,6 +233,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptDPA2DP.call.__wrapped__( diff --git a/deepmd/pt_expt/descriptor/se_e2_a.py b/deepmd/pt_expt/descriptor/se_e2_a.py index 61d611036e..11939c1499 100644 --- a/deepmd/pt_expt/descriptor/se_e2_a.py +++ b/deepmd/pt_expt/descriptor/se_e2_a.py @@ -139,6 +139,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptSeADP.call.__wrapped__( diff --git a/deepmd/pt_expt/descriptor/se_r.py b/deepmd/pt_expt/descriptor/se_r.py index 22302f54e6..5879084e39 100644 --- a/deepmd/pt_expt/descriptor/se_r.py +++ b/deepmd/pt_expt/descriptor/se_r.py @@ -128,6 +128,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptSeRDP.call.__wrapped__( diff --git a/deepmd/pt_expt/descriptor/se_t.py b/deepmd/pt_expt/descriptor/se_t.py index 061306f281..e3e69f7245 100644 --- a/deepmd/pt_expt/descriptor/se_t.py +++ b/deepmd/pt_expt/descriptor/se_t.py @@ -139,6 +139,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptSeTDP.call.__wrapped__( diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py index c0ae308971..34b7c5bde6 100644 --- a/deepmd/pt_expt/descriptor/se_t_tebd.py +++ b/deepmd/pt_expt/descriptor/se_t_tebd.py @@ -166,6 +166,7 @@ def call( nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> Any: if not self.compress: return DescrptSeTTebdDP.call.__wrapped__( diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 039981c4eb..ea015dc487 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -138,10 +138,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: "aparam", _model.get_dim_aparam(), atomic=True, must=True ) ) - has_chg_spin_ebd = getattr(_model, "has_chg_spin_ebd", False) - if callable(has_chg_spin_ebd): - has_chg_spin_ebd = has_chg_spin_ebd() - if has_chg_spin_ebd: + if _model.has_chg_spin_ebd(): has_default_cs = _model.has_default_chg_spin() if has_default_cs: default_cs = _model.get_default_chg_spin() diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 1232026a76..41cf33136a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4541,6 +4541,40 @@ def gen_json_schema(multi_task: bool = False) -> str: return json.dumps(generate_json_schema(arg)) +def _check_dpa3_chg_spin_migration(data: dict[str, Any]) -> None: + """Detect legacy DPA3 charge/spin configs that packed charge_spin into fparam. + + Prior to the charge_spin decoupling, DPA3 models with add_chg_spin_ebd=True + required numb_fparam=2 on the fitting net (and optionally default_fparam). + After the decoupling, charge/spin is a first-class input and fparam is no + longer used for that purpose. Raise a clear error so users update configs. + """ + model = data.get("model", {}) if isinstance(data, dict) else {} + if not isinstance(model, dict): + return + submodels = ( + [model] if "descriptor" in model else list(model.get("model_dict", {}).values()) + ) + for m in submodels: + if not isinstance(m, dict): + continue + desc = m.get("descriptor", {}) + fitting = m.get("fitting_net", {}) + if not isinstance(desc, dict) or not isinstance(fitting, dict): + continue + if desc.get("type") != "dpa3": + continue + if not desc.get("add_chg_spin_ebd", False): + continue + if fitting.get("numb_fparam", 0) or fitting.get("default_fparam") is not None: + raise ValueError( + "DPA3 `add_chg_spin_ebd=True` no longer uses `fparam` for " + "charge/spin. Remove `numb_fparam`/`default_fparam` from " + "`fitting_net` and provide charge/spin via the new " + "`charge_spin` input or the descriptor's `default_chg_spin`." + ) + + def normalize( data: dict[str, Any], multi_task: bool = False, *, check: bool = True ) -> dict[str, Any]: @@ -4550,6 +4584,7 @@ def normalize( if check: base.check_value(data, strict=True) validate_full_validation_config(data, multi_task=multi_task) + _check_dpa3_chg_spin_migration(data) return data diff --git a/source/tests/consistent/descriptor/common.py b/source/tests/consistent/descriptor/common.py index 16f6e0e1ca..078db4829f 100644 --- a/source/tests/consistent/descriptor/common.py +++ b/source/tests/consistent/descriptor/common.py @@ -119,14 +119,10 @@ def eval_dp_descriptor( dp_obj.get_sel(), distinguish_types=(not mixed_types), ) - return dp_obj( - ext_coords, - ext_atype, - nlist=nlist, - mapping=mapping, - fparam=fparam, - charge_spin=charge_spin, - ) + kwargs = {"nlist": nlist, "mapping": mapping, "fparam": fparam} + if hasattr(dp_obj, "get_dim_chg_spin") and dp_obj.get_dim_chg_spin() > 0: + kwargs["charge_spin"] = charge_spin + return dp_obj(ext_coords, ext_atype, **kwargs) def eval_pt_descriptor( self, @@ -161,16 +157,12 @@ def eval_pt_descriptor( if charge_spin is not None else None ) + kwargs = {"nlist": nlist, "mapping": mapping, "fparam": fparam_pt} + if hasattr(pt_obj, "get_dim_chg_spin") and pt_obj.get_dim_chg_spin() > 0: + kwargs["charge_spin"] = charge_spin_pt return [ x.detach().cpu().numpy() if torch.is_tensor(x) else x - for x in pt_obj( - ext_coords, - ext_atype, - nlist=nlist, - mapping=mapping, - fparam=fparam_pt, - charge_spin=charge_spin_pt, - ) + for x in pt_obj(ext_coords, ext_atype, **kwargs) ] def eval_pt_expt_descriptor( @@ -206,16 +198,15 @@ def eval_pt_expt_descriptor( if charge_spin is not None else None ) + kwargs = {"nlist": nlist, "mapping": mapping, "fparam": fparam_pt} + if ( + hasattr(pt_expt_obj, "get_dim_chg_spin") + and pt_expt_obj.get_dim_chg_spin() > 0 + ): + kwargs["charge_spin"] = charge_spin_pt return [ x.detach().cpu().numpy() if torch.is_tensor(x) else x - for x in pt_expt_obj( - ext_coords, - ext_atype, - nlist=nlist, - mapping=mapping, - fparam=fparam_pt, - charge_spin=charge_spin_pt, - ) + for x in pt_expt_obj(ext_coords, ext_atype, **kwargs) ] def eval_jax_descriptor( @@ -245,16 +236,12 @@ def eval_jax_descriptor( ) fparam_jax = jnp.array(fparam) if fparam is not None else None charge_spin_jax = jnp.array(charge_spin) if charge_spin is not None else None + kwargs = {"nlist": nlist, "mapping": mapping, "fparam": fparam_jax} + if hasattr(jax_obj, "get_dim_chg_spin") and jax_obj.get_dim_chg_spin() > 0: + kwargs["charge_spin"] = charge_spin_jax return [ np.asarray(x) if isinstance(x, jnp.ndarray) else x - for x in jax_obj( - ext_coords, - ext_atype, - nlist=nlist, - mapping=mapping, - fparam=fparam_jax, - charge_spin=charge_spin_jax, - ) + for x in jax_obj(ext_coords, ext_atype, **kwargs) ] def eval_pd_descriptor( @@ -323,16 +310,15 @@ def eval_array_api_strict_descriptor( charge_spin_array_api = ( array_api_strict.asarray(charge_spin) if charge_spin is not None else None ) + kwargs = {"nlist": nlist, "mapping": mapping, "fparam": fparam_array_api} + if ( + hasattr(array_api_strict_obj, "get_dim_chg_spin") + and array_api_strict_obj.get_dim_chg_spin() > 0 + ): + kwargs["charge_spin"] = charge_spin_array_api return [ to_numpy_array(x) if hasattr(x, "__array_namespace__") else x - for x in array_api_strict_obj( - ext_coords, - ext_atype, - nlist=nlist, - mapping=mapping, - fparam=fparam_array_api, - charge_spin=charge_spin_array_api, - ) + for x in array_api_strict_obj(ext_coords, ext_atype, **kwargs) ] From 621641d51ad9e498d36bb971145abfa9382da786 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 22:32:04 +0800 Subject: [PATCH 04/24] fix ut --- deepmd/dpmodel/descriptor/dpa3.py | 4 ++-- deepmd/jax/model/base_model.py | 4 ++++ deepmd/jax/model/dp_model.py | 2 ++ deepmd/jax/model/dp_zbl_model.py | 2 ++ 4 files changed, 10 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 1cce82baf4..add6da597c 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -415,11 +415,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.use_econf_tebd = use_econf_tebd self.add_chg_spin_ebd = add_chg_spin_ebd - self.default_chg_spin = default_chg_spin - if self.default_chg_spin is not None and len(self.default_chg_spin) != 2: + if default_chg_spin is not None and len(default_chg_spin) != 2: raise ValueError( "default_chg_spin must have exactly 2 values [charge, spin]" ) + self.default_chg_spin = default_chg_spin self.use_tebd_bias = use_tebd_bias self.use_loc_mapping = use_loc_mapping self.type_map = type_map diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py index 4522e25586..a9be0e2f16 100644 --- a/deepmd/jax/model/base_model.py +++ b/deepmd/jax/model/base_model.py @@ -26,6 +26,7 @@ def forward_common_atomic( aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, extended_coord_corr: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, @@ -34,6 +35,7 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) atomic_output_def = self.atomic_output_def() model_predict = {} @@ -75,6 +77,7 @@ def eval_output( mapping=mapping[None, ...] if mapping is not None else None, fparam=fparam[None, ...] if fparam is not None else None, aparam=aparam[None, ...] if aparam is not None else None, + charge_spin=charge_spin, ) return jnp.sum(atomic_ret[_kk][0], axis=_atom_axis) @@ -137,6 +140,7 @@ def eval_ce( mapping=mapping[None, ...] if mapping is not None else None, fparam=fparam[None, ...] if fparam is not None else None, aparam=aparam[None, ...] if aparam is not None else None, + charge_spin=charge_spin, ) nloc = nlist.shape[0] cc_loc = jax.lax.stop_gradient(cc_ext)[:nloc, ...] diff --git a/deepmd/jax/model/dp_model.py b/deepmd/jax/model/dp_model.py index 3e96eb6689..c4c2a2fafe 100644 --- a/deepmd/jax/model/dp_model.py +++ b/deepmd/jax/model/dp_model.py @@ -56,6 +56,7 @@ def forward_common_atomic( aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, extended_coord_corr: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return forward_common_atomic( self, @@ -67,6 +68,7 @@ def forward_common_atomic( aparam=aparam, do_atomic_virial=do_atomic_virial, extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, ) def format_nlist( diff --git a/deepmd/jax/model/dp_zbl_model.py b/deepmd/jax/model/dp_zbl_model.py index 7751d22a1f..cd1bbe342d 100644 --- a/deepmd/jax/model/dp_zbl_model.py +++ b/deepmd/jax/model/dp_zbl_model.py @@ -38,6 +38,7 @@ def forward_common_atomic( aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, extended_coord_corr: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return forward_common_atomic( self, @@ -49,6 +50,7 @@ def forward_common_atomic( aparam=aparam, do_atomic_virial=do_atomic_virial, extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, ) def format_nlist( From 8e5899b0c07c8f8cea8f3dfff0a68cb6a7d469d1 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 22:43:17 +0800 Subject: [PATCH 05/24] Update dp_atomic_model.py --- deepmd/jax/atomic_model/dp_atomic_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index 7227839f1f..8ee1403825 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -57,6 +57,7 @@ def forward_common_atomic( mapping: jnp.ndarray | None = None, fparam: jnp.ndarray | None = None, aparam: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, @@ -65,6 +66,7 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) return jax_atomic_model From adc550f23bfbfcd444dad08f25aab62e33f3d611 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 23:06:24 +0800 Subject: [PATCH 06/24] Update argcheck.py --- deepmd/utils/argcheck.py | 37 +++++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 41cf33136a..edec6be473 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4542,12 +4542,21 @@ def gen_json_schema(multi_task: bool = False) -> str: def _check_dpa3_chg_spin_migration(data: dict[str, Any]) -> None: - """Detect legacy DPA3 charge/spin configs that packed charge_spin into fparam. - - Prior to the charge_spin decoupling, DPA3 models with add_chg_spin_ebd=True - required numb_fparam=2 on the fitting net (and optionally default_fparam). - After the decoupling, charge/spin is a first-class input and fparam is no - longer used for that purpose. Raise a clear error so users update configs. + """Warn on likely legacy DPA3 configs that packed charge/spin into fparam. + + Before the charge_spin decoupling, enabling ``add_chg_spin_ebd`` on DPA3 + required ``numb_fparam=2`` on the fitting net so that charge/spin could be + carried via ``fparam``. After the decoupling, ``charge_spin`` is a + first-class input that is fully independent of ``fparam``, so users may + legitimately combine ``add_chg_spin_ebd`` with any ``numb_fparam`` for + genuine frame parameters. + + We cannot determine from the config alone whether a user's ``numb_fparam`` + is legacy (charge/spin in disguise) or genuine (real frame parameters). + But the combination ``add_chg_spin_ebd=True`` together with + ``numb_fparam=2`` is the strongest heuristic for the legacy pattern, since + that is exactly what the old code required. Emit a warning — not an error + — so users can audit their setup without breaking legitimate combinations. """ model = data.get("model", {}) if isinstance(data, dict) else {} if not isinstance(model, dict): @@ -4566,12 +4575,16 @@ def _check_dpa3_chg_spin_migration(data: dict[str, Any]) -> None: continue if not desc.get("add_chg_spin_ebd", False): continue - if fitting.get("numb_fparam", 0) or fitting.get("default_fparam") is not None: - raise ValueError( - "DPA3 `add_chg_spin_ebd=True` no longer uses `fparam` for " - "charge/spin. Remove `numb_fparam`/`default_fparam` from " - "`fitting_net` and provide charge/spin via the new " - "`charge_spin` input or the descriptor's `default_chg_spin`." + if fitting.get("numb_fparam", 0) == 2: + warnings.warn( + "DPA3 `add_chg_spin_ebd=True` with `numb_fparam=2` matches the " + "pre-decoupling pattern where charge/spin was carried via " + "`fparam`. `charge_spin` is now an independent input, so " + "`numb_fparam=2` will be treated as two genuine frame " + "parameters. If you intended to feed charge/spin, remove the " + "charge/spin part of `fparam` and use the `charge_spin` input " + "or the descriptor's `default_chg_spin` instead.", + stacklevel=2, ) From 34ad36dc680bd88d2dc07f2851fec2c3a80de6f0 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 7 May 2026 23:22:19 +0800 Subject: [PATCH 07/24] fix ut --- deepmd/pt_expt/model/spin_model.py | 4 ++ deepmd/pt_expt/utils/serialization.py | 72 ++++++++++++++++++++++----- 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/deepmd/pt_expt/model/spin_model.py b/deepmd/pt_expt/model/spin_model.py index e69ee29f5a..b7ab2c5f24 100644 --- a/deepmd/pt_expt/model/spin_model.py +++ b/deepmd/pt_expt/model/spin_model.py @@ -57,6 +57,7 @@ def forward_common_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, **make_fx_kwargs: Any, ) -> torch.nn.Module: @@ -96,6 +97,7 @@ def fn( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, ) -> dict[str, torch.Tensor]: extended_coord = extended_coord.detach().requires_grad_(True) nlist = _pad_nlist_for_export(nlist) @@ -107,6 +109,7 @@ def fn( mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) @@ -130,6 +133,7 @@ def fn( mapping, fparam, aparam, + charge_spin, ) finally: backbone.need_sorted_nlist_for_lower = _orig_need_sort diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 7b2559db4f..9448237333 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -119,8 +119,9 @@ def _make_sample_inputs( Returns ------- tuple - (ext_coord, ext_atype, nlist, mapping, fparam, aparam) or - (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) when has_spin. + (ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin) or + (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam, + charge_spin) when has_spin. """ rcut = model.get_rcut() sel = model.get_sel() @@ -187,14 +188,31 @@ def _make_sample_inputs( else: aparam = None + dim_chg_spin = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + if dim_chg_spin > 0: + charge_spin = torch.zeros( + nframes, dim_chg_spin, dtype=torch.float64, device=_env.DEVICE + ) + else: + charge_spin = None + if has_spin: nall = extended_coord.shape[1] ext_spin = torch.zeros( nframes, nall, 3, dtype=torch.float64, device=_env.DEVICE ) - return ext_coord, ext_atype, ext_spin, nlist_t, mapping_t, fparam, aparam + return ( + ext_coord, + ext_atype, + ext_spin, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, + ) - return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin def _build_dynamic_shapes( @@ -224,9 +242,10 @@ def _build_dynamic_shapes( nnei_dim = torch.export.Dim("nnei", min=max(1, model_nnei)) if has_spin: - # (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) + # (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam, charge_spin) fparam = sample_inputs[5] aparam = sample_inputs[6] + charge_spin = sample_inputs[7] return ( {0: nframes_dim, 1: nall_dim}, # extended_coord: (nframes, nall, 3) {0: nframes_dim, 1: nall_dim}, # extended_atype: (nframes, nall) @@ -239,11 +258,13 @@ def _build_dynamic_shapes( {0: nframes_dim, 1: nall_dim}, # mapping: (nframes, nall) {0: nframes_dim} if fparam is not None else None, # fparam {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, # aparam + {0: nframes_dim} if charge_spin is not None else None, # charge_spin ) else: - # (ext_coord, ext_atype, nlist, mapping, fparam, aparam) + # (ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin) fparam = sample_inputs[4] aparam = sample_inputs[5] + charge_spin = sample_inputs[6] return ( {0: nframes_dim, 1: nall_dim}, # extended_coord: (nframes, nall, 3) {0: nframes_dim, 1: nall_dim}, # extended_atype: (nframes, nall) @@ -255,6 +276,7 @@ def _build_dynamic_shapes( {0: nframes_dim, 1: nall_dim}, # mapping: (nframes, nall) {0: nframes_dim} if fparam is not None else None, # fparam {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, # aparam + {0: nframes_dim} if charge_spin is not None else None, # charge_spin ) @@ -487,11 +509,26 @@ def _trace_and_export( _env.DEVICE = _orig_device if is_spin: - ext_coord, ext_atype, ext_spin, nlist_t, mapping_t, fparam, aparam = ( - sample_inputs - ) + ( + ext_coord, + ext_atype, + ext_spin, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, + ) = sample_inputs else: - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = sample_inputs + ( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, + ) = sample_inputs # 4. Trace via make_fx on CPU. # This decomposes torch.autograd.grad into aten ops so the resulting @@ -505,13 +542,21 @@ def _trace_and_export( mapping_t, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) # 5. Extract output keys from the CPU-traced module. sample_out = traced( - ext_coord, ext_atype, ext_spin, nlist_t, mapping_t, fparam, aparam + ext_coord, + ext_atype, + ext_spin, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, ) else: traced = model.forward_common_lower_exportable( @@ -521,12 +566,15 @@ def _trace_and_export( mapping_t, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) # 5. Extract output keys from the CPU-traced module. - sample_out = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + sample_out = traced( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) output_keys = list(sample_out.keys()) From c9101d1df2375dcbc57c72793deb87fa5254c819 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 8 May 2026 21:15:02 +0800 Subject: [PATCH 08/24] fix ut --- deepmd/dpmodel/model/dipole_model.py | 4 ++++ deepmd/dpmodel/model/dos_model.py | 4 ++++ deepmd/dpmodel/model/dp_zbl_model.py | 4 ++++ deepmd/dpmodel/model/polar_model.py | 4 ++++ deepmd/dpmodel/model/property_model.py | 4 ++++ deepmd/dpmodel/model/spin_model.py | 8 ++++++++ deepmd/pt/model/model/dipole_model.py | 4 ++++ deepmd/pt/model/model/dos_model.py | 4 ++++ deepmd/pt/model/model/dp_linear_model.py | 4 ++++ deepmd/pt/model/model/dp_zbl_model.py | 4 ++++ deepmd/pt/model/model/make_hessian_model.py | 2 ++ deepmd/pt/model/model/make_model.py | 2 ++ deepmd/pt/model/model/polar_model.py | 4 ++++ deepmd/pt/model/model/property_model.py | 4 ++++ deepmd/pt/model/model/spin_model.py | 8 ++++++++ 15 files changed, 64 insertions(+) diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index fa5a76e0af..9e85403e9f 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -44,6 +44,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -52,6 +53,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["dipole"] = model_ret["dipole"] @@ -75,6 +77,7 @@ def call_lower( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, @@ -84,6 +87,7 @@ def call_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["dipole"] = model_ret["dipole"] diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index b75c9a2bcc..4d854ae007 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -44,6 +44,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -52,6 +53,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["atom_dos"] = model_ret["dos"] @@ -69,6 +71,7 @@ def call_lower( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, @@ -78,6 +81,7 @@ def call_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["atom_dos"] = model_ret["dos"] diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index d864b3b61e..d389c0b7b2 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -46,6 +46,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -54,6 +55,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["atom_energy"] = model_ret["energy"] @@ -77,6 +79,7 @@ def call_lower( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, @@ -86,6 +89,7 @@ def call_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["atom_energy"] = model_ret["energy"] diff --git a/deepmd/dpmodel/model/polar_model.py b/deepmd/dpmodel/model/polar_model.py index 5031166a5e..83c4b62f01 100644 --- a/deepmd/dpmodel/model/polar_model.py +++ b/deepmd/dpmodel/model/polar_model.py @@ -44,6 +44,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -52,6 +53,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["polar"] = model_ret["polarizability"] @@ -69,6 +71,7 @@ def call_lower( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, @@ -78,6 +81,7 @@ def call_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict["polar"] = model_ret["polarizability"] diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index bc1657f0bd..d3153b92f9 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -51,6 +51,7 @@ def call( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -59,6 +60,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) var_name = self.get_var_name() model_predict = {} @@ -77,6 +79,7 @@ def call_lower( fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, @@ -86,6 +89,7 @@ def call_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) var_name = self.get_var_name() model_predict = {} diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index be6566e303..f537672777 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -579,6 +579,7 @@ def call_common( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction with raw internal keys. @@ -624,6 +625,7 @@ def call_common( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, ) @@ -673,6 +675,7 @@ def call( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction with translated user-facing keys. @@ -710,6 +713,7 @@ def call( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_output_type = self.backbone_model.model_output_type() @@ -747,6 +751,7 @@ def call_common_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction with raw internal keys. Lower interface that takes @@ -798,6 +803,7 @@ def call_common_lower( mapping=mapping_updated, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, extended_coord_corr=extended_coord_corr, ) @@ -851,6 +857,7 @@ def call_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + charge_spin: Array | None = None, do_atomic_virial: bool = False, ) -> dict[str, Array]: """Return model prediction with translated user-facing keys. Lower interface. @@ -889,6 +896,7 @@ def call_lower( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_output_type = self.backbone_model.model_output_type() diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 9bd52dd428..7301659958 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -60,6 +60,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -68,6 +69,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} @@ -97,6 +99,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -108,6 +111,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index d28487ed9c..daebc2fb9f 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -52,6 +52,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -60,6 +61,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} @@ -89,6 +91,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -100,6 +103,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index b95f568cb1..c3004b5a5d 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -65,6 +65,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -73,6 +74,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} @@ -101,6 +103,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -112,6 +115,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) model_predict = {} diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index ea2cd17f38..0de2d7fadd 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -62,6 +62,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -70,6 +71,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} @@ -98,6 +100,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -109,6 +112,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) model_predict = {} diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py index 1b1bc3feba..c9e8eea078 100644 --- a/deepmd/pt/model/model/make_hessian_model.py +++ b/deepmd/pt/model/model/make_hessian_model.py @@ -68,6 +68,7 @@ def forward_common( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -101,6 +102,7 @@ def forward_common( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) vdef = self.atomic_output_def() hess_yes = [vdef[kk].r_hessian for kk in vdef.keys()] diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 0eacf85037..80c78a438a 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -689,6 +689,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: # directly call the forward_common method when no specific transform rule return self.forward_common( @@ -698,6 +699,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) return CM diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 7c9550dc3a..78bbd069a8 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -55,6 +55,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -63,6 +64,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} @@ -86,6 +88,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -97,6 +100,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) if self.get_fitting_net() is not None: model_predict = {} diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index c24ca7fa64..0c0e76d6c0 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -55,6 +55,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -63,6 +64,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, ) model_predict = {} model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()] @@ -97,6 +99,7 @@ def forward_lower( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, @@ -108,6 +111,7 @@ def forward_lower( do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), + charge_spin=charge_spin, ) model_predict = {} model_predict[f"atom_{self.get_var_name()}"] = model_ret[self.get_var_name()] diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 91c6e2ea71..5fe16b59a0 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -526,6 +526,7 @@ def forward_common( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: nframes, nloc = atype.shape @@ -540,6 +541,7 @@ def forward_common( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, ) @@ -578,6 +580,7 @@ def forward_common_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, @@ -601,6 +604,7 @@ def forward_common_lower( mapping=mapping_updated, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=extra_nlist_sort, @@ -695,6 +699,7 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -704,6 +709,7 @@ def forward( box, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, ) model_predict = {} @@ -729,6 +735,7 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: @@ -740,6 +747,7 @@ def forward_lower( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, extra_nlist_sort=self.backbone_model.need_sorted_nlist_for_lower(), From e619fa8f44d9a9eddbb77c82d73fa72c0e8b8702 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 8 May 2026 21:56:19 +0800 Subject: [PATCH 09/24] Update test_deep_eval.py --- source/tests/pt_expt/infer/test_deep_eval.py | 47 +++++++++++++++----- 1 file changed, 37 insertions(+), 10 deletions(-) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index f77b882b7c..7537575f1a 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -244,12 +244,18 @@ def test_dynamic_shapes(self) -> None: exported_mod = exported.module() for nloc in [2, 5, 10]: - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = ( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( _make_sample_inputs(self.model, nloc=nloc) ) pte_ret = exported_mod( - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, ) ec = ext_coord.detach().requires_grad_(True) @@ -261,6 +267,7 @@ def test_dynamic_shapes(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_c"): @@ -296,8 +303,8 @@ def test_oversized_nlist(self) -> None: nnei = sum(self.sel) # model's expected neighbor count nloc = 5 - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = _make_sample_inputs( - self.model, nloc=nloc + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( + _make_sample_inputs(self.model, nloc=nloc) ) # Pad nlist with -1 columns, then shuffle column order so real @@ -331,11 +338,18 @@ def test_oversized_nlist(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) # Exported model with same shuffled oversized nlist pte_ret = exported_mod( - ext_coord, ext_atype, nlist_shuffled, mapping_t, fparam, aparam + ext_coord, + ext_atype, + nlist_shuffled, + mapping_t, + fparam, + aparam, + charge_spin, ) for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_c"): @@ -362,6 +376,7 @@ def test_oversized_nlist(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) # The truncated result MUST differ from the correctly sorted result, # proving that naive truncation discards real neighbors. @@ -382,7 +397,7 @@ def test_serialize_round_trip(self) -> None: model2.eval() for nloc in [3, 7]: - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = ( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( _make_sample_inputs(self.model, nloc=nloc) ) ec1 = ext_coord.detach().requires_grad_(True) @@ -396,6 +411,7 @@ def test_serialize_round_trip(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) ret2 = model2.forward_common_lower( ec2, @@ -405,6 +421,7 @@ def test_serialize_round_trip(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_c"): @@ -943,8 +960,8 @@ def test_oversized_nlist(self) -> None: nnei = sum(self.sel) # model's expected neighbor count nloc = 5 - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = _make_sample_inputs( - self.model, nloc=nloc + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( + _make_sample_inputs(self.model, nloc=nloc) ) # Pad nlist with -1 columns, then shuffle column order so real @@ -977,10 +994,17 @@ def test_oversized_nlist(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) pte_ret = exported_mod( - ext_coord, ext_atype, nlist_shuffled, mapping_t, fparam, aparam + ext_coord, + ext_atype, + nlist_shuffled, + mapping_t, + fparam, + aparam, + charge_spin, ) for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_c"): @@ -1004,6 +1028,7 @@ def test_oversized_nlist(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) e_ref = ref_ret["energy_redu"].detach().cpu().numpy() e_trunc = trunc_ret["energy_redu"].detach().cpu().numpy() @@ -1022,7 +1047,7 @@ def test_serialize_round_trip(self) -> None: model2.eval() for nloc in [3, 7]: - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = ( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( _make_sample_inputs(self.model, nloc=nloc) ) ec1 = ext_coord.detach().requires_grad_(True) @@ -1036,6 +1061,7 @@ def test_serialize_round_trip(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) ret2 = model2.forward_common_lower( ec2, @@ -1045,6 +1071,7 @@ def test_serialize_round_trip(self) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) for key in ("energy", "energy_redu", "energy_derv_r", "energy_derv_c"): From e98b6a812f5691d1301b75e49f6bfac4f85fdb56 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sat, 9 May 2026 16:01:41 +0800 Subject: [PATCH 10/24] fix ut --- deepmd/pt/modifier/base_modifier.py | 1 + deepmd/pt_expt/infer/deep_eval.py | 4 +-- source/tests/pt/test_data_modifier.py | 3 ++ source/tests/pt_expt/export_helpers.py | 29 ++++++++++++--- .../pt_expt/model/test_export_pipeline.py | 36 +++++++++++++++---- 5 files changed, 60 insertions(+), 13 deletions(-) diff --git a/deepmd/pt/modifier/base_modifier.py b/deepmd/pt/modifier/base_modifier.py index 5a8c6538b0..957e7a81dc 100644 --- a/deepmd/pt/modifier/base_modifier.py +++ b/deepmd/pt/modifier/base_modifier.py @@ -83,6 +83,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Compute energy, force, and virial corrections.""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index bd586e19b4..604a76270a 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1137,7 +1137,7 @@ def _eval_model( mapping_t, fparam_t, aparam_t, - charge_spin=charge_spin_t, + charge_spin_t, ) # Apply communicate_extended_output to map extended atoms → local atoms @@ -1317,7 +1317,7 @@ def _eval_model_spin( mapping_t, fparam_t, aparam_t, - charge_spin=charge_spin_t, + charge_spin_t, ) # Apply communicate_extended_output to map extended atoms → local atoms diff --git a/source/tests/pt/test_data_modifier.py b/source/tests/pt/test_data_modifier.py index 18d66ef2ff..7b8f8096ce 100644 --- a/source/tests/pt/test_data_modifier.py +++ b/source/tests/pt/test_data_modifier.py @@ -116,6 +116,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Implementation of abstractmethod.""" return {} @@ -158,6 +159,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Implementation of abstractmethod.""" return {} @@ -205,6 +207,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Take scaled model prediction as data modification.""" model_pred = self.model( diff --git a/source/tests/pt_expt/export_helpers.py b/source/tests/pt_expt/export_helpers.py index ae4db82ddb..3f97192a67 100644 --- a/source/tests/pt_expt/export_helpers.py +++ b/source/tests/pt_expt/export_helpers.py @@ -118,6 +118,7 @@ def model_forward_lower_export_round_trip( fparam, aparam, output_keys: tuple[str, ...], + charge_spin=None, rtol: float = 1e-10, atol: float = 1e-10, ): @@ -141,6 +142,9 @@ def model_forward_lower_export_round_trip( Frame and atom parameters. output_keys : tuple of str Output dictionary keys to verify. + charge_spin : torch.Tensor or None + Charge/spin parameter for descriptors that consume it (e.g. DPA3 + with ``add_chg_spin_ebd=True``). rtol, atol : float Tolerances for np.testing.assert_allclose. """ @@ -156,6 +160,7 @@ def model_forward_lower_export_round_trip( mapping_t, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) # 2. Concrete trace @@ -166,21 +171,24 @@ def model_forward_lower_export_round_trip( mapping_t, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) assert isinstance(traced, torch.nn.Module) # 3. Basic export (no dynamic shapes) exported = torch.export.export( traced, - (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin), strict=False, ) assert exported is not None # 4. Compare traced and exported vs eager - ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_traced = traced( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) ret_exported = exported.module()( - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin ) for key in output_keys: np.testing.assert_allclose( @@ -201,7 +209,15 @@ def model_forward_lower_export_round_trip( # 5. Symbolic trace + dynamic shapes + .pte round-trip inputs_2f = tuple( torch.cat([t, t], dim=0) if t is not None else None - for t in (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + for t in ( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam, + aparam, + charge_spin, + ) ) traced_sym = md_pt.forward_lower_exportable( inputs_2f[0], @@ -210,6 +226,7 @@ def model_forward_lower_export_round_trip( inputs_2f[3], fparam=inputs_2f[4], aparam=inputs_2f[5], + charge_spin=inputs_2f[6], tracing_mode="symbolic", _allow_non_fake_inputs=True, ) @@ -226,7 +243,9 @@ def model_forward_lower_export_round_trip( loaded = torch.export.load(f.name).module() # 6. Compare loaded vs eager (nf=1 — different shapes) - ret_loaded_1f = loaded(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + ret_loaded_1f = loaded( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) for key in output_keys: np.testing.assert_allclose( ret_eager[key].detach().cpu().numpy(), diff --git a/source/tests/pt_expt/model/test_export_pipeline.py b/source/tests/pt_expt/model/test_export_pipeline.py index 23e0a62a98..478298c92f 100644 --- a/source/tests/pt_expt/model/test_export_pipeline.py +++ b/source/tests/pt_expt/model/test_export_pipeline.py @@ -121,7 +121,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: inputs_trace = _make_sample_inputs(model2, nframes=5, nloc=7) finally: _env.DEVICE = orig_device - ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = inputs_trace + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = ( + inputs_trace + ) # 4. Eager reference eager_out = model2.forward_common_lower( @@ -132,6 +134,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: fparam=fparam, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) # 5. Trace with symbolic mode (same as dp freeze) @@ -142,6 +145,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: mapping_t, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, do_atomic_virial=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, @@ -155,11 +159,12 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: mapping_t, fparam, aparam, + charge_spin, model_nnei=sum(model2.get_sel()), ) exported = torch.export.export( traced, - (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam), + (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin), dynamic_shapes=dynamic_shapes, strict=False, prefer_deferred_runtime_asserts_over_guards=True, @@ -171,7 +176,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: loaded = torch.export.load(tmp.name).module() # 8. Verify: traced output matches eager (same shapes as trace) - traced_out = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + traced_out = traced( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) for key in eager_out: np.testing.assert_allclose( eager_out[key].detach().cpu().numpy(), @@ -182,7 +189,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: ) # 9. Verify: loaded (.pte) output matches eager (same shapes) - loaded_out = loaded(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam) + loaded_out = loaded( + ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + ) for key in eager_out: np.testing.assert_allclose( eager_out[key].detach().cpu().numpy(), @@ -206,6 +215,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: mapping_t2, fparam2, aparam2, + charge_spin2, ) = inputs_infer eager_out2 = model2.forward_common_lower( @@ -216,9 +226,16 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: fparam=fparam2, aparam=aparam2, do_atomic_virial=True, + charge_spin=charge_spin2, ) loaded_out2 = loaded( - ext_coord2, ext_atype2, nlist_t2, mapping_t2, fparam2, aparam2 + ext_coord2, + ext_atype2, + nlist_t2, + mapping_t2, + fparam2, + aparam2, + charge_spin2, ) for key in eager_out2: np.testing.assert_allclose( @@ -248,9 +265,16 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None: fparam=fparam_ones, aparam=aparam, do_atomic_virial=True, + charge_spin=charge_spin, ) loaded_out_fp1 = loaded( - ext_coord, ext_atype, nlist_t, mapping_t, fparam_ones, aparam + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam_ones, + aparam, + charge_spin, ) # Loaded with fparam=1 should match eager with fparam=1 for key in eager_out_fp1: From 17b4c8a3b8659b2ccc72da6a7173c809b8f0f907 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sat, 9 May 2026 23:23:26 +0800 Subject: [PATCH 11/24] make charge_spin last-input --- deepmd/dpmodel/model/ener_model.py | 4 ++-- deepmd/dpmodel/model/spin_model.py | 8 ++++---- deepmd/pt/infer/deep_eval.py | 8 ++++---- deepmd/pt/model/model/spin_model.py | 8 ++++---- deepmd/pt_expt/model/dipole_model.py | 6 +++--- deepmd/pt_expt/model/dos_model.py | 6 +++--- deepmd/pt_expt/model/dp_linear_model.py | 6 +++--- deepmd/pt_expt/model/dp_zbl_model.py | 6 +++--- deepmd/pt_expt/model/ener_model.py | 6 +++--- deepmd/pt_expt/model/make_model.py | 6 +++--- deepmd/pt_expt/model/polar_model.py | 6 +++--- deepmd/pt_expt/model/property_model.py | 6 +++--- deepmd/pt_expt/model/spin_ener_model.py | 6 +++--- deepmd/pt_expt/model/spin_model.py | 2 +- deepmd/pt_expt/train/training.py | 4 ++-- deepmd/pt_expt/train/wrapper.py | 2 +- 16 files changed, 45 insertions(+), 45 deletions(-) diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 63a2344c5c..c8c75d3cca 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -86,8 +86,8 @@ def call( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common( coord, @@ -121,8 +121,8 @@ def call_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: model_ret = self.call_common_lower( extended_coord, diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index f537672777..6bd5b97edc 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -579,8 +579,8 @@ def call_common( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction with raw internal keys. @@ -675,8 +675,8 @@ def call( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction with translated user-facing keys. @@ -751,8 +751,8 @@ def call_common_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction with raw internal keys. Lower interface that takes extended atomic coordinates, types and spins, nlist, and mapping @@ -857,8 +857,8 @@ def call_lower( mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - charge_spin: Array | None = None, do_atomic_virial: bool = False, + charge_spin: Array | None = None, ) -> dict[str, Array]: """Return model prediction with translated user-facing keys. Lower interface. diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 6128e2633b..3a44bde4fc 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -408,7 +408,7 @@ def eval( request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, fparam, aparam, charge_spin, request_defs + coords, cells, atom_types, fparam, aparam, request_defs, charge_spin ) else: out = self._eval_func(self._eval_model_spin, numb_test, natoms)( @@ -418,8 +418,8 @@ def eval( np.array(kwargs["spin"]), fparam, aparam, - charge_spin, request_defs, + charge_spin, ) return dict( zip( @@ -520,8 +520,8 @@ def _eval_model( atom_types: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, - charge_spin: np.ndarray | None, request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None, ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION]] @@ -604,8 +604,8 @@ def _eval_model_spin( spins: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, - charge_spin: np.ndarray | None, request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None, ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 5fe16b59a0..e0e1002bf0 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -526,8 +526,8 @@ def forward_common( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc = atype.shape coord_updated, atype_updated, coord_corr_for_virial = self.process_spin_input( @@ -580,10 +580,10 @@ def forward_common_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc = nlist.shape[:2] ( @@ -699,8 +699,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( coord, @@ -735,9 +735,9 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, diff --git a/deepmd/pt_expt/model/dipole_model.py b/deepmd/pt_expt/model/dipole_model.py index a4598d4020..4d664adb89 100644 --- a/deepmd/pt_expt/model/dipole_model.py +++ b/deepmd/pt_expt/model/dipole_model.py @@ -44,8 +44,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -77,8 +77,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -129,8 +129,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/dos_model.py b/deepmd/pt_expt/model/dos_model.py index ed961efe37..125522b889 100644 --- a/deepmd/pt_expt/model/dos_model.py +++ b/deepmd/pt_expt/model/dos_model.py @@ -44,8 +44,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -71,8 +71,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -109,8 +109,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 08515960a5..4a29251932 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -47,8 +47,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -80,8 +80,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -134,8 +134,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/dp_zbl_model.py b/deepmd/pt_expt/model/dp_zbl_model.py index 9a65346782..9ca9fe6fff 100644 --- a/deepmd/pt_expt/model/dp_zbl_model.py +++ b/deepmd/pt_expt/model/dp_zbl_model.py @@ -44,8 +44,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -77,8 +77,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -131,8 +131,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 684bc5b07e..1fdef5eaad 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -57,8 +57,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -92,8 +92,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -148,8 +148,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_lower`` into an exportable module. diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index aa195631c1..ac28563500 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -61,8 +61,8 @@ def _cal_hessian_ext( mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None = None, create_graph: bool = False, + charge_spin: torch.Tensor | None = None, ) -> torch.Tensor: """Compute hessian of reduced output w.r.t. extended coordinates. @@ -285,9 +285,9 @@ def forward_common_atomic( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, extended_coord_corr: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, @@ -338,8 +338,8 @@ def forward_common_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower`` into an exportable module. diff --git a/deepmd/pt_expt/model/polar_model.py b/deepmd/pt_expt/model/polar_model.py index 75beb31afb..60fb004bd4 100644 --- a/deepmd/pt_expt/model/polar_model.py +++ b/deepmd/pt_expt/model/polar_model.py @@ -44,8 +44,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -71,8 +71,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -109,8 +109,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/property_model.py b/deepmd/pt_expt/model/property_model.py index 629f6042ad..6d8470f142 100644 --- a/deepmd/pt_expt/model/property_model.py +++ b/deepmd/pt_expt/model/property_model.py @@ -48,8 +48,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -76,8 +76,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -116,8 +116,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: model = self diff --git a/deepmd/pt_expt/model/spin_ener_model.py b/deepmd/pt_expt/model/spin_ener_model.py index f24770f19e..54d0cbb411 100644 --- a/deepmd/pt_expt/model/spin_ener_model.py +++ b/deepmd/pt_expt/model/spin_ener_model.py @@ -48,8 +48,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common( coord, @@ -83,8 +83,8 @@ def forward_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: model_ret = self.call_common_lower( extended_coord, @@ -123,8 +123,8 @@ def forward_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_lower`` into an exportable module. diff --git a/deepmd/pt_expt/model/spin_model.py b/deepmd/pt_expt/model/spin_model.py index b7ab2c5f24..a5f062c8f8 100644 --- a/deepmd/pt_expt/model/spin_model.py +++ b/deepmd/pt_expt/model/spin_model.py @@ -57,8 +57,8 @@ def forward_common_lower_exportable( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``call_common_lower`` into an exportable module. diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index ea015dc487..1059af0be6 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -212,8 +212,8 @@ def _trace_and_compile( mapping: torch.Tensor, fparam: torch.Tensor | None, aparam: torch.Tensor | None, - charge_spin: torch.Tensor | None = None, compile_opts: dict[str, Any] | None = None, + charge_spin: torch.Tensor | None = None, ) -> torch.nn.Module: """Symbolic-trace ``forward_lower`` and compile with inductor + dynamic=True. @@ -366,8 +366,8 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py index 2a301a92e3..6fd68b8edc 100644 --- a/deepmd/pt_expt/train/wrapper.py +++ b/deepmd/pt_expt/train/wrapper.py @@ -170,11 +170,11 @@ def forward( box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, cur_lr: float | torch.Tensor | None = None, label: dict[str, torch.Tensor] | None = None, task_key: str | None = None, do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None, dict | None]: if not self.multi_task: task_key = "Default" From 026707a1dd463afe562e436516f4709076bd4b41 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 10 May 2026 17:21:25 +0800 Subject: [PATCH 12/24] fix ut --- .../jax/atomic_model/linear_atomic_model.py | 2 ++ .../jax/atomic_model/pairtab_atomic_model.py | 2 ++ source/tests/pd/model/test_dpa3.py | 2 +- source/tests/pt_expt/model/test_ener_model.py | 32 +++++++++++++++++-- .../pt_expt/model/test_spin_ener_model.py | 30 +++++++++++++---- source/tests/pt_expt/test_training.py | 6 ++++ 6 files changed, 65 insertions(+), 9 deletions(-) diff --git a/deepmd/jax/atomic_model/linear_atomic_model.py b/deepmd/jax/atomic_model/linear_atomic_model.py index 1c183db7ac..a94194bf32 100644 --- a/deepmd/jax/atomic_model/linear_atomic_model.py +++ b/deepmd/jax/atomic_model/linear_atomic_model.py @@ -61,6 +61,7 @@ def forward_common_atomic( mapping: jnp.ndarray | None = None, fparam: jnp.ndarray | None = None, aparam: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, @@ -69,4 +70,5 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) diff --git a/deepmd/jax/atomic_model/pairtab_atomic_model.py b/deepmd/jax/atomic_model/pairtab_atomic_model.py index 7f18a6403c..96b7bb2842 100644 --- a/deepmd/jax/atomic_model/pairtab_atomic_model.py +++ b/deepmd/jax/atomic_model/pairtab_atomic_model.py @@ -46,6 +46,7 @@ def forward_common_atomic( mapping: jnp.ndarray | None = None, fparam: jnp.ndarray | None = None, aparam: jnp.ndarray | None = None, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, @@ -54,4 +55,5 @@ def forward_common_atomic( mapping=mapping, fparam=fparam, aparam=aparam, + charge_spin=charge_spin, ) diff --git a/source/tests/pd/model/test_dpa3.py b/source/tests/pd/model/test_dpa3.py index 2294b1810e..a582181085 100644 --- a/source/tests/pd/model/test_dpa3.py +++ b/source/tests/pd/model/test_dpa3.py @@ -66,7 +66,7 @@ def test_consistency( [1, 2], # n_multi_edge_message ["float64"], # precision [False], # use_econf_tebd - [False, True], # add_chg_spin_ebd + [False], # add_chg_spin_ebd (PD backend does not support charge_spin) ): dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) diff --git a/source/tests/pt_expt/model/test_ener_model.py b/source/tests/pt_expt/model/test_ener_model.py index b91653a260..79946221af 100644 --- a/source/tests/pt_expt/model/test_ener_model.py +++ b/source/tests/pt_expt/model/test_ener_model.py @@ -182,6 +182,12 @@ def test_forward_lower_exportable(self) -> None: dtype=torch.float64, device=self.device, ) + charge_spin_zero = torch.zeros( + nframes, + 2, + dtype=torch.float64, + device=self.device, + ) # --- eager reference with zero params --- ret_eager_zero = md.forward_lower( @@ -204,13 +210,22 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam=fparam_zero, aparam=aparam_zero, + charge_spin=charge_spin_zero, do_atomic_virial=True, ) self.assertIsInstance(traced, torch.nn.Module) exported = torch.export.export( traced, - (ext_coord, ext_atype, nlist_t, mapping_t, fparam_zero, aparam_zero), + ( + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam_zero, + aparam_zero, + charge_spin_zero, + ), strict=False, ) self.assertIsNotNone(exported) @@ -223,6 +238,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_zero, aparam_zero, + charge_spin_zero, ) ret_exported_zero = exported.module()( ext_coord, @@ -231,6 +247,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_zero, aparam_zero, + charge_spin_zero, ) for key in output_keys: np.testing.assert_allclose( @@ -278,6 +295,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_nz, aparam_nz, + charge_spin_zero, ) ret_exported_nz = exported.module()( ext_coord, @@ -286,6 +304,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_nz, aparam_nz, + charge_spin_zero, ) for key in output_keys: np.testing.assert_allclose( @@ -321,6 +340,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_zero, aparam_nz, + charge_spin_zero, ) self.assertFalse( np.allclose( @@ -351,6 +371,7 @@ def test_forward_lower_exportable(self) -> None: mapping_t, fparam_zero, aparam_zero, + charge_spin_zero, ) ) @@ -361,6 +382,7 @@ def test_forward_lower_exportable(self) -> None: inputs_5f[3], fparam=inputs_5f[4], aparam=inputs_5f[5], + charge_spin=inputs_5f[6], do_atomic_virial=True, tracing_mode="symbolic", _allow_non_fake_inputs=True, @@ -390,7 +412,13 @@ def test_forward_lower_exportable(self) -> None: do_atomic_virial=True, ) ret_loaded_1f = loaded( - ext_coord, ext_atype, nlist_t, mapping_t, fparam_zero, aparam_zero + ext_coord, + ext_atype, + nlist_t, + mapping_t, + fparam_zero, + aparam_zero, + charge_spin_zero, ) for key in ret_common: np.testing.assert_allclose( diff --git a/source/tests/pt_expt/model/test_spin_ener_model.py b/source/tests/pt_expt/model/test_spin_ener_model.py index f7f96392d3..1f600934b5 100644 --- a/source/tests/pt_expt/model/test_spin_ener_model.py +++ b/source/tests/pt_expt/model/test_spin_ener_model.py @@ -495,14 +495,23 @@ def test_forward_lower_exportable(self) -> None: # --- export with torch.export --- exported = torch.export.export( traced, - (ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None), + ( + ext_coord_t, + ext_atype_t, + ext_spin_t, + nlist_t, + mapping_t, + None, + None, + None, + ), strict=False, ) self.assertIsNotNone(exported) # --- verify traced matches eager --- ret_traced = traced( - ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None + ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None, None ) for key in output_keys: np.testing.assert_allclose( @@ -515,7 +524,7 @@ def test_forward_lower_exportable(self) -> None: # --- verify exported matches eager --- ret_exported = exported.module()( - ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None + ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None, None ) for key in output_keys: np.testing.assert_allclose( @@ -538,6 +547,7 @@ def test_forward_lower_exportable(self) -> None: torch.cat([mapping_t, mapping_t], dim=0), None, None, + None, ) traced_sym = model.forward_lower_exportable( @@ -551,7 +561,7 @@ def test_forward_lower_exportable(self) -> None: ) # Build dynamic shapes for spin model - # (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) + # (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam, charge_spin) nframes_dim = torch.export.Dim("nframes", min=1) nall_dim = torch.export.Dim("nall", min=1) nloc_dim = torch.export.Dim("nloc", min=1) @@ -563,6 +573,7 @@ def test_forward_lower_exportable(self) -> None: {0: nframes_dim, 1: nall_dim}, # mapping None, # fparam None, # aparam + None, # charge_spin ) exported_dyn = torch.export.export( traced_sym, @@ -577,7 +588,7 @@ def test_forward_lower_exportable(self) -> None: loaded = torch.export.load(f.name).module() ret_loaded_1f = loaded( - ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None + ext_coord_t, ext_atype_t, ext_spin_t, nlist_t, mapping_t, None, None, None ) for key in output_keys: np.testing.assert_allclose( @@ -632,7 +643,14 @@ def test_oversized_nlist(self) -> None: mapping_t, ) ret_traced = traced( - ext_coord_t, ext_atype_t, ext_spin_t, nlist_shuffled, mapping_t, None, None + ext_coord_t, + ext_atype_t, + ext_spin_t, + nlist_shuffled, + mapping_t, + None, + None, + None, ) ec = ext_coord_t.detach().requires_grad_(True) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index bb3123b1ed..07bbf2c06a 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -574,6 +574,9 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> list[float]: return [0.0, 1.0] + def has_chg_spin_ebd(self) -> bool: + return False + reqs = get_additional_data_requirement(_M()) self.assertEqual(len(reqs), 1) fparam_req = reqs[0] @@ -605,6 +608,9 @@ def has_default_fparam(self) -> bool: def get_default_fparam(self) -> None: return None + def has_chg_spin_ebd(self) -> bool: + return False + reqs = get_additional_data_requirement(_M()) self.assertEqual(len(reqs), 1) fparam_req = reqs[0] From 545a664cbea393275deb9b542ad18939b8b870d0 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 10 May 2026 17:46:20 +0800 Subject: [PATCH 13/24] Update test_dpa3.py --- source/tests/pt/model/test_dpa3.py | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/source/tests/pt/model/test_dpa3.py b/source/tests/pt/model/test_dpa3.py index cf00561379..58c346f548 100644 --- a/source/tests/pt/model/test_dpa3.py +++ b/source/tests/pt/model/test_dpa3.py @@ -5,7 +5,6 @@ import numpy as np import torch -from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DPDescrptDPA3 from deepmd.dpmodel.descriptor.dpa3 import ( RepFlowArgs, ) @@ -121,12 +120,10 @@ def test_consistency( # Prepare charge_spin per mode. charge_spin = None - charge_spin_np = None if need_cs_input: charge_spin = torch.tensor( [[5, 1]], dtype=dtype, device=env.DEVICE ).expand(nf, -1) - charge_spin_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) rd0, _, _, _, _ = dd0( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), @@ -150,21 +147,9 @@ def test_consistency( rtol=rtol, atol=atol, ) - # dp impl - dd2 = DPDescrptDPA3.deserialize(dd0.serialize()) - rd2, _, _, _, _ = dd2.call( - self.coord_ext, - self.atype_ext, - self.nlist, - self.mapping, - charge_spin=charge_spin_np, - ) - np.testing.assert_allclose( - rd0.detach().cpu().numpy(), - rd2, - rtol=rtol, - atol=atol, - ) + # Cross-backend (dpmodel vs pt) numeric consistency for + # add_chg_spin_ebd is covered by + # source/tests/consistent/descriptor/test_dpa3.py. # default_chg_spin should match explicit when value is the same. if cs_mode == "default_chg_spin": From f68fc18e692d1f9282a6952ea39e5db5e52a1d49 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 10 May 2026 20:49:49 +0800 Subject: [PATCH 14/24] fix hlo --- deepmd/jax/jax2tf/tfmodel.py | 1 + deepmd/jax/model/hlo.py | 1 + 2 files changed, 2 insertions(+) diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 1c968c8f41..2d820810cc 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -187,6 +187,7 @@ def call_lower( fparam: jnp.ndarray | None = None, aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: if do_atomic_virial: call_lower = self._call_lower_atomic_virial diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index c79bc727cf..8c1e85c59c 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -183,6 +183,7 @@ def call_lower( fparam: jnp.ndarray | None = None, aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, + charge_spin: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: if extended_coord.shape[1] > nlist.shape[1]: if do_atomic_virial: From 03d21412ffd522c5118a91ac495eb81113a6c9aa Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 11 May 2026 19:55:02 +0800 Subject: [PATCH 15/24] add ut --- source/tests/consistent/descriptor/test_dpa3.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index 2c9ac9428d..c3b3a48c68 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -81,6 +81,7 @@ "n_multi_edge_message", "precision", "add_chg_spin_ebd", + "default_chg_spin", ) @@ -100,6 +101,7 @@ "n_multi_edge_message": 1, "precision": "float64", "add_chg_spin_ebd": False, + "default_chg_spin": None, } @@ -123,6 +125,7 @@ def dpa3_case(**overrides: Any) -> tuple: dpa3_case(exclude_types=[[0, 1]]), dpa3_case(use_loc_mapping=False), dpa3_case(add_chg_spin_ebd=True), + dpa3_case(add_chg_spin_ebd=True, default_chg_spin=[5.0, 1.0]), # Repflow compression branches. dpa3_case(a_compress_rate=1), dpa3_case(a_compress_e_rate=2), @@ -161,6 +164,7 @@ def dpa3_descriptor_api_case(**overrides: Any) -> tuple: dpa3_descriptor_api_case(use_loc_mapping=False), dpa3_descriptor_api_case(fix_stat_std=0.0), dpa3_descriptor_api_case(add_chg_spin_ebd=True), + dpa3_descriptor_api_case(add_chg_spin_ebd=True, default_chg_spin=[5.0, 1.0]), # Repflow compression branches. dpa3_descriptor_api_case(a_compress_rate=1), dpa3_descriptor_api_case(a_compress_e_rate=2), @@ -205,6 +209,7 @@ def data(self) -> dict: n_multi_edge_message, precision, add_chg_spin_ebd, + default_chg_spin, ) = self.param return { "ntypes": self.ntypes, @@ -246,6 +251,7 @@ def data(self) -> dict: "use_loc_mapping": use_loc_mapping, "trainable": False, "add_chg_spin_ebd": add_chg_spin_ebd, + "default_chg_spin": default_chg_spin, } @property @@ -266,6 +272,7 @@ def skip_pt(self) -> bool: _n_multi_edge_message, _precision, _add_chg_spin_ebd, + _default_chg_spin, ) = self.param return CommonTest.skip_pt @@ -287,6 +294,7 @@ def skip_pd(self) -> bool: _n_multi_edge_message, _precision, add_chg_spin_ebd, + _default_chg_spin, ) = self.param return True if add_chg_spin_ebd else CommonTest.skip_pd @@ -308,6 +316,7 @@ def skip_dp(self) -> bool: _n_multi_edge_message, _precision, _add_chg_spin_ebd, + _default_chg_spin, ) = self.param return CommonTest.skip_dp @@ -329,6 +338,7 @@ def skip_tf(self) -> bool: _n_multi_edge_message, _precision, _add_chg_spin_ebd, + _default_chg_spin, ) = self.param return True @@ -394,8 +404,8 @@ def setUp(self) -> None: _n_multi_edge_message, _precision, add_chg_spin_ebd, + _default_chg_spin, ) = self.param - # charge_spin for charge=5, spin=1 when add_chg_spin_ebd is True self.charge_spin = ( np.array([[5, 1]], dtype=GLOBAL_NP_FLOAT_PRECISION) if add_chg_spin_ebd @@ -500,6 +510,7 @@ def rtol(self) -> float: _n_multi_edge_message, precision, _add_chg_spin_ebd, + _default_chg_spin, ) = self.param if precision == "float64": return 1e-10 @@ -527,6 +538,7 @@ def atol(self) -> float: _n_multi_edge_message, precision, _add_chg_spin_ebd, + _default_chg_spin, ) = self.param if precision == "float64": return 1e-6 # need to fix in the future, see issue https://github.com/deepmodeling/deepmd-kit/issues/3786 @@ -563,6 +575,7 @@ def data(self) -> dict: n_multi_edge_message, precision, add_chg_spin_ebd, + default_chg_spin, ) = self.param return { "ntypes": self.ntypes, @@ -604,4 +617,5 @@ def data(self) -> dict: "use_loc_mapping": use_loc_mapping, "trainable": False, "add_chg_spin_ebd": add_chg_spin_ebd, + "default_chg_spin": default_chg_spin, } From 86922bc9706b95f04c644ed5c00b6e2656965fa6 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 17 May 2026 21:41:59 +0800 Subject: [PATCH 16/24] fix comments --- deepmd/dpmodel/descriptor/hybrid.py | 1 + deepmd/dpmodel/model/make_model.py | 28 ++++++++++++++++------------ deepmd/jax/model/base_model.py | 13 +++++++++++-- deepmd/pt/model/descriptor/dpa3.py | 18 ++++++++++++++---- deepmd/pt/model/descriptor/hybrid.py | 20 ++++++++++++++++---- deepmd/pt_expt/infer/deep_eval.py | 22 ++++++++++++++++++++-- 6 files changed, 78 insertions(+), 24 deletions(-) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 1059a9ed0e..c9020cd886 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -349,6 +349,7 @@ def call( atype_ext, nl, mapping, + fparam=fparam, comm_dict=comm_dict, charge_spin=charge_spin, ) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 0892122a28..5aa39daf06 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -298,10 +298,10 @@ def call_common( The keys are defined by the `ModelOutputDef`. """ - cc, bb, fp, ap, input_prec = self._input_type_cast( - coord, box=box, fparam=fparam, aparam=aparam + cc, bb, fp, ap, cs, input_prec = self._input_type_cast( + coord, box=box, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) - del coord, box, fparam, aparam + del coord, box, fparam, aparam, charge_spin model_predict = model_call_from_call_lower( call_lower=self.call_common_lower, rcut=self.get_rcut(), @@ -315,7 +315,7 @@ def call_common( aparam=ap, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, - charge_spin=charge_spin, + charge_spin=cs, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -377,10 +377,10 @@ def call_common_lower( nlist, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) - cc_ext, _, fp, ap, input_prec = self._input_type_cast( - extended_coord, fparam=fparam, aparam=aparam + cc_ext, _, fp, ap, cs, input_prec = self._input_type_cast( + extended_coord, fparam=fparam, aparam=aparam, charge_spin=charge_spin ) - del extended_coord, fparam, aparam + del extended_coord, fparam, aparam, charge_spin model_predict = self.forward_common_atomic( cc_ext, extended_atype, @@ -391,7 +391,7 @@ def call_common_lower( do_atomic_virial=do_atomic_virial, extended_coord_corr=extended_coord_corr, comm_dict=comm_dict, - charge_spin=charge_spin, + charge_spin=cs, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -482,7 +482,8 @@ def _input_type_cast( box: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, - ) -> tuple[Array, Array | None, Array | None, Array | None, Any]: + charge_spin: Array | None = None, + ) -> tuple[Array, Array | None, Array | None, Array | None, Array | None, Any]: """Cast the input data to global float type.""" xp = array_api_compat.array_namespace(coord) input_dtype = coord.dtype @@ -494,17 +495,20 @@ def _input_type_cast( ### _lst: list[Array | None] = [ xp.astype(vv, input_dtype) if vv is not None else None - for vv in [box, fparam, aparam] + for vv in [box, fparam, aparam, charge_spin] ] - box, fparam, aparam = _lst + box, fparam, aparam, charge_spin = _lst if input_dtype == global_dtype: - return coord, box, fparam, aparam, input_dtype + return coord, box, fparam, aparam, charge_spin, input_dtype else: return ( xp.astype(coord, global_dtype), xp.astype(box, global_dtype) if box is not None else None, xp.astype(fparam, global_dtype) if fparam is not None else None, xp.astype(aparam, global_dtype) if aparam is not None else None, + xp.astype(charge_spin, global_dtype) + if charge_spin is not None + else None, input_dtype, ) diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py index 330fd13627..481ca1656d 100644 --- a/deepmd/jax/model/base_model.py +++ b/deepmd/jax/model/base_model.py @@ -68,6 +68,7 @@ def eval_output( mapping: jnp.ndarray | None, fparam: jnp.ndarray | None, aparam: jnp.ndarray | None, + charge_spin_: jnp.ndarray | None, *, _kk: str = kk, _atom_axis: int = atom_axis, @@ -79,7 +80,9 @@ def eval_output( mapping=mapping[None, ...] if mapping is not None else None, fparam=fparam[None, ...] if fparam is not None else None, aparam=aparam[None, ...] if aparam is not None else None, - charge_spin=charge_spin, + charge_spin=charge_spin_[None, ...] + if charge_spin_ is not None + else None, ) return jnp.sum(atomic_ret[_kk][0], axis=_atom_axis) @@ -92,6 +95,7 @@ def eval_output( mapping, fparam, aparam, + charge_spin, ) # extended_force: [nf, nall, *def, 3] def_ndim = len(vdef.shape) @@ -109,6 +113,7 @@ def eval_output( mapping, fparam, aparam, + charge_spin, ) kk_hessian = get_hessian_name(kk) model_predict[kk_hessian] = hessian @@ -130,6 +135,7 @@ def eval_ce( mapping: jnp.ndarray | None, fparam: jnp.ndarray | None, aparam: jnp.ndarray | None, + charge_spin_: jnp.ndarray | None, *, _kk: str = kk, _atom_axis: int = atom_axis - 1, @@ -142,7 +148,9 @@ def eval_ce( mapping=mapping[None, ...] if mapping is not None else None, fparam=fparam[None, ...] if fparam is not None else None, aparam=aparam[None, ...] if aparam is not None else None, - charge_spin=charge_spin, + charge_spin=charge_spin_[None, ...] + if charge_spin_ is not None + else None, ) nloc = nlist.shape[0] cc_loc = jax.lax.stop_gradient(cc_ext)[:nloc, ...] @@ -160,6 +168,7 @@ def eval_ce( mapping, fparam, aparam, + charge_spin, ) # move the first 3 to the last # [nf, *def, nall, 3, 3] diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index 8a78ea4860..35d35eac28 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -204,14 +204,14 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: if self.add_chg_spin_ebd: self.act = ActivationFn(activation_function) - # -100 ~ 100 is a conservative bound + # charge range [-100, 99] mapped to indices [0, 199] self.chg_embedding = TypeEmbedNet( 200, self.tebd_dim, precision=precision, seed=child_seed(seed, 3), ) - # 100 is a conservative upper bound + # spin range [0, 99] mapped to indices [0, 99] self.spin_embedding = TypeEmbedNet( 100, self.tebd_dim, @@ -588,9 +588,19 @@ def forward( assert charge_spin is not None assert self.chg_embedding is not None assert self.spin_embedding is not None - charge = charge_spin[:, 0].to(dtype=torch.int64) + 100 + charge = charge_spin[:, 0].to(dtype=torch.int64) spin = charge_spin[:, 1].to(dtype=torch.int64) - chg_ebd = self.chg_embedding(charge) + # Validate charge range [-100, 99] (200 embedding entries) + if torch.any(charge < -100) or torch.any(charge > 99): + raise ValueError( + f"charge must be in range [-100, 99], got min={charge.min().item()}, max={charge.max().item()}" + ) + # Validate spin range [0, 99] (100 embedding entries) + if torch.any(spin < 0) or torch.any(spin >= 100): + raise ValueError( + f"spin must be in range [0, 99], got min={spin.min().item()}, max={spin.max().item()}" + ) + chg_ebd = self.chg_embedding(charge + 100) spin_ebd = self.spin_embedding(spin) sys_cs_embd = self.act( self.mix_cs_mlp(torch.cat((chg_ebd, spin_ebd), dim=-1)) diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index d79678e668..d0229f6df7 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -101,14 +101,20 @@ def __init__( def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" - return 0 + return max( + (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0 + ) def has_default_chg_spin(self) -> bool: """Returns whether the descriptor has a default charge_spin value.""" - return False + return any(descrpt.has_default_chg_spin() for descrpt in self.descrpt_list) - def get_default_chg_spin(self) -> None: + def get_default_chg_spin(self) -> list[float] | None: """Returns the default charge_spin value, or None.""" + for descrpt in self.descrpt_list: + default = descrpt.get_default_chg_spin() + if default is not None: + return default return None def get_rcut(self) -> float: @@ -346,7 +352,13 @@ def forward( :, :, self.nlist_cut_idx[ii].to(atype_ext.device) ] odescriptor, gr, g2, h2, sw = descrpt( - coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin + coord_ext, + atype_ext, + nl, + mapping, + comm_dict=comm_dict, + fparam=fparam, + charge_spin=charge_spin, ) out_descriptor.append(odescriptor) if gr is not None: diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 604a76270a..26fc89ac57 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -624,6 +624,12 @@ def eval( aparam The atomic parameter. The array should be of size nframes x natoms x dim_aparam. + charge_spin + The charge and spin values for each frame. + The array should be of size nframes x 2, where the first column is charge + and the second column is spin. If the model has add_chg_spin_ebd=True and + no default_chg_spin is set, this parameter is required. If default_chg_spin + is configured, this parameter is optional and will override the default. **kwargs Other parameters @@ -1056,8 +1062,14 @@ def _prepare_inputs( # charge_spin handling: dedicated input, separate from fparam. if charge_spin is not None: + charge_spin_arr = np.asarray(charge_spin) + if charge_spin_arr.shape != (nframes, 2): + raise ValueError( + f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. " + f"Expected ({nframes}, 2) for {nframes} frame(s)." + ) charge_spin_t = torch.tensor( - np.asarray(charge_spin).reshape(nframes, 2), + charge_spin_arr, dtype=torch.float64, device=DEVICE, ) @@ -1272,8 +1284,14 @@ def _eval_model_spin( # charge_spin handling: dedicated input, separate from fparam. if charge_spin is not None: + charge_spin_arr = np.asarray(charge_spin) + if charge_spin_arr.shape != (nframes, 2): + raise ValueError( + f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. " + f"Expected ({nframes}, 2) for {nframes} frame(s)." + ) charge_spin_t = torch.tensor( - np.asarray(charge_spin).reshape(nframes, 2), + charge_spin_arr, dtype=torch.float64, device=DEVICE, ) From 44f9cb9da71126d8afe60d861628935fccccb829 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 18 May 2026 17:06:12 +0800 Subject: [PATCH 17/24] fix jit --- deepmd/dpmodel/descriptor/hybrid.py | 17 +++++++++++++++++ deepmd/pt/model/descriptor/hybrid.py | 28 ++++++++++++++++++++-------- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index c9020cd886..a8358359cb 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -123,6 +123,23 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return np.max([descrpt.get_rcut() for descrpt in self.descrpt_list]).item() + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return max( + (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0 + ) + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return any(descrpt.has_default_chg_spin() for descrpt in self.descrpt_list) + + def get_default_chg_spin(self) -> list[float] | None: + """Returns the default charge_spin value, or None.""" + for descrpt in self.descrpt_list: + if descrpt.has_default_chg_spin(): + return descrpt.get_default_chg_spin() + return None + def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" # may not be a good idea... diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index d0229f6df7..b466f62cb2 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -101,20 +101,32 @@ def __init__( def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" - return max( - (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0 - ) + # JIT-compiled via DPAtomicModel.get_dim_chg_spin; avoid generator + # expressions and `max(..., default=...)` which TorchScript rejects. + dim: int = 0 + for descrpt in self.descrpt_list: + d = descrpt.get_dim_chg_spin() + if d > dim: + dim = d + return dim def has_default_chg_spin(self) -> bool: """Returns whether the descriptor has a default charge_spin value.""" - return any(descrpt.has_default_chg_spin() for descrpt in self.descrpt_list) + # JIT-compiled via DPAtomicModel.has_default_chg_spin; keep as an + # explicit loop instead of `any(generator)` for TorchScript. + for descrpt in self.descrpt_list: + if descrpt.has_default_chg_spin(): + return True + return False - def get_default_chg_spin(self) -> list[float] | None: + @torch.jit.export + def get_default_chg_spin(self) -> Optional[torch.Tensor]: # noqa: UP045 """Returns the default charge_spin value, or None.""" + # JIT-compiled via DPAtomicModel.get_default_chg_spin; the caller + # invokes `.unsqueeze(0)` on the result, so return a Tensor (not list). for descrpt in self.descrpt_list: - default = descrpt.get_default_chg_spin() - if default is not None: - return default + if descrpt.has_default_chg_spin(): + return descrpt.get_default_chg_spin() return None def get_rcut(self) -> float: From b47840ddc905b5390f2655de4e0ec237f44fa488 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 21 May 2026 11:23:51 +0800 Subject: [PATCH 18/24] fix: UT charge_spin para passing --- source/tests/pt_expt/model/test_export_with_comm.py | 10 ++++++---- .../tests/pt_expt/model/test_spin_export_with_comm.py | 3 +++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/source/tests/pt_expt/model/test_export_with_comm.py b/source/tests/pt_expt/model/test_export_with_comm.py index dcbc628e53..ec305f2ed0 100644 --- a/source/tests/pt_expt/model/test_export_with_comm.py +++ b/source/tests/pt_expt/model/test_export_with_comm.py @@ -349,9 +349,11 @@ def test_pte_with_comm_dict_traces_and_loads(tmp_path) -> None: assert os.path.exists(pte_path) loaded = torch.export.load(pte_path) # Sanity: the loaded program has the expected number of inputs - # (6 base + 8 comm = 14). + # (7 base + 8 comm = 15): extended_coord, extended_atype, nlist, + # mapping, fparam, aparam, charge_spin (added in 0505_reformat_chg_spin) + # + the 8 comm tensors. spec = loaded.module().graph.find_nodes(op="placeholder") - assert len(spec) == 14, ( - f"with-comm exported program must accept 14 positional inputs " - f"(6 base + 8 comm); got {len(spec)}" + assert len(spec) == 15, ( + f"with-comm exported program must accept 15 positional inputs " + f"(7 base + 8 comm); got {len(spec)}" ) diff --git a/source/tests/pt_expt/model/test_spin_export_with_comm.py b/source/tests/pt_expt/model/test_spin_export_with_comm.py index 0e403d2b42..971a96623a 100644 --- a/source/tests/pt_expt/model/test_spin_export_with_comm.py +++ b/source/tests/pt_expt/model/test_spin_export_with_comm.py @@ -108,6 +108,7 @@ def test_spin_forward_common_lower_exportable_with_comm_traces() -> None: mapping = torch.zeros(1, nall, dtype=torch.int64) fparam = None aparam = None + charge_spin = None comm_inputs, _keepalive = _build_self_comm_inputs(nloc=nloc, nghost=nall - nloc) @@ -125,6 +126,7 @@ def test_spin_forward_common_lower_exportable_with_comm_traces() -> None: mapping, fparam, aparam, + charge_spin, *comm_inputs, do_atomic_virial=True, tracing_mode="symbolic", @@ -142,6 +144,7 @@ def test_spin_forward_common_lower_exportable_with_comm_traces() -> None: mapping, fparam, aparam, + charge_spin, *comm_inputs, ) assert isinstance(out, dict) From 30e7b07a8ffee6721ae7583cf0f30d588a148c27 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 21 May 2026 19:53:58 +0800 Subject: [PATCH 19/24] Update test_descriptor.py --- .../dpmodel/descriptor/test_descriptor.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 2b778d41c3..3b80a8e520 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -565,11 +565,25 @@ def DescriptorParamDPA3( "env_protection": (0.0, 1e-8), "precision": ("float64",), "use_loc_mapping": (True, False), - "add_chg_spin_ebd": (False, True), - "default_chg_spin": (None, [5.0, 1.0]), } ), ) + + +def DescriptorParamDPA3DefaultChgSpin(ntypes, rcut, rcut_smth, sel, type_map, **kwargs): + return DescriptorParamDPA3( + ntypes, + rcut, + rcut_smth, + sel, + type_map, + **kwargs, + add_chg_spin_ebd=True, + default_chg_spin=[5.0, 1.0], + ) + + +DescriptorParamDPA3List.append(DescriptorParamDPA3DefaultChgSpin) # to get name for the default function DescriptorParamDPA3 = DescriptorParamDPA3List[0] From ea7a511ff918d59143d62348e7e1291ec9c99fc5 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 21 May 2026 20:25:38 +0800 Subject: [PATCH 20/24] fix comment --- deepmd/dpmodel/model/make_model.py | 4 ++++ deepmd/pt/model/model/make_model.py | 5 ++++ deepmd/pt_expt/infer/deep_eval.py | 34 ++++++++++++++------------- deepmd/pt_expt/utils/serialization.py | 26 ++++++++++++++++++++ 4 files changed, 53 insertions(+), 16 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 5aa39daf06..ebcc671f62 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -725,6 +725,10 @@ def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" return self.atomic_model.has_chg_spin_ebd() + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + return self.atomic_model.get_dim_chg_spin() + def has_default_chg_spin(self) -> bool: """Check if the model has default charge_spin values.""" return self.atomic_model.has_default_chg_spin() diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 80c78a438a..713eab3d8c 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -560,6 +560,11 @@ def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" return self.atomic_model.has_chg_spin_ebd() + @torch.jit.export + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + return self.atomic_model.get_dim_chg_spin() + @torch.jit.export def has_default_chg_spin(self) -> bool: """Check if the model has default charge_spin values.""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 26fc89ac57..8f40600ffc 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -59,6 +59,17 @@ import ase.neighborlist +def _reshape_charge_spin(charge_spin: np.ndarray, nframes: int) -> np.ndarray: + charge_spin_arr = np.asarray(charge_spin) + try: + return charge_spin_arr.reshape(nframes, 2) + except ValueError as err: + raise ValueError( + f"charge_spin must be reshape-compatible with ({nframes}, 2), " + f"got shape {charge_spin_arr.shape}." + ) from err + + class DeepEval(DeepEvalBackend): """PyTorch Exportable backend implementation of DeepEval. @@ -626,10 +637,11 @@ def eval( The array should be of size nframes x natoms x dim_aparam. charge_spin The charge and spin values for each frame. - The array should be of size nframes x 2, where the first column is charge - and the second column is spin. If the model has add_chg_spin_ebd=True and - no default_chg_spin is set, this parameter is required. If default_chg_spin - is configured, this parameter is optional and will override the default. + The array should be reshape-compatible with nframes x 2, where the first + column is charge and the second column is spin. If the model has + add_chg_spin_ebd=True and no default_chg_spin is set, this parameter is + required. If default_chg_spin is configured, this parameter is optional + and will override the default. **kwargs Other parameters @@ -1062,12 +1074,7 @@ def _prepare_inputs( # charge_spin handling: dedicated input, separate from fparam. if charge_spin is not None: - charge_spin_arr = np.asarray(charge_spin) - if charge_spin_arr.shape != (nframes, 2): - raise ValueError( - f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. " - f"Expected ({nframes}, 2) for {nframes} frame(s)." - ) + charge_spin_arr = _reshape_charge_spin(charge_spin, nframes) charge_spin_t = torch.tensor( charge_spin_arr, dtype=torch.float64, @@ -1284,12 +1291,7 @@ def _eval_model_spin( # charge_spin handling: dedicated input, separate from fparam. if charge_spin is not None: - charge_spin_arr = np.asarray(charge_spin) - if charge_spin_arr.shape != (nframes, 2): - raise ValueError( - f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. " - f"Expected ({nframes}, 2) for {nframes} frame(s)." - ) + charge_spin_arr = _reshape_charge_spin(charge_spin, nframes) charge_spin_t = torch.tensor( charge_spin_arr, dtype=torch.float64, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 4fd602276f..48d073ed99 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import ctypes import json +from typing import ( + Any, +) import numpy as np import torch @@ -98,6 +101,16 @@ def _json_to_numpy(model_obj: dict) -> dict: ) +def _metadata_value_to_json(value: Any) -> Any: + if value is None: + return None + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, np.ndarray): + return value.tolist() + return value + + def _needs_with_comm_artifact(model: torch.nn.Module) -> bool: """Return ``True`` if the model needs a "with-comm" AOTI artifact compiled. @@ -450,6 +463,19 @@ def _collect_metadata(model: torch.nn.Module, is_spin: bool = False) -> dict: "mixed_types": model.mixed_types(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), + "has_chg_spin_ebd": ( + model.has_chg_spin_ebd() if hasattr(model, "has_chg_spin_ebd") else False + ), + "has_default_chg_spin": ( + model.has_default_chg_spin() + if hasattr(model, "has_default_chg_spin") + else False + ), + "default_chg_spin": ( + _metadata_value_to_json(model.get_default_chg_spin()) + if hasattr(model, "get_default_chg_spin") + else None + ), "fitting_output_defs": fitting_output_defs, # sel_type enables `DeepEval.get_sel_type()` without a dpmodel # round-trip; required for dipole/polar/wfc models in metadata-only From 5918bed13f5ddc8d79730e17f2798bffe4e692a7 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 21 May 2026 21:40:52 +0800 Subject: [PATCH 21/24] fix ut --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 4 ++++ deepmd/pt/model/atomic_model/base_atomic_model.py | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 7cd4ed8fb8..cf59af94db 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -160,6 +160,10 @@ def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" return False + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + return 0 + def has_default_chg_spin(self) -> bool: """Check if the model has default charge_spin values.""" return False diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index dfd4ce60fc..8605db9359 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -193,6 +193,11 @@ def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" return False + @torch.jit.export + def get_dim_chg_spin(self) -> int: + """Get the dimension of charge_spin input.""" + return 0 + def has_default_chg_spin(self) -> bool: """Check if the model has default charge_spin values.""" return False From 7286d3716cbf8c07495214e3e568e210f5d191d3 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 22 May 2026 14:36:40 +0800 Subject: [PATCH 22/24] fix comment --- deepmd/dpmodel/descriptor/hybrid.py | 21 +++++++- deepmd/pt/model/descriptor/hybrid.py | 25 ++++++---- source/tests/pt/model/test_dpa3.py | 48 +++++++++++++++++++ .../dpmodel/descriptor/test_descriptor.py | 34 +++++++++++++ 4 files changed, 118 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index a8358359cb..110ae1fbe0 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -131,12 +131,29 @@ def get_dim_chg_spin(self) -> int: def has_default_chg_spin(self) -> bool: """Returns whether the descriptor has a default charge_spin value.""" - return any(descrpt.has_default_chg_spin() for descrpt in self.descrpt_list) + default_chg_spin = None + found_chg_spin = False + for descrpt in self.descrpt_list: + if descrpt.get_dim_chg_spin() == 0: + continue + found_chg_spin = True + if not descrpt.has_default_chg_spin(): + return False + child_default_chg_spin = descrpt.get_default_chg_spin() + if child_default_chg_spin is None: + return False + if default_chg_spin is None: + default_chg_spin = child_default_chg_spin + elif child_default_chg_spin != default_chg_spin: + return False + return found_chg_spin def get_default_chg_spin(self) -> list[float] | None: """Returns the default charge_spin value, or None.""" + if not self.has_default_chg_spin(): + return None for descrpt in self.descrpt_list: - if descrpt.has_default_chg_spin(): + if descrpt.get_dim_chg_spin() > 0: return descrpt.get_default_chg_spin() return None diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index b466f62cb2..97a49b750a 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -112,20 +112,29 @@ def get_dim_chg_spin(self) -> int: def has_default_chg_spin(self) -> bool: """Returns whether the descriptor has a default charge_spin value.""" - # JIT-compiled via DPAtomicModel.has_default_chg_spin; keep as an - # explicit loop instead of `any(generator)` for TorchScript. + default_chg_spin: torch.Tensor | None = None + found_chg_spin: bool = False for descrpt in self.descrpt_list: - if descrpt.has_default_chg_spin(): - return True - return False + if descrpt.get_dim_chg_spin() > 0: + found_chg_spin = True + if not descrpt.has_default_chg_spin(): + return False + child_default_chg_spin = descrpt.get_default_chg_spin() + if child_default_chg_spin is None: + return False + if default_chg_spin is None: + default_chg_spin = child_default_chg_spin + elif not torch.equal(default_chg_spin, child_default_chg_spin): + return False + return found_chg_spin @torch.jit.export def get_default_chg_spin(self) -> Optional[torch.Tensor]: # noqa: UP045 """Returns the default charge_spin value, or None.""" - # JIT-compiled via DPAtomicModel.get_default_chg_spin; the caller - # invokes `.unsqueeze(0)` on the result, so return a Tensor (not list). + if not self.has_default_chg_spin(): + return None for descrpt in self.descrpt_list: - if descrpt.has_default_chg_spin(): + if descrpt.get_dim_chg_spin() > 0: return descrpt.get_default_chg_spin() return None diff --git a/source/tests/pt/model/test_dpa3.py b/source/tests/pt/model/test_dpa3.py index ce926f08c9..f99111b8d7 100644 --- a/source/tests/pt/model/test_dpa3.py +++ b/source/tests/pt/model/test_dpa3.py @@ -10,6 +10,7 @@ ) from deepmd.pt.model.descriptor import ( DescrptDPA3, + DescrptHybrid, ) from deepmd.pt.utils import ( env, @@ -31,6 +32,23 @@ dtype = env.GLOBAL_PT_FLOAT_PRECISION +def _repflow_args() -> RepFlowArgs: + return RepFlowArgs( + n_dim=8, + e_dim=6, + a_dim=4, + nlayers=1, + e_rcut=4.0, + e_rcut_smth=0.5, + e_sel=12, + a_rcut=3.5, + a_rcut_smth=0.5, + a_sel=8, + axis_neuron=4, + update_angle=False, + ) + + class TestDescrptDPA3(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) @@ -193,6 +211,36 @@ def test_consistency( atol=atol, ) + def test_hybrid_default_chg_spin_semantics(self) -> None: + def make_dpa3(default_chg_spin: list[float] | None) -> DescrptDPA3: + return DescrptDPA3( + self.nt, + repflow=_repflow_args(), + precision="float64", + add_chg_spin_ebd=True, + default_chg_spin=default_chg_spin, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + shared_default = DescrptHybrid( + list=[make_dpa3([5.0, 1.0]), make_dpa3([5.0, 1.0])] + ) + self.assertTrue(shared_default.has_default_chg_spin()) + torch.testing.assert_close( + shared_default.get_default_chg_spin(), + torch.tensor([5.0, 1.0], dtype=torch.float64, device=env.DEVICE), + ) + + missing_default = DescrptHybrid(list=[make_dpa3([5.0, 1.0]), make_dpa3(None)]) + self.assertFalse(missing_default.has_default_chg_spin()) + self.assertIsNone(missing_default.get_default_chg_spin()) + + mismatched_default = DescrptHybrid( + list=[make_dpa3([5.0, 1.0]), make_dpa3([6.0, 1.0])] + ) + self.assertFalse(mismatched_default.has_default_chg_spin()) + self.assertIsNone(mismatched_default.get_default_chg_spin()) + def test_jit( self, ) -> None: diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 3b80a8e520..422fd3819c 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -659,3 +659,37 @@ def setUp(self) -> None: self.nt, self.rcut, self.rcut_smth, self.sel, ["O", "H"] ) self.module = Descrpt(**self.input_dict) + + +class TestHybridChgSpinDefaultDP(unittest.TestCase): + def _make_dpa3(self, default_chg_spin: list[float] | None) -> DescrptDPA3: + return DescrptDPA3( + **DescriptorParamDPA3( + 2, + 4.0, + 0.5, + [6, 6], + ["O", "H"], + add_chg_spin_ebd=True, + default_chg_spin=default_chg_spin, + ) + ) + + def test_shared_default_required_for_hybrid_default(self) -> None: + shared_default = DescrptHybrid( + list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3([5.0, 1.0])] + ) + self.assertTrue(shared_default.has_default_chg_spin()) + self.assertEqual(shared_default.get_default_chg_spin(), [5.0, 1.0]) + + missing_default = DescrptHybrid( + list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3(None)] + ) + self.assertFalse(missing_default.has_default_chg_spin()) + self.assertIsNone(missing_default.get_default_chg_spin()) + + mismatched_default = DescrptHybrid( + list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3([6.0, 1.0])] + ) + self.assertFalse(mismatched_default.has_default_chg_spin()) + self.assertIsNone(mismatched_default.get_default_chg_spin()) From 215a969e1e82af608c72eb659d883a2b88a832b7 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 24 May 2026 00:01:11 +0800 Subject: [PATCH 23/24] Update dp_atomic_model.py --- deepmd/pt/model/atomic_model/dp_atomic_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 9b932aae72..783ee9e766 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -278,6 +278,7 @@ def forward_atomic( if self.add_chg_spin_ebd and charge_spin is None: default_cs_tensor = self.descriptor.get_default_chg_spin() if default_cs_tensor is not None: + default_cs_tensor = default_cs_tensor.to(device=extended_coord.device) charge_spin = torch.tile(default_cs_tensor.unsqueeze(0), [nframes, 1]) descriptor, rot_mat, g2, h2, sw = self.descriptor( From 18faac5f50db4f5c341f98a1bb9930376eed1a4e Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 24 May 2026 16:17:17 +0800 Subject: [PATCH 24/24] Update test_descriptor.py --- source/tests/universal/dpmodel/descriptor/test_descriptor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 422fd3819c..010cf4dcd6 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -583,7 +583,6 @@ def DescriptorParamDPA3DefaultChgSpin(ntypes, rcut, rcut_smth, sel, type_map, ** ) -DescriptorParamDPA3List.append(DescriptorParamDPA3DefaultChgSpin) # to get name for the default function DescriptorParamDPA3 = DescriptorParamDPA3List[0]