diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index d5a3a3eabd..7050a4a62e 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -250,10 +250,6 @@ 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 - def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return None diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index c5cf0b157c..0406eb4b91 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -117,12 +117,9 @@ def __init__( super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting - if hasattr(self.fitting_net, "reinit_exclude"): - self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.fitting_net.reinit_exclude(self.atom_exclude_types) self.type_map = type_map - self.add_chg_spin_ebd: bool = getattr( - self.descriptor, "add_chg_spin_ebd", False - ) + self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0 # Structural capability: only descriptors with a native spin # conditioning mechanism (currently DPA4) accept a ``spin`` kwarg on # ``call_graph`` at all -- unlike ``charge_spin``, which every @@ -151,15 +148,9 @@ def get_dim_chg_spin(self) -> int: 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(): + if self.add_chg_spin_ebd: return self.descriptor.get_default_chg_spin() return None diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 4999bc15a7..bb42ec37ae 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -723,19 +723,11 @@ def _fparam_consumers(self) -> list: """Children that actually consume ``fparam``.""" return [m for m in self.models if m.get_dim_fparam() > 0] - def has_default_chg_spin(self) -> bool: - """Whether every active child shares one default charge/spin.""" - return self._agreed_default( - self._chg_spin_consumers(), - lambda m: m.has_default_chg_spin(), - lambda m: m.get_default_chg_spin(), - )[0] - def get_default_chg_spin(self) -> "Array | None": """The shared default charge/spin conditions, if the children agree.""" return self._agreed_default( self._chg_spin_consumers(), - lambda m: m.has_default_chg_spin(), + lambda m: m.get_default_chg_spin() is not None, lambda m: m.get_default_chg_spin(), )[1] diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 7118aa5d7a..f5deadde11 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -84,6 +84,14 @@ def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.get_nsel() + def get_pair_exclude_types(self) -> list[tuple[int, int]]: + """Return the excluded atom-type pairs of this atomic model. + + Always set by ``__init__`` (empty list when no exclusion is + configured); an empty return means the pair-exclusion mask is off. + """ + return self.pair_exclude_types + @abstractmethod def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" diff --git a/deepmd/dpmodel/descriptor/descriptor.py b/deepmd/dpmodel/descriptor/descriptor.py index 1dce8fb7d1..d03fbe1a08 100644 --- a/deepmd/dpmodel/descriptor/descriptor.py +++ b/deepmd/dpmodel/descriptor/descriptor.py @@ -37,6 +37,12 @@ class DescriptorBlock(ABC, make_plugin_registry("DescriptorBlock")): local_cluster = False + # Stat-behavior flags with concrete defaults so stat machinery can read + # them on any block without getattr probes; blocks that configure them + # assign instance attributes in __init__ (issue #5897). + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is DescriptorBlock: try: diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 27d73b1251..1781728951 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -477,6 +477,10 @@ def __init__( self.concat_output_tebd = concat_output_tebd self.trainable = trainable self.precision = precision + # tebd-compression slots: declared here so presence is a class + # property, not a runtime accident (issue #5897); populated by + # enable_compression()/deserialize(). + self.type_embd_data = None self.tebd_compress = False self.geo_compress = False self.compress = False @@ -1113,6 +1117,10 @@ def enable_compression( stacklevel=2, ) + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active.""" + return self.geo_compress + def serialize(self) -> dict: """Serialize the descriptor to dict.""" obj = self.se_atten @@ -1172,7 +1180,7 @@ def serialize(self) -> dict: if self.compress: type_embd_data = ( self.type_embd_data - if hasattr(self, "type_embd_data") + if self.type_embd_data is not None else obj.type_embd_data ) compress_dict: dict = { @@ -1514,6 +1522,9 @@ def __init__( self.mean = np.zeros(wanted_shape, dtype=PRECISION_DICT[self.precision]) self.stddev = np.ones(wanted_shape, dtype=PRECISION_DICT[self.precision]) self.orig_sel = self.sel + # tebd-compression slots: declared here so presence is a class + # property, not a runtime accident (issue #5897); populated by + # type_embedding_compression()/enable_compression(). self.tebd_compress = False self.geo_compress = False self.is_sorted = len(self.exclude_types) == 0 diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index d94bb5fc02..55566fee54 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -1421,6 +1421,10 @@ def _call_dense( g1 = xp.concat([g1, g1_inp], axis=-1) return g1, rot_mat, g2, h2, sw + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active.""" + return self.geo_compress + def serialize(self) -> dict: repinit = self.repinit repformers = self.repformers diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index aa1d10fc33..111dc7a520 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -543,10 +543,6 @@ 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 diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index fd21e7798a..f1a08123c4 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2286,10 +2286,6 @@ def get_dim_chg_spin(self) -> int: """Return the charge/spin condition width.""" return 2 if self.add_chg_spin_ebd else 0 - def has_default_chg_spin(self) -> bool: - """Return whether default charge/spin conditions are configured.""" - return self.default_chg_spin is not None - def get_default_chg_spin(self) -> list[float] | None: """Return default charge/spin conditions.""" return self.default_chg_spin diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 9709ce37f0..83a903c687 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -129,33 +129,26 @@ def get_dim_chg_spin(self) -> int: (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.""" + def get_default_chg_spin(self) -> list[float] | None: + """Returns the default charge_spin value, or None. + + ``None`` unless every sub-descriptor that supports charge_spin + (``get_dim_chg_spin() > 0``) agrees on the same default value. + """ 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 + return None 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.get_dim_chg_spin() > 0: - return descrpt.get_default_chg_spin() - return None + return None + return default_chg_spin if found_chg_spin else None def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 2d058a3398..b0a77e22fc 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -48,6 +48,15 @@ def make_base_descriptor( class BD(ABC, PluginVariant, make_plugin_registry("descriptor")): """Base descriptor provides the interfaces of descriptor.""" + # Stat-behavior flags with concrete defaults so stat machinery (e.g. + # ``merge_env_stat``, which accepts either a ``Descriptor`` or a + # ``DescriptorBlock``) can read them on any descriptor without + # getattr probes; descriptors that configure them assign instance + # attributes in __init__ (issue #5897). Mirrors the same defaults on + # ``DescriptorBlock``. + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is BD: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) @@ -100,14 +109,18 @@ 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 + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active. + + Concrete default ``False``; descriptor families with a + geometric compression path override from their own state. + """ + return False + @abstractmethod def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index ec7e3e4f77..9a7d8a82fe 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -259,7 +259,7 @@ def serialize(self) -> dict: if self.compress: type_embd_data = ( self.type_embd_data - if hasattr(self, "type_embd_data") + if self.type_embd_data is not None else obj.type_embd_data ) compress_dict: dict = { diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index 3d5321c0ef..b810334031 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -206,6 +206,11 @@ def __init__( self.trainable = trainable self.precision = precision self.compress = False + # tebd-compression slot: declared here so presence is a class + # property, not a runtime accident (issue #5897); optionally + # populated by deserialize() (only present in compressed models + # that carry type embedding compression data). + self.type_embd_data = None def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -463,7 +468,7 @@ def serialize(self) -> dict: "compress_info": [to_numpy_array(i) for i in self.compress_info], }, } - if hasattr(self, "type_embd_data"): + if self.type_embd_data is not None: compress_dict["@variables"]["type_embd_data"] = to_numpy_array( self.type_embd_data ) diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index cf8172bd03..7a595d4bc8 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -67,6 +67,30 @@ def compute_output_stats(self, merged: Any) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError + def reinit_exclude(self, exclude_types: list[int] = []) -> None: + """Reinitialize the per-type output exclusion list. + + Concrete default for fittings without exclusion support: an + empty list is a no-op; a non-empty list raises, because + silently ignoring a requested exclusion would degrade the + model without any signal. + + Parameters + ---------- + exclude_types + Atom types whose fitting output is excluded. + + Raises + ------ + NotImplementedError + If ``exclude_types`` is non-empty and this fitting does + not support atom-type exclusion. + """ + if exclude_types: + raise NotImplementedError( + "this fitting does not support atom-type exclusion" + ) + @abstractmethod def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 9c85fa7e26..0ee9763b1f 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -107,6 +107,45 @@ def has_spin(self) -> bool: """ return False + def has_chg_spin_ebd(self) -> bool: + """Return whether the model conditions on charge/spin embedding. + + Concrete default ``False``; models wrapping an atomic model + override to delegate. + """ + return False + + def get_dim_chg_spin(self) -> int: + """Return the charge/spin condition width (0 if unsupported).""" + return 0 + + def get_default_chg_spin(self) -> list | None: + """Return default charge/spin conditions, or ``None`` if none + are configured. ``is not None`` is the support predicate. + """ + return None + + def get_var_name(self) -> str | None: + """Return the fitted property's variable name, or ``None`` if + this is not a property model. ``is not None`` is the support + predicate. + """ + return None + + def get_task_dim(self) -> int: + """Return the property output dimension (property models only). + + Raises + ------ + NotImplementedError + If the model is not a property model. + """ + raise NotImplementedError("get_task_dim: property models only") + + def get_intensive(self) -> bool: + """Return whether the fitted property is intensive.""" + return False + @abstractmethod def serialize(self) -> dict: """Serialize the model. diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 44967bba74..73b53ce6ea 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -428,7 +428,7 @@ def call_common( charge_spin=cs, neighbor_list=neighbor_list, # exclusion is a nlist-BUILD transform (decision #18/A4) - pair_excl=getattr(self.atomic_model, "pair_excl", None), + pair_excl=self.atomic_model.pair_excl, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -527,7 +527,7 @@ def _call_common_graph( # is constructed, so the graph lower / exported ``.pt2`` consumes an # already-excluded ``edge_mask`` and never re-applies it. Mirrors the # pt_expt eager path and the C++ ``applyPairExclusion`` at build. - pair_excl = getattr(self.atomic_model, "pair_excl", None) + pair_excl = self.atomic_model.pair_excl if method == "dense": ng = build_neighbor_graph( cc, atype, bb, self.get_rcut(), pair_excl=pair_excl @@ -1131,10 +1131,6 @@ 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() - def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return self.atomic_model.get_default_chg_spin() diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index faf0a8dacf..17c575a92b 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -214,12 +214,16 @@ def process_spin_input_lower( # pair exclusion in here (decision #18/A4 — the lower consumes a # pre-excluded nlist and never re-applies it). No-op when the backbone # has no pair_exclude_types. - pair_excl = getattr( - self.backbone_model.atomic_model + # ``backbone_model`` is either a full ``make_model``-wrapped model + # (exposes ``.atomic_model``) or, per the ``__init__`` annotation, + # a bare ``DPAtomicModel`` -- which always carries ``pair_excl`` + # (set unconditionally by ``BaseAtomicModel.__init__`` via + # ``reinit_pair_exclude``, ``None`` when no exclusion is + # configured). Direct access, not a defensive ``getattr`` probe. + pair_excl = ( + self.backbone_model.atomic_model.pair_excl if hasattr(self.backbone_model, "atomic_model") - else self.backbone_model, - "pair_excl", - None, + else self.backbone_model.pair_excl ) if pair_excl is not None: from deepmd.dpmodel.utils.nlist import ( diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 8288ede2e7..1bb123534d 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -69,9 +69,7 @@ def merge_env_stat( or getattr(link_obj, "stats", None) is None ): return - if getattr(base_obj, "set_stddev_constant", False) and getattr( - base_obj, "set_davg_zero", False - ): + if base_obj.set_stddev_constant and base_obj.set_davg_zero: return # Weighted merge of StatItem objects @@ -93,7 +91,7 @@ def merge_env_stat( xp = array_api_compat.array_namespace(current_stddev) device = array_api_compat.device(current_stddev) merged_mean = current_mean - if not getattr(base_obj, "set_davg_zero", False): + if not base_obj.set_davg_zero: merged_mean = xp.asarray(mean, dtype=current_mean.dtype, device=device) merged_stddev = xp.asarray( stddev, diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 915d216672..d20626e835 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -390,10 +390,11 @@ def _get_property_var_name(model: Any) -> str | None: """Return the property variable name of ``model``, or ``None``. Used by every backend's ``model_type`` to detect a property model. - ``get_var_name`` may be absent (dpmodel/pt live models expose it only on - property models) or present-but-unimplemented (jax/tf2 artifacts always - define it and raise ``NotImplementedError`` otherwise), so probe - defensively. + Live dpmodel-family models always expose ``get_var_name`` (the base + class's concrete default returns ``None`` for non-property models); + ``get_var_name`` is only absent on frozen pt legacy models, and only + present-but-unimplemented -- raising ``NotImplementedError`` -- on + jax/tf2 artifacts, so probe defensively. """ if not hasattr(model, "get_var_name"): return None @@ -441,6 +442,7 @@ def get_has_hessian(self) -> bool: def get_var_name(self) -> str: """Get the name of the fitting property (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_var_name"): return model.get_var_name() raise NotImplementedError @@ -448,6 +450,7 @@ def get_var_name(self) -> str: def get_task_dim(self) -> int: """Get the output dimension of the property (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_task_dim"): return model.get_task_dim() raise NotImplementedError @@ -455,6 +458,7 @@ def get_task_dim(self) -> int: def get_intensive(self) -> bool: """Whether the property is intensive (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_intensive"): return model.get_intensive() raise NotImplementedError diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 0e6c11ede6..d2930ff8c7 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -524,18 +524,17 @@ def has_default_fparam(self) -> bool: def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" - if hasattr(self.dp, "has_chg_spin_ebd"): - return self.dp.has_chg_spin_ebd() - return False + return self.dp.has_chg_spin_ebd() def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" - if hasattr(self.dp, "get_dim_chg_spin"): - return self.dp.get_dim_chg_spin() - return 0 + return self.dp.get_dim_chg_spin() def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - if hasattr(self.dp, "has_default_chg_spin"): - return self.dp.has_default_chg_spin() - return False + """Check if the model has default charge_spin values. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` on + the live-model interfaces; this wrapper method is kept for API + stability and computes the predicate directly. + """ + return self.dp.get_default_chg_spin() is not None diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 24ea448843..170cc516f5 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -264,6 +264,10 @@ def call( aparam: tf.Tensor | None = None, charge_spin: tf.Tensor | None = None, ) -> dict[str, tf.Tensor]: + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + am = getattr(model, "atomic_model", None) return model_call_from_call_lower( call_lower=call_lower, rcut=model.get_rcut(), @@ -277,12 +281,7 @@ def call( aparam=aparam, charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, - # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. Guard - # atomic_model too: test doubles (DummyModel) lack it. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + pair_excl=am.pair_excl if am is not None else None, ) return call @@ -418,9 +417,8 @@ def get_pair_exclude_types() -> tf.Tensor: # the LAMMPS nlist before the traced call_lower_* consumes it # (decision #18/A4). The upper ``call`` already pre-excludes its # freshly built nlist. Guard atomic_model: test doubles may lack it. - pet = getattr( - getattr(model, "atomic_model", None), "pair_exclude_types", [] - ) + am = getattr(model, "atomic_model", None) + pet = am.get_pair_exclude_types() if am is not None else [] flat = [int(t) for pair in (pet or []) for t in pair] return tf.constant(flat, dtype=tf.int64) @@ -470,7 +468,7 @@ def get_dim_chg_spin() -> tf.Tensor: @tf.function def has_default_chg_spin() -> tf.Tensor: - return tf.constant(model.has_default_chg_spin(), dtype=tf.bool) + return tf.constant(model.get_default_chg_spin() is not None, dtype=tf.bool) tf_model.has_default_chg_spin = has_default_chg_spin @@ -485,7 +483,8 @@ def get_default_chg_spin() -> tf.Tensor: # property models: persist the output name/dimension/intensiveness so # the evaluator can dispatch to DeepProperty and reshape the output. - if hasattr(model, "get_var_name"): + is_property_model = model.get_var_name() is not None + if is_property_model: @tf.function def get_var_name() -> tf.Tensor: diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 57efdb5c14..6347af2fc6 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -41,6 +41,32 @@ def decode_list_of_bytes(list_of_bytes: list[bytes]) -> list[str]: return [x.decode() for x in list_of_bytes] +def _decode_default_chg_spin( + has_default: bool, default_chg_spin: list[float] +) -> list[float] | None: + """Decode the exported default charge-spin values. + + The exporter always writes a ``get_default_chg_spin`` tensor next to + the boolean marker, using an EMPTY tensor when the model has no + default. Normalize the no-default case to ``None`` so the wrapper + upholds the live-model invariant + ``get_default_chg_spin() is None == no default``. + + Parameters + ---------- + has_default : bool + The boolean marker exported by ``has_default_chg_spin``. + default_chg_spin : list[float] + The decoded ``get_default_chg_spin`` tensor values. + + Returns + ------- + list[float] | None + The default values, or ``None`` when the artifact has none. + """ + return default_chg_spin if has_default else None + + class TFModelWrapper(tf.Module): def __init__( self, @@ -93,10 +119,13 @@ def __init__( if hasattr(self.model, "has_default_chg_spin") else False ) - self.default_chg_spin = ( - self.model.get_default_chg_spin().numpy().tolist() - if hasattr(self.model, "get_default_chg_spin") - else None + self.default_chg_spin = _decode_default_chg_spin( + self._has_default_chg_spin, + ( + self.model.get_default_chg_spin().numpy().tolist() + if hasattr(self.model, "get_default_chg_spin") + else [] + ), ) # property models only (absent for other model types). if hasattr(self.model, "get_var_name"): diff --git a/deepmd/jax/jax_md/__init__.py b/deepmd/jax/jax_md/__init__.py index 8450b96e5b..fa3859d6dc 100644 --- a/deepmd/jax/jax_md/__init__.py +++ b/deepmd/jax/jax_md/__init__.py @@ -328,7 +328,8 @@ def _eval_with_jax_md_neighbor( # re-applies it. The JAX-MD neighbor list is built without exclusion, so # fold it in at this ingestion seam -- otherwise excluded pairs would be # silently included (fail-open). - pair_excl = getattr(getattr(model, "atomic_model", None), "pair_excl", None) + am = getattr(model, "atomic_model", None) + pair_excl = am.pair_excl if am is not None else None if pair_excl is not None: from deepmd.dpmodel.utils.nlist import ( apply_pair_exclusion_nlist, diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index d78dbc38ed..a35de2a49b 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -973,7 +973,7 @@ def _prepare_batch( box=jax_data["box"] if jax_data["find_box"] else None, fparam=jax_data.get("fparam", None), aparam=jax_data.get("aparam", None), - pair_excl=getattr(model.atomic_model, "pair_excl", None), + pair_excl=model.atomic_model.pair_excl, ) return jax_data, extended_coord, extended_atype, nlist, mapping, fp, ap diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index a347c2dff8..62a6851160 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -335,6 +335,7 @@ def call_lower_with_fixed_do_atomic_virial( data["@variables"]["stablehlo_atomic_virial_no_ghost"] = np.void( serialized_atomic_virial_no_ghost ) + is_property_model = model.get_var_name() is not None data["constants"] = { "type_map": model.get_type_map(), "rcut": model.get_rcut(), @@ -352,15 +353,9 @@ def call_lower_with_fixed_do_atomic_virial( # property models: the output name/dimension/intensiveness cannot be # recovered from the StableHLO alone, so persist them for the # evaluator (None for non-property models). - "var_name": model.get_var_name() - if hasattr(model, "get_var_name") - else None, - "task_dim": model.get_task_dim() - if hasattr(model, "get_task_dim") - else None, - "intensive": model.get_intensive() - if hasattr(model, "get_intensive") - else False, + "var_name": model.get_var_name(), + "task_dim": model.get_task_dim() if is_property_model else None, + "intensive": model.get_intensive() if is_property_model else False, } save_dp_model(filename=model_file, model_dict=data) elif model_file.endswith(".savedmodel"): diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index ebfdd9981d..856fab953c 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -34,12 +34,12 @@ def canonical_model_eligible(model: Any) -> bool: fitting = getattr(atomic_model, "fitting_net", None) if descriptor is None or fitting is None: return False - if not bool(getattr(descriptor, "geo_compress", False)): + if not descriptor.get_geo_compress(): return False eligible = getattr(descriptor, "_fused_eligible", None) if not callable(eligible) or not bool(eligible("cuda")): return False - if getattr(atomic_model, "pair_excl", None) is not None: + if atomic_model.pair_excl is not None: return False if getattr(atomic_model, "atom_excl", None) is not None: return False diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 749bcdb76b..fd393d901d 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -239,12 +239,10 @@ def get_intensive(self) -> bool: def get_var_name(self) -> str: """Get the name of the property.""" - if hasattr(self.dp.model["Default"], "get_var_name") and callable( - getattr(self.dp.model["Default"], "get_var_name") - ): - return self.dp.model["Default"].get_var_name() - else: + var_name = self.dp.model["Default"].get_var_name() + if var_name is None: raise NotImplementedError + return var_name @property def model_type(self) -> type["DeepEvalWrapper"]: diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 50fc4b4a5f..5a1ff9d2cd 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -189,10 +189,6 @@ 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 - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values.""" return None diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 1b37184b18..334ebdc5f5 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -65,9 +65,7 @@ def __init__( self.sel = self.descriptor.get_sel() self.fitting_net = fitting super().init_out_stat() - self.add_chg_spin_ebd: bool = getattr( - self.descriptor, "add_chg_spin_ebd", False - ) + self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0 self.enable_eval_descriptor_hook = False self.enable_eval_fitting_last_layer_hook = False self.eval_descriptor_list = [] @@ -85,13 +83,12 @@ def _string_to_array(s: str | list[str]) -> list[int]: ), ) self.buffer_type_map.name = "buffer_type_map" - if hasattr(self.descriptor, "has_message_passing"): - # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector) - self.register_buffer( - "buffer_has_message_passing", - paddle.to_tensor(self.descriptor.has_message_passing(), dtype="int32"), - ) - self.buffer_has_message_passing.name = "buffer_has_message_passing" + # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector) + self.register_buffer( + "buffer_has_message_passing", + paddle.to_tensor(self.descriptor.has_message_passing(), dtype="int32"), + ) + self.buffer_has_message_passing.name = "buffer_has_message_passing" # register 'ntypes' as buffer self.register_buffer( "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int32") @@ -488,14 +485,8 @@ def get_dim_chg_spin(self) -> int: 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) -> paddle.Tensor | None: """Get the default charge_spin values as a tensor.""" - if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): + if self.add_chg_spin_ebd: return self.descriptor.get_default_chg_spin() return None diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index 6071ab7d2b..824adaad8d 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -47,6 +47,12 @@ class DescriptorBlock(paddle.nn.Layer, ABC, make_plugin_registry("DescriptorBloc local_cluster = False + # Stat-behavior flags with concrete defaults so stat machinery can read + # them on any block without getattr probes; blocks that configure them + # assign instance attributes in __init__ (issue #5897). + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Self: if cls is DescriptorBlock: try: @@ -145,8 +151,7 @@ def share_params( # link buffers if hasattr(self, "mean"): if not resume and ( - not getattr(self, "set_stddev_constant", False) - or not getattr(self, "set_davg_zero", False) + not self.set_stddev_constant or not self.set_davg_zero ): # in case of change params during resume base_env = EnvMatStatSe(base_class) diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index ccd8944cc8..8db2792117 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -370,10 +370,6 @@ 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 diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index b976ed6dd5..8ca36e2091 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -337,10 +337,6 @@ 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 diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 95fe69422b..9a83fe56c4 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -458,10 +458,6 @@ 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) -> paddle.Tensor | None: """Get the default charge_spin values as a tensor.""" if self.default_chg_spin is None: diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 0098ca2186..918b8be224 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -123,10 +123,6 @@ 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 diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index d61294650a..c83e22561d 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -194,10 +194,6 @@ 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 diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index 7db800bf22..f5beb501fc 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -562,10 +562,6 @@ 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() - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values.""" return self.atomic_model.get_default_chg_spin() diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 53f768fded..688f0c81c0 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -1336,10 +1336,9 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ] additional_data_requirement += spin_requirement_items 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 - ) + default_cs = _model.get_default_chg_spin() + has_default_cs = default_cs is not None + cs_default = default_cs.cpu().numpy() if has_default_cs else 0.0 additional_data_requirement.append( DataRequirementItem( "charge_spin", diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 5e0cc1f50e..93092ecd9c 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -338,6 +338,35 @@ class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer return False, value +def register_buffer_replacing_slot( + obj: torch.nn.Module, name: str, tensor: torch.Tensor +) -> None: + """Register a buffer, replacing a same-named plain attribute if present. + + Some dpmodel base ``__init__``s declare a capability slot such as + ``self.type_embd_data = None`` so its presence is a class property + rather than a runtime accident (issue #5897). A torch-native + compression path may later want to store the real value as a + persistent buffer via ``torch.nn.Module.register_buffer`` directly + (bypassing ``__setattr__``/``dpmodel_setattr``, e.g. because the + value is already a computed ``torch.Tensor`` rather than a + ``np.ndarray``). ``register_buffer`` raises ``KeyError`` if the name + already exists as a plain (non-buffer) attribute, so remove it first. + + Parameters + ---------- + obj : torch.nn.Module + The pt_expt wrapper object to register the buffer on. + name : str + The buffer name. + tensor : torch.Tensor + The tensor value to store as a buffer. + """ + if hasattr(obj, name) and name not in obj._buffers: + delattr(obj, name) + torch.nn.Module.register_buffer(obj, name, tensor) + + # --------------------------------------------------------------------------- # Utility # --------------------------------------------------------------------------- diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 47145d96ec..91266b8788 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -56,6 +56,7 @@ triton_infer_level, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -529,7 +530,7 @@ def _store_type_embd_data(self) -> None: self.se_atten.embeddings_strip[0].call(two_side_embd).detach() ) - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index 6d94f5f5cc..46bc774cc9 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -21,6 +21,7 @@ remap_atype_to_padding, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -284,7 +285,7 @@ def _store_type_embd_data(self) -> None: self.repinit.embeddings_strip[0].call(two_side_embd).detach() ) - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/descriptor/repflows.py b/deepmd/pt_expt/descriptor/repflows.py index dacab9f464..2b88da3562 100644 --- a/deepmd/pt_expt/descriptor/repflows.py +++ b/deepmd/pt_expt/descriptor/repflows.py @@ -57,7 +57,7 @@ def _exchange_ghosts( # entirely, so combining it with comm_dict is contradictory. # Surface this as a clear error rather than producing silently # wrong results. - if getattr(self, "use_loc_mapping", False): + if self.use_loc_mapping: raise RuntimeError( "DescrptBlockRepflows._exchange_ghosts: comm_dict is " "set but use_loc_mapping=True. Multi-rank parallel " diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py index 512e0ebe7a..69b786552d 100644 --- a/deepmd/pt_expt/descriptor/se_t_tebd.py +++ b/deepmd/pt_expt/descriptor/se_t_tebd.py @@ -16,6 +16,7 @@ remap_atype_to_padding, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -159,7 +160,7 @@ def _store_type_embd_data(self) -> None: ) # Run through the strip embedding network embd_tensor = self.se_ttebd.embeddings_strip[0].call(two_side).detach() - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 5d06b8b8e6..b084f3fdef 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -689,27 +689,13 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "sel": model.get_sel(), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), - "dim_chg_spin": model.get_dim_chg_spin() - if hasattr(model, "get_dim_chg_spin") - else 0, + "dim_chg_spin": model.get_dim_chg_spin(), "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 - ), + "has_chg_spin_ebd": model.has_chg_spin_ebd(), + "has_default_chg_spin": model.get_default_chg_spin() is not None, + "default_chg_spin": model.get_default_chg_spin(), "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } @@ -907,14 +893,19 @@ def get_dim_aparam(self) -> int: 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"): + if self._dpmodel is not None: return bool(self._dpmodel.has_chg_spin_ebd()) return bool(self.metadata.get("has_chg_spin_ebd", self.get_dim_chg_spin() > 0)) 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()) + """Check whether the model has a default charge_spin fallback. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` on + the live-model interfaces; this DeepEval wrapper method is kept for + API stability and computes the predicate directly. + """ + if self._dpmodel is not None: + return self._dpmodel.get_default_chg_spin() is not None return bool( self.metadata.get( "has_default_chg_spin", @@ -924,7 +915,7 @@ def has_default_chg_spin(self) -> bool: def get_dim_chg_spin(self) -> int: """Get the width of charge/spin condition inputs.""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_dim_chg_spin"): + if self._dpmodel is not None: return self._dpmodel.get_dim_chg_spin() return int(self.metadata.get("dim_chg_spin", 0)) @@ -965,6 +956,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: """The evaluator of the model type.""" if self._dpmodel is not None: model_output_type = self._dpmodel.model_output_type() + var_name = self._dpmodel.get_var_name() else: # Metadata-only mode: derive the output-type set from the # fitting_output_defs names. `model_output_type()` on a @@ -973,6 +965,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: model_output_type = [ d.name for d in self._model_output_def.def_outp.get_data().values() ] + var_name = None if "energy" in model_output_type: return DeepPot elif "dos" in model_output_type: @@ -983,11 +976,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar elif "wfc" in model_output_type: return DeepWFC - elif ( - self._dpmodel is not None - and hasattr(self._dpmodel, "get_var_name") - and self._dpmodel.get_var_name() in model_output_type - ): + elif var_name is not None and var_name in model_output_type: return DeepProperty else: raise RuntimeError("Unknown model type") @@ -1012,7 +1001,7 @@ def get_numb_dos(self) -> int: def get_var_name(self) -> str: """Get the name of the property (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_var_name"): + if self._dpmodel is not None and self._dpmodel.get_var_name() is not None: return self._dpmodel.get_var_name() raise NotImplementedError( "get_var_name is only available for property models with the " @@ -1021,7 +1010,7 @@ def get_var_name(self) -> str: def get_task_dim(self) -> int: """Get the output dimension of the property (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_task_dim"): + if self._dpmodel is not None: return self._dpmodel.get_task_dim() raise NotImplementedError( "get_task_dim is only available for property models with the " @@ -1030,7 +1019,7 @@ def get_task_dim(self) -> int: def get_intensive(self) -> bool: """Whether the property is intensive (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_intensive"): + if self._dpmodel is not None: return self._dpmodel.get_intensive() raise NotImplementedError( "get_intensive is only available for property models with the " @@ -2424,7 +2413,7 @@ def _model_pair_excl(self) -> "PairExcludeMask | None": ) if self._dpmodel is not None: - pe = getattr(self._dpmodel.atomic_model, "pair_excl", None) + pe = self._dpmodel.atomic_model.pair_excl pet = pe.get_exclude_types() if pe is not None else [] else: pet = self.metadata.get("pair_exclude_types", []) @@ -2618,9 +2607,7 @@ def eval_descriptor( ext_atype_t, nlist_t, mapping=mapping_t, - charge_spin=charge_spin_t - if getattr(dp_am, "add_chg_spin_ebd", False) - else None, + charge_spin=charge_spin_t if dp_am.has_chg_spin_ebd() else None, ) return descriptor.detach().cpu().numpy() @@ -2689,9 +2676,7 @@ def eval_fitting_last_layer( ext_atype_t, nlist_t, mapping=mapping_t, - charge_spin=charge_spin_t - if getattr(dp_am, "add_chg_spin_ebd", False) - else None, + charge_spin=charge_spin_t if dp_am.has_chg_spin_ebd() else None, ) atype = ext_atype_t[:, :natoms] fitting_net = dp_am.fitting_net diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index c6def8f136..f59d507727 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -780,9 +780,10 @@ def _call_common_graph( with_csr = ( not self.training and cuda_infer_level() >= 1 - and bool(getattr(_desc, "geo_compress", False)) + and _desc is not None + and _desc.get_geo_compress() ) - pair_excl = getattr(self.atomic_model, "pair_excl", None) + pair_excl = self.atomic_model.pair_excl ng = build_neighbor_graph_for_method( method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr ) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index a489ec2a75..805e55cbe9 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -383,9 +383,9 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ) if _model.has_chg_spin_ebd(): - has_default_cs = _model.has_default_chg_spin() + default_cs = _model.get_default_chg_spin() + has_default_cs = default_cs is not None if has_default_cs: - default_cs = _model.get_default_chg_spin() if hasattr(default_cs, "cpu"): default_cs = default_cs.cpu().numpy() else: @@ -1178,7 +1178,7 @@ def forward( distinguish_types=False, # model-level pair exclusion is a nlist-BUILD transform (decision # #18/A4); the compiled dense lower consumes a pre-excluded nlist. - pair_excl=getattr(self.original_model.atomic_model, "pair_excl", None), + pair_excl=self.original_model.atomic_model.pair_excl, ) ext_coord = ext_coord.reshape(nframes, -1, 3) @@ -1215,9 +1215,7 @@ def forward( .reshape(1, _dim_fparam) .expand(nframes, -1) ) - _dim_cs = ( - _model.get_dim_chg_spin() if hasattr(_model, "get_dim_chg_spin") else 0 - ) + _dim_cs = _model.get_dim_chg_spin() if charge_spin is None and _dim_cs > 0: _default_cs = _model.get_default_chg_spin() if _default_cs is not None: @@ -1390,9 +1388,7 @@ def _forward_graph( .reshape(1, _dim_fparam) .expand(nframes, -1) ) - _dim_cs = ( - _model.get_dim_chg_spin() if hasattr(_model, "get_dim_chg_spin") else 0 - ) + _dim_cs = _model.get_dim_chg_spin() if charge_spin is None and _dim_cs > 0: _default_cs = _model.get_default_chg_spin() if _default_cs is not None: @@ -1409,7 +1405,7 @@ def _forward_graph( # level pair_exclude is a graph-BUILD transform (decision #18): fold it # into edge_mask here so the compiled lower consumes a pre-excluded graph # (the lower no longer re-applies it), matching the eager path exactly. - pair_excl = getattr(_model.atomic_model, "pair_excl", None) + pair_excl = _model.atomic_model.pair_excl ng = build_neighbor_graph_for_method( getattr(_model, "neighbor_graph_method", "dense"), coord_3d, diff --git a/deepmd/pt_expt/utils/network.py b/deepmd/pt_expt/utils/network.py index 004ba94401..f3ca3b392d 100644 --- a/deepmd/pt_expt/utils/network.py +++ b/deepmd/pt_expt/utils/network.py @@ -99,7 +99,7 @@ def __setattr__(self, name: str, value: Any) -> None: self._buffers[name] = None return None return super().__setattr__(name, None) - if getattr(self, "trainable", False): + if self.trainable: param = ( value if isinstance(value, torch.nn.Parameter) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index e82ae24e0a..5d7be558bb 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -394,7 +394,7 @@ def _make_sample_inputs( else: aparam = None - dim_chg_spin = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + dim_chg_spin = model.get_dim_chg_spin() if dim_chg_spin > 0: charge_spin = torch.zeros( nframes, dim_chg_spin, dtype=torch.float64, device=_env.DEVICE @@ -515,7 +515,7 @@ def build_synthetic_graph_inputs( ntypes = len(model.get_type_map()) dim_fparam = model.get_dim_fparam() dim_aparam = model.get_dim_aparam() - dim_chg_spin = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + dim_chg_spin = model.get_dim_chg_spin() # Box large enough to avoid PBC degeneracy; centered coords. box_size = rcut * 3.0 @@ -1054,25 +1054,13 @@ def _collect_metadata( "nnei": sum(model.get_sel()), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), - "dim_chg_spin": ( - model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 - ), + "dim_chg_spin": model.get_dim_chg_spin(), "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 - ), + "has_chg_spin_ebd": model.has_chg_spin_ebd(), + "has_default_chg_spin": model.get_default_chg_spin() is not None, + "default_chg_spin": _metadata_value_to_json(model.get_default_chg_spin()), "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 @@ -1101,15 +1089,21 @@ def _collect_metadata( # ``atomic_model.has_message_passing()`` is important for composite # atomic models (e.g. ``LinearAtomicModel`` in DP-ZBL) which don't # expose a single ``.descriptor`` but do aggregate the flag across - # their sub-models. ``descriptor.has_message_passing()`` is the - # fallback for any future wrapper that lacks the higher-level - # methods. + # their sub-models. ``has_message_passing`` is declared on the base + # model/atomic-model/descriptor classes, so every concrete object at + # each level implements it; ``descriptor.has_message_passing()`` only + # matters as a fallback when an upstream level raises + # ``NotImplementedError`` (e.g. an atomic model that intentionally + # opts out), never for a missing method. def _probe_has_message_passing(obj: object) -> bool | None: - if obj is None or not hasattr(obj, "has_message_passing"): + # has_message_passing is @abstractmethod on the base descriptor, so + # every concrete descriptor implements it; a wrapper lacking it is a + # construction bug that must raise, not degrade silently. + if obj is None: return None try: return bool(obj.has_message_passing()) - except (AttributeError, NotImplementedError): + except NotImplementedError: return None result: bool | None = None @@ -1140,12 +1134,19 @@ def _probe_has_message_passing(obj: object) -> bool | None: # feeders (C++ ``DeepPotPTExpt::init``, metadata-only DeepEval) rebuild # the mask. Descriptor-level ``exclude_types`` needs NO metadata: it is # fully inside the compiled artifact. + from deepmd.dpmodel.atomic_model.base_atomic_model import ( + BaseAtomicModel, + ) + pair_exclude_types: list[list[int]] = [] for obj in ( getattr(model, "atomic_model", None), model, ): - pet = getattr(obj, "pair_exclude_types", None) + # `obj` may be the atomic model (the owner of pair_exclude_types) or + # the full model (e.g. the ``model`` fallback above); only the former + # implements the accessor, so gate on it instead of getattr-probing. + pet = obj.get_pair_exclude_types() if isinstance(obj, BaseAtomicModel) else None if pet: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break @@ -1542,7 +1543,7 @@ def _trace_and_export( ) _forbidden = forbidden_dims_from_model(model) - _dim_cs = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + _dim_cs = model.get_dim_chg_spin() if _dim_cs > 1: _forbidden.add(int(_dim_cs)) nframes_sample = next_safe_prime(5, _forbidden) diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 1326e1c881..bd69555dc4 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -87,6 +87,8 @@ def call_common( # ``_input_type_cast`` (dpmodel make_model) returns a ``spin`` slot # for the native-spin graph route; tf2 has no spin/graph lower, so # it is discarded here. + # Guard atomic_model too: test doubles may lack it. + am = getattr(self, "atomic_model", None) cc, bb, fp, ap, cs, _, input_prec = self._input_type_cast( to_tensorflow_array(coord), box=to_tensorflow_array(box), @@ -113,11 +115,8 @@ def call_common( # Model-level pair exclusion is a nlist-BUILD transform # (decision #18/A4): fold it into the freshly built nlist here so # the live/eager TF2 upper path matches the SavedModel export and - # the other backends. Identity when nothing is excluded. Guard - # atomic_model too: test doubles may lack it. - pair_excl=getattr( - getattr(self, "atomic_model", None), "pair_excl", None - ), + # the other backends. Identity when nothing is excluded. + pair_excl=am.pair_excl if am is not None else None, pass_lower_kwargs=True, ) return self._output_type_cast(model_predict, input_prec) diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 4c1e8af671..81a1139221 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -164,9 +164,10 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ) if _model.has_chg_spin_ebd(): - has_default_cs = _model.has_default_chg_spin() + default_chg_spin = _model.get_default_chg_spin() + has_default_cs = default_chg_spin is not None default_cs = ( - np.asarray(to_tf_tensor(_model.get_default_chg_spin()).numpy()) + np.asarray(to_tf_tensor(default_chg_spin).numpy()) if has_default_cs else 0.0 ) @@ -882,6 +883,8 @@ def compiled_prepare_lower_batch( aparam=to_tensorflow_array(aparam), charge_spin=to_tensorflow_array(charge_spin), ) + # Guard atomic_model for test doubles. + am = getattr(model, "atomic_model", None) return prepare_lower_inputs( rcut=model.get_rcut(), sel=model.get_sel(), @@ -895,10 +898,8 @@ def compiled_prepare_lower_batch( # Model-level pair exclusion is a nlist-BUILD transform # (decision #18/A4): the compiled lower consumes a pre-excluded # nlist, so fold exclusion in here at the compiled-training - # prepare seam. Guard atomic_model for test doubles. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + # prepare seam. + pair_excl=am.pair_excl if am is not None else None, ) return compiled_prepare_lower_batch diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index baedd5ba87..63bbdf6583 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -377,6 +377,10 @@ def call( fparam: tf.Tensor, aparam: tf.Tensor, ) -> dict[str, tf.Tensor]: + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # traced lower consumes a pre-excluded nlist. Guard atomic_model + # too: test doubles (DummyModel) lack it. + am = getattr(model, "atomic_model", None) return unwrap_value( model_call_from_call_lower( call_lower=call_lower, @@ -390,12 +394,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, - # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. Guard - # atomic_model too: test doubles (DummyModel) lack it. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + pair_excl=am.pair_excl if am is not None else None, ) ) @@ -517,7 +516,8 @@ def get_pair_exclude_types() -> tf.Tensor: # traced call_lower_* consumes it (decision #18/A4). The compiled ``call`` # already pre-excludes its freshly built nlist. Guard atomic_model: test # doubles may lack it. - pet = getattr(getattr(model, "atomic_model", None), "pair_exclude_types", []) + am = getattr(model, "atomic_model", None) + pet = am.get_pair_exclude_types() if am is not None else [] flat = [int(t) for pair in (pet or []) for t in pair] return tf.constant(flat, dtype=tf.int64) @@ -555,7 +555,12 @@ def get_default_fparam() -> tf.Tensor: # property models: persist the output name/dimension/intensiveness so the # evaluator can dispatch to DeepProperty and reshape the output. - if hasattr(model, "get_var_name"): + # ``get_var_name`` is declared on every model with a concrete default + # of ``None`` for non-property models (issue #5897), so a bare + # ``hasattr`` check is always true; gate on the actual return value + # instead, mirroring deepmd/jax/jax2tf/serialization.py. + is_property_model = model.get_var_name() is not None + if is_property_model: @tf.function def get_var_name() -> tf.Tensor: diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index 310a21e83f..2933facb13 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -9,6 +9,34 @@ There are several scenarios where one might want to adjust the output bias after such as zero-shot testing (similar to the procedure before the first step in fine-tuning) or manually setting the output bias. +## The two statistic modes, precisely + +The model energy decomposes as `E = E_model + E_bias`, where `E_model` is +whatever the model computes (a learned network, an analytical term such as +ZBL bridging, or a `linear_ener` combination of models) and `E_bias` is the +per-type output bias. + +- **`set` (`set-by-statistic`)** assigns `E_bias` directly: either the + user-given values (`-b`), or the per-type least-squares statistic of the + **raw data labels**. It is independent of `E_model` by definition — it + ignores a trained network, and it equally ignores an analytical + contribution such as the ZBL term of a bridged model. The result is + reproducible and idempotent for a given dataset, but it contains **no + compensation for `E_model`**: after `set`, the remaining error on the + calibration data is the configuration-dependent `E_model` itself, plus + any residual of the raw-label least-squares fit. +- **`change` (`change-by-statistic`)** assigns `E_bias` from the residual: + the per-type statistic of the labels **minus the complete model + prediction** (including any analytical bridging term), added to the + existing bias. Use this mode for a self-consistent calibration of a + trained (or bridged) model. + +For a bridged model — or any model whose `E_model` is significantly nonzero +on the calibration data — `set` leaves `E_model` uncompensated and can absorb +its composition-correlated component into `E_bias`, so the forward pass may add +that component again. Use `change` to fit the residual against the complete +model prediction. + The `dp change-bias` command supports the following methods for adjusting the bias: ::::{tab-set} diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 461b0b163d..3acee5ec35 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -341,6 +341,17 @@ When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the same value as `bridging_r_inner` so frames with shorter atom pairs are excluded from training. See `examples/water/dpa4/input-zbl.json` for a complete example. +> [!NOTE] +> Output-bias statistics and bridging: the model energy is +> `E = E_model + E_bias`, and the ZBL term belongs to `E_model`. The +> `set-by-statistic` bias mode (initial statistics, finetune with a +> random fitting, `dp change-bias --mode set`) fits `E_bias` to the raw +> data labels and by definition ignores `E_model` — the analytical ZBL +> contribution included. For a self-consistent calibration of a bridged +> model use `change-by-statistic`, which subtracts the complete bridged +> prediction. See [change-bias](change-bias.md) for the precise +> definitions. + ## Performance and precision ### Training-time settings diff --git a/source/jax2tf_tests/test_serialization.py b/source/jax2tf_tests/test_serialization.py index 27751b7c43..bc4cf02239 100644 --- a/source/jax2tf_tests/test_serialization.py +++ b/source/jax2tf_tests/test_serialization.py @@ -28,6 +28,22 @@ def _saved_model_ops(model_dir: Path) -> set[str]: return ops +def test_decode_default_chg_spin_preserves_the_no_default_marker() -> None: + """The exporter writes an EMPTY tensor when the model has no default + charge-spin; the wrapper must decode it back to ``None`` so the + live-model invariant ``get_default_chg_spin() is None == no default`` + holds at the artifact boundary too. + """ + pytest.importorskip("jax") + + from deepmd.jax.jax2tf.tfmodel import ( + _decode_default_chg_spin, + ) + + assert _decode_default_chg_spin(False, []) is None + assert _decode_default_chg_spin(True, [2.0, 1.0]) == [2.0, 1.0] + + def test_savedmodel_export_contains_xla_call_module(tmp_path, monkeypatch) -> None: pytest.importorskip("jax") pytest.importorskip("flax") @@ -122,6 +138,10 @@ def has_default_chg_spin(self) -> bool: def get_default_chg_spin(self) -> None: return None + def get_var_name(self) -> None: + # non-property model, matching the make_base_model default + return None + class DummyChargeSpinModel(DummyModel): dim_chg_spin = 2 diff --git a/source/op/pd/comm.cc b/source/op/pd/comm.cc index 548e5db83a..e5f0afd88f 100644 --- a/source/op/pd/comm.cc +++ b/source/op/pd/comm.cc @@ -12,6 +12,29 @@ #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) #include "device.h" + +template +static void copy_local_tensor_data(FPTYPE* dst, + const FPTYPE* src, + size_t count, + const paddle::Place& place) { + if (count == 0) { + return; + } + + // CUDA-aware MPI describes whether MPI can consume device pointers; it does + // not describe where this particular Paddle tensor lives. Self-swaps must + // select the copy primitive from the actual tensor place so CPU tensors also + // work in CUDA/ROCm-enabled builds. + if (phi::is_gpu_place(place)) { + gpuMemcpy(dst, src, count * sizeof(FPTYPE), gpuMemcpyDeviceToDevice); + } else { + // CPU and host-pinned tensors are both host-addressable. Defaulting + // non-GPU places to memcpy also avoids treating a future host place as a + // CUDA pointer merely because the operator was built with CUDA support. + memcpy(dst, src, count * sizeof(FPTYPE)); + } +} #endif #ifdef USE_MPI @@ -83,13 +106,16 @@ void Border_forward_t(const paddle::Tensor& sendlist_tensor, int tensor_size = g1.dims()[1]; + // nlocal and nghost are scalar protocol values, independent of the number + // of communication swaps. In particular, nswap == 0 still needs one slot + // for each value before the host dereference below. paddle::Tensor cpu_nlocal = - paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); cpu_nlocal.copy_(nlocal_tensor, paddle::CPUPlace(), true); int nlocal = *(cpu_nlocal.data()); paddle::Tensor cpu_nghost = - paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); cpu_nghost.copy_(nghost_tensor, paddle::CPUPlace(), true); int nghost = *(cpu_nghost.data()); @@ -175,20 +201,8 @@ void Border_forward_t(const paddle::Tensor& sendlist_tensor, #endif #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) -#ifdef USE_MPI - if (cuda_aware == 0) { - memcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE)); - } else { - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); - } -#else - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nsend * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); -#endif + copy_local_tensor_data(recv_g1, send_g1, (size_t)nsend * tensor_size, + recv_g1_tensor.place()); #else memcpy(recv_g1, send_g1, @@ -317,7 +331,6 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, cpu_recvnum.copy_(recvnum_tensor, paddle::CPUPlace(), true); int* sendnum = cpu_recvnum.data(); - FPTYPE* local_g1 = d_local_g1_tensor.data(); int tensor_size = d_local_g1_tensor.dims()[1]; paddle::Tensor cpu_nlocal = @@ -381,20 +394,8 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, #endif if (nrecv) { #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) -#ifdef USE_MPI - if (cuda_aware == 0) { - memcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); - } else { - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); - } -#else - gpuMemcpy(recv_g1, send_g1, - (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), - gpuMemcpyDeviceToDevice); -#endif + copy_local_tensor_data(recv_g1, send_g1, (size_t)nrecv * tensor_size, + d_local_g1_tensor.place()); #else memcpy(recv_g1, send_g1, (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); @@ -408,18 +409,34 @@ void Border_backward_t(const paddle::Tensor& sendlist_tensor, d_local_g1_tensor, irecvlist, recv_g1_tensor.slice(0, nrecv), 0); } } -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - gpuDeviceSynchronize(); -#endif -#ifdef USE_MPI + // Forward swaps overwrite every ghost row with owner data. Reverse + // communication accumulates each ghost output gradient into its owner, but + // the original ghost input has no path to the output and therefore must + // receive a zero gradient. With no swaps, forward is the identity and the + // upstream ghost gradient remains valid. + if (nswap > 0 && nghost > 0) { + FPTYPE* ghost_g1 = d_local_g1_tensor.data() + nlocal * tensor_size; + size_t ghost_size = (size_t)nghost * tensor_size; #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - if (cuda_aware == 0) { - recv_g1_tensor_grad.copy_(d_local_g1_tensor, recv_g1_tensor_grad.place(), - true); - } + if (phi::is_gpu_place(d_local_g1_tensor.place())) { + gpuMemset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); + } else { + memset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); + } +#else + memset(ghost_g1, 0, ghost_size * sizeof(FPTYPE)); #endif + } +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + gpuDeviceSynchronize(); #endif + + // The reverse exchange may update a temporary on the original device or on + // the CPU fallback used by non-CUDA-aware MPI. Always publish that result to + // the custom-op gradient output instead of relying on incidental aliasing. + recv_g1_tensor_grad.copy_(d_local_g1_tensor, recv_g1_tensor_grad.place(), + true); } void Border_backward(const paddle::Tensor& sendlist_tensor, diff --git a/source/tests/common/dpmodel/test_descriptor_block_defaults.py b/source/tests/common/dpmodel/test_descriptor_block_defaults.py new file mode 100644 index 0000000000..daef3c46aa --- /dev/null +++ b/source/tests/common/dpmodel/test_descriptor_block_defaults.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Class-level defaults for the ``set_davg_zero`` / ``set_stddev_constant`` +stat-behavior flags on ``DescriptorBlock`` (issue #5897): stat machinery +must be able to read these flags on any block without a ``getattr`` probe. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.descriptor.descriptor import ( + DescriptorBlock, +) +from deepmd.dpmodel.descriptor.make_base_descriptor import ( + make_base_descriptor, +) +from deepmd.dpmodel.utils.env_mat_stat import ( + merge_env_stat, +) + + +def test_block_stat_flags_have_class_defaults() -> None: + """Every DescriptorBlock answers the stat-behavior flags without a + getattr probe: class-level defaults False, blocks override in __init__. + """ + assert DescriptorBlock.set_davg_zero is False + assert DescriptorBlock.set_stddev_constant is False + + +def test_base_descriptor_stat_flags_have_class_defaults() -> None: + """The ``BD`` base in ``make_base_descriptor`` (the ``Descriptor``-side + twin of ``DescriptorBlock`` above) carries the same concrete class + defaults, so ``merge_env_stat`` -- which accepts either a ``Descriptor`` + or a ``DescriptorBlock`` as ``base_obj`` -- can read the flags on a bare + ``Descriptor`` without a ``getattr`` probe. + """ + bd = make_base_descriptor(np.ndarray, "call") + assert bd.set_davg_zero is False + assert bd.set_stddev_constant is False + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def test_merge_env_stat_on_bare_descriptor_no_attribute_error() -> None: + """``merge_env_stat`` reads ``base_obj.set_davg_zero`` / + ``set_stddev_constant`` unconditionally (no ``getattr`` probe). Pin that + this does not raise ``AttributeError`` when ``base_obj`` is a bare + se-family ``Descriptor`` (not a ``DescriptorBlock``) which never sets + those flags itself and instead relies on the ``BD`` base's class + defaults. + """ + base = DescrptSeA(6.0, 0.5, [10, 10]) + link = DescrptSeA(6.0, 0.5, [10, 10]) + sample = _sample() + base.compute_input_stats([sample]) + link.compute_input_stats([sample]) + # Would raise AttributeError before the BD-base class defaults existed + # if a descriptor never assigned instance attributes for these flags. + merge_env_stat(base, link) + + +def test_block_stat_flags_override_branch() -> None: + """A block constructed with ``set_davg_zero=True`` shadows the class + default with an instance attribute; a block constructed with the + default arguments keeps reading the class default (False). + """ + from deepmd.dpmodel.descriptor.dpa1 import ( + DescrptBlockSeAtten, + ) + + blk_default = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[6, 6], + ntypes=2, + ) + assert blk_default.set_davg_zero is False + + blk_override = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[6, 6], + ntypes=2, + set_davg_zero=True, + ) + assert blk_override.set_davg_zero is True diff --git a/source/tests/common/dpmodel/test_descriptor_dpa1.py b/source/tests/common/dpmodel/test_descriptor_dpa1.py index 1df473c087..371bd69c85 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa1.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa1.py @@ -5,6 +5,7 @@ from deepmd.dpmodel.descriptor import ( DescrptDPA1, + DescrptSeA, ) from ...seed import ( @@ -65,6 +66,27 @@ def test_lmax_two_serialization(self) -> None: for index in (0, 1, 4): np.testing.assert_allclose(actual[index], expected[index]) + def test_tebd_compression_slots_declared(self) -> None: + """Tebd-compression slots are class properties, not runtime accidents. + + ``type_embd_data``/``tebd_compress`` must be declared (defaulted) + in ``__init__`` of tebd-family descriptors so their presence does + not depend on whether compression was ever enabled. The jax + restore walker (``deepmd/jax/utils/serialization.py``) relies on + ``hasattr(obj, "tebd_compress")`` as a family-membership test, so + a non-tebd descriptor (e.g. ``DescrptSeA``) must never carry + either attribute (see issue #5897). + """ + em0 = DescrptDPA1(self.rcut, self.rcut_smth, self.sel, ntypes=2) + self.assertIsNone(em0.type_embd_data) + self.assertFalse(em0.tebd_compress) + self.assertIsNone(em0.se_atten.type_embd_data) + self.assertFalse(em0.se_atten.tebd_compress) + + se_a = DescrptSeA(self.rcut, self.rcut_smth, self.sel) + self.assertFalse(hasattr(se_a, "type_embd_data")) + self.assertFalse(hasattr(se_a, "tebd_compress")) + def test_multiple_frames(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape diff --git a/source/tests/common/dpmodel/test_make_base_fitting.py b/source/tests/common/dpmodel/test_make_base_fitting.py new file mode 100644 index 0000000000..0eacd2b117 --- /dev/null +++ b/source/tests/common/dpmodel/test_make_base_fitting.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +from deepmd.dpmodel.fitting.make_base_fitting import ( + make_base_fitting, +) + + +class MinimalFitting(make_base_fitting(np.ndarray)): + """Smallest concrete fitting: exercises the base-declared defaults.""" + + def output_def(self): + raise NotImplementedError + + def fwd(self, *args, **kwargs): + raise NotImplementedError + + def get_type_map(self): + return [] + + def change_type_map(self, type_map, model_with_new_type_stat=None): + raise NotImplementedError + + def serialize(self): + return {} + + @classmethod + def deserialize(cls, data): + return cls() + + +def test_reinit_exclude_default_noop_on_empty() -> None: + MinimalFitting().reinit_exclude([]) + + +def test_reinit_exclude_default_raises_on_nonempty() -> None: + with pytest.raises(NotImplementedError): + MinimalFitting().reinit_exclude([0]) diff --git a/source/tests/common/dpmodel/test_model_compression.py b/source/tests/common/dpmodel/test_model_compression.py index 29f5c281a8..8599d3e51b 100644 --- a/source/tests/common/dpmodel/test_model_compression.py +++ b/source/tests/common/dpmodel/test_model_compression.py @@ -344,6 +344,8 @@ def test_se_atten_enable_compression(self) -> None: self.assertTrue(compressed.compress) self.assertTrue(compressed.geo_compress) + self.assertTrue(compressed.get_geo_compress()) + self.assertFalse(descriptor.get_geo_compress()) serialized = compressed.serialize() self.assertEqual(serialized["@version"], 3) self.assertIn("compress", serialized) diff --git a/source/tests/common/dpmodel/test_pair_exclude_contract.py b/source/tests/common/dpmodel/test_pair_exclude_contract.py new file mode 100644 index 0000000000..ce2e21bdde --- /dev/null +++ b/source/tests/common/dpmodel/test_pair_exclude_contract.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Pin the ``pair_exclude_types``/``pair_excl`` construction-path contract +(issue #5897 task 5): ``BaseAtomicModel.__init__`` (via +``reinit_pair_exclude``) always sets both attributes, on every construction +path (direct ``__init__`` and ``deserialize``); ``get_pair_exclude_types()`` +is the public accessor, ``pair_excl`` a pinned direct-access attribute. +""" + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) + +RCUT = 2.2 +RCUT_SMTH = 0.4 +SEL = [5, 2] +NTYPES = 2 +TYPE_MAP = ["foo", "bar"] + + +def _make_minimal_atomic_model( + pair_exclude_types: list[tuple[int, int]], +) -> DPAtomicModel: + ds = DescrptSeA( + RCUT, + RCUT_SMTH, + SEL, + ) + ft = InvarFitting( + "energy", + NTYPES, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + return DPAtomicModel( + ds, + ft, + type_map=TYPE_MAP, + pair_exclude_types=pair_exclude_types, + ) + + +def test_pair_excl_exists_after_init_and_deserialize() -> None: + md0 = _make_minimal_atomic_model(pair_exclude_types=[(0, 1)]) + assert md0.get_pair_exclude_types() == [(0, 1)] + assert md0.pair_excl is not None + # the mask holds the SYMMETRIC CLOSURE of the accessor's pairs + assert {tuple(p) for p in md0.pair_excl.get_exclude_types()} == {(0, 1), (1, 0)} + md1 = type(md0).deserialize(md0.serialize()) + assert md1.get_pair_exclude_types() == [(0, 1)] + assert md1.pair_excl is not None + md2 = _make_minimal_atomic_model(pair_exclude_types=[]) + assert md2.get_pair_exclude_types() == [] + assert md2.pair_excl is None diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 0db497dd84..9d9ddba47d 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -396,7 +396,7 @@ def test_default_conditioning_accessors_are_forwarded(self) -> None: """``has_default_*`` must not fall through to the base either.""" bridged = self._model(bridging=True) plain = self._model(bridging=False) - assert bridged.has_default_chg_spin() == plain.has_default_chg_spin() + assert bridged.get_default_chg_spin() == plain.get_default_chg_spin() assert bridged.has_default_fparam() == plain.has_default_fparam() assert bridged.get_default_fparam() == plain.get_default_fparam() @@ -629,3 +629,49 @@ def test_forwarded_from_children(self) -> None: assert bridged.atomic_model.get_compute_stats_distinguish_types() == any( c.get_compute_stats_distinguish_types() for c in children ) + + +def test_set_by_statistic_fits_raw_labels_by_definition(): + """Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind. + + ``E = E_model + E_bias``; the set mode defines ``E_bias`` as the + per-type statistic of the raw labels, independent of ``E_model`` -- + the composition-level fit ignores the learned child and equally + ignores the analytical ZBL child. Children compute no output + statistics of their own (the composition is the one owner). Use + ``change-by-statistic`` for a calibration that compensates + ``E_model``. + """ + model = get_model(copy.deepcopy(ZBL_CONFIG)) + rng = np.random.default_rng(5) + coord = rng.uniform(1.0, 2.5, size=(1, 4, 3)) + box = (np.eye(3) * 8.0).reshape(1, 9) + samples, labels, counts_rows = [], [], [] + for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]): + counts = np.bincount(np.asarray(types[0]), minlength=2) + label = float(rng.normal()) + samples.append( + { + "coord": coord, + "atype": np.array(types), + "box": box, + "energy": np.array([[label]]), + "find_energy": np.float32(1.0), + "natoms": np.array([[4, 4, *counts]]), + } + ) + labels.append(label) + counts_rows.append(counts) + # Seed a nonzero bias: `set` must DISCARD it (an accidental additive + # implementation would shift the result by the seed). The dpmodel bias + # storage is separate from pt's, so the pin is mirrored here. + model.atomic_model.out_bias = np.ones_like(model.atomic_model.out_bias) + model.atomic_model.compute_or_load_out_stat(samples) + bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] + raw_fit = np.linalg.solve(np.array(counts_rows, dtype=np.float64), np.array(labels)) + np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) + # Idempotence: repeating the call from the fitted state must land on + # the same raw-label fit again. + model.atomic_model.compute_or_load_out_stat(samples) + repeated_bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] + np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) diff --git a/source/tests/infer/gen_dpa4_spin_chgspin.py b/source/tests/infer/gen_dpa4_spin_chgspin.py index 502767c6ac..e8aa53a359 100644 --- a/source/tests/infer/gen_dpa4_spin_chgspin.py +++ b/source/tests/infer/gen_dpa4_spin_chgspin.py @@ -125,7 +125,7 @@ def _build_model_dict() -> dict: f"expected the combined native-spin DPA4 to expose dim_chg_spin == 2, " f"got {model.get_dim_chg_spin()}" ) - assert model.has_default_chg_spin() + assert model.get_default_chg_spin() is not None model_dict = model.serialize() model_dict = jitter_zero_arrays(model_dict, np.random.default_rng(_JITTER_SEED)) return model_dict diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 0055a35b4b..913a851fe3 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -496,6 +496,11 @@ def __call__(self) -> tuple[np.ndarray, np.ndarray]: class _DescriptorWithStats: + # stat-behavior flags merge_env_stat reads directly on any + # Descriptor/DescriptorBlock (class defaults on the real bases) + set_davg_zero = False + set_stddev_constant = False + def __init__(self, stats: dict[str, StatItem]) -> None: self.stats = stats self.davg = np.asarray([0.0], dtype=np.float64) diff --git a/source/tests/pd/test_border_op.py b/source/tests/pd/test_border_op.py new file mode 100644 index 0000000000..adf0242cf4 --- /dev/null +++ b/source/tests/pd/test_border_op.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for Paddle border-exchange control and data tensors.""" + +import numpy as np +import paddle +import pytest + +deepmd_op_pd = pytest.importorskip( + "deepmd_op_pd", reason="the Paddle custom operator library is not built" +) + + +def _control_tensors(nswap: int) -> tuple[paddle.Tensor, ...]: + """Create the common CPU control tensors for a border exchange.""" + return ( + paddle.zeros([nswap], dtype="int32"), # sendproc + paddle.zeros([nswap], dtype="int32"), # recvproc + paddle.zeros([nswap], dtype="int32"), # sendnum + paddle.zeros([nswap], dtype="int32"), # recvnum + paddle.zeros([1], dtype="int64"), # unused communicator without MPI + ) + + +def test_border_op_accepts_no_swaps() -> None: + """Scalar atom counts must remain readable when ``nswap == 0``.""" + sendproc, recvproc, sendnum, recvnum, communicator = _control_tensors(0) + g1 = paddle.arange(6, dtype="float64").reshape([2, 3]) + + result = deepmd_op_pd.border_op( + paddle.zeros([0], dtype="int64"), + sendproc, + recvproc, + sendnum, + recvnum, + g1, + communicator, + paddle.to_tensor([2], dtype="int32"), + paddle.to_tensor([0], dtype="int32"), + ) + + np.testing.assert_array_equal(result.numpy(), g1.numpy()) + + +def test_border_op_self_copy_uses_cpu_place() -> None: + """A CUDA-enabled operator must not use a GPU copy for CPU tensors. + + NOTE: no CI job currently builds Paddle with CUDA (``test_cuda.yml`` + disables Paddle at the workflow level and ``test_python.yml`` installs the + CPU build), so ``copy_local_tensor_data`` is not compiled-and-executed by + any pipeline. This test therefore documents the intended CPU-branch + behavior rather than guarding it; it should gain real coverage once a CI + job builds Paddle with CUDA. + """ + # CUDA Paddle builds otherwise create tensors on the default GPU, which + # would leave the operator's CPU copy branch untested. + paddle.set_device("cpu") + sendproc, recvproc, sendnum, recvnum, communicator = _control_tensors(1) + sendnum = paddle.ones_like(sendnum) + recvnum = paddle.ones_like(recvnum) + + # The C++ operator receives the LAMMPS send lists as pointer-valued int64 + # entries. Keep this NumPy owner alive through the call so the pointed-to + # int32 index remains valid. + send_indices = np.array([1], dtype=np.int32) + sendlist = paddle.to_tensor([send_indices.ctypes.data], dtype="int64") + g1_leaf = paddle.to_tensor( + [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], stop_gradient=False + ) + # Paddle rejects an in-place custom op on an autograd leaf. This identity + # keeps a leaf for checking gradients while letting border_op update g1. + g1 = g1_leaf * 1.0 + + result = deepmd_op_pd.border_op( + sendlist, + sendproc, + recvproc, + sendnum, + recvnum, + g1, + communicator, + paddle.to_tensor([2], dtype="int32"), + paddle.to_tensor([1], dtype="int32"), + ) + + np.testing.assert_array_equal( + result.numpy(), np.array([[1.0, 2.0], [3.0, 4.0], [3.0, 4.0]]) + ) + # Backpropagation runs the reverse self-swap, which needs the same + # place-based CPU/GPU dispatch as the forward copy. + result.sum().backward() + expected_grad = np.array([[1.0, 1.0], [2.0, 2.0], [0.0, 0.0]]) + np.testing.assert_array_equal(g1_leaf.grad.numpy(), expected_grad) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 8855fe7aa4..d25b743fae 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -2259,6 +2259,64 @@ def test_change_out_bias_is_invariant_for_self_labels(self) -> None: ) ) + def test_set_by_statistic_fits_raw_labels_by_definition(self) -> None: + """Semantic pin (issue #5927): ``set-by-statistic`` is E_model-blind. + + The model energy decomposes as ``E = E_model + E_bias``. The set + mode DEFINES ``E_bias`` as the per-type statistic of the raw data + labels (or a user-given value), independent of ``E_model`` -- it + ignores a trained network and it equally ignores the analytical + ZBL term of a bridged model. Do NOT "fix" this by subtracting the + analytical contribution: that would silently turn the mode into a + third, model-dependent behavior. For a calibration that + compensates ``E_model``, use ``change-by-statistic`` (which since + #5910 uses the complete bridged predictor). + """ + params = self._build_model_params(bridging_method="ZBL") + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + model = get_sezm_model(params).to(self.device) + + rng = np.random.default_rng(5) + coord = torch.tensor( + rng.uniform(1.0, 2.5, size=(1, 4, 3)), + dtype=torch.float64, + device=self.device, + ) + box = torch.eye(3, dtype=torch.float64, device=self.device).reshape(1, 9) * 8.0 + samples, labels, counts_rows = [], [], [] + for types in ([[0, 0, 1, 1]], [[0, 1, 1, 1]]): + counts = np.bincount(np.asarray(types[0]), minlength=2) + label = float(rng.normal()) + samples.append( + { + "coord": coord, + "atype": torch.tensor(types, device=self.device), + "box": box, + "energy": torch.tensor( + [[label]], dtype=torch.float64, device=self.device + ), + "find_energy": np.float32(1.0), + "natoms": torch.tensor([[4, 4, *counts]], device=self.device), + } + ) + labels.append(label) + counts_rows.append(counts) + # Seed a nonzero bias: `set` must DISCARD it (an accidental + # additive implementation would shift the result by the seed). + model.set_out_bias(torch.ones_like(model.get_out_bias())) + model.change_out_bias(samples, bias_adjust_mode="set-by-statistic") + bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] + raw_fit = np.linalg.solve( + np.array(counts_rows, dtype=np.float64), np.array(labels) + ) + np.testing.assert_allclose(bias, raw_fit, atol=1.0e-8) + # Idempotence: repeating the call from the fitted state must land + # on the same raw-label fit again. + model.change_out_bias(samples, bias_adjust_mode="set-by-statistic") + repeated_bias = model.get_out_bias().detach().cpu().numpy().reshape(-1)[:2] + np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) + def test_zbl_respects_exclusions(self) -> None: """Excluded atoms and pairs contribute neither learned nor ZBL energy.""" coord = torch.tensor( diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index dd98f89059..34119b09dd 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -1374,7 +1374,7 @@ def test_training_smoke_combined(self, tmp_path) -> None: model = trainer.wrapper.model[DEFAULT_TASK_KEY] assert isinstance(model, NativeSpinEnergyModel) assert model.has_chg_spin_ebd() - assert model.has_default_chg_spin() + assert model.get_default_chg_spin() is not None tasks = trainer._make_training_tasks() task = trainer.select_task(tasks) diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py index 45b30c6454..fbf6600320 100644 --- a/source/tests/universal/common/cases/atomic_model/utils.py +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -97,6 +97,26 @@ def test_has_message_passing(self) -> None: module.has_message_passing(), self.expected_has_message_passing ) + def test_pair_exclude_contract(self) -> None: + """``pair_exclude_types``/``pair_excl`` are set by + ``BaseAtomicModel.__init__`` on every construction path; the + accessor is the public surface, the mask attribute is internal. + """ + pet = self.module.get_pair_exclude_types() + assert isinstance(pet, list) + # invariant pinned by reinit_pair_exclude: + assert (self.module.pair_excl is None) == (len(pet) == 0) + if self.module.pair_excl is not None: + # the mask must hold exactly the SYMMETRIC CLOSURE of the + # accessor's pairs (both backends symmetrize on reinit; + # compare as sets of tuples: pt stores a set, dpmodel a list) + symmetrized = { + pair for i, j in map(tuple, pet) for pair in ((i, j), (j, i)) + } + assert { + tuple(p) for p in self.module.pair_excl.get_exclude_types() + } == symmetrized + def test_forward(self) -> None: """Test forward.""" nf = 1 diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index 03dd5248da..ca5a750e1c 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -48,6 +48,20 @@ def test_capability_contract(self) -> None: # off -- a no-op on descriptors without one. self.module.disable_graph_lower() assert self.module.uses_graph_lower() is False + # chg-spin family: concrete base defaults, never probed (issue #5897). + assert isinstance(self.module.get_dim_chg_spin(), int) + dcs = self.module.get_default_chg_spin() + # Backend-agnostic: dpmodel returns a list, frozen pt returns a + # torch.Tensor -- pin the shape contract, not the container type. + assert dcs is None or len(dcs) == self.module.get_dim_chg_spin() + # has_default_chg_spin was merged into get_default_chg_spin: the + # predicate is ``get_default_chg_spin() is not None``. The absence + # of ``has_default_chg_spin`` on the dpmodel side is pinned in + # source/tests/universal/dpmodel/descriptor/test_descriptor.py -- + # pt is frozen and still declares the (now-redundant) method. + # Geometric-compression state query: base default False, the + # dpa1/dpa2 families override from their ``geo_compress`` attribute. + assert isinstance(self.module.get_geo_compress(), bool) def test_forward_consistency(self) -> None: ret = [] diff --git a/source/tests/universal/common/cases/fitting/utils.py b/source/tests/universal/common/cases/fitting/utils.py index de6b12c3a2..eae1cdde3c 100644 --- a/source/tests/universal/common/cases/fitting/utils.py +++ b/source/tests/universal/common/cases/fitting/utils.py @@ -104,6 +104,17 @@ def test_exclude_types( )[var_name] np.testing.assert_allclose(rd, rd_ex) + def test_reinit_exclude_contract(self) -> None: + """``reinit_exclude`` is declared on the base fitting: empty input + is always accepted; fittings with exclusion support apply it. + """ + self.module.reinit_exclude([]) # must never raise + if hasattr(self.module, "emask"): # override branch + self.module.reinit_exclude([0]) + assert self.module.exclude_types == [0] + self.module.reinit_exclude([]) + assert self.module.exclude_types == [] + def test_change_type_map(self) -> None: if not self.module.mixed_types: # skip if not mixed_types diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 5ff4254934..c51f9aefd1 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -131,6 +131,38 @@ def test_has_spin(self) -> None: expected = getattr(self, "test_spin", False) self.assertEqual(self.module.has_spin(), expected) + def test_chg_spin_capability_contract(self) -> None: + """chg-spin queries are declared on the base model with concrete + defaults (False/0/None) -- direct calls, never ``hasattr`` probes. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (predicate: ``get_default_chg_spin() is not None``). Its absence on + the dpmodel side is pinned in + source/tests/universal/dpmodel/model/test_model.py -- pt is frozen + and still declares the (now-redundant) method. + """ + assert isinstance(self.module.has_chg_spin_ebd(), bool) + assert isinstance(self.module.get_dim_chg_spin(), int) + dcs = self.module.get_default_chg_spin() + # Backend-agnostic: dpmodel returns a list, frozen pt returns a + # torch.Tensor -- pin the shape contract, not the container type. + assert dcs is None or len(dcs) == self.module.get_dim_chg_spin() + + def test_property_capability_contract(self) -> None: + """Property queries are declared on the base model with concrete + defaults: ``get_var_name`` returns ``None`` for non-property models + (the support predicate), ``get_intensive`` defaults ``False``, + ``get_task_dim`` raises for non-property models. + """ + vn = self.module.get_var_name() + assert vn is None or isinstance(vn, str) + assert isinstance(self.module.get_intensive(), bool) + if vn is None: + with self.assertRaises(NotImplementedError): + self.module.get_task_dim() + else: + assert isinstance(self.module.get_task_dim(), int) + def test_forward(self) -> None: """Test forward and forward_lower.""" test_spin = getattr(self, "test_spin", False) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 6999c0b779..85faf46ed9 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -19,6 +19,9 @@ DescrptSeT, DescrptSeTTebd, ) +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) from deepmd.dpmodel.descriptor.dpa2 import ( RepformerArgs, RepinitArgs, @@ -26,6 +29,9 @@ from deepmd.dpmodel.descriptor.dpa3 import ( RepFlowArgs, ) +from deepmd.dpmodel.descriptor.make_base_descriptor import ( + make_base_descriptor, +) from deepmd.dpmodel.descriptor.repflows import ( DescrptBlockRepflows, ) @@ -999,17 +1005,47 @@ 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.assertIsNotNone(shared_default.get_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()) + + +class TestHasDefaultChgSpinAbsentDP(unittest.TestCase): + """Pin the dpmodel-side half of the ``has_default_chg_spin`` merge. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (issue #5897): the predicate is ``get_default_chg_spin() is not None``. + The shared universal descriptor case only asserts the concrete + replacement (``get_default_chg_spin``), since the frozen pt backend + still declares the (now-redundant) ``has_default_chg_spin`` method on + several descriptors. This test pins that the method is gone from the + dpmodel base-descriptor family, where the merge is authoritative. + """ + + def test_absent_from_base_descriptor(self) -> None: + assert not hasattr(BaseDescriptor, "has_default_chg_spin") + assert not hasattr( + make_base_descriptor(np.ndarray, "call"), "has_default_chg_spin" + ) + + def test_absent_from_concrete_descriptors(self) -> None: + for cls in ( + DescrptSeA, + DescrptSeR, + DescrptSeT, + DescrptSeTTebd, + DescrptDPA1, + DescrptDPA2, + DescrptDPA3, + DescrptHybrid, + ): + assert not hasattr(cls, "has_default_chg_spin") diff --git a/source/tests/universal/dpmodel/model/test_model.py b/source/tests/universal/dpmodel/model/test_model.py index ece94a58ad..a51611069d 100644 --- a/source/tests/universal/dpmodel/model/test_model.py +++ b/source/tests/universal/dpmodel/model/test_model.py @@ -18,6 +18,10 @@ EnergyModel, SpinModel, ) +from deepmd.dpmodel.model.base_model import ( + BaseModel, + make_base_model, +) from deepmd.utils.spin import ( Spin, ) @@ -271,3 +275,24 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() cls.skip_test_autodiff = True + + +class TestHasDefaultChgSpinAbsentDP(unittest.TestCase): + """Pin the dpmodel-side half of the ``has_default_chg_spin`` merge. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (issue #5897): the predicate is ``get_default_chg_spin() is not None``. + The shared universal model case only asserts the concrete replacement + (``get_default_chg_spin``), since the frozen pt backend still declares + the (now-redundant) ``has_default_chg_spin`` method on ``make_model``. + This test pins that the method is gone from the dpmodel base-model + family, where the merge is authoritative. + """ + + def test_absent_from_base_model(self) -> None: + assert not hasattr(BaseModel, "has_default_chg_spin") + assert not hasattr(make_base_model(), "has_default_chg_spin") + + def test_absent_from_concrete_models(self) -> None: + for cls in (EnergyModel, SpinModel): + assert not hasattr(cls, "has_default_chg_spin")