diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 24bc56074d..3a7922e8bd 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -47,6 +47,12 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + check_preset_out_bias, + normalize_preset_out_bias, + preset_out_bias_rows, + remap_preset_out_bias, +) from .make_base_atomic_model import ( make_base_atomic_model, @@ -71,7 +77,7 @@ def __init__( atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, - preset_out_bias: dict[str, Array] | None = None, + preset_out_bias: dict | None = None, data_stat_protect: float = 1e-2, ) -> None: super().__init__() @@ -79,7 +85,7 @@ def __init__( self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) self.rcond = rcond - self.preset_out_bias = preset_out_bias + self.preset_out_bias = normalize_preset_out_bias(preset_out_bias, type_map) self.data_stat_protect = data_stat_protect self._observed_type: list[str] | None = None @@ -131,6 +137,7 @@ def init_out_stat(self) -> None: """Initialize the output bias.""" ntypes = self.get_ntypes() self.bias_keys: list[str] = list(self.fitting_output_def().keys()) + check_preset_out_bias(self.preset_out_bias, self.bias_keys) self.max_out_size = max( [self.atomic_output_def()[kk].size for kk in self.bias_keys] ) @@ -172,6 +179,10 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference into the fitting bias; nothing to fold without a fitting network.""" + return + def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" return False @@ -313,6 +324,7 @@ def change_type_map( self.reinit_pair_exclude( map_pair_exclude_types(self.pair_exclude_types, remap_index) ) + self.preset_out_bias = remap_preset_out_bias(self.preset_out_bias, remap_index) if has_new_type: xp = array_api_compat.array_namespace(self.out_bias) extend_shape = [ @@ -725,6 +737,7 @@ def compute_or_load_out_stat( merged, stat_file_path=stat_file_path, bias_adjust_mode="set-by-statistic", + observed_type=self.observed_type, ) def _make_wrapped_sampler( @@ -786,6 +799,7 @@ def change_out_bias( sample_merged: Callable[[], list[dict]] | list[dict], stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", + observed_type: list[str] | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -803,60 +817,84 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + An output assigned in `preset_out_bias` is fixed by the preset in both modes: + every element in the data must be assigned, the assigned types take the + preset value, and the types absent from the data keep zero + ('set-by-statistic') or their stored bias ('change-by-statistic'); no + statistics are computed, read or written for such an output, whose + output std keeps its stored value. stat_file_path : Optional[DPPath] The path to the stat file. + observed_type : list[str], optional + The elements that occur in the data; derived from the sample when + not given. """ from deepmd.dpmodel.utils.stat import ( + collect_observed_types, compute_output_stats, ) - if bias_adjust_mode == "change-by-statistic": - delta_bias, out_std = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=self.bias_keys, - stat_file_path=stat_file_path, - model_forward=self._get_forward_wrapper_func(), - rcond=self.rcond, - preset_bias=self.preset_out_bias, - stats_distinguish_types=self.get_compute_stats_distinguish_types(), - intensive=self.get_intensive(), - ) - self._store_out_stat(delta_bias, out_std, add=True) - elif bias_adjust_mode == "set-by-statistic": - bias_out, std_out = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=self.bias_keys, - stat_file_path=stat_file_path, - rcond=self.rcond, - preset_bias=self.preset_out_bias, - stats_distinguish_types=self.get_compute_stats_distinguish_types(), - intensive=self.get_intensive(), - ) - self._store_out_stat(bias_out, std_out) - else: + if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + change = bias_adjust_mode == "change-by-statistic" + distinguish_types = self.get_compute_stats_distinguish_types() + check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + sampled = sample_merged + # === Step 1. Outputs fixed by the preset === + fitted_keys = self.bias_keys + if self.preset_out_bias: + if observed_type is None: + sampled = sample_merged() if callable(sample_merged) else sample_merged + observed_type = collect_observed_types(sampled, self.type_map) + rows = preset_out_bias_rows( + self.preset_out_bias, + self.type_map, + observed_type, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + keep_unassigned=change, + excluded_types=self.atom_exclude_types, + ) + self._store_out_stat(rows) + fitted_keys = [kk for kk in self.bias_keys if kk not in rows] + if not fitted_keys: + return + # === Step 2. Outputs fitted from the data === + # The fitted statistics are absolute biases in 'set-by-statistic' mode and + # shifts of the stored bias in 'change-by-statistic' mode. + forward = self._get_forward_wrapper_func() if change else None + out_bias, out_std = compute_output_stats( + sampled, + self.get_ntypes(), + keys=fitted_keys, + stat_file_path=stat_file_path, + model_forward=forward, + rcond=self.rcond, + stats_distinguish_types=distinguish_types, + intensive=self.get_intensive(), + ) + self._store_out_stat(out_bias, out_std, add=change) def _store_out_stat( self, out_bias: dict[str, np.ndarray], - out_std: dict[str, np.ndarray], + out_std: dict[str, np.ndarray] | None = None, add: bool = False, ) -> None: - """Store output bias and std into the model.""" + """Store the output bias, and the output std when given, into the model.""" ntypes = self.get_ntypes() out_bias_data = np.array(to_numpy_array(self.out_bias)) out_std_data = np.array(to_numpy_array(self.out_std)) for kk in out_bias.keys(): - assert kk in out_std.keys() idx = self._get_bias_index(kk) size = self._varsize(self.atomic_output_def()[kk].shape) if not add: out_bias_data[idx, :, :size] = out_bias[kk].reshape(ntypes, size) else: out_bias_data[idx, :, :size] += out_bias[kk].reshape(ntypes, size) - out_std_data[idx, :, :size] = out_std[kk].reshape(ntypes, size) + if out_std is not None: + out_std_data[idx, :, :size] = out_std[kk].reshape(ntypes, size) self.out_bias = out_bias_data self.out_std = out_std_data diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index c78eb44d5d..7ce0e566e7 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -50,6 +50,15 @@ def __init__( "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel" ) super().__init__(descriptor, fitting, type_map, **kwargs) + if self.preset_out_bias is not None and any( + entry is not None + for entries in self.preset_out_bias.values() + for entry in entries + ): + raise ValueError( + "Dipole models do not apply an output bias and cannot use " + "assigned preset_out_bias values." + ) def apply_out_stat( self, diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 01ea0faf78..c134ddc249 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -12,6 +12,8 @@ NeighborGraph, ) +import array_api_compat + from deepmd.dpmodel.array_api import ( Array, xp_take_first_n, @@ -28,6 +30,11 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.vacuum_reference import ( + reference_charge_spin, + reference_spin, + resolve_vacuum_ref, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -107,6 +114,8 @@ class DPAtomicModel(BaseAtomicModel): """ + CONFIG_DERIVED_ARRAYS = ("vacuum_charge_spin", "vacuum_spin") + def __init__( self, descriptor: BaseDescriptor, @@ -136,7 +145,49 @@ def __init__( # declares it; other descriptors' ``call_graph`` would ``TypeError`` # on an unconditional ``charge_spin=`` kwarg. self.supports_charge_spin: bool = self.descriptor.supports_charge_spin() + self.add_spin_ebd: bool = ( + self._supports_native_spin and self.descriptor.use_spin is not None + ) super().init_out_stat() + # === Reference of the isolated atoms === + # The fitting references its output only when the preset fixes the + # bias of that output; the conditions of the reference atoms follow. + resolve_vacuum_ref(self.fitting_net, self.preset_out_bias) + self.init_vacuum_conditions() + + def init_vacuum_conditions(self) -> None: + """Build the conditioning inputs of the reference atoms for the type map. + + A table exists for every condition the descriptor takes when the + fitting references the isolated atoms; the type names must then be + element symbols. + """ + vacuum_ref = self.fitting_net.vacuum_ref + self.vacuum_charge_spin = ( + reference_charge_spin(self.type_map) + if vacuum_ref and self.add_chg_spin_ebd + else None + ) + self.vacuum_spin = ( + reference_spin(self.type_map) if vacuum_ref and self.add_spin_ebd else None + ) + + def vacuum_conditions(self) -> dict[str, Array]: + """Conditioning inputs of one isolated neutral ground-state atom per type. + + Returns + ------- + dict[str, Array] + The charge/spin conditions with shape (ntypes, 2) under + ``charge_spin`` and the spin vectors with shape (ntypes, 3) under + ``spin``, for the conditions the descriptor takes. + """ + conditions = {} + if self.vacuum_charge_spin is not None: + conditions["charge_spin"] = self.vacuum_charge_spin + if self.vacuum_spin is not None: + conditions["spin"] = self.vacuum_spin + return conditions def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" @@ -306,10 +357,6 @@ def forward_atomic( if self.add_chg_spin_ebd and charge_spin is None: default_cs = self.descriptor.get_default_chg_spin() if default_cs is not None: - from deepmd.dpmodel.array_api import ( - array_api_compat, - ) - xp = array_api_compat.array_namespace(extended_coord) cs_array = xp.asarray( default_cs, @@ -326,6 +373,11 @@ def forward_atomic( comm_dict=comm_dict, charge_spin=charge_spin if self.add_chg_spin_ebd else None, ) + # The vacuum descriptor of every type is handed to a fitting that + # references its atoms; other fittings do not take the keyword. + vacuum_kwargs = {} + if self.fitting_net.needs_vacuum_descriptor(): + vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor() ret = self.fitting_net( descriptor, atype, @@ -334,9 +386,141 @@ def forward_atomic( h2=h2, fparam=fparam, aparam=aparam, + **vacuum_kwargs, ) return ret + def vacuum_descriptor(self) -> Array: + """Descriptor of an isolated atom of every type. + + Every type is evaluated as a single-atom frame without neighbors, + conditioned as the neutral ground-state atom: zero charge with the + ground-state multiplicity when the descriptor takes the charge/spin + condition, and a spin vector of one Bohr magneton per unpaired + electron when it takes the native spin. + + Returns + ------- + Array + The vacuum descriptor with shape (ntypes, dim_descrpt), in the + array namespace, precision and device of the fitting bias. + """ + bias = self.fitting_net.bias_atom_e + xp = array_api_compat.array_namespace(bias) + device = array_api_compat.device(bias) + ntypes = self.get_ntypes() + coord = xp.zeros((ntypes, 1, 3), dtype=bias.dtype, device=device) + atype = xp.reshape( + xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1) + ) + nlist = xp.full( + (ntypes, 1, self.descriptor.get_nnei()), + -1, + dtype=xp.int64, + device=device, + ) + mapping = xp.zeros((ntypes, 1), dtype=xp.int64, device=device) + conditions = { + name: xp.asarray(table, dtype=bias.dtype, device=device) + for name, table in self.vacuum_conditions().items() + } + if "spin" in conditions: + conditions["spin"] = conditions["spin"][:, None, :] + descriptor = self.descriptor( + coord, atype, nlist, mapping=mapping, **conditions + )[0] + return xp.reshape(descriptor, (ntypes, -1)) + + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the fitting so the forward carries no reference atoms. + + The vacuum descriptor of every type is evaluated once with the current + parameters and handed to the fitting, which folds the reference into + its bias or stores the table (see its ``fold_vacuum_reference``). + """ + fitting = self.fitting_net + if fitting.needs_vacuum_descriptor(): + fitting.fold_vacuum_reference(self.vacuum_descriptor()) + + def append_vacuum_frames( + self, + graph: "NeighborGraph", + atype: Array, + charge_spin: Array | None, + spin: Array | None, + ) -> tuple["NeighborGraph", Array, Array | None, Array | None]: + """Append one isolated atom of every type as a single-node frame. + + The reference nodes follow the node axis of the real frames, padding + included, as single-node frames: one node of every type without + neighbors, conditioned as the neutral ground-state atom (zero charge + and the ground-state multiplicity for the charge/spin conditioning, a + spin vector of as many Bohr magnetons as unpaired electrons for the + native spin). Node-wise operations of the descriptor leave the real + nodes unaffected, so the descriptor rows of the reference nodes are the + vacuum descriptor of every type. + + Parameters + ---------- + graph : NeighborGraph + Neighbor graph of the real frames. + atype : Array + Flat node types with shape (N,). + charge_spin : Array, optional + Charge/spin condition of the real frames with shape (nf, 2) or + (1, 2); None takes the descriptor default. + spin : Array, optional + Per-node spin vectors with shape (N, 3), or None. + + Returns + ------- + tuple + The extended graph, node types, charge/spin conditions and spin + vectors, each with the reference entries appended. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + append_isolated_frames, + ) + + xp = array_api_compat.array_namespace(atype, graph.edge_vec) + device = array_api_compat.device(atype) + dtype = graph.edge_vec.dtype + ntypes = self.get_ntypes() + nf = graph.n_node.shape[0] + conditions = self.vacuum_conditions() + graph = append_isolated_frames(graph, atype.shape[0], ntypes) + atype = xp.concat( + [atype, xp.arange(ntypes, dtype=atype.dtype, device=device)], axis=0 + ) + if self.add_chg_spin_ebd: + if charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is None: + raise ValueError("`charge_spin` is required for this descriptor.") + charge_spin = xp.reshape( + xp.asarray(default_cs, dtype=dtype, device=device), (1, 2) + ) + # one row per real frame, one for the padding frame, one per type + charge_spin = xp.concat( + [ + xp.broadcast_to( + xp.astype(xp.reshape(charge_spin, (-1, 2)), dtype), (nf, 2) + ), + xp.zeros((1, 2), dtype=dtype, device=device), + xp.asarray(conditions["charge_spin"], dtype=dtype, device=device), + ], + axis=0, + ) + if spin is not None: + spin = xp.concat( + [ + xp.reshape(spin, (-1, 3)), + xp.asarray(conditions["spin"], dtype=spin.dtype, device=device), + ], + axis=0, + ) + return graph, atype, charge_spin, spin + def forward_atomic_graph( self, graph: "NeighborGraph", @@ -385,13 +569,39 @@ def forward_atomic_graph( the result dict on the flat node axis, defined by the `FittingOutputDef`. """ - import array_api_compat - from deepmd.dpmodel.utils.neighbor_graph import ( frame_id_from_n_node, ) xp = array_api_compat.array_namespace(graph.edge_vec) + n_real = atype.shape[0] + # === Step 1. Conditioning of the real nodes === + fparam_node = None + if fparam is not None: + # Pass the STATIC flat node count (``atype.shape[0] == N``) so the + # helper does not fall back to ``int(sum(n_node))``: that int() on a + # traced tensor breaks make_fx / torch.export + # (``GuardOnDataDependentSymNode``) for the graph .pt2 export and + # compiled-training paths when ``numb_fparam > 0``. + frame_id = frame_id_from_n_node(graph.n_node, n_total=atype.shape[0]) + fparam_node = xp.take(fparam, frame_id, axis=0) # (N, ndf) + aparam_node = aparam + if aparam is not None and graph.n_local is not None and aparam.ndim == 3: + aparam_node = _extend_graph_aparam( + aparam, + graph.n_node, + graph.n_local, + atype.shape[0], + ) + # === Step 2. Vacuum reference nodes === + # The reference atoms are appended as single-node frames, so the same + # descriptor call yields the vacuum descriptor of every type. + vacuum_ref = self.fitting_net.needs_vacuum_descriptor() + if vacuum_ref: + graph, atype, charge_spin, spin = self.append_vacuum_frames( + graph, atype, charge_spin, spin + ) + # === Step 3. Descriptor === # Descriptor-owned: dpa1/dpa2 hand out their full tebd table; DPA4 # embeds types internally from ``atype`` and returns None. type_embedding = self.descriptor.graph_type_embedding_table() @@ -411,23 +621,13 @@ def forward_atomic_graph( **spin_kwargs, **charge_spin_kwargs, ) - fparam_node = None - if fparam is not None: - # Pass the STATIC flat node count (``atype.shape[0] == N``) so the - # helper does not fall back to ``int(sum(n_node))``: that int() on a - # traced tensor breaks make_fx / torch.export - # (``GuardOnDataDependentSymNode``) for the graph .pt2 export and - # compiled-training paths when ``numb_fparam > 0``. - frame_id = frame_id_from_n_node(graph.n_node, n_total=atype.shape[0]) - fparam_node = xp.take(fparam, frame_id, axis=0) # (N, ndf) - aparam_node = aparam - if aparam is not None and graph.n_local is not None and aparam.ndim == 3: - aparam_node = _extend_graph_aparam( - aparam, - graph.n_node, - graph.n_local, - atype.shape[0], - ) + vacuum_kwargs = {} + if vacuum_ref: + vacuum_kwargs["vacuum_descriptor"] = gg[n_real:] + gg = gg[:n_real] + atype = atype[:n_real] + rot_mat = None if rot_mat is None else rot_mat[:n_real] + # === Step 4. Fitting === return self.fitting_net.call_graph( gg, atype, @@ -436,6 +636,7 @@ def forward_atomic_graph( h2=None, fparam=fparam_node, aparam=aparam_node, + **vacuum_kwargs, ) def compute_or_load_stat( @@ -467,12 +668,11 @@ def compute_or_load_stat( self.fitting_net.compute_input_stats( wrapped_sampler, stat_file_path=stat_file_path ) - if compute_or_load_out_stat: - self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) def change_type_map( self, type_map: list[str], model_with_new_type_stat: Any | None = None @@ -491,6 +691,7 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + self.init_vacuum_conditions() def compute_fitting_input_stat( self, diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 6d9a93ac65..cd42e983c4 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -219,6 +219,11 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of every sub-model into its bias.""" + for model in self.models: + model.fold_vacuum_reference() + def change_type_map( self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 8a466b9702..41a1741e14 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -232,18 +232,17 @@ def compute_or_load_stat( compute_or_load_out_stat : bool Whether to compute the output statistics. """ - if compute_or_load_out_stat: - wrapped_sampler = self._make_wrapped_sampler(sampled_func) - self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - + observed_stat_path = stat_file_path if stat_file_path is not None and self.type_map is not None: - stat_file_path /= " ".join(self.type_map) - + observed_stat_path = stat_file_path / " ".join(self.type_map) self._collect_and_set_observed_type( sampled_func if callable(sampled_func) else lambda: sampled_func, - stat_file_path, + observed_stat_path, preset_observed_type, ) + if compute_or_load_out_stat: + wrapped_sampler = self._make_wrapped_sampler(sampled_func) + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) def forward_atomic( self, diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index f1106cd395..f19165b87c 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -65,6 +65,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, + frame_id_from_n_node, graph_from_dense_quartet, ) from deepmd.dpmodel.utils.seed import ( @@ -1373,7 +1374,6 @@ def call( x_scalar, _ = self._run_graph( graph, atype_flat, - nf=nf, n_out_nodes=nf * nloc, force_embedding=force_embedding, charge_spin=charge_spin, @@ -1455,7 +1455,7 @@ def call_graph( ref=graph.edge_vec, ) x_scalar, _ = self._run_graph( - graph, atype, nf=nf, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict + graph, atype, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict ) # ``_run_graph`` returns the read-out with its SO(3) singleton # axes still attached, shape (n_nodes, 1, 1, channels); flatten to the @@ -1469,7 +1469,6 @@ def _run_graph( graph: NeighborGraph, atype_flat: Array, *, - nf: int = 1, n_out_nodes: int | None = None, force_embedding: Array | None = None, charge_spin: Array | None = None, @@ -1505,8 +1504,6 @@ def _run_graph( geometry/autograd leaf, ``edge_mask`` flags valid edges. atype_flat Flat node types with shape (N,). - nf - Frame count (only consumed by the charge/spin FiLM conditioning). n_out_nodes Leading node count kept for the read-out (owned atoms). ``None`` keeps all nodes (``atype_flat.shape[0]``). @@ -1555,10 +1552,7 @@ def _run_graph( ) # (N, C) if self.charge_spin_embedding is not None: type_ebed = self._apply_charge_spin_embedding( - type_ebed, - charge_spin, - nf=nf, - nloc=n_out_nodes // nf, + type_ebed, charge_spin, graph.n_node ) n_nodes = type_ebed.shape[0] @@ -2012,9 +2006,7 @@ def _apply_charge_spin_embedding( self, type_ebed: Array, charge_spin: Array, - *, - nf: int, - nloc: int, + n_node: Array, ) -> Array: """ Add frame-level charge and spin conditions to scalar type features. @@ -2022,23 +2014,22 @@ def _apply_charge_spin_embedding( Parameters ---------- type_ebed - Flattened type embeddings with shape (nf * nloc, channels). + Flattened type embeddings with shape (N, channels). charge_spin Frame-level charge and spin conditions with shape (nf, 2). - nf - Number of frames. - nloc - Number of local atoms. + n_node + Node count of every frame with shape (nf,); the frames occupy + consecutive blocks of the node axis. Returns ------- Array - Conditioned type embeddings with shape (nf * nloc, channels). + Conditioned type embeddings with shape (N, channels). """ xp = array_api_compat.array_namespace(type_ebed, charge_spin) condition = self.charge_spin_embedding(xp.astype(charge_spin, type_ebed.dtype)) - condition = xp.broadcast_to(condition[:, None, :], (nf, nloc, self.channels)) - return type_ebed + xp.reshape(condition, type_ebed.shape) + frame_id = frame_id_from_n_node(n_node, n_total=type_ebed.shape[0]) + return type_ebed + xp.take(condition, frame_id, axis=0) def _apply_spin_embedding( self, diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 5e39622bb9..0dd0c7f6ef 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -126,6 +126,7 @@ def __init__( spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], + vacuum_ref: bool = False, r_differentiable: bool = True, c_differentiable: bool = True, type_map: list[str] | None = None, @@ -163,6 +164,7 @@ def __init__( spin=spin, mixed_types=mixed_types, exclude_types=exclude_types, + vacuum_ref=vacuum_ref, type_map=type_map, seed=seed, default_fparam=default_fparam, @@ -183,7 +185,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) var_name = data.pop("var_name", None) assert var_name == "dipole" return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index b15935602a..61c9a65e9b 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -55,6 +55,7 @@ def __init__( precision: str = DEFAULT_PRECISION, mixed_types: bool = False, exclude_types: list[int] = [], + vacuum_ref: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None, default_fparam: list | None = None, @@ -80,6 +81,7 @@ def __init__( precision=precision, mixed_types=mixed_types, exclude_types=exclude_types, + vacuum_ref=vacuum_ref, type_map=type_map, seed=seed, default_fparam=default_fparam, @@ -102,7 +104,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data["numb_dos"] = data.pop("dim_out") data.pop("tot_ener_zero", None) data.pop("var_name", None) diff --git a/deepmd/dpmodel/fitting/dpa4_ener.py b/deepmd/dpmodel/fitting/dpa4_ener.py index ef3597af16..266d07430f 100644 --- a/deepmd/dpmodel/fitting/dpa4_ener.py +++ b/deepmd/dpmodel/fitting/dpa4_ener.py @@ -464,7 +464,7 @@ def deserialize(cls, data: dict) -> "SeZMEnergyFittingNet": data = data.copy() variables = data.pop("@variables") nets = data.pop("nets") - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("@class", None) data.pop("type", None) data.pop("var_name") diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 44a5e2d0c4..0f154076b8 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -41,6 +41,7 @@ def __init__( tot_ener_zero: bool = False, trainable: list[bool] | None = None, atom_ener: list[float] | None = None, + vacuum_ref: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: list[str | None] | None = None, @@ -66,6 +67,7 @@ def __init__( tot_ener_zero=tot_ener_zero, trainable=trainable, atom_ener=atom_ener, + vacuum_ref=vacuum_ref, activation_function=activation_function, precision=precision, layer_name=layer_name, @@ -81,7 +83,7 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 4474f9e6db..fe7b6f3d17 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -106,6 +106,14 @@ class GeneralFitting(NativeOP, BaseFitting): Remove vacuum contribution before the bias is added. The list assigned each type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same length as `ntypes` signaling if or not removing the vacuum contribution for the atom types in the list. + vacuum_ref: bool + Reference the network output of every atom to the output of the same + network for an isolated atom of the same type under the same frame + parameters, atomic parameters and case embedding. The output of an atom + without neighbors is then exactly ``bias_atom_e`` of its type. The + call takes the vacuum descriptor of every type until + :meth:`fold_vacuum_reference` has folded the reference into the bias + or stored the table. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. seed: Optional[Union[int, list[int]]] @@ -115,6 +123,11 @@ class GeneralFitting(NativeOP, BaseFitting): this value will be used as the default value for the frame parameter in the fitting net. """ + # A deployment constant kept out of checkpoints; see + # :meth:`fold_vacuum_reference`. A subclass extends this tuple rather than + # replacing it. + CONFIG_DERIVED_ARRAYS = ("vacuum_table",) + def __init__( self, var_name: str, @@ -137,6 +150,7 @@ def __init__( mixed_types: bool = True, exclude_types: list[int] = [], remove_vaccum_contribution: list[bool] | None = None, + vacuum_ref: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None, default_fparam: list[float] | None = None, @@ -174,6 +188,7 @@ def __init__( if self.spin is not None: raise NotImplementedError("spin is not supported") self.remove_vaccum_contribution = remove_vaccum_contribution + self.vacuum_ref = vacuum_ref self.eval_return_middle_output = False net_dim_out = self._net_out_dim() @@ -199,6 +214,7 @@ def __init__( self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec) else: self.case_embd = None + self.vacuum_table = None if self.default_fparam is not None: if self.numb_fparam > 0: @@ -510,6 +526,8 @@ def change_type_map( ) self.bias_atom_e = xp.concat([self.bias_atom_e, extend_bias_atom_e], axis=0) self.bias_atom_e = self.bias_atom_e[remap_index] + # the stored references belong to the old type map + self.vacuum_table = None def __setitem__(self, key: str, value: Any) -> None: if key in ["bias_atom_e"]: @@ -562,7 +580,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 4, + "@version": 5, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -577,6 +595,7 @@ def serialize(self) -> dict: "precision": self.precision, "mixed_types": self.mixed_types, "exclude_types": self.exclude_types, + "vacuum_ref": self.vacuum_ref, "nets": self.nets.serialize(), "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), @@ -608,6 +627,295 @@ def deserialize(cls, data: dict) -> "GeneralFitting": obj.nets = NetworkCollection.deserialize(nets) return obj + def needs_vacuum_descriptor(self) -> bool: + """Whether the call takes the vacuum descriptor of every type from the descriptor. + + A referencing fitting takes it until :meth:`fold_vacuum_reference` has + folded the reference into the bias or stored the table. + """ + return self.vacuum_ref and self.vacuum_table is None + + def uniform_conditioning(self) -> bool: + """Whether every atom receives the same conditioning columns. + + Frame parameters vary between frames and atomic parameters between + atoms, while the case embedding is shared by all atoms of a call. + """ + return self.numb_fparam == 0 and ( + self.numb_aparam == 0 or self.use_aparam_as_mask + ) + + def conditioning_columns( + self, + descriptor: Array, + fparam: Array | None, + aparam: Array | None, + ) -> Array | None: + """Normalized conditioning columns appended to the descriptor of every atom. + + Parameters + ---------- + descriptor : Array + Descriptor with shape (nf, nloc, nd); it supplies the frame and + atom counts, the dtype and the device of the columns. + fparam : Array, optional + Frame parameters with shape (nf, numb_fparam). The default frame + parameter is used when omitted. + aparam : Array, optional + Atomic parameters with shape (nf, nloc, numb_aparam). + + Returns + ------- + Array or None + The frame parameters, atomic parameters and case embedding of every + atom, normalized and concatenated, with shape (nf, nloc, ncond); + None when the descriptor alone is the fitting input. + """ + xp = array_api_compat.array_namespace(descriptor) + nf, nloc, _ = descriptor.shape + device = array_api_compat.device(descriptor) + # Fitting statistics remain NumPy for portable serialization; the + # constants are materialized next to the runtime descriptor instead of + # asking a backend namespace to reshape a foreign array directly. + columns = [] + if self.numb_fparam > 0: + if fparam is None: + assert self.default_fparam_tensor is not None + default_fparam = xp.asarray( + self.default_fparam_tensor, dtype=descriptor.dtype, device=device + ) + fparam = xp.tile( + xp.reshape(default_fparam, (1, self.numb_fparam)), (nf, 1) + ) + try: + fparam = xp.reshape(fparam, (nf, 1, self.numb_fparam)) + except (ValueError, RuntimeError) as e: + raise ValueError( + f"input fparam: cannot reshape {fparam.shape} " + f"into ({nf}, {self.numb_fparam})." + ) from e + fparam_device = array_api_compat.device(fparam) + fparam_avg = xp.asarray( + self.fparam_avg, dtype=fparam.dtype, device=fparam_device + ) + fparam_inv_std = xp.asarray( + self.fparam_inv_std, dtype=fparam.dtype, device=fparam_device + ) + fparam = (fparam - fparam_avg) * fparam_inv_std + columns.append(xp.broadcast_to(fparam, (nf, nloc, self.numb_fparam))) + if self.numb_aparam > 0 and not self.use_aparam_as_mask: + assert aparam is not None, "aparam should not be None" + try: + aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam)) + except (ValueError, RuntimeError) as e: + raise ValueError( + f"input aparam: cannot reshape {aparam.shape} " + f"into ({nf}, {nloc}, {self.numb_aparam})." + ) from e + aparam_device = array_api_compat.device(aparam) + aparam_avg = xp.asarray( + self.aparam_avg, dtype=aparam.dtype, device=aparam_device + ) + aparam_inv_std = xp.asarray( + self.aparam_inv_std, dtype=aparam.dtype, device=aparam_device + ) + columns.append((aparam - aparam_avg) * aparam_inv_std) + if self.dim_case_embd > 0: + assert self.case_embd is not None + case_embd = xp.asarray( + self.case_embd, dtype=descriptor.dtype, device=device + ) + columns.append( + xp.broadcast_to( + xp.reshape(case_embd, (1, 1, -1)), (nf, nloc, self.dim_case_embd) + ) + ) + if len(columns) == 0: + return None + return xp.concat(columns, axis=-1) + + def vacuum_input( + self, + vacuum_descriptor: Array | None, + atype: Array, + cond: Array | None, + ) -> Array | None: + """Fitting input rows of the vacuum references. + + The vacuum reference of an atom is an isolated atom of its type under + the conditioning of the atom itself. With frame or atomic parameters + the conditioning differs between atoms and the rows follow the atoms, + with shape (nf, nloc, in_dim). Otherwise one row per type suffices, + with shape (ntypes, in_dim), and :meth:`vacuum_output` gathers the + network output by type. + + Parameters + ---------- + vacuum_descriptor : Array, optional + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt); None takes the table stored by + :meth:`fold_vacuum_reference`. + atype : Array + Atom types with shape (nf, nloc). + cond : Array, optional + Conditioning columns of the atoms with shape (nf, nloc, ncond). + + Returns + ------- + Array or None + The reference rows, or None when ``vacuum_ref`` is off. + + Raises + ------ + ValueError + If ``vacuum_ref`` is on and the vacuum descriptor is missing or + has the wrong shape. + """ + if not self.vacuum_ref: + return None + xp = array_api_compat.array_namespace(atype) + if vacuum_descriptor is None: + if self.vacuum_table is None: + raise ValueError( + "vacuum_ref requires the vacuum descriptor of every atom type" + ) + vacuum_descriptor = xp.asarray( + self.vacuum_table, + dtype=get_xp_precision(xp, self.precision), + device=array_api_compat.device(atype), + ) + if tuple(vacuum_descriptor.shape) != (self.ntypes, self.dim_descrpt): + raise ValueError( + f"vacuum descriptor of shape {tuple(vacuum_descriptor.shape)} " + f"does not match ({self.ntypes}, {self.dim_descrpt})" + ) + if not self.uniform_conditioning(): + assert cond is not None + nf, nloc = atype.shape + x_vac = xp.reshape( + xp.take(vacuum_descriptor, xp.reshape(atype, (-1,)), axis=0), + (nf, nloc, self.dim_descrpt), + ) + return xp.concat([x_vac, cond], axis=-1) + if self.dim_case_embd == 0: + return vacuum_descriptor + assert self.case_embd is not None + case_embd = xp.asarray( + self.case_embd, + dtype=vacuum_descriptor.dtype, + device=array_api_compat.device(vacuum_descriptor), + ) + return xp.concat( + [ + vacuum_descriptor, + xp.broadcast_to( + xp.reshape(case_embd, (1, -1)), (self.ntypes, self.dim_case_embd) + ), + ], + axis=-1, + ) + + def vacuum_output(self, vacuum_property: Array, atype: Array) -> Array: + """Network output of the vacuum reference of every atom. + + Parameters + ---------- + vacuum_property : Array + Network output on the rows of :meth:`vacuum_input`. + atype : Array + Atom types with shape (nf, nloc). + + Returns + ------- + Array + The reference output of every atom with shape (nf, nloc, dim_out). + """ + if not self.uniform_conditioning(): + return vacuum_property + xp = array_api_compat.array_namespace(vacuum_property, atype) + nf, nloc = atype.shape + return xp.reshape( + xp.take(vacuum_property, xp.reshape(atype, (-1,)), axis=0), + (nf, nloc, vacuum_property.shape[-1]), + ) + + def vacuum_property(self, vacuum_descriptor: Array | None) -> Array: + """Network output on the vacuum descriptor of every type. + + Defined under uniform conditioning, where the output of a reference + depends on its type alone. + + Parameters + ---------- + vacuum_descriptor : Array + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt). + + Returns + ------- + Array + The reference output of every type with shape (ntypes, dim_out), + in the precision of ``vacuum_descriptor``. + """ + if vacuum_descriptor is None: + vacuum_descriptor = self.vacuum_table + if vacuum_descriptor is None: + raise ValueError( + "vacuum_ref requires the vacuum descriptor of every atom type" + ) + xp = array_api_compat.array_namespace(vacuum_descriptor) + atype = xp.arange( + self.ntypes, + dtype=xp.int64, + device=array_api_compat.device(vacuum_descriptor), + ) + xx_vac = self.vacuum_input( + xp.astype(vacuum_descriptor, get_xp_precision(xp, self.precision)), + atype, + None, + ) + if self.mixed_types: + out = self.nets[()](xx_vac) + else: + out = xp.concat( + [ + self.nets[(type_i,)](xx_vac[type_i : type_i + 1, :]) + for type_i in range(self.ntypes) + ], + axis=0, + ) + return xp.astype(out, vacuum_descriptor.dtype) + + def fold_vacuum_reference(self, vacuum_descriptor: Array) -> None: + """Fold the vacuum reference into the fitting so the call needs no reference atoms. + + Under uniform conditioning the reference output of an atom is a + constant of its type, so subtracting it from ``bias_atom_e`` yields + the same outputs as the referenced call and the option is switched + off. With frame or atomic parameters the reference output varies + between atoms, so the vacuum descriptor is stored instead and the call + evaluates the references from the stored table. The table is a + deployment constant: an exported model bakes it, checkpoints leave it + out, and a fitting loaded from a checkpoint takes the reference from + the descriptor again. + + Parameters + ---------- + vacuum_descriptor : Array + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt). + """ + if not self.vacuum_ref: + return + if not self.uniform_conditioning(): + self.vacuum_table = to_numpy_array(vacuum_descriptor) + return + reference = to_numpy_array(self.vacuum_property(vacuum_descriptor)) + self["bias_atom_e"] = to_numpy_array(self.bias_atom_e) - reference.astype( + GLOBAL_NP_FLOAT_PRECISION + ) + self.vacuum_ref = False + def _call_common( self, descriptor: Array, @@ -617,6 +925,7 @@ def _call_common( h2: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + vacuum_descriptor: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. @@ -639,6 +948,9 @@ def _call_common( The frame parameter. shape: nf x nfp. nfp being `numb_fparam` aparam The atomic parameter. shape: nf x nloc x nap. nap being `numb_aparam` + vacuum_descriptor + The descriptor of an isolated atom of every type, required by + ``vacuum_ref``. shape: ntypes x nd """ xp = array_api_compat.array_namespace(descriptor, atype) @@ -650,117 +962,25 @@ def _call_common( "get an input descriptor of dim {nd}," "which is not consistent with {self.dim_descrpt}." ) - xx = descriptor - if self.remove_vaccum_contribution is not None: - # TODO: comput the input for vacuum when setting remove_vaccum_contribution - # Ideally, the input for vacuum should be computed; - # we consider it as always zero for convenience. - # Needs a compute_input_stats for vacuum passed from the - # descriptor. - xx_zeros = xp.zeros_like(xx) - else: - xx_zeros = None - - if self.numb_fparam > 0 and fparam is None: - # use default fparam - assert self.default_fparam_tensor is not None - # Fitting statistics remain NumPy for portable serialization. - # Materialize constants next to the runtime descriptor instead of - # asking a backend namespace to reshape a foreign array directly. - default_fparam_tensor = xp.asarray( - self.default_fparam_tensor, - dtype=descriptor.dtype, - device=array_api_compat.device(descriptor), - ) - fparam = xp.tile( - xp.reshape(default_fparam_tensor, (1, self.numb_fparam)), (nf, 1) - ) - - # check fparam dim, concate to input descriptor - if self.numb_fparam > 0: - assert fparam is not None, "fparam should not be None" - try: - fparam = xp.reshape(fparam, (nf, self.numb_fparam)) - except (ValueError, RuntimeError) as e: - raise ValueError( - f"input fparam: cannot reshape {fparam.shape} " - f"into ({nf}, {self.numb_fparam})." - ) from e - fparam_device = array_api_compat.device(fparam) - fparam_avg = xp.asarray( - self.fparam_avg, - dtype=fparam.dtype, - device=fparam_device, - ) - fparam_inv_std = xp.asarray( - self.fparam_inv_std, - dtype=fparam.dtype, - device=fparam_device, - ) - fparam = (fparam - fparam_avg) * fparam_inv_std - fparam = xp.tile( - xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1) - ) - xx = xp.concat( - [xx, fparam], - axis=-1, - ) - if xx_zeros is not None: - xx_zeros = xp.concat( - [xx_zeros, fparam], - axis=-1, - ) - # check aparam dim, concate to input descriptor - if self.numb_aparam > 0 and not self.use_aparam_as_mask: - assert aparam is not None, "aparam should not be None" - try: - aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam)) - except (ValueError, RuntimeError) as e: - raise ValueError( - f"input aparam: cannot reshape {aparam.shape} " - f"into ({nf}, {nloc}, {self.numb_aparam})." - ) from e - aparam_device = array_api_compat.device(aparam) - aparam_avg = xp.asarray( - self.aparam_avg, - dtype=aparam.dtype, - device=aparam_device, - ) - aparam_inv_std = xp.asarray( - self.aparam_inv_std, - dtype=aparam.dtype, - device=aparam_device, - ) - aparam = (aparam - aparam_avg) * aparam_inv_std - xx = xp.concat( - [xx, aparam], - axis=-1, - ) - if xx_zeros is not None: - xx_zeros = xp.concat( - [xx_zeros, aparam], - axis=-1, - ) - if self.dim_case_embd > 0: - assert self.case_embd is not None - case_embd_buffer = xp.asarray( - self.case_embd, - dtype=descriptor.dtype, - device=array_api_compat.device(descriptor), - ) - case_embd = xp.tile(xp.reshape(case_embd_buffer, (1, 1, -1)), (nf, nloc, 1)) - xx = xp.concat( - [xx, case_embd], - axis=-1, - ) + # === Step 1. Assemble the fitting input === + # The conditioning columns are shared by the atoms and by their vacuum + # references, so that the reference of an atom differs from the atom + # in its descriptor only. + xx = descriptor + cond = self.conditioning_columns(descriptor, fparam, aparam) + # ``remove_vaccum_contribution`` subtracts the network output for a zero + # descriptor under the same conditioning columns. + xx_zeros = ( + None if self.remove_vaccum_contribution is None else xp.zeros_like(xx) + ) + if cond is not None: + xx = xp.concat([xx, cond], axis=-1) if xx_zeros is not None: - xx_zeros = xp.concat( - [xx_zeros, case_embd], - axis=-1, - ) + xx_zeros = xp.concat([xx_zeros, cond], axis=-1) + xx_vac = self.vacuum_input(vacuum_descriptor, atype, cond) - # calculate the prediction + # === Step 2. Evaluate the fitting networks === results: dict[str, Array] = {} if not self.mixed_types: outs = xp.zeros( @@ -786,6 +1006,17 @@ def _call_common( ): assert xx_zeros is not None atom_property -= self.nets[(type_i,)](xx_zeros) + if xx_vac is not None: + # The mask below keeps the atoms of type ``type_i`` alone, so + # the network of the type runs on its own reference row when + # the references are per type, and on the per-atom rows + # otherwise. + reference = ( + xx_vac[type_i : type_i + 1] + if self.uniform_conditioning() + else xx_vac + ) + atom_property = atom_property - self.nets[(type_i,)](reference) atom_property = xp.where( mask, atom_property, xp.zeros_like(atom_property) ) @@ -802,6 +1033,8 @@ def _call_common( outs = self.nets[()](xx) if xx_zeros is not None: outs -= self.nets[()](xx_zeros) + if xx_vac is not None: + outs = outs - self.vacuum_output(self.nets[()](xx_vac), atype) if self.eval_return_middle_output and len(self.neuron) > 0: middle_outs = self.nets[()].call_until_last(xx) bias_atom_e = xp.asarray( @@ -836,6 +1069,7 @@ def call_graph( h2: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + vacuum_descriptor: Array | None = None, ) -> dict[str, Array]: """Graph-native (flat node axis) fitting forward. @@ -861,6 +1095,9 @@ def call_graph( NODE-level frame parameter (already gathered by frame_id). N x nfp aparam atomic parameter. N x nap + vacuum_descriptor + the descriptor of an isolated atom of every type, required by + ``vacuum_ref``. ntypes x nd Returns ------- @@ -887,5 +1124,20 @@ def call_graph( ap1 = None if aparam is None else xp.reshape(aparam, (n, 1, aparam.shape[-1])) # fparam: dense API expects (nf, nfp); here nf'=N single-atom frames, so the # node-level (N, nfp) IS the per-(pseudo)frame param -- tiled over nloc'=1. - ret = self.__call__(d1, a1, gr=g1, g2=g2, h2=h2, fparam=fparam, aparam=ap1) + # Only referencing fittings take the keyword; it travels with a table. + vacuum_kwargs = ( + {} + if vacuum_descriptor is None + else {"vacuum_descriptor": vacuum_descriptor} + ) + ret = self.__call__( + d1, + a1, + gr=g1, + g2=g2, + h2=h2, + fparam=fparam, + aparam=ap1, + **vacuum_kwargs, + ) return {kk: xp.reshape(vv, (n, *vv.shape[2:])) for kk, vv in ret.items()} diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index f7daf6d98f..57c9eec185 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -93,6 +93,11 @@ class InvarFitting(GeneralFitting): this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. atom_ener Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descriptor should be set. + vacuum_ref + Reference the network output of every atom to the output of the same + network for an isolated atom of the same type under the same + conditioning, so that an atom without neighbors contributes exactly + its output bias. activation_function The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision @@ -131,6 +136,7 @@ def __init__( tot_ener_zero: bool = False, trainable: list[bool] | None = None, atom_ener: list[float] | None = None, + vacuum_ref: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: list[str | None] | None = None, @@ -151,6 +157,15 @@ def __init__( self.dim_out = dim_out self.atom_ener = atom_ener + if ( + vacuum_ref + and atom_ener is not None + and any(x is not None for x in atom_ener) + ): + raise ValueError( + "atom_ener and vacuum_ref are exclusive; vacuum_ref references every " + "atom to the isolated atom of its type by itself" + ) super().__init__( var_name=var_name, ntypes=ntypes, @@ -174,6 +189,7 @@ def __init__( remove_vaccum_contribution=None if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 else [x is not None for x in atom_ener], + vacuum_ref=vacuum_ref, type_map=type_map, seed=seed, default_fparam=default_fparam, @@ -189,7 +205,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) return super().deserialize(data) def _net_out_dim(self) -> int: @@ -224,6 +240,7 @@ def call( h2: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, + vacuum_descriptor: Array | None = None, ) -> dict[str, Array]: """Calculate the fitting. @@ -246,6 +263,18 @@ def call( The frame parameter. shape: nf x nfp. nfp being `numb_fparam` aparam The atomic parameter. shape: nf x nloc x nap. nap being `numb_aparam` + vacuum_descriptor + The descriptor of an isolated atom of every type, required by + ``vacuum_ref``. shape: ntypes x nd """ - return self._call_common(descriptor, atype, gr, g2, h2, fparam, aparam) + return self._call_common( + descriptor, + atype, + gr, + g2, + h2, + fparam, + aparam, + vacuum_descriptor=vacuum_descriptor, + ) diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index 7a595d4bc8..0933d0e281 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -39,6 +39,17 @@ def make_base_fitting( class BF(ABC, PluginVariant, make_plugin_registry("fitting")): """Base fitting provides the interfaces of fitting net.""" + vacuum_ref: bool = False + """Whether every atom is referenced to the isolated atom of its type. + + A fitting that sets it takes the vacuum descriptor of every type in + its forward; the default holds for fittings without the option. + """ + + def needs_vacuum_descriptor(self) -> bool: + """Whether the forward takes the vacuum descriptor of every type from the descriptor.""" + return False + def __new__(cls: type, *args: Any, **kwargs: Any) -> Any: if cls is BF: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index f4e504a7aa..95e2648cdf 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -145,6 +145,7 @@ def __init__( spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], + vacuum_ref: bool = False, fit_diag: bool = True, scale: list[float] | None = None, shift_diag: bool = True, @@ -196,6 +197,7 @@ def __init__( spin=spin, mixed_types=mixed_types, exclude_types=exclude_types, + vacuum_ref=vacuum_ref, type_map=type_map, seed=seed, default_fparam=default_fparam, @@ -224,7 +226,7 @@ def __getitem__(self, key: str) -> Array: def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 5 + data["@version"] = 6 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -235,7 +237,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 5, 1) + check_version_compatibility(data.pop("@version", 1), 6, 1) var_name = data.pop("var_name", None) assert var_name == "polar" return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 78082ad4b9..03882f1f1a 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -88,6 +88,7 @@ def __init__( precision: str = DEFAULT_PRECISION, mixed_types: bool = True, exclude_types: list[int] = [], + vacuum_ref: bool = False, type_map: list[str] | None = None, default_fparam: list | None = None, distinguish_types: bool = True, @@ -114,6 +115,7 @@ def __init__( precision=precision, mixed_types=mixed_types, exclude_types=exclude_types, + vacuum_ref=vacuum_ref, type_map=type_map, default_fparam=default_fparam, ) @@ -136,7 +138,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version"), 6, 1) + check_version_compatibility(data.pop("@version"), 7, 1) data.setdefault("distinguish_types", False) data.pop("dim_out") data["property_name"] = data.pop("var_name") @@ -158,7 +160,7 @@ def serialize(self) -> dict: "intensive": self.intensive, "distinguish_types": self.distinguish_types, } - dd["@version"] = 6 + dd["@version"] = 7 return dd diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index bee5b79215..edafba7efc 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1238,6 +1238,10 @@ def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the atomic model into its fitting bias.""" + self.atomic_model.fold_vacuum_reference() + def compute_or_load_stat( self, sampled_func: Callable[[], Any], diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index d570e12eb5..d49ad1d9df 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -86,6 +86,7 @@ def get_standard_model( type_map=data["type_map"], atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), + preset_out_bias=data.get("preset_out_bias"), ) @@ -124,6 +125,7 @@ def get_zbl_model( smin_alpha=data.get("smin_alpha", 0.1), atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), + preset_out_bias=data.get("preset_out_bias"), ) @@ -199,6 +201,7 @@ def get_linear_atomic_model( data = copy.deepcopy(data) type_map = data["type_map"] children = data["models"] + preset_out_bias = data.get("preset_out_bias") inner_indices = [ i for i, sub in enumerate(children) if sub.get("type") == "inner_potential" ] @@ -279,6 +282,7 @@ def get_linear_atomic_model( learned_descriptor["inner_clamp_r_inner"] = float(inner_cfg.get("r_inner", 0.5)) learned_descriptor["inner_clamp_r_outer"] = float(inner_cfg.get("r_outer", 0.8)) route_canonical_learned_options(data, children[learned_indices[0]]) + preset_out_bias = children[learned_indices[0]].get("preset_out_bias") built: dict[int, Any] = {} for i, sub in enumerate(children): @@ -336,6 +340,8 @@ def get_linear_atomic_model( # graph, so "excluded" must cover the analytical term too. atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), + # The composition computes the output bias, so the preset belongs to it. + preset_out_bias=preset_out_bias, ) diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index c7daebef71..4e2321b412 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -499,6 +499,10 @@ def get_ntypes(self) -> int: """Returns the number of element types.""" return len(self.get_type_map()) + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the backbone fitting into its bias.""" + self.backbone_model.fold_vacuum_reference() + def get_rcut(self) -> float: """Get the cut-off radius.""" return self.backbone_model.get_rcut() diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 32c0905a2d..6895ce3e1d 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -43,6 +43,7 @@ from .graph import ( GraphLayout, NeighborGraph, + append_isolated_frames, apply_pair_exclusion, compact_nodes, expand_node_values, @@ -68,6 +69,7 @@ "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", + "append_isolated_frames", "apply_pair_exclusion", "attach_angles", "attach_edge_csr", diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 1a98504005..7aca89f1b6 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -12,6 +12,7 @@ annotations, ) +import dataclasses from dataclasses import ( dataclass, field, @@ -235,6 +236,54 @@ def frame_id_from_n_node(n_node: Array, n_total: int | None = None) -> Array: return xp.minimum(frame_id, xp.astype(last_frame, xp.int64)) +def append_isolated_frames( + graph: NeighborGraph, n_total: int, count: int +) -> NeighborGraph: + """Append ``count`` single-node frames without edges after the node axis. + + The new nodes take the positions ``[n_total, n_total + count)`` of the + flat node axis, each as a frame of one owned node and no neighbors, so + every edge-indexed field of the graph stays valid unchanged. A padding + suffix of the compact-prefix layout, ``[sum(n_node), n_total)``, becomes + one frame of its own without owned nodes, so that the frame bookkeeping + keeps assigning every position of the extended axis to its frame. The + compressed-sparse-row offsets, when present, gain empty rows. + + Parameters + ---------- + graph + Neighbor graph of the real frames. + n_total + Size of the flat node axis of ``graph``, padding included. + count + Number of isolated nodes to append. + + Returns + ------- + NeighborGraph + A ``dataclasses.replace`` copy with the extended frame bookkeeping. + """ + xp = array_api_compat.array_namespace(graph.n_node) + device = array_api_compat.device(graph.n_node) + n_node = graph.n_node + padding = xp.reshape(n_total - xp.sum(n_node), (1,)) + ones = xp.ones((count,), dtype=n_node.dtype, device=device) + fields = {"n_node": xp.concat([n_node, xp.astype(padding, n_node.dtype), ones])} + if graph.n_local is not None: + fields["n_local"] = xp.concat( + [ + graph.n_local, + xp.zeros((1,), dtype=graph.n_local.dtype, device=device), + xp.ones((count,), dtype=graph.n_local.dtype, device=device), + ] + ) + for name in ("destination_row_ptr", "source_row_ptr"): + row_ptr = getattr(graph, name) + if row_ptr is not None: + fields[name] = xp.concat([row_ptr, xp.broadcast_to(row_ptr[-1:], (count,))]) + return dataclasses.replace(graph, **fields) + + def node_ownership_mask(n_node: Array, n_local: Array, n_total: int) -> Array: """Return the owned-node mask for a local-plus-halo graph. diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py index f62e80defe..855d8b0bd7 100644 --- a/deepmd/dpmodel/utils/stat.py +++ b/deepmd/dpmodel/utils/stat.py @@ -25,6 +25,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + make_preset_out_bias, + override_assigned_bias, +) from deepmd.utils.stat_file import ( load_paired_items, load_required_items, @@ -168,31 +172,6 @@ def _post_process_stat( return out_bias, new_std -def _make_preset_out_bias( - ntypes: int, - ibias: list[np.ndarray | None], -) -> np.ndarray | None: - """Make preset out bias. - - output: - a np array of shape [ntypes, *(odim0, odim1, ...)] is any item is not None - None if all items are None. - """ - if len(ibias) != ntypes: - raise ValueError("the length of preset bias list should be ntypes") - if all(ii is None for ii in ibias): - return None - for refb in ibias: - if refb is not None: - break - refb = np.array(refb) - nbias = [ - np.full_like(refb, np.nan, dtype=np.float64) if ii is None else ii - for ii in ibias - ] - return np.array(nbias) - - def _fill_stat_with_global( atomic_stat: np.ndarray | None, global_stat: np.ndarray, @@ -288,9 +267,12 @@ def compute_output_stats( rcond : float, optional The condition number for the regression of atomic energy. preset_bias : dict[str, list[Optional[np.ndarray]]], optional - Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. + Assigned values of the returned bias, given by key:value pairs. + The value is a list with one element per type: None leaves the type to the + statistics, an np.ndarray of output shape assigns the type. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The values live in the frame of the returned bias: absolute biases without + `model_forward`, shifts of the model's stored bias with `model_forward`. The `set_davg_zero` key in the descriptor should be set. model_forward : Callable, optional The wrapped forward function of atomic model. @@ -307,6 +289,11 @@ def compute_output_stats( assert isinstance(keys, list) requested_keys = list(keys) + # Model residuals depend on parameters not recorded in the statistics cache. + # Neither reuse nor persist them as absolute output statistics. + if model_forward is not None: + stat_file_path = None + # try to restore the bias from stat file bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) @@ -380,13 +367,22 @@ def compute_output_stats( else None ) + # assigned bias of every output as a (ntypes, ...) array, NaN where a + # type is left to the statistics + assigned_bias = { + kk: make_preset_out_bias(ntypes, preset_bias[kk]) + if preset_bias is not None and kk in preset_bias + else None + for kk in keys + } + # compute stat bias_atom_g, std_atom_g = _compute_output_stats_global( sampled, ntypes, keys, rcond, - preset_bias, + assigned_bias, global_sampled_idx, stats_distinguish_types, intensive, @@ -398,6 +394,7 @@ def compute_output_stats( keys, atomic_sampled_idx, model_pred_a, + assigned_bias, ) # merge global/atomic bias @@ -435,7 +432,7 @@ def _compute_output_stats_global( ntypes: int, keys: list[str], rcond: float | None = None, - preset_bias: dict[str, list[np.ndarray | None]] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, global_sampled_idx: dict | None = None, stats_distinguish_types: bool = True, intensive: bool = False, @@ -482,15 +479,8 @@ def _compute_output_stats_global( } nf = {kk: merged_natoms[kk].shape[0] for kk in keys if kk in merged_natoms} - if preset_bias is not None: - assigned_atom_ener = { - kk: _make_preset_out_bias(ntypes, preset_bias[kk]) - if kk in preset_bias.keys() - else None - for kk in keys - } - else: - assigned_atom_ener = dict.fromkeys(keys) + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) if model_pred is None: stats_input = merged_output @@ -511,7 +501,6 @@ def _compute_output_stats_global( compute_stats_do_not_distinguish_types( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], intensive=intensive, ) ) @@ -519,7 +508,7 @@ def _compute_output_stats_global( bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], + assigned_bias=assigned_bias[kk], rcond=rcond, intensive=intensive, ) @@ -561,6 +550,7 @@ def _compute_output_stats_atomic( keys: list[str], atomic_sampled_idx: dict | None = None, model_pred: dict[str, np.ndarray] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: """Compute output statistics from atomic labels.""" # return directly if no atomic samples @@ -568,6 +558,8 @@ def _compute_output_stats_atomic( len(v) == 0 for v in atomic_sampled_idx.values() ): return {}, {} + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { @@ -640,6 +632,9 @@ def _compute_output_stats_atomic( nan_padding.fill(np.nan) bias_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0) std_atom_e[kk] = np.concatenate([std_atom_e[kk], nan_padding], axis=0) + # the per-type means are independent, so an assigned type is + # overridden exactly + bias_atom_e[kk] = override_assigned_bias(bias_atom_e[kk], assigned_bias[kk]) else: # this key does not have atomic labels, skip it. continue diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 5a1ff9d2cd..8a8aaefd48 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -9,7 +9,6 @@ Optional, ) -import numpy as np import paddle from deepmd.dpmodel.atomic_model import ( @@ -19,6 +18,9 @@ FittingOutputDef, OutputVariableDef, ) +from deepmd.dpmodel.utils.stat import ( + collect_observed_types, +) from deepmd.pd.utils import ( AtomExcludeMask, PairExcludeMask, @@ -42,6 +44,12 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + check_preset_out_bias, + normalize_preset_out_bias, + preset_out_bias_rows, + remap_preset_out_bias, +) log = logging.getLogger(__name__) dtype = env.GLOBAL_PD_FLOAT_PRECISION @@ -65,11 +73,13 @@ class BaseAtomicModel(paddle.nn.Layer, BaseAtomicModel_): of the atomic model. Implemented by removing the pairs from the nlist. rcond : float, optional The condition number for the regression of atomic energy. - preset_out_bias : dict[str, list[Optional[np.ndarray]]], optional - Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. - For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] - The `set_davg_zero` key in the descriptor should be set. + preset_out_bias : dict, optional + Preset output bias, typically the atomic energy in vacuum, keyed by output + name. An assigned output is fixed by the preset: every type that occurs in + the data must be assigned and no statistics are computed for it. The value + is a list with one entry per type (None leaves a type unassigned), a dict + keyed by element name, or the name of a bundled table or the path of a + JSON file holding such a dict. """ @@ -79,7 +89,7 @@ def __init__( atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, - preset_out_bias: dict[str, np.ndarray] | None = None, + preset_out_bias: dict | None = None, data_stat_protect: float = 1e-2, ) -> None: paddle.nn.Layer.__init__(self) @@ -97,12 +107,13 @@ def __init__( self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) self.rcond = rcond - self.preset_out_bias = preset_out_bias + self.preset_out_bias = normalize_preset_out_bias(preset_out_bias, type_map) self.data_stat_protect = data_stat_protect def init_out_stat(self) -> None: """Initialize the output bias.""" self.bias_keys: list[str] = list(self.fitting_output_def().keys()) + check_preset_out_bias(self.preset_out_bias, self.bias_keys) self.max_out_size = max( [self.atomic_output_def()[kk].size for kk in self.bias_keys] ) @@ -356,6 +367,7 @@ def change_type_map( self.reinit_pair_exclude( map_pair_exclude_types(self.pair_exclude_types, remap_index) ) + self.preset_out_bias = remap_preset_out_bias(self.preset_out_bias, remap_index) if has_new_type: extend_shape = [ self.out_bias.shape[0], @@ -488,6 +500,7 @@ def change_out_bias( sample_merged: Callable[[], list[dict]] | list[dict], stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", + observed_type: list[str] | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -505,34 +518,59 @@ def change_out_bias( 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. + An output assigned in `preset_out_bias` is fixed by the preset in both modes: + every element in the data must be assigned, the assigned types take the + preset value, and the types absent from the data keep zero + ('set-by-statistic') or their stored bias ('change-by-statistic'); no + statistics are computed, read or written for such an output, whose + output std keeps its stored value. stat_file_path : Optional[DPPath] The path to the stat file. + observed_type : list[str], optional + The elements that occur in the data; derived from the sample when + not given. """ - if bias_adjust_mode == "change-by-statistic": - delta_bias, out_std = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=list(self.atomic_output_def().keys()), - stat_file_path=stat_file_path, - model_forward=self._get_forward_wrapper_func(), - rcond=self.rcond, - preset_bias=self.preset_out_bias, - ) - self._store_out_stat(delta_bias, out_std, add=True) - elif bias_adjust_mode == "set-by-statistic": - bias_out, std_out = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=list(self.atomic_output_def().keys()), - stat_file_path=stat_file_path, - rcond=self.rcond, - preset_bias=self.preset_out_bias, - stats_distinguish_types=self.get_compute_stats_distinguish_types(), - intensive=self.get_intensive(), - ) - self._store_out_stat(bias_out, std_out) - else: + if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + change = bias_adjust_mode == "change-by-statistic" + distinguish_types = self.get_compute_stats_distinguish_types() + check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + sampled = sample_merged + # === Step 1. Outputs fixed by the preset === + fitted_keys = self.bias_keys + if self.preset_out_bias: + if observed_type is None: + sampled = sample_merged() if callable(sample_merged) else sample_merged + observed_type = collect_observed_types(sampled, self.type_map) + rows = preset_out_bias_rows( + self.preset_out_bias, + self.type_map, + observed_type, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + keep_unassigned=change, + excluded_types=self.atom_exclude_types, + ) + self._store_out_stat({kk: to_paddle_tensor(vv) for kk, vv in rows.items()}) + fitted_keys = [kk for kk in self.bias_keys if kk not in rows] + if not fitted_keys: + return + # === Step 2. Outputs fitted from the data === + # The fitted statistics are absolute biases in 'set-by-statistic' mode and + # shifts of the stored bias in 'change-by-statistic' mode. + forward = self._get_forward_wrapper_func() if change else None + out_bias, out_std = compute_output_stats( + sampled, + self.get_ntypes(), + keys=fitted_keys, + stat_file_path=stat_file_path, + model_forward=forward, + rcond=self.rcond, + stats_distinguish_types=distinguish_types, + intensive=self.get_intensive(), + ) + self._store_out_stat(out_bias, out_std, add=change) def compute_fitting_input_stat( self, @@ -626,21 +664,22 @@ def _get_bias_index( def _store_out_stat( self, out_bias: dict[str, paddle.Tensor], - out_std: dict[str, paddle.Tensor], + out_std: dict[str, paddle.Tensor] | None = None, add: bool = False, ) -> None: + """Store the output bias, and the output std when given, into the model.""" ntypes = self.get_ntypes() out_bias_data = paddle.clone(self.out_bias) out_std_data = paddle.clone(self.out_std) for kk in out_bias.keys(): - assert kk in out_std.keys() idx = self._get_bias_index(kk) size = self._varsize(self.atomic_output_def()[kk].shape) if not add: out_bias_data[idx, :, :size] = out_bias[kk].reshape([ntypes, size]) else: out_bias_data[idx, :, :size] += out_bias[kk].reshape([ntypes, size]) - out_std_data[idx, :, :size] = out_std[kk].reshape([ntypes, size]) + if out_std is not None: + out_std_data[idx, :, :size] = out_std[kk].reshape([ntypes, size]) paddle.assign(out_bias_data, self.out_bias) paddle.assign(out_std_data, self.out_std) diff --git a/deepmd/pd/model/model/__init__.py b/deepmd/pd/model/model/__init__.py index 8348ac039e..59dce4f89b 100644 --- a/deepmd/pd/model/model/__init__.py +++ b/deepmd/pd/model/model/__init__.py @@ -13,11 +13,6 @@ import copy import json -from typing import ( - Any, -) - -import numpy as np from deepmd.pd.model.descriptor.base_descriptor import ( BaseDescriptor, @@ -68,40 +63,6 @@ def _get_standard_model_components( return descriptor, fitting, fitting_net["type"] -def _can_be_converted_to_float(value: Any) -> bool: - try: - float(value) - return True - except (TypeError, ValueError): - # return false for any failure... - return False - - -def _convert_preset_out_bias_to_array( - preset_out_bias: dict | None, type_map: list[str] -) -> dict | None: - if preset_out_bias is not None: - for kk in preset_out_bias: - if len(preset_out_bias[kk]) != len(type_map): - raise ValueError( - "length of the preset_out_bias should be the same as the type_map" - ) - for jj in range(len(preset_out_bias[kk])): - if preset_out_bias[kk][jj] is not None: - if isinstance(preset_out_bias[kk][jj], list): - bb = preset_out_bias[kk][jj] - elif _can_be_converted_to_float(preset_out_bias[kk][jj]): - bb = [float(preset_out_bias[kk][jj])] - else: - raise ValueError( - f"unsupported type/value of the {jj}th element of " - f"preset_out_bias['{kk}'] " - f"{type(preset_out_bias[kk][jj])}" - ) - preset_out_bias[kk][jj] = np.array(bb) - return preset_out_bias - - def get_standard_model(model_params: dict) -> BaseModel: model_params_old = model_params model_params = copy.deepcopy(model_params) @@ -111,10 +72,6 @@ def get_standard_model(model_params: dict) -> BaseModel: ) atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) - preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_array( - preset_out_bias, model_params["type_map"] - ) if fitting_net_type in ["ener", "direct_force_ener"]: modelcls = EnergyModel @@ -127,7 +84,7 @@ def get_standard_model(model_params: dict) -> BaseModel: type_map=model_params["type_map"], atom_exclude_types=atom_exclude_types, pair_exclude_types=pair_exclude_types, - preset_out_bias=preset_out_bias, + preset_out_bias=model_params.get("preset_out_bias"), ) model.model_def_script = json.dumps(model_params_old) return model diff --git a/deepmd/pd/model/task/ener.py b/deepmd/pd/model/task/ener.py index c11124df8e..3cbc4a49c3 100644 --- a/deepmd/pd/model/task/ener.py +++ b/deepmd/pd/model/task/ener.py @@ -20,9 +20,6 @@ from deepmd.pd.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.utils.version import ( - check_version_compatibility, -) dtype = env.GLOBAL_PD_FLOAT_PRECISION device = env.DEVICE @@ -71,7 +68,6 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/pd/model/task/fitting.py b/deepmd/pd/model/task/fitting.py index 01ba370c2b..b3a8acbdc2 100644 --- a/deepmd/pd/model/task/fitting.py +++ b/deepmd/pd/model/task/fitting.py @@ -401,7 +401,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 4, + "@version": 5, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -410,6 +410,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "vacuum_ref": False, "default_fparam": self.default_fparam, "activation_function": self.activation_function, "precision": self.precision, diff --git a/deepmd/pd/model/task/invar_fitting.py b/deepmd/pd/model/task/invar_fitting.py index 04c66befc1..4709cf33c4 100644 --- a/deepmd/pd/model/task/invar_fitting.py +++ b/deepmd/pd/model/task/invar_fitting.py @@ -146,7 +146,12 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 4, 1) + version = data.pop("@version", 1) + check_version_compatibility(version, 5, 1) + if version >= 5 and data.pop("vacuum_ref"): + raise NotImplementedError( + "vacuum_ref is not supported by the Paddle backend" + ) return super().deserialize(data) def output_def(self) -> FittingOutputDef: diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py index 654a28c85d..3cc2599fad 100644 --- a/deepmd/pd/utils/stat.py +++ b/deepmd/pd/utils/stat.py @@ -41,6 +41,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + make_preset_out_bias, + override_assigned_bias, +) from deepmd.utils.stat_file import ( load_output_stat_full_scan, save_output_stat_full_scan, @@ -298,31 +302,6 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> paddle.Tensor: return model_predict -def _make_preset_out_bias( - ntypes: int, - ibias: list[np.ndarray | None], -) -> np.ndarray | None: - """Make preset out bias. - - output: - a np array of shape [ntypes, *(odim0, odim1, ...)] is any item is not None - None if all items are None. - """ - if len(ibias) != ntypes: - raise ValueError("the length of preset bias list should be ntypes") - if all(ii is None for ii in ibias): - return None - for refb in ibias: - if refb is not None: - break - refb = np.array(refb) - nbias = [ - np.full_like(refb, np.nan, dtype=np.float64) if ii is None else ii - for ii in ibias - ] - return np.array(nbias) - - def _fill_stat_with_global( atomic_stat: np.ndarray | None, global_stat: np.ndarray, @@ -378,10 +357,13 @@ def compute_output_stats( The path to the stat file. rcond : float, optional The condition number for the regression of atomic energy. - preset_bias : dict[str, list[Optional[paddle.Tensor]]], optional - Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. + preset_bias : dict[str, list[Optional[np.ndarray]]], optional + Assigned values of the returned bias, given by key:value pairs. + The value is a list with one element per type: None leaves the type to the + statistics, an np.ndarray of output shape assigns the type. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The values live in the frame of the returned bias: absolute biases without + `model_forward`, shifts of the model's stored bias with `model_forward`. The `set_davg_zero` key in the descriptor should be set. model_forward : Callable[..., paddle.Tensor], optional The wrapped forward function of atomic model. @@ -405,6 +387,11 @@ def compute_output_stats( ) redu_scanner = None + # Model residuals depend on parameters not recorded in the statistics cache. + # Neither reuse nor persist them as absolute output statistics. + if model_forward is not None: + stat_file_path = None + # try to restore the bias from stat file bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) if ( @@ -500,13 +487,22 @@ def compute_output_stats( else None ) + # assigned bias of every output as a (ntypes, ...) array, NaN where a + # type is left to the statistics + assigned_bias = { + kk: make_preset_out_bias(ntypes, preset_bias[kk]) + if preset_bias is not None and kk in preset_bias + else None + for kk in keys + } + # compute stat bias_atom_g, std_atom_g = _compute_output_stats_global( sampled, ntypes, keys, rcond, - preset_bias, + assigned_bias, global_sampled_idx, stats_distinguish_types, intensive, @@ -519,6 +515,7 @@ def compute_output_stats( keys, atomic_sampled_idx, model_pred_a, + assigned_bias, ) # merge global/atomic bias @@ -557,7 +554,7 @@ def _compute_output_stats_global( ntypes: int, keys: list[str], rcond: float | None = None, - preset_bias: dict[str, list[paddle.Tensor | None]] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, global_sampled_idx: dict | None = None, stats_distinguish_types: bool = True, intensive: bool = False, @@ -609,15 +606,8 @@ def _compute_output_stats_global( if len(input_natoms[kk]) > 0 } nf = {kk: merged_natoms[kk].shape[0] for kk in keys if kk in merged_natoms} - if preset_bias is not None: - assigned_atom_ener = { - kk: _make_preset_out_bias(ntypes, preset_bias[kk]) - if kk in preset_bias.keys() - else None - for kk in keys - } - else: - assigned_atom_ener = dict.fromkeys(keys) + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) if model_pred is None: stats_input = merged_output @@ -651,7 +641,7 @@ def _compute_output_stats_global( for kk in keys: if scan is not None and kk in scan.stats: bias_atom_e[kk], std_atom_e[kk] = scan.stats[kk].solve( - assigned_bias=assigned_atom_ener[kk], + assigned_bias=assigned_bias[kk], rcond=rcond, type_mask=type_mask, ) @@ -662,7 +652,6 @@ def _compute_output_stats_global( compute_stats_do_not_distinguish_types( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], intensive=intensive, ) ) @@ -670,7 +659,7 @@ def _compute_output_stats_global( bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], + assigned_bias=assigned_bias[kk], rcond=rcond, intensive=intensive, ) @@ -726,6 +715,7 @@ def _compute_output_stats_atomic( keys: list[str], atomic_sampled_idx: dict | None = None, model_pred: dict[str, np.ndarray] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: """Compute output statistics from atomic labels.""" # return directly if no atomic samples @@ -733,6 +723,8 @@ def _compute_output_stats_atomic( len(v) == 0 for v in atomic_sampled_idx.values() ): return {}, {} + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { @@ -804,6 +796,9 @@ def _compute_output_stats_atomic( nan_padding.fill(np.nan) bias_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0) std_atom_e[kk] = np.concatenate([std_atom_e[kk], nan_padding], axis=0) + # the per-type means are independent, so an assigned type is + # overridden exactly + bias_atom_e[kk] = override_assigned_bias(bias_atom_e[kk], assigned_bias[kk]) else: # this key does not have atomic labels, skip it. continue diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 62693d547d..b3cd33798a 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -985,7 +985,17 @@ def _freeze_sezm_to_pt2( model = get_model(params) is_spin = _model_has_spin(model) ModelWrapper(model).load_state_dict(state_dict) + if model.get_active_mode() == "dens": + raise ValueError( + "SeZM freeze supports only the `ener` mode: the DeNS head serves " + "training alone and is not exported." + ) model.eval() + # The vacuum reference is resolved on the target device, folded into the + # fitting bias or stored as a per-type table, so the exported graph + # carries no reference atoms. + model.to(target_device) + model.fold_vacuum_reference() model.to("cpu") # Device-dependent Python branches resolve on the CPU tracing inputs, so diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index d4f4b4efc7..d18e780cd3 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -10,7 +10,6 @@ Optional, ) -import numpy as np import torch from deepmd.dpmodel.atomic_model import ( @@ -20,6 +19,9 @@ FittingOutputDef, OutputVariableDef, ) +from deepmd.dpmodel.utils.stat import ( + collect_observed_types, +) from deepmd.pt.utils import ( AtomExcludeMask, PairExcludeMask, @@ -46,6 +48,12 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + check_preset_out_bias, + normalize_preset_out_bias, + preset_out_bias_rows, + remap_preset_out_bias, +) log = logging.getLogger(__name__) dtype = env.GLOBAL_PT_FLOAT_PRECISION @@ -69,11 +77,13 @@ class BaseAtomicModel(torch.nn.Module, BaseAtomicModel_): of the atomic model. Implemented by removing the pairs from the nlist. rcond : float, optional The condition number for the regression of atomic energy. - preset_out_bias : dict[str, list[Optional[np.ndarray]]], optional - Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. - For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] - The `set_davg_zero` key in the descriptor should be set. + preset_out_bias : dict, optional + Preset output bias, typically the atomic energy in vacuum, keyed by output + name. An assigned output is fixed by the preset: every type that occurs in + the data must be assigned and no statistics are computed for it. The value + is a list with one entry per type (None leaves a type unassigned), a dict + keyed by element name, or the name of a bundled table or the path of a + JSON file holding such a dict. """ @@ -83,7 +93,7 @@ def __init__( atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, - preset_out_bias: dict[str, np.ndarray] | None = None, + preset_out_bias: dict | None = None, data_stat_protect: float = 1e-2, ) -> None: torch.nn.Module.__init__(self) @@ -92,7 +102,7 @@ def __init__( self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) self.rcond = rcond - self.preset_out_bias = preset_out_bias + self.preset_out_bias = normalize_preset_out_bias(preset_out_bias, type_map) self.data_stat_protect = data_stat_protect self._observed_type: list[str] | None = None @@ -144,6 +154,7 @@ def init_out_stat(self) -> None: """Initialize the output bias.""" ntypes = self.get_ntypes() self.bias_keys: list[str] = list(self.fitting_output_def().keys()) + check_preset_out_bias(self.preset_out_bias, self.bias_keys) self.max_out_size = max( [self.atomic_output_def()[kk].size for kk in self.bias_keys] ) @@ -182,6 +193,10 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference into the fitting bias; nothing to fold without a fitting network.""" + return + def get_compute_stats_distinguish_types(self) -> bool: """Get whether the fitting net computes stats which are not distinguished between different types of atoms.""" return True @@ -477,6 +492,7 @@ def change_type_map( self.reinit_pair_exclude( map_pair_exclude_types(self.pair_exclude_types, remap_index) ) + self.preset_out_bias = remap_preset_out_bias(self.preset_out_bias, remap_index) if has_new_type: extend_shape = [ self.out_bias.shape[0], @@ -580,6 +596,7 @@ def compute_or_load_out_stat( merged, stat_file_path=stat_file_path, bias_adjust_mode="set-by-statistic", + observed_type=self.observed_type, ) def apply_out_stat( @@ -610,6 +627,8 @@ def change_out_bias( sample_merged: Callable[[], list[dict]] | list[dict], stat_file_path: DPPath | None = None, bias_adjust_mode: str = "change-by-statistic", + model_forward: Callable[..., dict[str, torch.Tensor]] | None = None, + observed_type: list[str] | None = None, ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -622,41 +641,74 @@ def change_out_bias( - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. bias_adjust_mode : str The mode for changing output bias : ['change-by-statistic', 'set-by-statistic'] 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. - stat_file_path : Optional[DPPath] - The path to the stat file. + An output assigned in `preset_out_bias` is fixed by the preset in both modes: + every element in the data must be assigned, the assigned types take the + preset value, and the types absent from the data keep zero + ('set-by-statistic') or their stored bias ('change-by-statistic'); no + statistics are computed, read or written for such an output, whose + output std keeps its stored value. + model_forward : Callable[..., dict[str, torch.Tensor]], optional + Predictor of the complete atomic outputs used by 'change-by-statistic'. + Defaults to the forward of this atomic model; a model that adds + contributions on top of the atomic model passes its complete forward. + observed_type : list[str], optional + The elements that occur in the data; derived from the sample when + not given. """ - if bias_adjust_mode == "change-by-statistic": - delta_bias, out_std = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=self.bias_keys, - stat_file_path=stat_file_path, - model_forward=self._get_forward_wrapper_func(), - rcond=self.rcond, - preset_bias=self.preset_out_bias, - stats_distinguish_types=self.get_compute_stats_distinguish_types(), - intensive=self.get_intensive(), + if bias_adjust_mode not in ("change-by-statistic", "set-by-statistic"): + raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + change = bias_adjust_mode == "change-by-statistic" + distinguish_types = self.get_compute_stats_distinguish_types() + check_preset_out_bias(self.preset_out_bias, self.bias_keys, distinguish_types) + sampled = sample_merged + # === Step 1. Outputs fixed by the preset === + fitted_keys = self.bias_keys + if self.preset_out_bias: + if observed_type is None: + sampled = sample_merged() if callable(sample_merged) else sample_merged + observed_type = collect_observed_types(sampled, self.type_map) + rows = preset_out_bias_rows( + self.preset_out_bias, + self.type_map, + observed_type, + to_numpy_array(self.out_bias), + self.bias_keys, + [self.atomic_output_def()[kk].size for kk in self.bias_keys], + keep_unassigned=change, + excluded_types=self.atom_exclude_types, ) - self._store_out_stat(delta_bias, out_std, add=True) - elif bias_adjust_mode == "set-by-statistic": - bias_out, std_out = compute_output_stats( - sample_merged, - self.get_ntypes(), - keys=self.bias_keys, - stat_file_path=stat_file_path, - rcond=self.rcond, - preset_bias=self.preset_out_bias, - stats_distinguish_types=self.get_compute_stats_distinguish_types(), - intensive=self.get_intensive(), + self._store_out_stat({kk: to_torch_tensor(vv) for kk, vv in rows.items()}) + fitted_keys = [kk for kk in self.bias_keys if kk not in rows] + if not fitted_keys: + return + # === Step 2. Outputs fitted from the data === + # The fitted statistics are absolute biases in 'set-by-statistic' mode and + # shifts of the stored bias in 'change-by-statistic' mode. + forward = None + if change: + forward = ( + self._get_forward_wrapper_func() + if model_forward is None + else model_forward ) - self._store_out_stat(bias_out, std_out) - else: - raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + out_bias, out_std = compute_output_stats( + sampled, + self.get_ntypes(), + keys=fitted_keys, + stat_file_path=stat_file_path, + model_forward=forward, + rcond=self.rcond, + stats_distinguish_types=distinguish_types, + intensive=self.get_intensive(), + ) + self._store_out_stat(out_bias, out_std, add=change) def compute_fitting_input_stat( self, @@ -760,21 +812,22 @@ def _get_bias_index( def _store_out_stat( self, out_bias: dict[str, torch.Tensor], - out_std: dict[str, torch.Tensor], + out_std: dict[str, torch.Tensor] | None = None, add: bool = False, ) -> None: + """Store the output bias, and the output std when given, into the model.""" ntypes = self.get_ntypes() out_bias_data = torch.clone(self.out_bias) out_std_data = torch.clone(self.out_std) for kk in out_bias.keys(): - assert kk in out_std.keys() idx = self._get_bias_index(kk) size = self._varsize(self.atomic_output_def()[kk].shape) if not add: out_bias_data[idx, :, :size] = out_bias[kk].view(ntypes, size) else: out_bias_data[idx, :, :size] += out_bias[kk].view(ntypes, size) - out_std_data[idx, :, :size] = out_std[kk].view(ntypes, size) + if out_std is not None: + out_std_data[idx, :, :size] = out_std[kk].view(ntypes, size) self.out_bias.copy_(out_bias_data) self.out_std.copy_(out_std_data) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index c9badefcad..48ae314e17 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -23,6 +23,15 @@ def __init__( "fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel" ) super().__init__(descriptor, fitting, type_map, **kwargs) + if self.preset_out_bias is not None and any( + entry is not None + for entries in self.preset_out_bias.values() + for entry in entries + ): + raise ValueError( + "Dipole models do not apply an output bias and cannot use " + "assigned preset_out_bias values." + ) def apply_out_stat( self, diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 35983ef086..8adc0abbf8 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -48,6 +48,8 @@ class DPAtomicModel(BaseAtomicModel): For example `type_map[1]` gives the name of the type 1. """ + _supports_vacuum_ref: bool = False + def __init__( self, descriptor: BaseDescriptor, @@ -55,6 +57,11 @@ def __init__( type_map: list[str], **kwargs: Any, ) -> None: + if fitting.vacuum_ref and not self._supports_vacuum_ref: + raise NotImplementedError( + "vacuum_ref is only supported by DPA4/SeZM models " + "(model.type='dpa4' or 'sezm') in the PyTorch backend" + ) super().__init__(type_map, **kwargs) ntypes = len(type_map) self.type_map = type_map @@ -395,12 +402,11 @@ def compute_or_load_stat( wrapped_sampler = self._make_wrapped_sampler(sampled_func) self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) self.compute_fitting_input_stat(wrapped_sampler, stat_file_path) - if compute_or_load_out_stat: - self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) - self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) def compute_fitting_input_stat( self, diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index ede4497df3..46b3a50599 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -148,6 +148,11 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of every sub-model into its bias.""" + for model in self.models: + model.fold_vacuum_reference() + def change_type_map( self, type_map: list[str], diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index eb70cc0e78..7496a2f6f8 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -59,8 +59,6 @@ class PairTabAtomicModel(BaseAtomicModel): For example `type_map[1]` gives the name of the type 1. rcond : float, optional The condition number for the regression of atomic energy. - atom_ener - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descriptor should be set. """ @@ -249,17 +247,16 @@ def compute_or_load_stat( If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ - if compute_or_load_out_stat: - self.compute_or_load_out_stat(sampled_func, stat_file_path) - + observed_stat_path = stat_file_path if stat_file_path is not None and self.type_map is not None: - stat_file_path /= " ".join(self.type_map) - + observed_stat_path = stat_file_path / " ".join(self.type_map) self._collect_and_set_observed_type( sampled_func if callable(sampled_func) else lambda: sampled_func, - stat_file_path, + observed_stat_path, preset_observed_type, ) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(sampled_func, stat_file_path) def forward_atomic( self, diff --git a/deepmd/pt/model/atomic_model/sezm_atomic_model.py b/deepmd/pt/model/atomic_model/sezm_atomic_model.py index e7014b6963..190b5960fe 100644 --- a/deepmd/pt/model/atomic_model/sezm_atomic_model.py +++ b/deepmd/pt/model/atomic_model/sezm_atomic_model.py @@ -35,12 +35,20 @@ from deepmd.pt.model.task.sezm_ener import ( SeZMEnergyFittingNet, ) +from deepmd.pt.utils import ( + env, +) from deepmd.pt.utils.utils import ( to_torch_tensor, ) from deepmd.utils.stat_file import ( load_required_items, ) +from deepmd.utils.vacuum_reference import ( + reference_charge_spin, + reference_spin, + resolve_vacuum_ref, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -78,6 +86,8 @@ class SeZMAtomicModel(DPAtomicModel): If fitting is not an energy fitting network. """ + _supports_vacuum_ref: bool = True + def __init__( self, descriptor: Any, @@ -106,6 +116,16 @@ def __init__( "dens_force_rmsd", self.out_std.new_tensor(1.0), ) + # === Reference of the isolated atoms === + # Each head references its output only when the preset fixes the + # bias of that output; the conditions of the reference atoms follow. + resolve_vacuum_ref(fitting, self.preset_out_bias) + if dens_fitting is not None: + resolve_vacuum_ref(dens_fitting, self.preset_out_bias) + self.add_spin_ebd: bool = self.descriptor.use_spin is not None + self.register_buffer("vacuum_charge_spin", None, persistent=False) + self.register_buffer("vacuum_spin", None, persistent=False) + self.init_vacuum_conditions() self.dens_fitting_net = dens_fitting # Start unlocked when `active_mode` is not provided. # The mode will be decided later by training setup (`loss.type`) @@ -115,6 +135,97 @@ def __init__( if active_mode is not None: self.set_active_mode(active_mode) + def init_vacuum_conditions(self) -> None: + """Build the conditioning inputs of the reference atoms for the type map. + + A table exists for every condition the descriptor takes when the + fitting references the isolated atoms; the type names must then be + element symbols. + """ + vacuum_ref = self.fitting_net.vacuum_ref + dtype = env.GLOBAL_PT_FLOAT_PRECISION + self.vacuum_charge_spin = ( + torch.as_tensor( + reference_charge_spin(self.type_map), dtype=dtype, device=env.DEVICE + ) + if vacuum_ref and self.add_chg_spin_ebd + else None + ) + self.vacuum_spin = ( + torch.as_tensor( + reference_spin(self.type_map), dtype=dtype, device=env.DEVICE + ) + if vacuum_ref and self.add_spin_ebd + else None + ) + + def vacuum_conditions(self) -> dict[str, torch.Tensor]: + """Conditioning inputs of one isolated neutral ground-state atom per type. + + Returns + ------- + dict[str, torch.Tensor] + The charge/spin conditions with shape (ntypes, 2) under + ``charge_spin`` and the spin vectors with shape (ntypes, 3) under + ``spin``, for the conditions the descriptor takes. + """ + conditions = {} + if self.vacuum_charge_spin is not None: + conditions["charge_spin"] = self.vacuum_charge_spin + if self.vacuum_spin is not None: + conditions["spin"] = self.vacuum_spin + return conditions + + def vacuum_descriptor(self) -> torch.Tensor: + """Descriptor of an isolated atom of every type. + + Every type is evaluated as a single-atom frame without neighbors under + the reference conditions. + + Returns + ------- + torch.Tensor + The vacuum descriptor with shape (ntypes, dim_descrpt) on the + device of the fitting bias. + """ + bias = self.fitting_net.bias_atom_e + dtype = env.GLOBAL_PT_FLOAT_PRECISION + ntypes = self.get_ntypes() + coord = torch.zeros((ntypes, 1, 3), dtype=dtype, device=bias.device) + atype = torch.arange(ntypes, dtype=torch.long, device=bias.device).view( + ntypes, 1 + ) + nlist = torch.full( + (ntypes, 1, self.descriptor.get_nnei()), + -1, + dtype=torch.long, + device=bias.device, + ) + mapping = torch.zeros((ntypes, 1), dtype=torch.long, device=bias.device) + conditions = { + name: table.to(dtype=dtype, device=bias.device) + for name, table in self.vacuum_conditions().items() + } + if "spin" in conditions: + conditions["spin"] = conditions["spin"][:, None, :] + descriptor = self.descriptor( + coord, atype, nlist, mapping=mapping, **conditions + )[0] + return descriptor.reshape(ntypes, -1) + + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the energy fitting so the forward carries no reference atoms. + + The vacuum descriptor of every type is evaluated once with the current + parameters and handed to the fitting, which folds the reference into + its bias or stores the table (see its ``fold_vacuum_reference``). The + DeNS head serves training alone and is not exported, so it keeps its + reference. + """ + fitting = self.fitting_net + if fitting.needs_vacuum_descriptor(): + fitting.fold_vacuum_reference(self.vacuum_descriptor()) + def _load_from_state_dict( self, state_dict: dict[str, torch.Tensor], @@ -545,6 +656,7 @@ def change_type_map( type_map=type_map, model_with_new_type_stat=ref_dens, ) + self.init_vacuum_conditions() def compute_or_load_stat( self, @@ -578,6 +690,11 @@ def compute_or_load_stat( wrapped_sampler = self._make_wrapped_sampler(sampled_func) self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) self.compute_fitting_input_stat(wrapped_sampler, stat_file_path) + self._collect_and_set_observed_type( + wrapped_sampler, + stat_file_path, + preset_observed_type, + ) if compute_or_load_out_stat: self.set_active_mode("ener") try: @@ -587,12 +704,6 @@ def compute_or_load_stat( if original_mode == "dens": self._compute_or_load_dens_force_stat(wrapped_sampler, stat_file_path) - self._collect_and_set_observed_type( - wrapped_sampler, - stat_file_path, - preset_observed_type, - ) - def apply_out_stat( self, ret: dict[str, torch.Tensor], @@ -728,6 +839,7 @@ def _build_ener_fitting_kwargs(self) -> dict[str, Any]: "exclude_types": copy.deepcopy(fitting.exclude_types), "trainable": copy.deepcopy(fitting.trainable), "atom_ener": copy.deepcopy(fitting.atom_ener), + "vacuum_ref": bool(fitting.vacuum_ref), "use_aparam_as_mask": bool(fitting.use_aparam_as_mask), } diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 3988cd29eb..999bbea447 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -1258,7 +1258,7 @@ def forward( dtype=extended_coord.dtype, device=extended_coord.device, ) - descriptor, _ = self.forward_with_edges( + descriptor, _, _ = self.forward_with_edges( extended_coord=extended_coord, extended_atype=extended_atype, edge_index=edge_index, @@ -1474,7 +1474,8 @@ def forward_with_edges( spin: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, nloc: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + vacuum_conditions: dict[str, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Compute the descriptor from a sparse edge list. @@ -1516,12 +1517,21 @@ def forward_with_edges( nloc Number of owned (local) atoms per frame. Required when ``comm_dict`` is provided; the final scalar read-out is restricted to these atoms. + vacuum_conditions + Conditioning inputs of one isolated atom per type, the neutral + ground-state atom, under ``charge_spin`` with shape (ntypes, 2) + and ``spin`` with shape (ntypes, 3) as the descriptor takes them. + When given, the reference atoms are carried through the same + forward as additional nodes and their vacuum descriptor is + returned. Returns ------- - tuple[torch.Tensor, torch.Tensor] - The scalar descriptor with shape ``(nf, nloc, channels)`` and the - final equivariant latent with shape ``(nf * nloc, D_final, 1, channels)``. + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] + The scalar descriptor with shape ``(nf, nloc, channels)``, the + final equivariant latent with shape ``(nf * nloc, D_final, 1, channels)`` + and, with ``vacuum_conditions``, the vacuum descriptor with shape + ``(ntypes, channels)``; ``None`` otherwise. """ # === Step 1. Setup dimensions === # ``n_per_frame`` is the per-frame node count: ``nloc`` in the @@ -1545,11 +1555,41 @@ def forward_with_edges( ensure_comm_registered() out_nloc = nloc if parallel else n_per_frame + n_real_nodes = nf * n_per_frame atype_flat = extended_atype.reshape(-1) # (N,) + # === Step 1b. Vacuum reference nodes === + # One isolated atom of every type follows the real nodes, conditioned + # as the neutral ground-state atom. Every node-wise operation leaves + # the real nodes unaffected, so the read-out rows of the reference + # nodes are the vacuum descriptor of every type. + vacuum_ref = vacuum_conditions is not None + if vacuum_conditions is not None: + atype_flat = torch.cat( + [ + atype_flat, + torch.arange( + self.ntypes, dtype=atype_flat.dtype, device=atype_flat.device + ), + ] + ) + if spin is not None and self.spin_embedding is not None: + spin = torch.cat( + [spin.reshape(-1, 3), vacuum_conditions["spin"].to(spin.dtype)] + ) + if force_embedding is not None: + force_embedding = torch.cat( + [ + force_embedding, + force_embedding.new_zeros( + (self.ntypes, *force_embedding.shape[1:]) + ), + ] + ) + # === Step 2. Type embedding (l=0) === with nvtx_range("type_embedding"): - type_ebed = self.type_embedding(extended_atype).reshape( + type_ebed = self.type_embedding(atype_flat).reshape( -1, self.channels ) # (N, C) if self.charge_spin_embedding is not None: @@ -1558,6 +1598,9 @@ def forward_with_edges( charge_spin, nf=nf, nloc=n_per_frame, + vacuum_reference=None + if vacuum_conditions is None + else vacuum_conditions["charge_spin"], ) n_nodes = type_ebed.shape[0] @@ -1701,19 +1744,36 @@ def forward_with_edges( # === Step 11. Keep the owned-atom rows for the read-out === # ``n_out_nodes`` is the owned-node count in the flattened layout # (``nf * nloc``). Single-domain: ``out_nloc == n_per_frame``, so this - # equals the whole node set and the slice is a no-op. Parallel + # equals the whole real node set and the slice is a no-op. Parallel # (single-frame): it drops the trailing ghost rows that only fed message - # passing -- LAMMPS orders owned atoms before ghosts, so they lead. + # passing -- LAMMPS orders owned atoms before ghosts, so they lead. The + # vacuum reference rows, when present, trail the real nodes and share + # the read-out with the owned rows. n_out_nodes = nf * out_nloc - x = x[:n_out_nodes] + latent = x[:n_out_nodes] + if vacuum_ref: + x = torch.cat([latent, x[n_real_nodes:]], dim=0) + n_readout = n_out_nodes + self.ntypes + else: + x = latent + n_readout = n_out_nodes # === Step 12. Final l=0 output mixing === with nvtx_range("output_ffn"): - x_scalar = self._apply_readout(x, n_out_nodes) + x_scalar = self._apply_readout(x, n_readout).to( + dtype=env.GLOBAL_PT_FLOAT_PRECISION + ) # === Step 13. Reshape to (nf, nloc, channels) and return === - descriptor = x_scalar.reshape(nf, out_nloc, self.channels) # (nf, nloc, C) - return descriptor.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), x.contiguous() + descriptor = x_scalar[:n_out_nodes].reshape( + nf, out_nloc, self.channels + ) # (nf, nloc, C) + vacuum = ( + x_scalar[n_out_nodes:].reshape(self.ntypes, self.channels) + if vacuum_ref + else None + ) + return descriptor, latent.contiguous(), vacuum def _forward_blocks( self, @@ -1994,6 +2054,7 @@ def _apply_charge_spin_embedding( *, nf: int, nloc: int, + vacuum_reference: torch.Tensor | None = None, ) -> torch.Tensor: """ Add frame-level charge and spin conditions to scalar type features. @@ -2001,22 +2062,32 @@ def _apply_charge_spin_embedding( Parameters ---------- type_ebed - Flattened type embeddings with shape (nf * nloc, channels). + Flattened type embeddings with shape (nf * nloc, channels), followed + by one row per type when ``vacuum_reference`` is given. charge_spin Frame-level charge and spin conditions with shape (nf, 2). nf Number of frames. nloc Number of local atoms. + vacuum_reference + Charge and spin conditions of the vacuum reference nodes that trail + the real nodes, with shape (ntypes, 2), or None. Returns ------- torch.Tensor - Conditioned type embeddings with shape (nf * nloc, channels). + Conditioned type embeddings with the shape of ``type_ebed``. """ condition = self.charge_spin_embedding(charge_spin.to(dtype=type_ebed.dtype)) condition = condition[:, None, :].expand(nf, nloc, self.channels) - return type_ebed + condition.reshape_as(type_ebed) + condition = condition.reshape(nf * nloc, self.channels) + if vacuum_reference is not None: + reference = self.charge_spin_embedding( + vacuum_reference.to(dtype=type_ebed.dtype) + ) + condition = torch.cat([condition, reference], dim=0) + return type_ebed + condition def _apply_spin_embedding( self, @@ -2362,6 +2433,10 @@ def get_ntypes(self) -> int: def get_type_map(self) -> list[str]: return self.type_map if self.type_map is not None else [] + def supports_native_spin(self) -> bool: + """SeZM accepts per-atom ``spin`` vectors (native magnetic conditioning).""" + return True + def get_dim_chg_spin(self) -> int: """Return the charge/spin condition width.""" return 2 if self.add_chg_spin_ebd else 0 diff --git a/deepmd/pt/model/descriptor/sezm_nn/dens.py b/deepmd/pt/model/descriptor/sezm_nn/dens.py index a5f084a0a8..10cbc2fb86 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/dens.py +++ b/deepmd/pt/model/descriptor/sezm_nn/dens.py @@ -418,6 +418,9 @@ class SeZMDeNSFittingNet(torch.nn.Module): Whether the `dens` fitting parameters are trainable. atom_ener Optional vacuum atomic energy contribution for the scalar energy branch. + vacuum_ref + Whether the scalar energy branch references every atom to the isolated + atom of its type. use_aparam_as_mask Whether atomic parameters act as masks in the scalar energy branch. """ @@ -448,6 +451,7 @@ def __init__( exclude_types: list[int] | None = None, trainable: bool | list[bool] = True, atom_ener: list[torch.Tensor | None] | None = None, + vacuum_ref: bool = False, use_aparam_as_mask: bool = False, ) -> None: super().__init__() @@ -505,6 +509,7 @@ def __init__( exclude_types=self.exclude_types, trainable=self.trainable, atom_ener=self.atom_ener, + vacuum_ref=bool(vacuum_ref), use_aparam_as_mask=self.use_aparam_as_mask, ) @@ -575,6 +580,24 @@ def get_default_fparam(self) -> torch.Tensor | None: """Return default frame parameters of the energy branch.""" return self.energy_head.get_default_fparam() + def needs_vacuum_descriptor(self) -> bool: + """Whether the energy head takes the vacuum descriptor of every type from the descriptor.""" + return self.energy_head.needs_vacuum_descriptor() + + @property + def vacuum_ref(self) -> bool: + """Whether the scalar energy branch references every atom to the isolated atom of its type.""" + return self.energy_head.vacuum_ref + + @vacuum_ref.setter + def vacuum_ref(self, value: bool) -> None: + self.energy_head.vacuum_ref = bool(value) + + @property + def var_name(self) -> str: + """Output name of the scalar energy branch.""" + return self.energy_head.var_name + def get_dim_aparam(self) -> int: """Return the atomic-parameter width of the energy branch.""" return self.energy_head.get_dim_aparam() @@ -644,6 +667,7 @@ def forward( noise_mask: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, return_components: bool = False, ) -> dict[str, torch.Tensor]: """ @@ -663,6 +687,9 @@ def forward( Optional frame parameters. aparam Optional atomic parameters. + vacuum_descriptor + Descriptor of an isolated atom of every type with shape + `(ntypes, dim_descrpt)`, required by ``vacuum_ref``. return_components If true, also return the clean-force and denoising branches. @@ -681,6 +708,7 @@ def forward( atype, fparam=fparam, aparam=aparam, + vacuum_descriptor=vacuum_descriptor, ) clean_force = self.direct_force_head(latent).view(nf, nloc, 3) denoising_force = self.denoising_head(latent).view(nf, nloc, 3) @@ -730,6 +758,7 @@ def serialize(self) -> dict[str, Any]: "exclude_types": self.exclude_types.copy(), "trainable": self.trainable, "atom_ener": self.atom_ener, + "vacuum_ref": self.vacuum_ref, "use_aparam_as_mask": self.use_aparam_as_mask, }, "@variables": {key: np_safe(value) for key, value in state.items()}, diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 18225dbdc8..27c2d05af4 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -273,6 +273,7 @@ def get_linear_model(model_params: dict) -> BaseModel: weights=weights, atom_exclude_types=atom_exclude_types, pair_exclude_types=pair_exclude_types, + preset_out_bias=model_params.get("preset_out_bias"), ) model.shared_links = shared_links if shared_links: @@ -388,45 +389,12 @@ def get_zbl_model(model_params: dict) -> DPZBLModel: type_map=model_params["type_map"], atom_exclude_types=atom_exclude_types, pair_exclude_types=pair_exclude_types, + preset_out_bias=model_params.get("preset_out_bias"), ) model.model_def_script = json.dumps(model_params) return model -def _can_be_converted_to_float(value: Any) -> bool | None: - try: - float(value) - return True - except (TypeError, ValueError): - # return false for any failure... - return False - - -def _convert_preset_out_bias_to_array( - preset_out_bias: dict | None, type_map: list[str] -) -> dict | None: - if preset_out_bias is not None: - for kk in preset_out_bias: - if len(preset_out_bias[kk]) != len(type_map): - raise ValueError( - "length of the preset_out_bias should be the same as the type_map" - ) - for jj in range(len(preset_out_bias[kk])): - if preset_out_bias[kk][jj] is not None: - if isinstance(preset_out_bias[kk][jj], list): - bb = preset_out_bias[kk][jj] - elif _can_be_converted_to_float(preset_out_bias[kk][jj]): - bb = [float(preset_out_bias[kk][jj])] - else: - raise ValueError( - f"unsupported type/value of the {jj}th element of " - f"preset_out_bias['{kk}'] " - f"{type(preset_out_bias[kk][jj])}" - ) - preset_out_bias[kk][jj] = np.array(bb) - return preset_out_bias - - def get_standard_model(model_params: dict) -> BaseModel: bridging_method = str(model_params.get("bridging_method", "none")) if bridging_method.lower() not in ("none", ""): @@ -446,9 +414,6 @@ def get_standard_model(model_params: dict) -> BaseModel: atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_array( - preset_out_bias, model_params["type_map"] - ) data_stat_protect = model_params.get("data_stat_protect", 1e-2) if fitting_net_type == "dipole": @@ -582,9 +547,6 @@ def get_sezm_model(model_params: dict) -> BaseModel: ) atom_exclude_types = model_params.get("atom_exclude_types", []) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_array( - preset_out_bias, model_params["type_map"] - ) data_stat_protect = model_params.get("data_stat_protect", 1e-2) use_compile = bool(model_params.get("use_compile", False)) enable_tf32 = bool(model_params.get("enable_tf32", True)) @@ -681,9 +643,6 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel: fitting = SeZMEnergyFittingNet(**fitting_net) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_array( - preset_out_bias, model_params["type_map"] - ) data_stat_protect = model_params.get("data_stat_protect", 1e-2) use_compile = bool(model_params.get("use_compile", False)) enable_tf32 = bool(model_params.get("enable_tf32", True)) @@ -764,9 +723,6 @@ def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel: fitting_net["dim_descrpt"] = descriptor.get_dim_out() fitting = SeZMEnergyFittingNet(**fitting_net) preset_out_bias = model_params.get("preset_out_bias") - preset_out_bias = _convert_preset_out_bias_to_array( - preset_out_bias, model_params["type_map"] - ) data_stat_protect = model_params.get("data_stat_protect", 1e-2) use_compile = bool(model_params.get("use_compile", False)) enable_tf32 = bool(model_params.get("enable_tf32", True)) diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 494772d453..f81b3a5b93 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -37,9 +37,6 @@ extend_input_and_build_neighbor_list, nlist_distinguish_types, ) -from deepmd.pt.utils.stat import ( - compute_output_stats, -) from deepmd.utils.path import ( DPPath, ) @@ -328,27 +325,6 @@ def predict_atomic_outputs_for_stat( spin=spin, ) - def _change_out_bias_with_model_forward( - self, - merged: Callable[[], list[dict]] | list[dict], - model_forward: Callable[..., dict[str, torch.Tensor]], - ) -> None: - """Fit a residual output-bias shift from a complete model predictor.""" - atomic_model = self.atomic_model - delta_bias, out_std = compute_output_stats( - merged, - atomic_model.get_ntypes(), - keys=atomic_model.bias_keys, - model_forward=model_forward, - rcond=atomic_model.rcond, - preset_bias=atomic_model.preset_out_bias, - stats_distinguish_types=( - atomic_model.get_compute_stats_distinguish_types() - ), - intensive=atomic_model.get_intensive(), - ) - atomic_model._store_out_stat(delta_bias, out_std, add=True) - def change_out_bias( self, merged: Any, @@ -356,6 +332,9 @@ def change_out_bias( ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. + The residual fit of 'change-by-statistic' evaluates the complete model + prediction through ``predict_atomic_outputs_for_stat``. + Parameters ---------- merged : Union[Callable[[], list[dict]], list[dict]] @@ -371,15 +350,10 @@ def change_out_bias( and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. """ - if bias_adjust_mode == "change-by-statistic": - self._change_out_bias_with_model_forward( - merged, - self.predict_atomic_outputs_for_stat, - ) - return self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, + model_forward=self.predict_atomic_outputs_for_stat, ) def forward_common_lower( @@ -778,6 +752,10 @@ def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the atomic model into its fitting bias.""" + self.atomic_model.fold_vacuum_reference() + def compute_or_load_stat( self, sampled_func: Callable[[], Any], diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 1bb46cdcfd..72d728046b 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -655,7 +655,10 @@ def _sezm_structure_key(model: SeZMModel) -> tuple[Any, ...]: descriptor.inner_clamp_r_outer, int(descriptor.get_dim_chg_spin()), ) - fitting_state = (_int_tuple(fitting.exclude_types),) + fitting_state = ( + _int_tuple(fitting.exclude_types), + bool(fitting.needs_vacuum_descriptor()), + ) atomic_state = (_int_tuple(atomic_model.atom_exclude_types),) model_state = ( str(model.bridging_method), @@ -1550,8 +1553,9 @@ def core_compute( # either way. ``comm_dict`` (possibly ``None``) and ``nloc`` are # forwarded unconditionally -- ``forward_with_edges`` ignores ``nloc`` # without ``comm_dict``, and ``extended_coord`` only supplies the device. + fitting_net = self.atomic_model.fitting_net with nvtx_range("SeZM/descriptor"): - descriptor, _ = descriptor_model.forward_with_edges( + descriptor, _, vacuum = descriptor_model.forward_with_edges( extended_coord=coord, extended_atype=descriptor_atype, edge_index=edge_index, @@ -1561,17 +1565,21 @@ def core_compute( spin=spin, comm_dict=comm_dict, nloc=nloc, + vacuum_conditions=self.atomic_model.vacuum_conditions() + if fitting_net.needs_vacuum_descriptor() + else None, ) # === Step 3. Fitting net === # The same fitting forward serves both modes; ``embedding_only`` only asks # it to also return the last hidden activation. with nvtx_range("SeZM/fitting_net"): - fit_ret = self.atomic_model.fitting_net( + fit_ret = fitting_net( descriptor, atype, fparam=fparam, aparam=aparam, + vacuum_descriptor=vacuum, return_atomic_feature=embedding_only, ) @@ -1751,7 +1759,7 @@ def core_compute_dens( # === Step 3. Descriptor forward with force embedding === with nvtx_range("SeZM/descriptor_dens"): - descriptor, latent = descriptor_model.forward_with_edges( + descriptor, latent, vacuum = descriptor_model.forward_with_edges( extended_coord=extended_coord[:, :nloc, :], extended_atype=atype, edge_index=edge_index, @@ -1759,6 +1767,9 @@ def core_compute_dens( edge_mask=edge_mask, force_embedding=force_embedding, charge_spin=charge_spin, + vacuum_conditions=self.atomic_model.vacuum_conditions() + if dens_fitting.needs_vacuum_descriptor() + else None, ) # === Step 4. Dens fitting net === @@ -1770,6 +1781,7 @@ def core_compute_dens( noise_mask=noise_mask, fparam=fparam, aparam=aparam, + vacuum_descriptor=vacuum, return_components=True, ) return torch.cat( @@ -3211,6 +3223,19 @@ def reset_head_for_mode(self, mode: str) -> None: Target mode to reset. """ self.atomic_model.reset_head_for_mode(mode) + self.drop_compiled_graphs(mode) + + def drop_compiled_graphs(self, mode: str) -> None: + """ + Drop the compiled graphs of one head so the next forward retraces. + + Parameters + ---------- + mode + ``"dens"`` for the DeNS head; any other value for the energy head, + whose embedding graph reads the same fitting head and is dropped + together with it. + """ if mode == "dens": self._dens_compiled = False self._dens_pending_compile_t0 = None @@ -3218,13 +3243,21 @@ def reset_head_for_mode(self, mode: str) -> None: else: self._core_compute_pending_compile_t0 = None self._core_compute_pending_compile_key = None - # Drop every compile slot so the next forward retraces against the - # reinitialised fitting head. The embedding graph reads the same - # fitting head, so it is invalidated together with the energy graph. self.compiled_core_compute_cache.clear() object.__setattr__(self, "compiled_embedding", None) object.__setattr__(self, "_embedding_task_buf_order", None) + def fold_vacuum_reference(self) -> None: + """ + Fold the vacuum reference into the energy fitting and drop its compiled graphs. + + A traced graph bakes in whether reference nodes trail the real nodes, + so the energy head retraces after the fold. The DeNS head serves + training alone and is not exported, so it keeps its reference. + """ + self.atomic_model.fold_vacuum_reference() + self.drop_compiled_graphs("ener") + # ========================================================================= # Bridging Helpers # ========================================================================= diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 029f6c1dd0..02efb3fcb3 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -386,6 +386,10 @@ def get_ntypes(self) -> int: """Returns the number of element types.""" return len(self.get_type_map()) + def fold_vacuum_reference(self) -> None: + """Fold the vacuum reference of the backbone fitting into its bias.""" + self.backbone_model.fold_vacuum_reference() + @torch.jit.export def get_rcut(self) -> float: """Get the cut-off radius.""" diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 245903a9c5..a3a2c89d1e 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -138,7 +138,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name", None) return super().deserialize(data) diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index b11bfd7c7f..1af4b3820b 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -51,6 +51,7 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, exclude_types: list[int] = [], + vacuum_ref: bool = False, mixed_types: bool = True, type_map: list[str] | None = None, default_fparam: list | None = None, @@ -78,6 +79,7 @@ def __init__( rcond=rcond, seed=seed, exclude_types=exclude_types, + vacuum_ref=vacuum_ref, trainable=trainable, type_map=type_map, default_fparam=default_fparam, @@ -99,7 +101,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "DOSFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("@class", None) data.pop("var_name", None) data.pop("tot_ener_zero", None) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 2aec091e65..ce55464272 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -81,7 +81,7 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index d4918149c0..c59a336d71 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -400,6 +400,14 @@ class GeneralFitting(Fitting): Remove vacuum contribution before the bias is added. The list assigned each type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same length as `ntypes` signaling if or not removing the vacuum contribution for the atom types in the list. + vacuum_ref : bool + Reference the network output of every atom to the output of the same + network for an isolated atom of the same type under the same frame + parameters, atomic parameters and case embedding. The output of an atom + without neighbors is then exactly ``bias_atom_e`` of its type. The + forward takes the vacuum descriptor of every type until + :meth:`fold_vacuum_reference` has folded the reference into the bias + or stored the table. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool @@ -428,6 +436,7 @@ def __init__( exclude_types: list[int] = [], trainable: bool | list[bool] = True, remove_vaccum_contribution: list[bool] | None = None, + vacuum_ref: bool = False, type_map: list[str] | None = None, use_aparam_as_mask: bool = False, default_fparam: list[float] | None = None, @@ -459,6 +468,7 @@ def __init__( all(self.trainable) if isinstance(self.trainable, list) else self.trainable ) self.remove_vaccum_contribution = remove_vaccum_contribution + self.vacuum_ref = vacuum_ref net_dim_out = self._net_out_dim() # init constants @@ -471,6 +481,8 @@ def __init__( if not self.mixed_types: assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" self.register_buffer("bias_atom_e", bias_atom_e) + # A deployment constant; see :meth:`fold_vacuum_reference`. + self.register_buffer("vacuum_table", None, persistent=False) if self.numb_fparam > 0: self.register_buffer( @@ -580,12 +592,14 @@ def change_type_map( ) self.bias_atom_e = torch.cat([self.bias_atom_e, extend_bias_atom_e], dim=0) self.bias_atom_e = self.bias_atom_e[remap_index] + # the stored references belong to the old type map + self.vacuum_table = None def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 4, + "@version": 5, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -601,6 +615,7 @@ def serialize(self) -> dict: "nets": self.filter_layers.serialize(), "rcond": self.rcond, "exclude_types": self.exclude_types, + "vacuum_ref": self.vacuum_ref, "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), "case_embd": to_numpy_array(self.case_embd), @@ -612,7 +627,6 @@ def serialize(self) -> dict: "type_map": self.type_map, # "tot_ener_zero": self.tot_ener_zero , # "trainable": self.trainable , - # "atom_ener": self.atom_ener , # "layer_name": self.layer_name , # "spin": self.spin , ## NOTICE: not supported by far @@ -725,11 +739,234 @@ def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" pass - def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: - return torch.tile(xx.view([1, self.numb_fparam]), [nb, 1]) + def needs_vacuum_descriptor(self) -> bool: + """Whether the forward takes the vacuum descriptor of every type from the descriptor. - def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: - return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) + A referencing fitting takes it until :meth:`fold_vacuum_reference` has + folded the reference into the bias or stored the table. + """ + return self.vacuum_ref and self.vacuum_table is None + + def uniform_conditioning(self) -> bool: + """Whether every atom receives the same conditioning columns. + + Frame parameters vary between frames and atomic parameters between + atoms, while the case embedding is shared by all atoms of a forward. + """ + return self.numb_fparam == 0 and ( + self.numb_aparam == 0 or self.use_aparam_as_mask + ) + + def conditioning_columns( + self, + nf: int, + nloc: int, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + case_embd: torch.Tensor | None, + ) -> torch.Tensor | None: + """Normalized conditioning columns appended to the descriptor of every atom. + + Parameters + ---------- + nf : int + Number of frames. + nloc : int + Number of local atoms per frame. + fparam : torch.Tensor, optional + Frame parameters with shape (nf, numb_fparam). The default frame + parameter is used when omitted. + aparam : torch.Tensor, optional + Atomic parameters with shape (nf, nloc, numb_aparam). + case_embd : torch.Tensor, optional + Case embedding with shape (dim_case_embd,) appended to every atom, + or None when the case embedding is not concatenated. + + Returns + ------- + torch.Tensor or None + The frame parameters, atomic parameters and case embedding of every + atom, normalized and concatenated, with shape (nf, nloc, ncond); + None when the descriptor alone is the fitting input. + """ + columns: list[torch.Tensor] = [] + if self.numb_fparam > 0: + if fparam is None: + assert self.default_fparam_tensor is not None + fparam = self.default_fparam_tensor.unsqueeze(0).expand(nf, -1) + if fparam.numel() != nf * self.numb_fparam: + raise ValueError( + f"input fparam: cannot reshape {list(fparam.shape)} " + f"into ({nf}, {self.numb_fparam})." + ) + assert self.fparam_avg is not None + assert self.fparam_inv_std is not None + fparam = fparam.to(self.prec).view([nf, 1, self.numb_fparam]) + fparam = (fparam - self.fparam_avg) * self.fparam_inv_std + columns.append(fparam.expand(nf, nloc, self.numb_fparam)) + if self.numb_aparam > 0 and not self.use_aparam_as_mask: + assert aparam is not None, "aparam should not be None" + if aparam.numel() != nf * nloc * self.numb_aparam: + raise ValueError( + f"input aparam: cannot reshape {list(aparam.shape)} " + f"into ({nf}, {nloc}, {self.numb_aparam})." + ) + assert self.aparam_avg is not None + assert self.aparam_inv_std is not None + aparam = aparam.to(self.prec).view([nf, nloc, self.numb_aparam]) + columns.append((aparam - self.aparam_avg) * self.aparam_inv_std) + if case_embd is not None: + columns.append(case_embd.view([1, 1, -1]).expand(nf, nloc, -1)) + if len(columns) == 0: + return None + return torch.cat(columns, dim=-1) + + def vacuum_input( + self, + vacuum_descriptor: torch.Tensor | None, + atype: torch.Tensor, + cond: torch.Tensor | None, + case_embd: torch.Tensor | None, + ) -> torch.Tensor | None: + """Fitting input rows of the vacuum references. + + The vacuum reference of an atom is an isolated atom of its type under + the conditioning of the atom itself. With frame or atomic parameters + the conditioning differs between atoms and the rows follow the atoms, + with shape (nf, nloc, in_dim). Otherwise one row per type suffices, + with shape (ntypes, in_dim), and :meth:`vacuum_output` gathers the + network output by type. + + Parameters + ---------- + vacuum_descriptor : torch.Tensor, optional + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt); None takes the table stored by + :meth:`fold_vacuum_reference`. + atype : torch.Tensor + Atom types with shape (nf, nloc). + cond : torch.Tensor, optional + Conditioning columns of the atoms with shape (nf, nloc, ncond). + case_embd : torch.Tensor, optional + Case embedding with shape (dim_case_embd,) appended to every row, + or None when the case embedding is not concatenated. + + Returns + ------- + torch.Tensor or None + The reference rows, or None when ``vacuum_ref`` is off. + + Raises + ------ + ValueError + If ``vacuum_ref`` is on and the vacuum descriptor is missing or + has the wrong shape. + """ + if not self.vacuum_ref: + return None + if vacuum_descriptor is None: + vacuum_descriptor = self.vacuum_table + if vacuum_descriptor is None: + raise ValueError( + "vacuum_ref requires the vacuum descriptor of every atom type" + ) + if list(vacuum_descriptor.shape) != [self.ntypes, self.dim_descrpt]: + raise ValueError( + f"vacuum descriptor of shape {list(vacuum_descriptor.shape)} " + f"does not match ({self.ntypes}, {self.dim_descrpt})" + ) + x_vac = vacuum_descriptor.to(device=atype.device, dtype=self.prec) + if not self.uniform_conditioning(): + assert cond is not None + return torch.cat([x_vac[atype], cond], dim=-1) + if case_embd is None: + return x_vac + return torch.cat( + [x_vac, case_embd.view([1, -1]).expand(self.ntypes, -1)], dim=-1 + ) + + def vacuum_output( + self, vacuum_property: torch.Tensor, atype: torch.Tensor + ) -> torch.Tensor: + """Network output of the vacuum reference of every atom. + + Parameters + ---------- + vacuum_property : torch.Tensor + Network output on the rows of :meth:`vacuum_input`. + atype : torch.Tensor + Atom types with shape (nf, nloc). + + Returns + ------- + torch.Tensor + The reference output of every atom with shape (nf, nloc, dim_out). + """ + if self.uniform_conditioning(): + return vacuum_property[atype] + return vacuum_property + + def vacuum_property(self, vacuum_descriptor: torch.Tensor) -> torch.Tensor: + """Network output on the vacuum descriptor of every type. + + Defined under uniform conditioning, where the output of a reference + depends on its type alone. + + Parameters + ---------- + vacuum_descriptor : torch.Tensor + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt). + + Returns + ------- + torch.Tensor + The reference output of every type with shape (ntypes, dim_out). + """ + atype = torch.arange( + self.ntypes, dtype=torch.long, device=vacuum_descriptor.device + ) + xx_vac = self.vacuum_input(vacuum_descriptor, atype, None, self.case_embd) + if self.mixed_types: + return self.filter_layers.networks[0](xx_vac) + return torch.cat( + [ + ll(xx_vac[type_i : type_i + 1]) + for type_i, ll in enumerate(self.filter_layers.networks) + ], + dim=0, + ) + + def fold_vacuum_reference(self, vacuum_descriptor: torch.Tensor) -> None: + """Fold the vacuum reference into the fitting so the forward needs no reference atoms. + + Under uniform conditioning the reference output of an atom is a + constant of its type, so subtracting it from ``bias_atom_e`` yields + the same outputs as the referenced forward and the option is switched + off. With frame or atomic parameters the reference output varies + between atoms, so the vacuum descriptor is stored instead and the + forward evaluates the references from the stored table. The table is + a deployment constant: an exported model bakes it, checkpoints leave + it out, and a fitting loaded from a checkpoint takes the reference + from the descriptor again. + + Parameters + ---------- + vacuum_descriptor : torch.Tensor + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt). + """ + if not self.vacuum_ref: + return + with torch.no_grad(): + if self.uniform_conditioning(): + reference = self.vacuum_property(vacuum_descriptor) + self.bias_atom_e = self.bias_atom_e - reference.to( + self.bias_atom_e.dtype + ) + self.vacuum_ref = False + else: + self.vacuum_table = vacuum_descriptor.detach().to(self.prec).clone() def _forward_common( self, @@ -740,99 +977,36 @@ def _forward_common( h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, ) -> dict[str, torch.Tensor]: - # cast the input to internal precsion + # cast the input to internal precision xx = descriptor.to(self.prec) nf, nloc, nd = xx.shape - - if self.numb_fparam > 0 and fparam is None: - # use default fparam - assert self.default_fparam_tensor is not None - fparam = torch.tile(self.default_fparam_tensor.unsqueeze(0), [nf, 1]) - - fparam = fparam.to(self.prec) if fparam is not None else None - aparam = aparam.to(self.prec) if aparam is not None else None - - if self.remove_vaccum_contribution is not None: - # TODO: compute the input for vaccm when remove_vaccum_contribution is set - # Ideally, the input for vacuum should be computed; - # we consider it as always zero for convenience. - # Needs a compute_input_stats for vacuum passed from the - # descriptor. - xx_zeros = torch.zeros_like(xx) - else: - xx_zeros = None - net_dim_out = self._net_out_dim() - if nd != self.dim_descrpt: raise ValueError( f"get an input descriptor of dim {nd}," f"which is not consistent with {self.dim_descrpt}." ) - # check fparam dim, concate to input descriptor - if self.numb_fparam > 0: - assert fparam is not None, "fparam should not be None" - assert self.fparam_avg is not None - assert self.fparam_inv_std is not None - if fparam.numel() != nf * self.numb_fparam: - raise ValueError( - f"input fparam: cannot reshape {list(fparam.shape)} " - f"into ({nf}, {self.numb_fparam})." - ) - fparam = fparam.view([nf, self.numb_fparam]) - nb, _ = fparam.shape - t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) - t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) - fparam = (fparam - t_fparam_avg) * t_fparam_inv_std - fparam = torch.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) - xx = torch.cat( - [xx, fparam], - dim=-1, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, fparam], - dim=-1, - ) - # check aparam dim, concate to input descriptor - if self.numb_aparam > 0 and not self.use_aparam_as_mask: - assert aparam is not None, "aparam should not be None" - assert self.aparam_avg is not None - assert self.aparam_inv_std is not None - if aparam.numel() % (nf * self.numb_aparam) != 0: - raise ValueError( - f"input aparam: cannot reshape {list(aparam.shape)} " - f"into ({nf}, nloc, {self.numb_aparam})." - ) - aparam = aparam.view([nf, -1, self.numb_aparam]) - nb, nloc, _ = aparam.shape - t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) - t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) - aparam = (aparam - t_aparam_avg) * t_aparam_inv_std - xx = torch.cat( - [xx, aparam], - dim=-1, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, aparam], - dim=-1, - ) + net_dim_out = self._net_out_dim() - if self.dim_case_embd > 0: - assert self.case_embd is not None - case_embd = torch.tile(self.case_embd.reshape([1, 1, -1]), [nf, nloc, 1]) - xx = torch.cat( - [xx, case_embd], - dim=-1, - ) + # === Step 1. Assemble the fitting input === + # The conditioning columns are shared by the atoms and by their vacuum + # references, so that the reference of an atom differs from the atom + # in its descriptor only. + cond = self.conditioning_columns(nf, nloc, fparam, aparam, self.case_embd) + # ``remove_vaccum_contribution`` subtracts the network output for a zero + # descriptor under the same conditioning columns. + xx_zeros = ( + None if self.remove_vaccum_contribution is None else torch.zeros_like(xx) + ) + if cond is not None: + xx = torch.cat([xx, cond], dim=-1) if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, case_embd], - dim=-1, - ) + xx_zeros = torch.cat([xx_zeros, cond], dim=-1) + xx_vac = self.vacuum_input(vacuum_descriptor, atype, cond, self.case_embd) + # === Step 2. Evaluate the fitting networks === outs = torch.zeros( (nf, nloc, net_dim_out), dtype=self.prec, @@ -847,7 +1021,11 @@ def _forward_common( 0 ].call_until_last(xx) if xx_zeros is not None: - atom_property -= self.filter_layers.networks[0](xx_zeros) + atom_property = atom_property - self.filter_layers.networks[0](xx_zeros) + if xx_vac is not None: + atom_property = atom_property - self.vacuum_output( + self.filter_layers.networks[0](xx_vac), atype + ) outs = ( outs + atom_property + self.bias_atom_e[atype].to(self.prec) ) # Shape is [nframes, natoms[0], net_dim_out] @@ -885,7 +1063,16 @@ def _forward_common( len(self.remove_vaccum_contribution) > type_i and not self.remove_vaccum_contribution[type_i] ): - atom_property -= ll(xx_zeros) + atom_property = atom_property - ll(xx_zeros) + if xx_vac is not None: + # The mask below keeps the atoms of type ``type_i`` alone, so + # the network of the type runs on its own reference row when + # the references are per type, and on the per-atom rows + # otherwise. + if self.uniform_conditioning(): + atom_property = atom_property - ll(xx_vac[type_i : type_i + 1]) + else: + atom_property = atom_property - ll(xx_vac) atom_property = atom_property + self.bias_atom_e[type_i].to(self.prec) atom_property = torch.where(mask, atom_property, 0.0) outs = ( diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 584915321b..8137772a3a 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -75,6 +75,11 @@ class InvarFitting(GeneralFitting): The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descriptor should be set. + vacuum_ref : bool + Reference the network output of every atom to the output of the same + network for an isolated atom of the same type under the same + conditioning, so that an atom without neighbors contributes exactly + its output bias. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool @@ -103,6 +108,7 @@ def __init__( seed: int | list[int] | None = None, exclude_types: list[int] = [], atom_ener: list[torch.Tensor | None] | None = None, + vacuum_ref: bool = False, type_map: list[str] | None = None, use_aparam_as_mask: bool = False, default_fparam: list[float] | None = None, @@ -110,6 +116,15 @@ def __init__( ) -> None: self.dim_out = dim_out self.atom_ener = atom_ener + if ( + vacuum_ref + and atom_ener is not None + and any(x is not None for x in atom_ener) + ): + raise ValueError( + "atom_ener and vacuum_ref are exclusive; vacuum_ref references every " + "atom to the isolated atom of its type by itself" + ) super().__init__( var_name=var_name, ntypes=ntypes, @@ -129,6 +144,7 @@ def __init__( remove_vaccum_contribution=None if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 else [x is not None for x in atom_ener], + vacuum_ref=vacuum_ref, type_map=type_map, use_aparam_as_mask=use_aparam_as_mask, default_fparam=default_fparam, @@ -149,7 +165,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) return super().deserialize(data) def output_def(self) -> FittingOutputDef: @@ -174,6 +190,7 @@ def forward( h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, ) -> dict[str, torch.Tensor]: """Based on embedding net output, alculate total energy. @@ -181,6 +198,8 @@ def forward( Args: - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. + - vacuum_descriptor: descriptor of an isolated atom of every type with + shape [ntypes, self.dim_descrpt], required by ``vacuum_ref``. - return_atomic_feature: also return the last hidden activation under the ``atomic_feature`` key. @@ -196,6 +215,7 @@ def forward( h2, fparam, aparam, + vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, ) result = {self.var_name: out[self.var_name].to(env.GLOBAL_PT_FLOAT_PRECISION)} diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 3d463a7723..a39564ddf8 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -198,7 +198,7 @@ def change_type_map( def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 5 + data["@version"] = 6 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -209,7 +209,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 5, 1) + check_version_compatibility(data.pop("@version", 1), 6, 1) data.pop("var_name", None) return super().deserialize(data) diff --git a/deepmd/pt/model/task/population.py b/deepmd/pt/model/task/population.py index bc542934da..45cae27b4d 100644 --- a/deepmd/pt/model/task/population.py +++ b/deepmd/pt/model/task/population.py @@ -123,7 +123,7 @@ def output_def(self) -> FittingOutputDef: def deserialize(cls, data: dict) -> "PopulationFittingNet": """Deserialize the fitting from a dict.""" data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) # var_name and dim_out are hardcoded in __init__; remove them so they # don't conflict with the positional arguments passed by super().__init__. data.pop("var_name", None) @@ -137,7 +137,7 @@ def serialize(self) -> dict: **InvarFitting.serialize(self), "type": "population", } - dd["@version"] = 4 + dd["@version"] = 5 return dd # make jit happy with torch 2.0.0 diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index d181acfaa1..deb5d5ee4f 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -144,7 +144,7 @@ def get_distinguish_types(self) -> bool: @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 6, 1) + check_version_compatibility(data.pop("@version", 1), 7, 1) data.setdefault("distinguish_types", False) data.pop("dim_out") data["property_name"] = data.pop("var_name") @@ -161,7 +161,7 @@ def serialize(self) -> dict: "intensive": self.intensive, "distinguish_types": self.distinguish_types, } - dd["@version"] = 6 + dd["@version"] = 7 return dd diff --git a/deepmd/pt/model/task/sezm_ener.py b/deepmd/pt/model/task/sezm_ener.py index 0932ec7086..2358c1942e 100644 --- a/deepmd/pt/model/task/sezm_ener.py +++ b/deepmd/pt/model/task/sezm_ener.py @@ -645,6 +645,7 @@ def _forward_common( h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, ) -> dict[str, torch.Tensor]: """Run the SeZM fitting path with optional case FiLM.""" @@ -657,6 +658,7 @@ def _forward_common( h2, fparam, aparam, + vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, ) return self._forward_case_film( @@ -664,6 +666,7 @@ def _forward_common( atype, fparam, aparam, + vacuum_descriptor=vacuum_descriptor, return_atomic_feature=return_atomic_feature, ) @@ -673,6 +676,7 @@ def _forward_case_film( atype: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, return_atomic_feature: bool = False, ) -> dict[str, torch.Tensor]: """ @@ -688,6 +692,9 @@ def _forward_case_film( Frame parameters with shape (nf, numb_fparam). aparam Atomic parameters with shape (nf, nloc, numb_aparam). + vacuum_descriptor + Descriptor of an isolated atom of every type with shape + (ntypes, dim_descrpt), required by ``vacuum_ref``. return_atomic_feature When True, also return the last hidden activation under the ``atomic_feature`` key. @@ -699,88 +706,65 @@ def _forward_case_film( """ xx = descriptor.to(self.prec) nf, nloc, nd = xx.shape - if self.numb_fparam > 0 and fparam is None: - assert self.default_fparam_tensor is not None - fparam = torch.tile(self.default_fparam_tensor.unsqueeze(0), [nf, 1]) - fparam = fparam.to(self.prec) if fparam is not None else None - aparam = aparam.to(self.prec) if aparam is not None else None - - if self.remove_vaccum_contribution is not None: - xx_zeros = torch.zeros_like(xx) - else: - xx_zeros = None - net_dim_out = self._net_out_dim() - if nd != self.dim_descrpt: raise ValueError( f"get an input descriptor of dim {nd}," f"which is not consistent with {self.dim_descrpt}." ) - - if self.numb_fparam > 0: - assert fparam is not None, "fparam should not be None" - assert self.fparam_avg is not None - assert self.fparam_inv_std is not None - if fparam.numel() != nf * self.numb_fparam: - raise ValueError( - f"input fparam: cannot reshape {list(fparam.shape)} " - f"into ({nf}, {self.numb_fparam})." - ) - fparam = fparam.view([nf, self.numb_fparam]) - nb, _ = fparam.shape - t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) - t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) - fparam = (fparam - t_fparam_avg) * t_fparam_inv_std - fparam = torch.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) - xx = torch.cat([xx, fparam], dim=-1) - if xx_zeros is not None: - xx_zeros = torch.cat([xx_zeros, fparam], dim=-1) - - if self.numb_aparam > 0 and not self.use_aparam_as_mask: - assert aparam is not None, "aparam should not be None" - assert self.aparam_avg is not None - assert self.aparam_inv_std is not None - if aparam.numel() % (nf * self.numb_aparam) != 0: - raise ValueError( - f"input aparam: cannot reshape {list(aparam.shape)} " - f"into ({nf}, nloc, {self.numb_aparam})." - ) - aparam = aparam.view([nf, -1, self.numb_aparam]) - nb, nloc, _ = aparam.shape - t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) - t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) - aparam = (aparam - t_aparam_avg) * t_aparam_inv_std - xx = torch.cat([xx, aparam], dim=-1) + # The case embedding modulates the hidden features through FiLM and is + # not concatenated to the input; the atoms and their vacuum references + # share the remaining conditioning columns. + cond = self.conditioning_columns(nf, nloc, fparam, aparam, None) + # ``remove_vaccum_contribution`` subtracts the network output for a zero + # descriptor under the same conditioning columns. + xx_zeros = ( + None if self.remove_vaccum_contribution is None else torch.zeros_like(xx) + ) + if cond is not None: + xx = torch.cat([xx, cond], dim=-1) if xx_zeros is not None: - xx_zeros = torch.cat([xx_zeros, aparam], dim=-1) + xx_zeros = torch.cat([xx_zeros, cond], dim=-1) + xx_vac = self.vacuum_input(vacuum_descriptor, atype, cond, None) assert self.case_embd is not None - outs = torch.zeros( - (nf, nloc, net_dim_out), - dtype=self.prec, - device=descriptor.device, - ) - results = {} - fitting = self.filter_layers.networks[0] + results = {} atom_property = fitting(xx, self.case_embd) if return_atomic_feature: results["atomic_feature"] = fitting.call_until_last(xx, self.case_embd) if xx_zeros is not None: - atom_property -= fitting(xx_zeros, self.case_embd) - outs = outs + atom_property + self.bias_atom_e[atype].to(self.prec) + atom_property = atom_property - fitting(xx_zeros, self.case_embd) + if xx_vac is not None: + atom_property = atom_property - self.vacuum_output( + fitting(xx_vac, self.case_embd), atype + ) + outs = atom_property + self.bias_atom_e[atype].to(self.prec) mask = self.emask(atype).to(torch.bool) outs = torch.where(mask[:, :, None], outs, 0.0) results.update({self.var_name: outs}) return results + def vacuum_property(self, vacuum_descriptor: torch.Tensor) -> torch.Tensor: + """Network output on the vacuum descriptor of every type. + + With case FiLM the case embedding modulates the hidden features + instead of extending the input rows. + """ + if not self.case_film_embd: + return super().vacuum_property(vacuum_descriptor) + atype = torch.arange( + self.ntypes, dtype=torch.long, device=vacuum_descriptor.device + ) + xx_vac = self.vacuum_input(vacuum_descriptor, atype, None, None) + return self.filter_layers.networks[0](xx_vac, self.case_embd) + @classmethod def deserialize(cls, data: dict) -> GeneralFitting: data = data.copy() variables = data.pop("@variables") nets = data.pop("nets") - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name") data.pop("dim_out") obj = cls(**data) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index ea5dbf1ebe..6252b56c27 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -40,6 +40,10 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.preset_out_bias import ( + make_preset_out_bias, + override_assigned_bias, +) from deepmd.utils.stat_file import ( load_output_stat_full_scan, load_paired_items, @@ -466,31 +470,6 @@ def _reduce_model_prediction( return reduced -def _make_preset_out_bias( - ntypes: int, - ibias: list[np.ndarray | None], -) -> np.ndarray | None: - """Make preset out bias. - - output: - a np array of shape [ntypes, *(odim0, odim1, ...)] is any item is not None - None if all items are None. - """ - if len(ibias) != ntypes: - raise ValueError("the length of preset bias list should be ntypes") - if all(ii is None for ii in ibias): - return None - for refb in ibias: - if refb is not None: - break - refb = np.array(refb) - nbias = [ - np.full_like(refb, np.nan, dtype=np.float64) if ii is None else ii - for ii in ibias - ] - return np.array(nbias) - - def _fill_stat_with_global( atomic_stat: np.ndarray | None, global_stat: np.ndarray, @@ -549,9 +528,12 @@ def compute_output_stats( rcond : float, optional The condition number for the regression of atomic energy. preset_bias : dict[str, list[Optional[np.ndarray]]], optional - Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. + Assigned values of the returned bias, given by key:value pairs. + The value is a list with one element per type: None leaves the type to the + statistics, an np.ndarray of output shape assigns the type. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] + The values live in the frame of the returned bias: absolute biases without + `model_forward`, shifts of the model's stored bias with `model_forward`. The `set_davg_zero` key in the descriptor should be set. model_forward : Callable[..., dict[str, torch.Tensor]], optional The wrapped forward function of atomic model. @@ -593,6 +575,11 @@ def compute_output_stats( ) redu_scanner = None + # Model residuals depend on parameters not recorded in the statistics cache. + # Neither reuse nor persist them as absolute output statistics. + if model_forward is not None: + stat_file_path = None + # try to restore the bias from stat file bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) if ( @@ -700,13 +687,22 @@ def compute_output_stats( else None ) + # assigned bias of every output as a (ntypes, ...) array, NaN where a + # type is left to the statistics + assigned_bias = { + kk: make_preset_out_bias(ntypes, preset_bias[kk]) + if preset_bias is not None and kk in preset_bias + else None + for kk in keys + } + # compute stat bias_atom_g, std_atom_g = _compute_output_stats_global( sampled, ntypes, keys, rcond, - preset_bias, + assigned_bias, global_sampled_idx, stats_distinguish_types, intensive, @@ -719,6 +715,7 @@ def compute_output_stats( keys, atomic_sampled_idx, model_pred_a, + assigned_bias, ) # merge global/atomic bias @@ -762,7 +759,7 @@ def _compute_output_stats_global( ntypes: int, keys: list[str], rcond: float | None = None, - preset_bias: dict[str, list[np.ndarray | None]] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, global_sampled_idx: dict | None = None, stats_distinguish_types: bool = True, intensive: bool = False, @@ -814,15 +811,8 @@ def _compute_output_stats_global( if len(input_natoms[kk]) > 0 } nf = {kk: merged_natoms[kk].shape[0] for kk in keys if kk in merged_natoms} - if preset_bias is not None: - assigned_atom_ener = { - kk: _make_preset_out_bias(ntypes, preset_bias[kk]) - if kk in preset_bias.keys() - else None - for kk in keys - } - else: - assigned_atom_ener = dict.fromkeys(keys) + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) if model_pred is None: stats_input = merged_output @@ -856,7 +846,7 @@ def _compute_output_stats_global( for kk in keys: if scan is not None and kk in scan.stats: bias_atom_e[kk], std_atom_e[kk] = scan.stats[kk].solve( - assigned_bias=assigned_atom_ener[kk], + assigned_bias=assigned_bias[kk], rcond=rcond, type_mask=type_mask, ) @@ -867,7 +857,6 @@ def _compute_output_stats_global( compute_stats_do_not_distinguish_types( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], intensive=intensive, ) ) @@ -875,7 +864,7 @@ def _compute_output_stats_global( bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( stats_input[kk], merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], + assigned_bias=assigned_bias[kk], rcond=rcond, intensive=intensive, ) @@ -926,6 +915,7 @@ def _compute_output_stats_atomic( keys: list[str], atomic_sampled_idx: dict | None = None, model_pred: dict[str, np.ndarray] | None = None, + assigned_bias: dict[str, np.ndarray | None] | None = None, ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: """Compute output statistics from atomic labels.""" # return directly if no atomic samples @@ -933,6 +923,8 @@ def _compute_output_stats_atomic( len(v) == 0 for v in atomic_sampled_idx.values() ): return {}, {} + if assigned_bias is None: + assigned_bias = dict.fromkeys(keys) # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { @@ -1004,6 +996,9 @@ def _compute_output_stats_atomic( nan_padding.fill(np.nan) bias_atom_e[kk] = np.concatenate([bias_atom_e[kk], nan_padding], axis=0) std_atom_e[kk] = np.concatenate([std_atom_e[kk], nan_padding], axis=0) + # the per-type means are independent, so an assigned type is + # overridden exactly + bias_atom_e[kk] = override_assigned_bias(bias_atom_e[kk], assigned_bias[kk]) else: # this key does not have atomic labels, skip it. continue diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 076b9b6a04..42080a9142 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -248,9 +248,9 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, 1. **numpy arrays → torch buffers**: State such as statistics (davg, dstd) that is saved in state_dict and moved with .to(device). An array the owning class lists in ``CONFIG_DERIVED_ARRAYS`` becomes a NON-persistent buffer - instead: being a pure function of the configuration it is rebuilt by - ``__init__``, so adopting a stored copy would let a checkpoint whose - configuration differs override the built value. + instead: it is either rebuilt by ``__init__`` from the configuration or + resolved at export time, so adopting a stored copy would let a + checkpoint override the value the model builds for itself. 2. **None values → clear buffers**: Setting an existing buffer to None. 3. **dpmodel objects → pt_expt modules**: Nested dpmodel objects like AtomExcludeMaskDP or NetworkCollectionDP are converted to their pt_expt diff --git a/deepmd/pt_expt/fitting/ener_fitting.py b/deepmd/pt_expt/fitting/ener_fitting.py index 2a5c3c6992..bd3faf28ac 100644 --- a/deepmd/pt_expt/fitting/ener_fitting.py +++ b/deepmd/pt_expt/fitting/ener_fitting.py @@ -42,6 +42,7 @@ def call_graph( h2: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, + vacuum_descriptor: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Graph-native fitting forward, fused when the backend supports it. @@ -50,7 +51,9 @@ def call_graph( routes through the fused operator of the backend device; anything else keeps the dpmodel reference. Routing resolves against the backend device rather than a traced tensor, because every export traces on CPU - and moves the program afterwards. + and moves the program afterwards. A vacuum reference is a per-type + constant on the eligible configuration and enters the operator through + the bias. """ if ( not self.training @@ -60,7 +63,12 @@ def call_graph( and fused_fitting_available() and fitting_eligible(self) ): - return graph_fitting(self, descriptor, atype) + atom_bias = self.bias_atom_e + if self.vacuum_ref: + atom_bias = atom_bias - self.vacuum_property(vacuum_descriptor).to( + atom_bias.dtype + ) + return graph_fitting(self, descriptor, atype, atom_bias) return EnergyFittingNetDP.call_graph( self, descriptor, @@ -70,4 +78,5 @@ def call_graph( h2=h2, fparam=fparam, aparam=aparam, + vacuum_descriptor=vacuum_descriptor, ) diff --git a/deepmd/pt_expt/kernels/graph_fitting.py b/deepmd/pt_expt/kernels/graph_fitting.py index 7408014d46..33c7209b59 100644 --- a/deepmd/pt_expt/kernels/graph_fitting.py +++ b/deepmd/pt_expt/kernels/graph_fitting.py @@ -401,6 +401,7 @@ def graph_fitting( fit: Any, descriptor: torch.Tensor, atype: torch.Tensor, + atom_bias: torch.Tensor, ) -> dict[str, torch.Tensor]: """Fused energy fitting on the flat node axis. @@ -417,6 +418,8 @@ def graph_fitting( Flat descriptor with shape (N, nd). atype : torch.Tensor Flat node atom types with shape (N,), int64. + atom_bias : torch.Tensor + Per-type energy bias added to the network output, shape (ntypes, 1). Returns ------- @@ -433,7 +436,7 @@ def graph_fitting( arguments.residuals, arguments.head_weight, arguments.head_bias, - fit.bias_atom_e.to(torch.float64).reshape(-1, 1)[:, 0].contiguous(), + atom_bias.to(torch.float64).reshape(-1, 1)[:, 0].contiguous(), arguments.activation, ) return {fit.var_name: e} diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py index c0932ba1c4..999430e4f1 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -2945,7 +2945,7 @@ def _mixing_stack_impl( ) x_local = torch.empty((n_edge, n_focus, row), device=u0.device, dtype=u0.dtype) if _has_no_edges(n_edge): - return x_local, z_all, u0 + return x_local, z_all, u0.clone() m0_config, m1_config, _ = stack_fp32_configs(focus_dim, lmax) m0_bm, m0_bn, m0_bk, m0_warps, m0_stages = m0_config diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 4b8230f2e7..c87ff016b5 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -93,8 +93,7 @@ def get_sezm_model(data: dict) -> BaseModel: this builder rejects the flag. Still unsupported here, each raising ``NotImplementedError``: the - virtual-atom (``deepspin``) spin scheme, ``lora``, and - ``preset_out_bias``. + virtual-atom (``deepspin``) spin scheme and ``lora``. Notes ----- @@ -134,10 +133,6 @@ def get_sezm_model(data: dict) -> BaseModel: "set `training.enable_compile` instead." ) _WARNED_ONCE.add("use_compile") - if data.get("preset_out_bias"): - raise NotImplementedError( - "`preset_out_bias` is not supported for DPA4/SeZM in the pt_expt backend." - ) data.pop("type", None) data.setdefault("descriptor", {}) data.setdefault("fitting_net", {}) @@ -179,6 +174,7 @@ def get_sezm_model(data: dict) -> BaseModel: type_map=data["type_map"], atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, + preset_out_bias=data.get("preset_out_bias"), ) @@ -235,8 +231,8 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: injected into the descriptor config (consumed by the descriptor's equivariant spin embedding). The non-spin backbone is built by the standard builder for the config's model type -- :func:`get_sezm_model` - for the DPA4/SeZM family (keeping its bridging/lora/compile/ - preset_out_bias rejections and ``exclude_types`` consistency check), + for the DPA4/SeZM family (keeping its bridging/lora/compile rejections + and ``exclude_types`` consistency check), else :func:`get_standard_model` -- then re-classed through the registered :class:`NativeSpinEnergyModel`. Eligibility is the atomic model's own ``supports_native_spin()`` capability, not a descriptor-type @@ -296,8 +292,8 @@ def _dpa4_family_child_builder(sub: dict) -> "BaseModel | None": A ``linear_ener`` child of the DPA4/SeZM model type must get exactly the semantics of a standalone ``type: "dpa4"`` model -- the descriptor/fitting type defaults, the exclusion consistency check, and - the loud rejections of unsupported options (``lora``, ``use_compile``, - ``preset_out_bias``) -- instead of the generic component build that + the loud rejections of unsupported options (``lora``, ``use_compile``) + -- instead of the generic component build that would silently ignore them. Returns ``None`` for non-DPA4-family children so the shared builder uses its generic path. diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index fa76e1e4fc..c61a5edca0 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -109,8 +109,16 @@ def _fused_energy_force_graph( fused = getattr(desc, "fused_energy_force_graph", None) if fused is None or fit is None: return None + if fit.vacuum_ref and not fit.uniform_conditioning(): + # The reference varies between atoms while the operators take a + # per-type bias only, so such a model uses the autograd lower. + return None graph, atype, output_mask = am._prepare_graph_inputs(graph, atype) atom_bias = fit.bias_atom_e[:, 0] + am.out_bias[0, :, 0] + if fit.vacuum_ref: + # The fused operators see the real atoms alone; the vacuum reference + # of every type is a constant here and enters through the bias. + atom_bias = atom_bias - fit.vacuum_property(am.vacuum_descriptor())[:, 0] out = fused( fit, graph, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 78df67fa33..dfe335d0d1 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -280,6 +280,11 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: per-node moment input and return additional magnetic outputs. A spin task must therefore never reuse a spin-free task's compiled graph even when the descriptor and fitting parameters are shared. + + The fitting's vacuum-reference state selects static branches in the + atomic model and fitting. Shared weights do not make those branches + interchangeable: tasks without a preset bias disable the reference, + while a folded reference no longer appends isolated atoms to the graph. """ descriptor_id: int = 0 try: @@ -293,14 +298,19 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: pass fitting_id: int = id(model) + vacuum_state = (0, 0) try: fitting = model.get_fitting_net() + vacuum_state = ( + int(fitting.vacuum_ref), + int(fitting.needs_vacuum_descriptor()), + ) for _, child in fitting.named_children(): fitting_id = id(child) break except AttributeError: pass - return (int(model.has_spin()), descriptor_id, fitting_id) + return (int(model.has_spin()), descriptor_id, fitting_id, *vacuum_state) # --------------------------------------------------------------------------- @@ -811,9 +821,9 @@ def _trace_and_compile_graph( nloc_trace += 1 trace_N = trace_nf * nloc_trace - # Shared with the .pt2 export trace (serialization.py) so the two graph - # traces can never desync on the input schema. Training uses the run-time - # float precision and device; optional tensors match the actual call. + # The positional input order is shared with .pt2 export (serialization.py). + # Training uses the runtime precision, device, and optional-input presence; + # its graph builders omit the CSR metadata carried by deployment inputs. from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, check_graph_trace_torch_version, @@ -854,6 +864,7 @@ def _trace_and_compile_graph( want_aparam=aparam is not None, want_charge_spin=charge_spin is not None, want_spin=spin is not None, + canonicalize=False, ) ( s_atype, @@ -880,10 +891,10 @@ def fn( edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, - destination_order: torch.Tensor, - destination_row_ptr: torch.Tensor, - source_order: torch.Tensor, - source_row_ptr: torch.Tensor, + destination_order: torch.Tensor | None, + destination_row_ptr: torch.Tensor | None, + source_order: torch.Tensor | None, + source_row_ptr: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 43998596b5..aeb11d798a 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -494,19 +494,20 @@ def build_synthetic_graph_inputs( want_aparam: bool = True, want_charge_spin: bool = True, want_spin: bool = False, + canonicalize: bool = True, ) -> tuple[torch.Tensor | None, ...]: """Build a synthetic carry-all ``NeighborGraph`` for graph-lower tracing. Single source of the trace-time graph inputs, shared by ``.pt2`` export (:func:`_trace_and_export`) and compiled training - (:func:`deepmd.pt_expt.train.training._trace_and_compile_graph`), so the two - traces can never desync on the graph input schema. Builds a small random - system, runs the carry-all + (:func:`deepmd.pt_expt.train.training._trace_and_compile_graph`). Builds a + small random system and runs the carry-all :func:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph` with a - padded ``GraphLayout(edge_capacity=e_max)`` trace sample, then canonicalizes - it to the destination-major deployment ABI. The exported edge axis remains - dynamic; the concrete capacity only supplies representative tensors to - ``make_fx``. Inputs follow the positional order expected by + padded ``GraphLayout(edge_capacity=e_max)`` trace sample. Export uses the + destination-major graph with CSR metadata; compiled training selects the + plain graph without CSR to match its runtime builders. The edge axis + remains dynamic; the concrete capacity only supplies representative tensors + to ``make_fx``. Inputs follow the positional order expected by ``forward_(common_)lower_graph``: ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, fparam, aparam, @@ -519,8 +520,7 @@ def build_synthetic_graph_inputs( (native spin rejects ``add_chg_spin_ebd`` at build). The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = - i % ntypes``) is identical for both callers; the only two former differences - are now parameters. + i % ntypes``) is identical for both callers. Parameters ---------- @@ -559,6 +559,11 @@ def build_synthetic_graph_inputs( all-zero spin leaf can hit degenerate branches, e.g. a ``norm(spin) == 0`` special case, in the equivariant spin embedding). + canonicalize : bool, optional + Sort edges by destination and include destination/source CSR metadata, + as required by the deployment ABI. If ``False``, preserve the builder's + edge order and leave all four CSR metadata inputs as ``None``, matching + the compiled-training graph contract. """ import deepmd.pt_expt.utils.env as _env from deepmd.dpmodel.utils.neighbor_graph import ( @@ -595,7 +600,7 @@ def build_synthetic_graph_inputs( box_t, rcut, layout=GraphLayout(edge_capacity=e_max), - canonicalize=True, + canonicalize=canonicalize, ) fparam = ( @@ -1827,8 +1832,13 @@ def _trace_and_export_impl( # Registry-dispatched (incl. native spin, type "native_spin"): the # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) - model.to("cpu") model.eval() + # The vacuum reference is resolved on the export object, folded into the + # fitting bias or stored as a per-type table, so the exported graph + # carries no reference atoms; the move to the tracing device follows, so + # the table lands there with the rest of the model. + model.fold_vacuum_reference() + model.to("cpu") # Device-dependent Python branches resolve on the CPU tracing inputs, so # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 990c0d41d5..d175c94179 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -429,7 +429,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "dipole", - "@version": 4, + "@version": 5, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, @@ -472,6 +472,7 @@ def serialize(self, suffix: str) -> dict: "rcond": None, "tot_ener_zero": False, "trainable": self.trainable, + "vacuum_ref": False, "layer_name": None, "use_aparam_as_mask": False, "spin": None, @@ -495,7 +496,12 @@ def deserialize(cls, data: dict, suffix: str) -> "DipoleFittingSeA": The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + version = data.pop("@version", 1) + check_version_compatibility(version, 5, 1) + if version >= 5 and data.pop("vacuum_ref"): + raise NotImplementedError( + "vacuum_ref is not supported by the TensorFlow backend" + ) exclude_types = data.pop("exclude_types", []) if len(exclude_types) > 0: data["sel_type"] = [ diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 3d99bcb64d..8a6d75d34b 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -699,7 +699,12 @@ def deserialize(cls, data: dict, suffix: str = "") -> "DOSFitting": The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + version = data.pop("@version", 1) + check_version_compatibility(version, 5, 1) + if version >= 5 and data.pop("vacuum_ref"): + raise NotImplementedError( + "vacuum_ref is not supported by the TensorFlow backend" + ) data["numb_dos"] = data.pop("dim_out") fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( @@ -726,7 +731,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "dos", - "@version": 4, + "@version": 5, "var_name": "dos", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -740,6 +745,7 @@ def serialize(self, suffix: str = "") -> dict: "default_fparam": self.default_fparam, "rcond": self.rcond, "trainable": self.trainable, + "vacuum_ref": False, "activation_function": self.activation_function, "precision": self.fitting_precision.name, "exclude_types": [], diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index a787e5249c..43faef5fb7 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -902,7 +902,12 @@ def deserialize(cls, data: dict, suffix: str = "") -> "EnerFitting": The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + version = data.pop("@version", 1) + check_version_compatibility(version, 5, 1) + if version >= 5 and data.pop("vacuum_ref"): + raise NotImplementedError( + "vacuum_ref is not supported by the TensorFlow backend" + ) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], @@ -928,7 +933,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "ener", - "@version": 4, + "@version": 5, "var_name": "energy", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt + self.tebd_dim, @@ -943,6 +948,7 @@ def serialize(self, suffix: str = "") -> dict: "rcond": self.rcond, "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, + "vacuum_ref": False, "atom_ener": self.atom_ener_v, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 3143f89c8b..220d87ce97 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -653,7 +653,7 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "polar", - "@version": 5, + "@version": 6, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, @@ -698,6 +698,7 @@ def serialize(self, suffix: str) -> dict: "rcond": None, "tot_ener_zero": False, "trainable": self.trainable, + "vacuum_ref": False, "layer_name": None, "use_aparam_as_mask": False, "spin": None, @@ -719,9 +720,12 @@ def deserialize(cls, data: dict, suffix: str) -> "PolarFittingSeA": The deserialized model """ data = data.copy() - check_version_compatibility( - data.pop("@version", 1), 5, 1 - ) # to allow PT version. + version = data.pop("@version", 1) + check_version_compatibility(version, 6, 1) # to allow PT version. + if version >= 6 and data.pop("vacuum_ref"): + raise NotImplementedError( + "vacuum_ref is not supported by the TensorFlow backend" + ) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index d8c3beb138..1cbfca040f 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -105,6 +105,14 @@ def supported_backends(*backends: str) -> str: doc_se_e2_r = "Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor." doc_se_e3 = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor." doc_se_a_tpe = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor." +doc_vacuum_ref = ( + "Reference the fitting network output of every atom to the output the same network gives " + "an isolated atom of the same type under the same frame parameters, atomic parameters and " + "case embedding, so that the energy of an atom without neighbors is exactly its output bias. " + "The reference applies to an output whose bias `preset_out_bias` fixes, which makes the preset value " + "the isolated-atom energy of every element; an output whose bias is fitted from the data keeps the " + "plain network output. It cannot be combined with `atom_ener`." +) doc_se_atten = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism will be used by this descriptor." doc_se_atten_v2 = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism with new modifications will be used by this descriptor." doc_se_a_mask = "Used by the smooth edition of Deep Potential. It can accept a variable number of atoms in a frame (Non-PBC system). *aparam* are required as an indicator matrix for the real/virtual sign of input atoms." @@ -2854,6 +2862,13 @@ def fitting_ener() -> list[Argument]: default=[], doc=doc_atom_ener, ), + Argument( + "vacuum_ref", + bool, + optional=True, + default=False, + doc=supported_backends("pt_expt") + doc_vacuum_ref, + ), Argument("layer_name", list[str], optional=True, doc=doc_layer_name), Argument( "use_aparam_as_mask", @@ -2950,6 +2965,13 @@ def fitting_sezm_ener() -> list[Argument]: default=[], doc=doc_atom_ener, ), + Argument( + "vacuum_ref", + bool, + optional=True, + default=False, + doc=supported_backends("pt", "pt_expt") + doc_vacuum_ref, + ), Argument("layer_name", list[str], optional=True, doc=doc_layer_name), Argument( "use_aparam_as_mask", @@ -3417,7 +3439,21 @@ def model_args( doc_spin = "The settings for systems with spin." doc_atom_exclude_types = "Exclude the atomic contribution of the listed atom types" doc_pair_exclude_types = "The atom pairs of the listed types are not treated to be neighbors, i.e. they do not see each other." - doc_preset_out_bias = "The preset bias of the atomic output. Note that the set_davg_zero should be set to true. The bias is provided as a dict. Taking the energy model that has three atom types for example, the `preset_out_bias` may be given as `{ 'energy': [null, 0., 1.] }`. In this case the energy bias of type 1 and 2 are set to 0. and 1., respectively. A dipole model with two atom types may set `preset_out_bias` as `{ 'dipole': [null, [0., 1., 2.]] }`" + doc_preset_out_bias = ( + "Fix the atomic output bias of chosen elements instead of fitting it from the data, keyed by output name. " + "For an energy model it is the energy of an isolated atom, so with `vacuum_ref` in the fitting net an atom " + "without neighbors gives exactly this energy. Four forms are accepted: a dict keyed by element symbol, " + "e.g. `{'energy': {'O': -430.1, 'H': -13.6}}`; the name of a bundled table of isolated-atom energies, e.g. " + "`{'energy': 'omat24'}`, one of `omat24`, `omol25`, `omc25` and `odac25`, which takes precedence " + "over a file of the same name; " + "the path of a JSON file holding such a dict, e.g. `{'energy': 'e0.json'}`, relative to the working directory; " + "or a list with one entry per type of the `type_map` (`null` leaves a type unassigned), which is also the form " + "for tensor outputs, e.g. `{'dipole': [null, [0., 1., 2.]]}`. " + "Elements outside the `type_map` are ignored and every element that occurs in the data must be assigned. " + "A table is resolved once when the input is processed and its values are stored in the model. An assigned " + "output is taken from the preset without statistics, both when a model is initialized and when its bias is " + "changed by fine-tuning or `dp change-bias`. Dipole models apply no output bias and take no preset." + ) doc_finetune_head = ( "The chosen fitting net to fine-tune on, when doing multi-task fine-tuning. " "If not set or set to 'RANDOM', the fitting net will be randomly initialized." @@ -3491,10 +3527,11 @@ def model_args( ), Argument( "preset_out_bias", - dict[str, list[float | list[float] | None]], + dict[str, list[float | list | None] | dict[str, float | list] | str], optional=True, default=None, - doc=supported_backends("pt", "pd") + doc_preset_out_bias, + doc=supported_backends("pt", "pd", "pt_expt", "jax") + + doc_preset_out_bias, ), Argument( "srtab_add_bias", diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index 590313186d..927a78ce69 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -21,6 +21,9 @@ from deepmd.utils.model_preset import ( expand_model_preset, ) +from deepmd.utils.preset_out_bias import ( + resolve_preset_out_bias_tables, +) def convert_input_v0_v1( @@ -461,7 +464,7 @@ def is_deepmd_v1_input(jdata: dict[str, Any]) -> bool: jdata = migrate_training_warmup(jdata, warning=warning) jdata = convert_optimizer_v31_to_v32(jdata, warning=warning) - jdata["model"] = expand_model_preset(jdata["model"]) + jdata["model"] = resolve_preset_out_bias_tables(expand_model_preset(jdata["model"])) return jdata diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index 29006f2580..e61acb8400 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -144,14 +144,14 @@ def compute_stats_from_atomic( def compute_stats_do_not_distinguish_types( output_redu: np.ndarray, natoms: np.ndarray, - assigned_bias: np.ndarray | None = None, intensive: bool = False, ) -> tuple[np.ndarray, np.ndarray]: """Compute element-independent statistics for property fitting. Computes mean and standard deviation of the output, treating all elements equally. For extensive properties, the output is normalized by the total number of atoms - before computing statistics. + before computing statistics. The statistics do not resolve atom types, so a + per-type preset bias cannot be assigned here. Parameters ---------- @@ -160,10 +160,6 @@ def compute_stats_do_not_distinguish_types( natoms The number of atoms for each atom, shape is [nframes, ntypes]. Used for normalization of extensive properties and generating uniform bias. - assigned_bias - The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)]. - Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias - of the type is not assigned. intensive Whether the output is intensive or extensive. If False, the output will be normalized by the total number of atoms before computing statistics. diff --git a/deepmd/utils/preset_out_bias.py b/deepmd/utils/preset_out_bias.py new file mode 100644 index 0000000000..48a53f644a --- /dev/null +++ b/deepmd/utils/preset_out_bias.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Preset output bias of atomic models. + +The ``preset_out_bias`` model option assigns the output bias of selected atom +types, typically their energy in vacuum. An assigned output is fixed by the +preset alone: every type that occurs in the training data must be assigned, +the assigned types take the preset value, and no statistics are computed for +the output. A string entry names a bundled table of isolated-atom energies or +the JSON file holding a table. The helpers of this module turn the configured +form into the canonical per-type form stored on the atomic model, carry it +through type-map changes, and expand it into the per-type rows written to the +model. + +The canonical form holds nested Python lists rather than arrays: it is written +to the serialized model as it is, and no array data lives on the atomic model +outside its registered variables, which the module wrappers of some backends +require. +""" + +import json +from collections.abc import ( + Iterable, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np + +BUNDLED_TABLES_FILE = Path(__file__).with_name("preset_out_bias_tables.json") + + +def bundled_preset_out_bias_tables() -> dict[str, dict[str, Any]]: + """Preset tables shipped with the package, keyed by table name. + + The database file holds one entry per table name with the ``source`` of + the table and its ``values`` keyed by element symbol; the bundled entries + are tables of isolated-atom energies in eV. + + Returns + ------- + dict[str, dict[str, Any]] + For every table name, its values keyed by element symbol. + """ + with open(BUNDLED_TABLES_FILE) as stream: + database = json.load(stream) + return {name: entry["values"] for name, entry in database.items()} + + +def load_preset_out_bias_table(spec: str) -> dict[str, Any]: + """Look up a preset table by its bundled name or read it from a JSON file. + + Parameters + ---------- + spec + The name of a bundled table, or the path of a JSON file keyed by + element name, relative to the working directory or absolute. A + bundled name takes precedence. + + Returns + ------- + dict[str, Any] + The table, mapping element names to preset values. + + Raises + ------ + ValueError + If ``spec`` is neither a bundled table nor an existing file, or if the + file does not hold a JSON object. + """ + tables = bundled_preset_out_bias_tables() + if spec in tables: + return tables[spec] + if not Path(spec).is_file(): + raise ValueError( + f"the preset_out_bias table {spec!r} is neither one of the bundled " + f"tables {sorted(tables)} nor an existing JSON file" + ) + with open(spec) as stream: + table = json.load(stream) + if not isinstance(table, dict): + raise ValueError( + f"the preset_out_bias table {spec!r} must be a JSON object keyed by " + f"element name, got {type(table).__name__}" + ) + return table + + +def resolve_preset_out_bias_tables(model_config: dict[str, Any]) -> dict[str, Any]: + """Replace the string entries of ``preset_out_bias`` by their tables. + + Parameters + ---------- + model_config + A model configuration, or a multi-task configuration whose branches + live under ``model_dict``; a preset next to ``model_dict`` is resolved + as well. + + Returns + ------- + dict[str, Any] + The configuration with every string entry of ``preset_out_bias`` + replaced by its bundled table or the content of its JSON file, so that + the configuration is self-contained. + """ + preset = model_config.get("preset_out_bias") + if preset and any(isinstance(spec, str) for spec in preset.values()): + resolved = { + key: load_preset_out_bias_table(spec) if isinstance(spec, str) else spec + for key, spec in preset.items() + } + model_config = {**model_config, "preset_out_bias": resolved} + if "model_dict" in model_config: + branches = { + name: resolve_preset_out_bias_tables(branch) + for name, branch in model_config["model_dict"].items() + } + model_config = {**model_config, "model_dict": branches} + return model_config + + +def normalize_preset_out_bias( + preset_out_bias: dict[str, Any] | None, + type_map: list[str], +) -> dict[str, list[list | None]] | None: + """Normalize the ``preset_out_bias`` model option into per-type entries. + + Parameters + ---------- + preset_out_bias + The preset bias of the atomic outputs, keyed by output name. Every value + is a sequence with one entry per type of ``type_map``, where None + leaves the type unassigned; a dict keyed by element name, whose + elements outside ``type_map`` are ignored; or a string naming a bundled + table or the path of a JSON file holding such a dict. An entry is a + number, or a nested list or array of the output shape. The normalized + form is accepted as well, so the function is idempotent. + type_map + Element names of the model, indexed by type. + + Returns + ------- + dict[str, list[list | None]] or None + For every output, one entry per type of ``type_map``: None for a type + without preset, otherwise the preset as a nested list of floats with + at least one dimension. None if no preset is given. + + Raises + ------ + ValueError + If a value is neither a sequence with one entry per type, a dict, a + bundled table name nor a file path, or if an entry is not entirely + finite and numeric. + """ + if preset_out_bias is None: + return None + normalized = {} + for key, spec in preset_out_bias.items(): + if isinstance(spec, str): + spec = load_preset_out_bias_table(spec) + if isinstance(spec, dict): + entries = [spec.get(name) for name in type_map] + elif isinstance(spec, (list, tuple, np.ndarray)) and len(spec) == len(type_map): + entries = list(spec) + else: + raise ValueError( + f"preset_out_bias['{key}'] must be a list with one entry per type " + f"of the type_map ({len(type_map)}), a dict keyed by element name, " + f"a bundled table name or the path of a JSON file, got {spec!r}" + ) + values = [] + for entry in entries: + if entry is None: + values.append(None) + continue + try: + value = np.atleast_1d(np.array(entry, dtype=np.float64)) + except (TypeError, ValueError) as err: + raise ValueError( + f"unsupported value {entry!r} in preset_out_bias['{key}']: " + "expected a number or a nested list of numbers" + ) from err + if not np.isfinite(value).all(): + raise ValueError( + f"preset_out_bias['{key}'] entry {entry!r} must be assigned " + "completely with finite values; a type is either preset " + "or left unassigned" + ) + values.append(value.tolist()) + normalized[key] = values + return normalized + + +def remap_preset_out_bias( + preset_out_bias: dict[str, list[list | None]] | None, + remap_index: list[int], +) -> dict[str, list[list | None]] | None: + """Reorder the per-type entries of a normalized preset bias onto a new type map. + + Parameters + ---------- + preset_out_bias + Normalized preset bias on the old type map. + remap_index + For every new type, the index of the type in the old type map, or a + negative index for a type absent from the old type map, as returned by + :func:`deepmd.utils.finetune.get_index_between_two_maps`. + + Returns + ------- + dict[str, list[list | None]] or None + The preset bias on the new type map; a new type is unassigned. + """ + if preset_out_bias is None: + return None + padding: list[list | None] = [None] * len(remap_index) + return { + key: [(entries + padding)[ii] for ii in remap_index] + for key, entries in preset_out_bias.items() + } + + +def check_preset_out_bias( + preset_out_bias: dict[str, list[list | None]] | None, + keys: list[str], + distinguish_types: bool | None = None, +) -> None: + """Validate a normalized preset bias against the outputs of a model. + + Parameters + ---------- + preset_out_bias + Normalized preset bias, or None. + keys + Names of the model outputs that carry a bias. + distinguish_types + Whether the output statistics of the model resolve atom types, or None + while this is not yet known. A preset assigns the bias of individual + types and therefore requires type-resolved statistics. + + Raises + ------ + ValueError + If the preset names an output the model does not produce, or assigns + a type while the statistics do not distinguish types. + """ + if preset_out_bias is None: + return + unknown = sorted(set(preset_out_bias) - set(keys)) + if unknown: + raise ValueError( + f"preset_out_bias names the outputs {unknown} which the model does " + f"not produce; its outputs are {keys}" + ) + assigned = any( + entry is not None for entries in preset_out_bias.values() for entry in entries + ) + if assigned and distinguish_types is False: + raise ValueError( + "preset_out_bias assigns the bias of individual atom types, but the " + "output statistics of this model do not distinguish atom types" + ) + + +def preset_assigns( + preset_out_bias: dict[str, list[list | None]] | None, key: str +) -> bool: + """Whether the normalized preset fixes the bias of the output ``key`` for at least one type. + + Parameters + ---------- + preset_out_bias + Normalized preset bias, or None. + key + Output name. + + Returns + ------- + bool + True if at least one type of the output has a preset entry. + """ + entries = (preset_out_bias or {}).get(key) + return entries is not None and any(entry is not None for entry in entries) + + +def make_preset_out_bias( + ntypes: int, + ibias: list[list | np.ndarray | None] | np.ndarray, +) -> np.ndarray | None: + """Assemble the preset bias of one output into a per-type array. + + Parameters + ---------- + ntypes + The number of atom types. + ibias + One entry per type: None for a type without preset, otherwise a nested + list or array of the output shape. An already assembled array of shape + (ntypes, ...) whose unassigned rows hold NaN is returned unchanged. + + Returns + ------- + np.ndarray or None + Array of shape (ntypes, *(odim0, odim1, ...)) with NaN for the types + without preset, or None if no type is assigned. + + Raises + ------ + ValueError + If ``ibias`` does not have one entry per type. + """ + if len(ibias) != ntypes: + raise ValueError("the length of preset bias list should be ntypes") + if all(ii is None for ii in ibias): + return None + for refb in ibias: + if refb is not None: + break + refb = np.array(refb) + nbias = [ + np.full_like(refb, np.nan, dtype=np.float64) if ii is None else ii + for ii in ibias + ] + return np.array(nbias) + + +def preset_out_bias_rows( + preset_out_bias: dict[str, list[list | None]], + type_map: list[str], + observed_types: list[str], + stored_bias: np.ndarray, + keys: list[str], + sizes: list[int], + keep_unassigned: bool, + excluded_types: Iterable[int] = (), +) -> dict[str, np.ndarray]: + """Expand the preset into the per-type bias rows of every assigned output. + + Parameters + ---------- + preset_out_bias + Normalized preset bias. + type_map + Element names of the model, indexed by type. + observed_types + Element names that occur in the data whose bias is being set; names + outside ``type_map`` are ignored. + stored_bias + Stored output bias with shape (n_out, ntypes, max_size); the output + ``keys[i]`` occupies ``stored_bias[i, :, :sizes[i]]``. + keys + Output names in the order of the first axis of ``stored_bias``. + sizes + Flattened size of every output. + keep_unassigned + Whether a type without preset keeps its stored bias; otherwise its + bias is zero. + excluded_types + Types whose atomic contribution is excluded from the outputs, such as + the virtual types of a spin model; they need no preset. + + Returns + ------- + dict[str, np.ndarray] + For every output with at least one assigned type, the rows with shape + (ntypes, size). + + Raises + ------ + ValueError + If a type that occurs in the data has no preset for an assigned + output. + """ + ntypes = len(type_map) + index = {name: ii for ii, name in enumerate(type_map)} + excluded = set(excluded_types) + required = [ + index[name] + for name in observed_types + if name in index and index[name] not in excluded + ] + rows = {} + for idx, (key, size) in enumerate(zip(keys, sizes, strict=True)): + preset = make_preset_out_bias(ntypes, preset_out_bias.get(key, [None] * ntypes)) + if preset is None: + continue + preset = preset.reshape(ntypes, size) + unassigned = np.isnan(preset).any(axis=1) + missing = [type_map[ii] for ii in required if unassigned[ii]] + if missing: + raise ValueError( + f"preset_out_bias['{key}'] does not assign the elements {missing} " + "that occur in the data; an assigned output needs a preset for " + "every element in the data" + ) + fill = stored_bias[idx, :, :size] if keep_unassigned else np.zeros_like(preset) + rows[key] = np.where(unassigned[:, None], fill, preset) + return rows + + +def override_assigned_bias( + bias: np.ndarray, + assigned_bias: np.ndarray | None, +) -> np.ndarray: + """Replace the rows of the assigned types by their assigned values. + + Parameters + ---------- + bias + Computed bias with shape (ntypes, ...). + assigned_bias + Assigned bias with the same number of elements, where the rows of the + unassigned types hold NaN, or None. + + Returns + ------- + np.ndarray + The bias with the assigned rows overridden, in the shape of ``bias``. + """ + if assigned_bias is None: + return bias + ntypes = bias.shape[0] + assigned = np.asarray(assigned_bias).reshape(ntypes, -1) + mask = ~np.isnan(assigned).any(axis=1) + result = bias.reshape(ntypes, -1).copy() + result[mask] = assigned[mask] + return result.reshape(bias.shape) diff --git a/deepmd/utils/preset_out_bias_tables.json b/deepmd/utils/preset_out_bias_tables.json new file mode 100644 index 0000000000..8b2eb6b43e --- /dev/null +++ b/deepmd/utils/preset_out_bias_tables.json @@ -0,0 +1,382 @@ +{ + "odac25": { + "source": "isolated-atom reference energies of the UMA odac task (ODAC25), in eV; fairchem configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml (MIT license), odac_elem_refs", + "values": { + "H": -1.11737936, + "He": -0.00011835, + "Li": -0.2941727, + "Be": -0.03868426, + "B": -0.34862832, + "C": -1.31552566, + "N": -3.12457285, + "O": -1.6052078, + "F": -0.49653389, + "Ne": -0.01137327, + "Na": -0.21957281, + "Mg": -0.0008343, + "Al": -0.2750172, + "Si": -0.88417265, + "P": -1.887378, + "S": -0.94903558, + "Cl": -0.31628167, + "Ar": -0.02014536, + "K": -0.15901053, + "Ca": -0.00731884, + "Sc": -1.96521355, + "Ti": -1.89045209, + "V": -2.53057428, + "Cr": -5.43600675, + "Mn": -5.09739336, + "Fe": -3.03088746, + "Co": -1.23786562, + "Ni": -0.40650749, + "Cu": -0.2416017, + "Zn": -0.01139188, + "Ga": -0.26282496, + "Ge": -0.82446455, + "As": -1.70237206, + "Se": -0.84245376, + "Br": -0.28544892, + "Kr": -0.02239991, + "Rb": -0.14115912, + "Sr": -0.02840799, + "Y": -2.09540994, + "Zr": -1.85863996, + "Nb": -1.12257399, + "Mo": -4.32965355, + "Tc": -3.30670045, + "Ru": -1.19460755, + "Rh": -1.26257601, + "Pd": -1.46832888, + "Ag": -0.19779414, + "Cd": -0.0144274, + "In": -0.23668767, + "Sn": -0.70836953, + "Sb": -1.43186113, + "Te": -0.71701186, + "I": -0.24883129, + "Xe": -0.01118184, + "Cs": -0.13173447, + "Ba": -0.0318395, + "La": -0.41195547, + "Ce": -1.23134873, + "Pr": -2.03082996, + "Nd": 0.1375954, + "Pm": -5.45866275, + "Sm": -7.59139905, + "Eu": -5.99965965, + "Gd": -8.43495767, + "Tb": -2.6578407, + "Dy": -7.77349787, + "Ho": -5.30762201, + "Er": -5.15109657, + "Tm": -4.41466995, + "Yb": -0.02995219, + "Lu": -0.2544495, + "Hf": -3.23821202, + "Ta": -3.45887214, + "W": -4.53635003, + "Re": -4.60979468, + "Os": -2.90707964, + "Ir": -1.28286153, + "Pt": -0.57716664, + "Au": -0.18337108, + "Hg": -0.01135944, + "Tl": -0.22045398, + "Pb": -0.66150479, + "Bi": -1.32506342, + "Po": -0.66500178, + "At": -0.22643927, + "Rn": -0.00728197, + "Fr": -0.11208472, + "Ra": -0.00757856, + "Ac": -0.21798637, + "Th": -0.91078787, + "Pa": -1.78187161, + "U": -3.89912261, + "Np": -3.94192659, + "Pu": -7.59026042 + } + }, + "omat24": { + "source": "isolated-atom reference energies of the UMA omat task (OMat24), in eV; fairchem configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml (MIT license), omat_elem_refs", + "values": { + "H": -1.11700253, + "He": 0.00079886, + "Li": -0.29731164, + "Be": -0.04129868, + "B": -0.29106192, + "C": -1.27751531, + "N": -3.12342715, + "O": -1.54797136, + "F": -0.43969356, + "Ne": -0.01250908, + "Na": -0.22855413, + "Mg": -0.00943179, + "Al": -0.21707638, + "Si": -0.82619133, + "P": -1.88667434, + "S": -0.89093583, + "Cl": -0.25816211, + "Ar": -0.02414768, + "K": -0.17662425, + "Ca": -0.02568319, + "Sc": -2.13001165, + "Ti": -2.38688845, + "V": -3.55934233, + "Cr": -5.44700879, + "Mn": -5.14749562, + "Fe": -3.30662847, + "Co": -1.42167737, + "Ni": -0.63181379, + "Cu": -0.23449167, + "Zn": -0.01146636, + "Ga": -0.21291259, + "Ge": -0.77939897, + "As": -1.70148487, + "Se": -0.78386705, + "Br": -0.22690657, + "Kr": -0.02245409, + "Rb": -0.16092396, + "Sr": -0.02798717, + "Y": -2.25685695, + "Zr": -2.23690495, + "Nb": -2.15347771, + "Mo": -4.60251809, + "Tc": -3.36416792, + "Ru": -2.23062607, + "Rh": -1.15550917, + "Pd": -1.47553527, + "Ag": -0.19918102, + "Cd": -0.01475888, + "In": -0.19767692, + "Sn": -0.68005773, + "Sb": -1.43073368, + "Te": -0.65790462, + "I": -0.18915279, + "Xe": -0.01179476, + "Cs": -0.13507902, + "Ba": -0.03056979, + "La": -0.36017439, + "Ce": -0.86279246, + "Pr": -0.20573327, + "Nd": -0.2734463, + "Pm": -0.20046965, + "Sm": -0.25444338, + "Eu": -8.37972664, + "Gd": -9.58424928, + "Tb": -0.19466184, + "Dy": -0.24860115, + "Ho": -0.19531288, + "Er": -0.15401392, + "Tm": -0.14577898, + "Yb": -0.19655747, + "Lu": -0.15645898, + "Hf": -3.49380556, + "Ta": -3.5317097, + "W": -4.57108006, + "Re": -4.63425205, + "Os": -2.88247063, + "Ir": -1.45679675, + "Pt": -0.50290184, + "Au": -0.18521704, + "Hg": -0.01123956, + "Tl": -0.17483649, + "Pb": -0.63132037, + "Bi": -1.3248562, + "Ac": -0.24135757, + "Th": -1.04601971, + "Pa": -2.04574044, + "U": -3.84544799, + "Np": -7.28626119, + "Pu": -7.3136314 + } + }, + "omc25": { + "source": "isolated-atom reference energies of the UMA omc task (OMC25), in eV; fairchem configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml (MIT license), omc_elem_refs", + "values": { + "H": -0.02831808, + "He": 4.512e-05, + "Li": -0.03227157, + "Be": -0.03842519, + "B": -0.05829283, + "C": -0.0845041, + "N": -0.08806738, + "O": -0.09021346, + "F": -0.06669846, + "Ne": -0.01218631, + "Na": -0.03650269, + "Mg": -0.00059093, + "Al": -0.05787736, + "Si": -0.08730952, + "P": -0.0975534, + "S": -0.09264199, + "Cl": -0.07124762, + "Ar": -0.02374602, + "K": -0.05299112, + "Ca": -0.02631476, + "Sc": -1.7772147, + "Ti": -1.25083444, + "V": -0.79579447, + "Cr": -0.49099317, + "Mn": -0.31414986, + "Fe": -0.20292182, + "Co": -0.14011632, + "Ni": -0.09929659, + "Cu": -0.03771207, + "Zn": -0.01117902, + "Ga": -0.06168715, + "Ge": -0.08873364, + "As": -0.09512942, + "Se": -0.09035978, + "Br": -0.06910849, + "Kr": -0.02244872, + "Rb": -0.05303651, + "Sr": -0.02871903, + "Y": -1.94805417, + "Zr": -1.33379896, + "Nb": -0.69169331, + "Mo": -0.26184306, + "Tc": -0.20631599, + "Ru": -0.48251608, + "Rh": -0.96911893, + "Pd": -1.47569462, + "Ag": -0.03845194, + "Cd": -0.0142445, + "In": -0.07118991, + "Sn": -0.09940292, + "Sb": -0.09235056, + "Te": -0.08755943, + "I": -0.06544925, + "Xe": -0.01246646, + "Cs": -0.04692937, + "Ba": -0.03225123, + "La": -0.26086039, + "Ce": -27.20024339, + "Pr": -0.08412926, + "Nd": -0.08225924, + "Pm": -0.07799715, + "Sm": -0.07806185, + "Eu": 0.00043759, + "Gd": -0.07459766, + "Dy": -0.06842841, + "Ho": -0.07758266, + "Er": -0.07025152, + "Tm": -0.08055003, + "Yb": -0.07118177, + "Lu": -0.07159568, + "Hf": -2.69202862, + "Ta": -2.21926765, + "W": -1.679756, + "Re": -1.06135075, + "Os": -0.4554231, + "Ir": -0.14488432, + "Pt": -0.18377098, + "Au": -0.03603118, + "Hg": -0.01076585, + "Tl": -0.06381411, + "Pb": -0.0905623, + "Bi": -0.10095787, + "Po": -0.09501217, + "At": -0.0574478, + "Rn": -0.00599173, + "Fr": -0.04134751, + "Ra": -0.0082683, + "Ac": -0.08704692, + "Th": -0.49656425, + "Pa": -5.24233138, + "U": -2.32542606, + "Np": -4.3376616, + "Pu": -5.96430676, + "Bk": -0.03842519 + } + }, + "omol25": { + "source": "isolated-atom reference energies of the UMA omol task (OMol25), neutral atoms, in eV; fairchem configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml (MIT license), omol_elem_refs", + "values": { + "H": -13.445584057373376, + "He": -78.82028064861346, + "Li": -203.32566089610248, + "Be": -398.94745406238457, + "B": -670.7528161616635, + "C": -1029.854128936652, + "N": -1485.5420268052571, + "O": -2042.9785186843171, + "F": -2714.2404072902364, + "Ne": -3508.7435104433293, + "Na": -4415.24245175699, + "Mg": -5443.89764689278, + "Al": -6594.618979783701, + "Si": -7873.6885552583635, + "P": -9285.660195840774, + "S": -10832.622363522705, + "Cl": -12520.669728186529, + "Ar": -14354.279383593506, + "K": -16323.548287099005, + "Ca": -18436.4802235745, + "Sc": -20696.184432436938, + "Ti": -23110.540829573696, + "V": -25682.996766371485, + "Cr": -28418.38077740936, + "Mn": -31317.9261898231, + "Fe": -34383.42850178591, + "Co": -37623.47198155286, + "Ni": -41039.92808633622, + "Cu": -44637.390644693216, + "Zn": -48417.15330794875, + "Ga": -52373.88353905847, + "Ge": -56512.7749670792, + "As": -60836.15457735998, + "Se": -65344.29463090149, + "Br": -70041.24926492578, + "Kr": -74929.57000138272, + "Rb": -653.6478304509429, + "Sr": -833.3193042712877, + "Y": -1038.0281969578277, + "Zr": -1273.9680039726502, + "Nb": -1542.4549613021854, + "Mo": -1850.7417570878108, + "Tc": -2193.9167529178317, + "Ru": -2577.187592352399, + "Rh": -3004.13632479042, + "Pd": -3477.528290952186, + "Ag": -3997.318637706032, + "Cd": -4563.758476480422, + "In": -5171.823431250868, + "Sn": -5828.85389738702, + "Sb": -6535.615918982279, + "Te": -7291.548625930308, + "I": -8099.879925553251, + "Xe": -8962.18002388331, + "Cs": -546.0321954061023, + "Ba": -690.6089668328505, + "La": -854.1124475883443, + "Ce": -12923.042210466992, + "Pr": -14064.262597328401, + "Nd": -15272.688366565537, + "Pm": -16550.207110957523, + "Sm": -17900.366871587325, + "Eu": -19323.23592276476, + "Gd": -20829.090486783083, + "Tb": -22428.734746941605, + "Dy": -24078.6824036755, + "Ho": -25794.42345644991, + "Er": -27616.684560203954, + "Tm": -29523.55544372604, + "Yb": -31526.683164730683, + "Lu": -33615.3810306281, + "Hf": -1300.1780354638831, + "Ta": -1544.4093860120772, + "W": -1818.6231567944906, + "Re": -2123.144371335989, + "Os": -2461.7605162840764, + "Ir": -2833.7631469210337, + "Pt": -3242.7992608176864, + "Au": -3690.363355421551, + "Hg": -4174.9981270766675, + "Tl": -4691.757194525804, + "Pb": -5245.360635715531, + "Bi": -5838.120614817622 + } + } +} diff --git a/deepmd/utils/vacuum_reference.py b/deepmd/utils/vacuum_reference.py new file mode 100644 index 0000000000..2647c613c6 --- /dev/null +++ b/deepmd/utils/vacuum_reference.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Conditions of the isolated neutral ground-state atom of every type. + +The isolated-atom energy reference (``vacuum_ref``) pins the energy of an +atom without neighbors to its preset bias. The reference atom is the neutral +atom in its ground state: zero charge, the ground-state spin multiplicity as +the charge/spin condition of the descriptor, and a spin vector whose magnitude +is the number of unpaired electrons for the native-spin descriptors. Every +quantity derives from the spin-resolved orbital occupation table of the +elements. +""" + +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.econf_embd import ( + electronic_configuration_embedding, +) +from deepmd.utils.preset_out_bias import ( + preset_assigns, +) + + +def unpaired_electrons(type_map: list[str]) -> np.ndarray: + """Number of unpaired electrons of the neutral ground-state atom of every type. + + The occupation table encodes every orbital as a pair of entries: ``[1, 1]`` + doubly occupied, ``[-1, 1]`` singly occupied and ``[-1, -1]`` empty. The + singly occupied orbitals are counted. + + Parameters + ---------- + type_map : list[str] + Element symbol of every type. + + Returns + ------- + np.ndarray + The unpaired-electron count of every type with shape (ntypes,). + + Raises + ------ + ValueError + If a type name is not an element symbol. + """ + unknown = [ + name for name in type_map if name not in electronic_configuration_embedding + ] + if unknown: + raise ValueError( + "the isolated-atom reference requires element symbols as type names; " + f"unknown names: {unknown}" + ) + counts = [] + for name in type_map: + occupation = electronic_configuration_embedding[name].reshape(-1, 2) + counts.append(np.sum((occupation[:, 0] == -1) & (occupation[:, 1] == 1))) + return np.array(counts, dtype=np.int64) + + +def reference_charge_spin(type_map: list[str]) -> np.ndarray: + """Charge and spin condition of the isolated neutral ground-state atom of every type. + + Parameters + ---------- + type_map : list[str] + Element symbol of every type. + + Returns + ------- + np.ndarray + ``[charge, multiplicity]`` of every type with shape (ntypes, 2): the + charge is zero and the multiplicity is the unpaired-electron count + plus one. + """ + unpaired = unpaired_electrons(type_map) + return np.stack([np.zeros_like(unpaired), unpaired + 1], axis=-1).astype(np.float64) + + +def reference_spin(type_map: list[str]) -> np.ndarray: + """Spin vector of the isolated neutral ground-state atom of every type. + + Parameters + ---------- + type_map : list[str] + Element symbol of every type. + + Returns + ------- + np.ndarray + Spin vectors with shape (ntypes, 3) in Bohr magnetons, of magnitude + equal to the unpaired-electron count and directed along ``z``; the + energy of an isolated atom does not depend on the direction. + """ + spin = np.zeros((len(type_map), 3), dtype=np.float64) + spin[:, 2] = unpaired_electrons(type_map) + return spin + + +def resolve_vacuum_ref( + fitting: Any, preset_out_bias: dict[str, list[list | None]] | None +) -> None: + """Keep the vacuum reference of a fitting only on an output whose bias the preset fixes. + + The reference makes an atom without neighbors contribute exactly its bias, + which is the isolated-atom value only when ``preset_out_bias`` fixes it. An + output whose bias is fitted from the data keeps the plain network output: + in multi-task training a branch without a table then differs from the + referenced branches that share its network by a constant per type, which + the case conditioning represents, instead of by an offset that would have + to vanish on the isolated atoms. + + Parameters + ---------- + fitting + A fitting network; its ``vacuum_ref`` option is switched off in place + when the preset does not fix the bias of its output ``var_name``. + preset_out_bias + Normalized preset bias of the atomic model, or None. + """ + if fitting.vacuum_ref and not preset_assigns(preset_out_bias, fitting.var_name): + fitting.vacuum_ref = False diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index c3fb7fe54d..575fd273e9 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -645,7 +645,11 @@ dp --pt freeze -c model.ckpt.pt -o frozen_model ``` The PyTorch backend detects DPA4/SeZM and writes `frozen_model.pt2`. The -pt_expt backend uses the same kernel-level policy for a DPA4/SeZM `.pt2`. +pt_expt backend uses the same kernel-level policy for a DPA4/SeZM `.pt2`. A +fitting with the [isolated-atom energy reference](train-energy.md#isolated-atom-energy-reference) +has its reference resolved at this point: folded into the fitting bias, or +stored as a per-type table when frame or atomic parameters make it vary +between atoms. Unless the environment says otherwise, a CUDA archive is built at `DP_TRITON_INFER=2` and `DP_CUDA_INFER=1`, the fastest all-float32 combination; set either variable to override, for instance `DP_CUDA_INFER=2` on a part with diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index 8a68885e9e..f448017ae6 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -106,6 +106,110 @@ The construction of the fitting net is given by section {ref}`fitting_net ` is set to `true`, then a timestep is used in the ResNet. - {ref}`seed ` gives the random seed that is used to generate random numbers when initializing the model parameters. +## Isolated-atom energy reference + +An atom without neighbors should contribute a known energy, for example the +isolated-atom energy of the reference calculation, so that the dissociation +limit of the model is pinned. Two model options provide this together: + +- `preset_out_bias` fixes the energy bias of every element to a given value + instead of fitting it from the data. The value may name a bundled table of + isolated-atom energies, be given as a dictionary keyed by element symbol, be + read from a JSON file holding such a dictionary, or be listed per type of + the `type_map`. Every element that occurs in the training data must be + assigned; elements of a table that are not in the `type_map` are ignored. + The elements that occur in the data are collected from the batches sampled + for the statistics, or from every frame with `data_stat_full`. An assigned + output takes no statistics: elements absent from the data get a zero bias + when the model is initialized and keep their bias when it is fine-tuned, + and the output std keeps its stored value. +- `vacuum_ref` in the fitting network references the network output of every + atom to the output the same network gives an isolated atom of the same type, + under the same frame parameters, atomic parameters and case embedding. With + it, the energy of an atom without neighbors is exactly its bias, whatever the + network parameters are, so the preset value is the isolated-atom energy. The + reference applies to an output whose bias `preset_out_bias` fixes; an output + whose bias is fitted from the data keeps the plain network output. It cannot + be combined with the fitting option `atom_ener`. + +```json +{ + "model": { + "type_map": [ + "O", + "H" + ], + "preset_out_bias": { + "energy": "omat24" + }, + "fitting_net": { + "vacuum_ref": true + } + } +} +``` + +`omat24` is one of the bundled tables of isolated-atom energies in eV: + +| Name | Reference calculation | Elements | +| -------- | --------------------- | ----------------------------------------- | +| `omat24` | OMat24 | 89, H to Pu without Po, At, Rn, Fr and Ra | +| `omol25` | OMol25, neutral atoms | 83, H to Bi | +| `omc25` | OMC25 | 94, H to Bk without Tb, Am and Cm | +| `odac25` | ODAC25 | 94, H to Pu | + +These four tables are the isolated-atom reference energies of the UMA training +tasks published with fairchem +(`configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml`, MIT +license). Each table is on the energy scale of its own reference calculation +(the OMat24 table gives H = -1.117 eV, the OMol25 table H = -13.446 eV), so a +model takes the table of the calculation that produced its training data. A +table of another reference calculation is given as a JSON file +mapping element symbols to energies in eV, for instance `"energy": "e0.json"` +with the file holding `{"O": -432.0, "H": -13.6}`; the path is relative to the +working directory. In every form the table is resolved once when the input is +processed and its values are stored in the model, so a trained model depends +neither on the file nor on the bundled data. + +In multi-task training `preset_out_bias` is given in each branch, so a branch +trained on another reference calculation names its own table; written once +next to `model_dict`, it applies to every branch that does not set its own. A +branch without a table takes its bias from the statistics and keeps the plain +network output even when it shares a fitting network with `vacuum_ref`: the +shared network then represents the same binding energy in every branch, the +branches differ by a constant per element that the case embedding carries, and +the isolated-atom energies of the branches with a table stay exactly their +presets. +`examples/water/dpa4/input_e0.json` and +`examples/water/dpa4/input_multitask_e0.json` show the single-task and the +multi-task setup on the water example, whose reference calculation is not one +of the bundled tables: the first reads the table from `e0.json` next to the +input, the second gives every branch an explicit dictionary. + +The reference atom is the neutral atom in its ground state. With charge/spin +conditioning (`add_chg_spin_ebd`) the reference carries zero charge and the +ground-state spin multiplicity of the element, and with +[native spin](dpa4.md#native-scheme) it carries a spin vector of one Bohr +magneton per unpaired electron; charged, excited or differently magnetized +isolated atoms keep the deviation the network learns for them. The type names +must then be element symbols. + +The reference follows the network parameters at every training step. Freezing +evaluates the vacuum descriptor of every type once and removes the reference +atoms from the exported model: a fitting without frame or atomic parameters +folds the reference into its bias and runs at the speed of a model without the +option, and a fitting with such parameters, whose reference output varies +between atoms, stores the table and evaluates its references from it. The +fused inference operators of the pt_expt backend take a per-type reference +through the bias; a fitting whose reference varies between atoms is served by +the autograd route instead. + +`vacuum_ref` is available for DPA4/SeZM on the PyTorch backend and for the +graph-native models (DPA1, DPA2, DPA4 and DPA4C) on the pt_expt backend, whose +flat node axis carries the reference atoms. Fine-tuning and `dp change-bias` +keep the reference; a `preset_out_bias` given to them fixes the bias again +without statistics. + ## Loss The loss function $L$ for training energy is given by diff --git a/examples/water/dpa4/e0.json b/examples/water/dpa4/e0.json new file mode 100644 index 0000000000..d16e092265 --- /dev/null +++ b/examples/water/dpa4/e0.json @@ -0,0 +1,4 @@ +{ + "O": -432.0, + "H": -13.6 +} diff --git a/examples/water/dpa4/input_e0.json b/examples/water/dpa4/input_e0.json new file mode 100644 index 0000000000..82469e2510 --- /dev/null +++ b/examples/water/dpa4/input_e0.json @@ -0,0 +1,123 @@ +{ + "_comment": "DPA4-Mini energy-training example for the water dataset with the isolated-atom energy reference: preset_out_bias reads the energy bias of O and H from the JSON table e0.json (representative isolated-atom energies on the energy scale of the example data; a production run uses the values of its own reference calculation or a bundled table such as omat24), and vacuum_ref pins the energy of an atom without neighbors to that bias.", + "model": { + "type": "dpa4", + "type_map": [ + "O", + "H" + ], + "preset_out_bias": { + "energy": "e0.json" + }, + "descriptor": { + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 + }, + "fitting_net": { + "neuron": [ + 0 + ], + "precision": "float32", + "seed": 42, + "vacuum_ref": true + }, + "use_compile": false, + "enable_tf32": true + }, + "learning_rate": { + "type": "wsd", + "start_lr": 0.00045, + "stop_lr": 1e-06, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.65, + "decay_type": "cosine" + }, + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4_e0.hdf5", + "stat_file_mode": "update", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1 + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_batch": 1 + }, + "numb_steps": 2000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "save_dir": "ckpt", + "max_ckpt_keep": 3, + "enable_ema": true, + "ema_decay": 0.999, + "ema_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_avg": true, + "disp_training": true, + "time_training": true, + "tensorboard": false, + "enable_profiler": false, + "tensorboard_freq": 1000, + "tensorboard_log_dir": "tb_log", + "profiling": false, + "profiling_file": "timeline.json", + "zero_stage": 1, + "seed": 42 + }, + "validating": { + "compiled_infer": false, + "tf32_infer": false, + "amp_infer": false, + "save_best_dir": "ckpt_best" + } +} diff --git a/examples/water/dpa4/input_multitask_e0.json b/examples/water/dpa4/input_multitask_e0.json new file mode 100644 index 0000000000..826e122542 --- /dev/null +++ b/examples/water/dpa4/input_multitask_e0.json @@ -0,0 +1,191 @@ +{ + "_comment": "DPA4-Mini multitask example with a shared descriptor and a case-conditioned shared fitting network with vacuum_ref; preset_out_bias fixes the energy bias of each branch, so a branch trained on another reference calculation names its own table.", + "model": { + "shared_dict": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "dpa4", + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 + }, + "shared_fit_with_id": { + "type": "dpa4_ener", + "neuron": [ + 0 + ], + "precision": "float32", + "dim_case_embd": 2, + "case_film_embd": true, + "seed": 42, + "vacuum_ref": true + } + }, + "model_dict": { + "water_1": { + "use_compile": false, + "enable_tf32": true, + "type": "dpa4", + "type_map": "type_map", + "descriptor": "descriptor", + "fitting_net": "shared_fit_with_id", + "preset_out_bias": { + "energy": { + "O": -432.0, + "H": -13.6 + } + }, + "model_branch_alias": [ + "Default", + "Water" + ], + "info": { + "description": "Water branch with shared DPA4/SeZM descriptor and case-embedded shared fitting net" + } + }, + "water_2": { + "use_compile": false, + "enable_tf32": true, + "type": "dpa4", + "type_map": "type_map", + "descriptor": "descriptor", + "fitting_net": "shared_fit_with_id", + "preset_out_bias": { + "energy": { + "O": -432.0, + "H": -13.6 + } + }, + "model_branch_alias": [ + "Water2" + ], + "info": { + "description": "Second water branch with shared DPA4/SeZM descriptor and case-embedded shared fitting net" + } + } + } + }, + "learning_rate": { + "type": "wsd", + "start_lr": 0.00045, + "stop_lr": 1e-06, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.65, + "decay_type": "cosine" + }, + "loss_dict": { + "water_1": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "water_2": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + } + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "enable_tf32": true, + "model_prob": { + "water_1": 0.5, + "water_2": 0.5 + }, + "data_dict": { + "water_1": { + "stat_file": "./dpa4_e0_water_1.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1 + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_batch": 1 + } + }, + "water_2": { + "stat_file": "./dpa4_e0_water_2.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1 + } + } + }, + "numb_steps": 2000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "enable_ema": true, + "ema_decay": 0.999, + "ema_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_avg": true, + "disp_training": true, + "time_training": true, + "tensorboard": false, + "enable_profiler": false, + "tensorboard_freq": 1000, + "tensorboard_log_dir": "tb_log", + "profiling": false, + "profiling_file": "timeline.json", + "zero_stage": 1, + "seed": 42 + } +} diff --git a/source/tests/common/dpmodel/test_atomic_model_global_stat.py b/source/tests/common/dpmodel/test_atomic_model_global_stat.py index dbca577524..0059374e0a 100644 --- a/source/tests/common/dpmodel/test_atomic_model_global_stat.py +++ b/source/tests/common/dpmodel/test_atomic_model_global_stat.py @@ -1,631 +1,645 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import tempfile -import unittest -from pathlib import ( - Path, -) -from typing import ( - NoReturn, -) - -import h5py -import numpy as np - -from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel -from deepmd.dpmodel.common import ( - NativeOP, -) -from deepmd.dpmodel.descriptor import ( - DescrptSeA, -) -from deepmd.dpmodel.fitting import ( - InvarFitting, -) -from deepmd.dpmodel.fitting.base_fitting import ( - BaseFitting, -) -from deepmd.dpmodel.output_def import ( - FittingOutputDef, - OutputVariableDef, -) -from deepmd.utils.path import ( - DPPath, -) - -from .case_single_frame_with_nlist import ( - TestCaseSingleFrameWithNlist, -) - - -class FooFitting(NativeOP, BaseFitting): - """Test fitting with multiple outputs for testing global statistics.""" - - def __init__(self): - pass - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "foo", - [1], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - OutputVariableDef( - "pix", - [1], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - OutputVariableDef( - "bar", - [1, 2], - reducible=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - - def serialize(self) -> dict: - return { - "@class": "Fitting", - "type": "foo", - "@version": 1, - } - - @classmethod - def deserialize(cls, data: dict): - return cls() - - def get_dim_fparam(self) -> int: - return 0 - - def get_dim_aparam(self) -> int: - return 0 - - def get_sel_type(self) -> list[int]: - return [] - - def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None - ) -> None: - pass - - def get_type_map(self) -> list[str]: - return [] - - def call( - self, - descriptor, - atype, - gr=None, - g2=None, - h2=None, - fparam=None, - aparam=None, - ): - nf, nloc, _ = descriptor.shape - ret = {} - ret["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *self.output_def()["foo"].shape]) - ret["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *self.output_def()["pix"].shape]) - ret["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *self.output_def()["bar"].shape]) - return ret - - -class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.merged_output_stat = [ - { - "coord": np.zeros([2, 3, 3]), - "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - "box": np.zeros([2, 3, 3]), - "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - # bias of foo: 1, 3 - "foo": np.array([5.0, 7.0]).reshape(2, 1), - # no bias of pix - # bias of bar: [1, 5], [3, 2] - "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - "find_foo": np.float32(1.0), - "find_bar": np.float32(1.0), - } - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_output_stat(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = FooFitting() - type_map = ["foo", "bar"] - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - # nf x nloc - at = self.atype_ext[:, :nloc] - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - expected_std = np.array( - [[[0, 1], [0, 1]], [[1, 1], [1, 1]], [[0, 0], [0, 0]]] - ) # 3 keys, 2 atypes, 2 max dims. - # nt x odim - foo_bias = np.array([1.0, 3.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - np.testing.assert_almost_equal(md0.out_std, expected_std) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - np.testing.assert_almost_equal(md0.out_std, expected_std) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - self.coord_ext, - np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), - self.nlist, - ] - ret3 = md0.forward_common_atomic(*args) - ## model output on foo: [[2, 3, 6], [5, 8, 9]] given bias [1, 3] - ## foo sumed: [11, 22] compared with [5, 7], fit target is [-6, -15] - ## fit bias is [1, -8] - ## old bias + fit bias [2, -5] - ## new model output is [[3, 4, -2], [6, 0, 1]], which sumed to [5, 7] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. - np.testing.assert_almost_equal(md0.out_std, expected_std) - - def test_preset_bias(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = FooFitting() - type_map = ["foo", "bar"] - preset_out_bias = { - "foo": [None, 2], - "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), - } - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - preset_out_bias=preset_out_bias, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - # nf x nloc - at = self.atype_ext[:, :nloc] - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - # foo sums: [5, 7], - # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] - # the solution of type 0 is 1.8 - foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) - bar_bias = preset_out_bias["bar"] - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - - # 4. test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - self.coord_ext, - np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), - self.nlist, - ] - ret3 = md0.forward_common_atomic(*args) - ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] - ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] - ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) - ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] - ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( - 2, 3, 1 - ) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. - - def test_preset_bias_all_none(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = FooFitting() - type_map = ["foo", "bar"] - preset_out_bias = { - "foo": [None, None], - } - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - preset_out_bias=preset_out_bias, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - # nf x nloc - at = self.atype_ext[:, :nloc] - - # 1. test run without bias - # nf x na x odim - ret0 = md0.forward_common_atomic(*args) - expected_ret0 = {} - expected_ret0["foo"] = np.array( - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) - expected_ret0["pix"] = np.array( - [ - [3.0, 2.0, 1.0], - [6.0, 5.0, 4.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) - expected_ret0["bar"] = np.array( - [ - [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], - [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], - ] - ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied (all None preset = same as no preset) - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - # nt x odim - foo_bias = np.array([1.0, 3.0]).reshape(2, 1) - bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) - expected_ret1 = {} - expected_ret1["foo"] = ret0["foo"] + foo_bias[at] - expected_ret1["pix"] = ret0["pix"] - expected_ret1["bar"] = ret0["bar"] + bar_bias[at] - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - def test_serialize(self) -> None: - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = InvarFitting( - "foo", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - ) - type_map = ["A", "B"] - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret0 = md0.forward_common_atomic(*args) - md1 = DPDPAtomicModel.deserialize(md0.serialize()) - ret1 = md1.forward_common_atomic(*args) - - for kk in ["foo"]: - np.testing.assert_almost_equal(ret0[kk], ret1[kk]) - - -class TestChangeByStatMixedLabels(unittest.TestCase, TestCaseSingleFrameWithNlist): - """Test change-by-statistic with mixed atomic and global labels.""" - - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.merged_output_stat = [ - { - "coord": np.zeros([2, 3, 3]), - "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - "box": np.zeros([2, 3, 3]), - "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - # foo: atomic label - "atom_foo": np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape( - 2, 3, 1 - ), - # pix: global label - "pix": np.array([5.0, 12.0]).reshape(2, 1), - # bar: global label - "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), - "find_atom_foo": np.float32(1.0), - "find_pix": np.float32(1.0), - "find_bar": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_change_by_statistic(self) -> None: - """Test change-by-statistic with atomic foo + global pix + global bar.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = FooFitting() - type_map = ["foo", "bar"] - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - - ret0 = md0.forward_common_atomic(*args) - - # set initial bias - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - - # change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - # use atype_ext from merged_output_stat for inference - args = [ - self.coord_ext, - np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), - self.nlist, - ] - ret3 = md0.forward_common_atomic(*args) - # foo: atomic label, bias after set-by-stat: [5, 6] - # model output with bias [5,6], atype [[0,0,1],[0,1,1]]: - # [[6, 7, 9], [9, 11, 12]] - # atom_foo labels: [[5, 5, 5], [5, 6, 7]] - # per-atom delta: [[-1, -2, -4], [-4, -5, -5]] - # delta bias (mean per type): type0=-7/3, type1=-14/3 - # new bias = [5-7/3, 6-14/3] = [8/3, 4/3] - # new output: [[11/3, 14/3, 13/3], [20/3, 19/3, 22/3]] - expected_ret3 = {} - expected_ret3["foo"] = np.array( - [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] - ).reshape(2, 3, 1) - # pix: global label, bias after set-by-stat: [-2/3, 19/3] - # model pix with bias, atype [[0,0,1],[0,1,1]]: - # [[7/3, 4/3, 22/3], [16/3, 34/3, 31/3]], sums [11, 27] - # labels [5, 12], delta [-6, -15] - # lstsq: delta bias [1, -8], new bias [1/3, -5/3] - # new output: [[10/3, 7/3, -2/3], [19/3, 10/3, 7/3]] - expected_ret3["pix"] = np.array( - [[3.3333, 2.3333, -0.6667], [6.3333, 3.3333, 2.3333]] - ).reshape(2, 3, 1) - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) - # bar is too complicated to be manually computed. - - -class TestEnergyModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - """Test statistics computation with real energy fitting net.""" - - def tearDown(self) -> None: - self.tempdir.cleanup() - - def setUp(self) -> None: - TestCaseSingleFrameWithNlist.setUp(self) - self.merged_output_stat = [ - { - "coord": np.zeros([2, 3, 3]), - "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), - "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), - "box": np.zeros([2, 3, 3]), - "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), - # energy data - "energy": np.array([10.0, 20.0]).reshape(2, 1), - "find_energy": np.float32(1.0), - }, - ] - self.tempdir = tempfile.TemporaryDirectory() - h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) - with h5py.File(h5file, "w") as f: - pass - self.stat_file_path = DPPath(h5file, "a") - - def test_energy_stat(self) -> None: - """Test energy statistics computation with real energy fitting net.""" - nf, nloc, nnei = self.nlist.shape - ds = DescrptSeA( - self.rcut, - self.rcut_smth, - self.sel, - ) - ft = InvarFitting( - "energy", - self.nt, - ds.get_dim_out(), - 1, - mixed_types=ds.mixed_types(), - ) - type_map = ["foo", "bar"] - md0 = DPDPAtomicModel( - ds, - ft, - type_map=type_map, - ) - args = [self.coord_ext, self.atype_ext, self.nlist] - - # test run without bias - ret0 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret0) - - # compute statistics - md0.compute_or_load_out_stat( - self.merged_output_stat, stat_file_path=self.stat_file_path - ) - ret1 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret1) - - # Check that bias was computed (out_bias should be non-zero) - self.assertFalse(np.all(md0.out_bias == 0)) - - # test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - np.testing.assert_allclose( - ret1["energy"], - ret2["energy"], - ) - - # test change bias - md0.change_out_bias( - self.merged_output_stat, bias_adjust_mode="change-by-statistic" - ) - ret3 = md0.forward_common_atomic(*args) - self.assertIn("energy", ret3) - - -if __name__ == "__main__": - unittest.main() +# SPDX-License-Identifier: LGPL-3.0-or-later +import tempfile +import unittest +from pathlib import ( + Path, +) +from typing import ( + NoReturn, +) + +import h5py +import numpy as np + +from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.path import ( + DPPath, +) + +from .case_single_frame_with_nlist import ( + TestCaseSingleFrameWithNlist, +) + + +class FooFitting(NativeOP, BaseFitting): + """Test fitting with multiple outputs for testing global statistics.""" + + def __init__(self): + pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "foo", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "pix", + [1], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + OutputVariableDef( + "bar", + [1, 2], + reducible=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def serialize(self) -> dict: + return { + "@class": "Fitting", + "type": "foo", + "@version": 1, + } + + @classmethod + def deserialize(cls, data: dict): + return cls() + + def get_dim_fparam(self) -> int: + return 0 + + def get_dim_aparam(self) -> int: + return 0 + + def get_sel_type(self) -> list[int]: + return [] + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat=None + ) -> None: + pass + + def get_type_map(self) -> list[str]: + return [] + + def call( + self, + descriptor, + atype, + gr=None, + g2=None, + h2=None, + fparam=None, + aparam=None, + ): + nf, nloc, _ = descriptor.shape + ret = {} + ret["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *self.output_def()["foo"].shape]) + ret["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *self.output_def()["pix"].shape]) + ret["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *self.output_def()["bar"].shape]) + return ret + + +class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # bias of foo: 1, 3 + "foo": np.array([5.0, 7.0]).reshape(2, 1), + # no bias of pix + # bias of bar: [1, 5], [3, 2] + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_foo": np.float32(1.0), + "find_bar": np.float32(1.0), + } + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_output_stat(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + expected_std = np.array( + [[[0, 1], [0, 1]], [[1, 1], [1, 1]], [[0, 0], [0, 0]]] + ) # 3 keys, 2 atypes, 2 max dims. + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + np.testing.assert_almost_equal(md0.out_std, expected_std) + + # 3. test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], ret2[kk]) + np.testing.assert_almost_equal(md0.out_std, expected_std) + + # 4. test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + ## model output on foo: [[2, 3, 6], [5, 8, 9]] given bias [1, 3] + ## foo sumed: [11, 22] compared with [5, 7], fit target is [-6, -15] + ## fit bias is [1, -8] + ## old bias + fit bias [2, -5] + ## new model output is [[3, 4, -2], [6, 0, 1]], which sumed to [5, 7] + expected_ret3 = {} + expected_ret3["foo"] = np.array([[3, 4, -2], [6, 0, 1]]).reshape(2, 3, 1) + expected_ret3["pix"] = ret0["pix"] + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) + # bar is too complicated to be manually computed. + np.testing.assert_almost_equal(md0.out_std, expected_std) + + def test_preset_bias(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + # both types occur in the data, so an assigned output assigns both + preset_out_bias = { + "foo": [3, 2], + "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), + } + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. the preset fixes foo and bar; pix has no label and keeps a zero bias + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + foo_bias = np.array(preset_out_bias["foo"], dtype=np.float64).reshape(2, 1) + bar_bias = preset_out_bias["bar"] + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + # 3. change-by-statistic keeps every assigned type at its preset + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + ret3 = md0.forward_common_atomic(*args) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret3[kk], ret1[kk]) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(out_bias["foo"], foo_bias) + np.testing.assert_almost_equal(out_bias["bar"], bar_bias) + + # 4. a preset that leaves an observed type unassigned is rejected + md1 = DPDPAtomicModel( + ds, + FooFitting(), + type_map=type_map, + preset_out_bias={"foo": [None, 2]}, + ) + with self.assertRaisesRegex(ValueError, "foo"): + md1.compute_or_load_out_stat(self.merged_output_stat) + + def test_preset_bias_all_none(self) -> None: + nf, nloc, nnei = self.nlist.shape + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + preset_out_bias = { + "foo": [None, None], + } + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + preset_out_bias=preset_out_bias, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + # nf x nloc + at = self.atype_ext[:, :nloc] + + # 1. test run without bias + # nf x na x odim + ret0 = md0.forward_common_atomic(*args) + expected_ret0 = {} + expected_ret0["foo"] = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["foo"].shape]) + expected_ret0["pix"] = np.array( + [ + [3.0, 2.0, 1.0], + [6.0, 5.0, 4.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["pix"].shape]) + expected_ret0["bar"] = np.array( + [ + [1.0, 2.0, 3.0, 7.0, 8.0, 9.0], + [4.0, 5.0, 6.0, 10.0, 11.0, 12.0], + ] + ).reshape([nf, nloc, *md0.fitting_output_def()["bar"].shape]) + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) + + # 2. test bias is applied (all None preset = same as no preset) + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + # nt x odim + foo_bias = np.array([1.0, 3.0]).reshape(2, 1) + bar_bias = np.array([1.0, 5.0, 3.0, 2.0]).reshape(2, 1, 2) + expected_ret1 = {} + expected_ret1["foo"] = ret0["foo"] + foo_bias[at] + expected_ret1["pix"] = ret0["pix"] + expected_ret1["bar"] = ret0["bar"] + bar_bias[at] + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + + def test_serialize(self) -> None: + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = InvarFitting( + "foo", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + type_map = ["A", "B"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret0 = md0.forward_common_atomic(*args) + md1 = DPDPAtomicModel.deserialize(md0.serialize()) + ret1 = md1.forward_common_atomic(*args) + + for kk in ["foo"]: + np.testing.assert_almost_equal(ret0[kk], ret1[kk]) + + +class TestChangeByStatMixedLabels(unittest.TestCase, TestCaseSingleFrameWithNlist): + """Test change-by-statistic with mixed atomic and global labels.""" + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # foo: atomic label + "atom_foo": np.array([[5.0, 5.0, 5.0], [5.0, 6.0, 7.0]]).reshape( + 2, 3, 1 + ), + # pix: global label + "pix": np.array([5.0, 12.0]).reshape(2, 1), + # bar: global label + "bar": np.array([5.0, 12.0, 7.0, 9.0]).reshape(2, 1, 2), + "find_atom_foo": np.float32(1.0), + "find_pix": np.float32(1.0), + "find_bar": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_change_by_statistic(self) -> None: + """Test change-by-statistic with atomic foo + global pix + global bar.""" + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = FooFitting() + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + + # set initial bias + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + + # change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + # use atype_ext from merged_output_stat for inference + args = [ + self.coord_ext, + np.array(self.merged_output_stat[0]["atype_ext"], dtype=np.int64), + self.nlist, + ] + ret3 = md0.forward_common_atomic(*args) + # foo: atomic label, bias after set-by-stat: [5, 6] + # model output with bias [5,6], atype [[0,0,1],[0,1,1]]: + # [[6, 7, 9], [9, 11, 12]] + # atom_foo labels: [[5, 5, 5], [5, 6, 7]] + # per-atom delta: [[-1, -2, -4], [-4, -5, -5]] + # delta bias (mean per type): type0=-7/3, type1=-14/3 + # new bias = [5-7/3, 6-14/3] = [8/3, 4/3] + # new output: [[11/3, 14/3, 13/3], [20/3, 19/3, 22/3]] + expected_ret3 = {} + expected_ret3["foo"] = np.array( + [[3.6667, 4.6667, 4.3333], [6.6667, 6.3333, 7.3333]] + ).reshape(2, 3, 1) + # pix: global label, bias after set-by-stat: [-2/3, 19/3] + # model pix with bias, atype [[0,0,1],[0,1,1]]: + # [[7/3, 4/3, 22/3], [16/3, 34/3, 31/3]], sums [11, 27] + # labels [5, 12], delta [-6, -15] + # lstsq: delta bias [1, -8], new bias [1/3, -5/3] + # new output: [[10/3, 7/3, -2/3], [19/3, 10/3, 7/3]] + expected_ret3["pix"] = np.array( + [[3.3333, 2.3333, -0.6667], [6.3333, 3.3333, 2.3333]] + ).reshape(2, 3, 1) + for kk in ["foo", "pix"]: + np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk], decimal=4) + # bar is too complicated to be manually computed. + + def test_preset_with_atomic_labels(self) -> None: + """An assigned output ignores the atomic labels in both modes.""" + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + md0 = DPDPAtomicModel( + ds, + FooFitting(), + type_map=["foo", "bar"], + preset_out_bias={"foo": {"foo": 5.0, "bar": 2.0}}, + ) + # atom_foo labels [[5, 5, 5], [5, 6, 7]] with atype [[0, 0, 1], [0, 1, 1]] + # are never fitted: both types take their preset + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + out_bias, _ = md0._fetch_out_stat(["foo"]) + np.testing.assert_almost_equal(out_bias["foo"], np.array([[5.0], [2.0]])) + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + out_bias, _ = md0._fetch_out_stat(["foo"]) + np.testing.assert_almost_equal(out_bias["foo"], np.array([[5.0], [2.0]])) + # a preset that leaves the observed type foo unassigned is rejected + md1 = DPDPAtomicModel( + ds, + FooFitting(), + type_map=["foo", "bar"], + preset_out_bias={"foo": {"bar": 2.0}}, + ) + with self.assertRaisesRegex(ValueError, "foo"): + md1.compute_or_load_out_stat(self.merged_output_stat) + + +class TestEnergyModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): + """Test statistics computation with real energy fitting net.""" + + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.merged_output_stat = [ + { + "coord": np.zeros([2, 3, 3]), + "atype": np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32), + "atype_ext": np.array([[0, 0, 1, 0], [0, 1, 1, 0]], dtype=np.int32), + "box": np.zeros([2, 3, 3]), + "natoms": np.array([[3, 3, 2, 1], [3, 3, 1, 2]], dtype=np.int32), + # energy data + "energy": np.array([10.0, 20.0]).reshape(2, 1), + "find_energy": np.float32(1.0), + }, + ] + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + + def test_energy_stat(self) -> None: + """Test energy statistics computation with real energy fitting net.""" + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + type_map = ["foo", "bar"] + md0 = DPDPAtomicModel( + ds, + ft, + type_map=type_map, + ) + args = [self.coord_ext, self.atype_ext, self.nlist] + + # test run without bias + ret0 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret0) + + # compute statistics + md0.compute_or_load_out_stat( + self.merged_output_stat, stat_file_path=self.stat_file_path + ) + ret1 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret1) + + # Check that bias was computed (out_bias should be non-zero) + self.assertFalse(np.all(md0.out_bias == 0)) + + # test bias load from file + def raise_error() -> NoReturn: + raise RuntimeError + + md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) + ret2 = md0.forward_common_atomic(*args) + np.testing.assert_allclose( + ret1["energy"], + ret2["energy"], + ) + + # test change bias + md0.change_out_bias( + self.merged_output_stat, bias_adjust_mode="change-by-statistic" + ) + ret3 = md0.forward_common_atomic(*args) + self.assertIn("energy", ret3) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/dpmodel/test_fitting_call_graph.py b/source/tests/common/dpmodel/test_fitting_call_graph.py index 2e143046eb..b2f72ecd54 100644 --- a/source/tests/common/dpmodel/test_fitting_call_graph.py +++ b/source/tests/common/dpmodel/test_fitting_call_graph.py @@ -13,16 +13,28 @@ ) +@pytest.mark.parametrize("vacuum_ref", [False, True]) # vacuum references per type @pytest.mark.parametrize("ndf", [0, 3]) # numb_fparam: no-fparam AND fparam -def test_call_graph_matches_dense_raveled(ndf): +def test_call_graph_matches_dense_raveled(ndf, vacuum_ref): rng = np.random.default_rng(0) nf, nloc, nd, ntypes, ng = 2, 4, 8, 2, 5 - ft = InvarFitting("energy", ntypes, nd, 1, mixed_types=True, numb_fparam=ndf) + ft = InvarFitting( + "energy", + ntypes, + nd, + 1, + mixed_types=True, + numb_fparam=ndf, + vacuum_ref=vacuum_ref, + ) desc = rng.normal(size=(nf, nloc, nd)) atype = rng.integers(0, ntypes, size=(nf, nloc)) gr = rng.normal(size=(nf, nloc, ng, 3)) fparam = rng.normal(size=(nf, ndf)) if ndf else None - dense = ft(desc, atype, gr=gr, fparam=fparam)["energy"] # (nf, nloc, 1) + vacuum = rng.normal(size=(ntypes, nd)) if vacuum_ref else None + dense = ft(desc, atype, gr=gr, fparam=fparam, vacuum_descriptor=vacuum)[ + "energy" + ] # (nf, nloc, 1) N = nf * nloc frame_id = np.repeat(np.arange(nf), nloc) fparam_node = fparam[frame_id] if ndf else None # (N, ndf) @@ -31,6 +43,7 @@ def test_call_graph_matches_dense_raveled(ndf): atype.reshape(N), gr=gr.reshape(N, ng, 3), fparam=fparam_node, + vacuum_descriptor=vacuum, )["energy"] # (N, 1) assert flat.shape == (N, 1) np.testing.assert_allclose(flat, dense.reshape(N, 1), rtol=1e-12, atol=1e-12) diff --git a/source/tests/common/dpmodel/test_fitting_invar_fitting.py b/source/tests/common/dpmodel/test_fitting_invar_fitting.py index 0590feae07..2c63a6dac3 100644 --- a/source/tests/common/dpmodel/test_fitting_invar_fitting.py +++ b/source/tests/common/dpmodel/test_fitting_invar_fitting.py @@ -227,3 +227,154 @@ def test_runtime_buffers_follow_torch_descriptor(self) -> None: self.assertEqual(result.dtype, torch.float64) self.assertEqual(result.device.type, "cpu") np.testing.assert_allclose(result.detach().cpu().numpy(), expected) + + +VACUUM_CONDITIONING = [(0, 0), (2, 0), (0, 1), (2, 1)] + + +class TestVacuumRef(unittest.TestCase): + """``vacuum_ref`` references every atom to the isolated atom of its type.""" + + ntypes, nd, nf, nloc = 3, 8, 2, 5 + + def setUp(self) -> None: + self.rng = np.random.default_rng(GLOBAL_SEED) + self.descriptor = self.rng.normal(size=(self.nf, self.nloc, self.nd)) + self.vacuum = self.rng.normal(size=(self.ntypes, self.nd)) + self.atype = self.rng.integers(0, self.ntypes, size=(self.nf, self.nloc)) + self.atype[0, : self.ntypes] = np.arange(self.ntypes) + self.bias = self.rng.normal(size=(self.ntypes, 1)) + + def build(self, vacuum_ref: bool, **kwargs) -> InvarFitting: + ft = InvarFitting( + "energy", + self.ntypes, + self.nd, + 1, + neuron=[6, 6], + bias_atom=self.bias, + vacuum_ref=vacuum_ref, + seed=GLOBAL_SEED, + **kwargs, + ) + if ft.dim_case_embd > 0: + ft.set_case_embd(1) + return ft + + def params(self, nfp: int, nap: int) -> dict: + return { + "fparam": self.rng.normal(size=(self.nf, nfp)) if nfp else None, + "aparam": self.rng.normal(size=(self.nf, self.nloc, nap)) if nap else None, + } + + def test_isolated_atom_gives_bias(self) -> None: + for mixed_types, (nfp, nap), ncase, mask in itertools.product( + [True, False], VACUUM_CONDITIONING, [0, 2], [False, True] + ): + ft = self.build( + True, + mixed_types=mixed_types, + numb_fparam=nfp, + numb_aparam=nap, + dim_case_embd=ncase, + use_aparam_as_mask=mask, + ) + out = ft( + self.vacuum[self.atype], + self.atype, + vacuum_descriptor=self.vacuum, + **self.params(nfp, nap), + )["energy"] + np.testing.assert_allclose( + out, self.bias[self.atype], rtol=1e-10, atol=1e-10 + ) + + def test_matches_reference_subtraction(self) -> None: + for mixed_types, (nfp, nap), ncase in itertools.product( + [True, False], VACUUM_CONDITIONING, [0, 2] + ): + ft_ref = self.build( + False, + mixed_types=mixed_types, + numb_fparam=nfp, + numb_aparam=nap, + dim_case_embd=ncase, + ) + ft_vac = InvarFitting.deserialize( + {**ft_ref.serialize(), "vacuum_ref": True} + ) + params = self.params(nfp, nap) + out = ft_vac( + self.descriptor, self.atype, vacuum_descriptor=self.vacuum, **params + )["energy"] + expected = ( + ft_ref(self.descriptor, self.atype, **params)["energy"] + - ft_ref(self.vacuum[self.atype], self.atype, **params)["energy"] + + self.bias[self.atype] + ) + np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + def test_vacuum_descriptor_required(self) -> None: + ft = self.build(True) + with self.assertRaises(ValueError): + ft(self.descriptor, self.atype) + with self.assertRaises(ValueError): + ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum[:, :-1]) + + def test_fold_vacuum_reference(self) -> None: + for mixed_types, ncase in itertools.product([True, False], [0, 2]): + ft = self.build(True, mixed_types=mixed_types, dim_case_embd=ncase) + expected = ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum) + ft.fold_vacuum_reference(self.vacuum) + self.assertFalse(ft.vacuum_ref) + np.testing.assert_allclose( + ft(self.descriptor, self.atype)["energy"], + expected["energy"], + rtol=1e-10, + atol=1e-10, + ) + # a conditioned fitting stores the table and references from it; the + # table is a deployment constant that serialization leaves out and a + # type-map change drops + ft = self.build(True, numb_aparam=1, type_map=["O", "H", "B"]) + aparam = self.rng.normal(size=(self.nf, self.nloc, 1)) + expected = ft( + self.descriptor, self.atype, aparam=aparam, vacuum_descriptor=self.vacuum + ) + ft.fold_vacuum_reference(self.vacuum) + self.assertTrue(ft.vacuum_ref) + self.assertFalse(ft.needs_vacuum_descriptor()) + np.testing.assert_allclose( + ft(self.descriptor, self.atype, aparam=aparam)["energy"], + expected["energy"], + rtol=1e-10, + atol=1e-10, + ) + restored = InvarFitting.deserialize(ft.serialize()) + self.assertTrue(restored.needs_vacuum_descriptor()) + np.testing.assert_allclose( + restored( + self.descriptor, + self.atype, + aparam=aparam, + vacuum_descriptor=self.vacuum, + )["energy"], + expected["energy"], + rtol=1e-10, + atol=1e-10, + ) + ft.change_type_map(["B", "O", "H"]) + self.assertTrue(ft.needs_vacuum_descriptor()) + + def test_serialization(self) -> None: + data = self.build(True).serialize() + self.assertTrue(data["vacuum_ref"]) + self.assertTrue(InvarFitting.deserialize(data).vacuum_ref) + # a dictionary of the previous version carries no key + older = {k: v for k, v in data.items() if k != "vacuum_ref"} + self.assertFalse(InvarFitting.deserialize({**older, "@version": 4}).vacuum_ref) + + def test_atom_ener_is_exclusive(self) -> None: + self.assertTrue(self.build(True, atom_ener=[None] * self.ntypes).vacuum_ref) + with self.assertRaises(ValueError): + self.build(True, atom_ener=[1.0] + [None] * (self.ntypes - 1)) diff --git a/source/tests/common/dpmodel/test_model_factory.py b/source/tests/common/dpmodel/test_model_factory.py index 2807f4d932..c4582f0424 100644 --- a/source/tests/common/dpmodel/test_model_factory.py +++ b/source/tests/common/dpmodel/test_model_factory.py @@ -227,6 +227,7 @@ def test_standard_construction_is_shared_and_non_mutating(self) -> None: "fitting_net": {"type": "dipole", "custom": 5}, "atom_exclude_types": [1], "pair_exclude_types": [[0, 1]], + "preset_out_bias": {"dipole": {"H": [0.0, 1.0, 2.0]}}, } expected = { "type_map": ["O", "H"], @@ -234,6 +235,7 @@ def test_standard_construction_is_shared_and_non_mutating(self) -> None: "fitting_net": {"type": "dipole", "custom": 5}, "atom_exclude_types": [1], "pair_exclude_types": [[0, 1]], + "preset_out_bias": {"dipole": {"H": [0.0, 1.0, 2.0]}}, } model = get_standard_model( data, @@ -251,6 +253,9 @@ def test_standard_construction_is_shared_and_non_mutating(self) -> None: self.assertEqual(model.kwargs["fitting"].kwargs["embedding_width"], 7) self.assertEqual(model.kwargs["atom_exclude_types"], [1]) self.assertEqual(model.kwargs["pair_exclude_types"], [[0, 1]]) + self.assertEqual( + model.kwargs["preset_out_bias"], {"dipole": {"H": [0.0, 1.0, 2.0]}} + ) def test_model_level_type_embedding_is_rejected(self) -> None: """Cover the shared validation used by every dpmodel-driven backend.""" @@ -299,6 +304,7 @@ def test_zbl_forwards_nondefault_softmin_and_descriptor_cutoff(self) -> None: "sw_rmin": 0.2, "sw_rmax": 4.0, "smin_alpha": 0.37, + "preset_out_bias": {"energy": [None, 1.0]}, }, descriptor_base=_DescriptorBase, fitting_base=_FittingBase, @@ -311,6 +317,9 @@ def test_zbl_forwards_nondefault_softmin_and_descriptor_cutoff(self) -> None: self.assertEqual(model.kwargs["smin_alpha"], 0.37) self.assertEqual(model.pairtab.rcut, 5.0) self.assertEqual(model.pairtab.sel, [4, 8]) + # the composition computes the bias and therefore owns the preset + self.assertEqual(model.kwargs["preset_out_bias"], {"energy": [None, 1.0]}) + self.assertNotIn("preset_out_bias", model.dp_model.kwargs) if __name__ == "__main__": diff --git a/source/tests/common/dpmodel/test_vacuum_ref_model.py b/source/tests/common/dpmodel/test_vacuum_ref_model.py new file mode 100644 index 0000000000..8289337b50 --- /dev/null +++ b/source/tests/common/dpmodel/test_vacuum_ref_model.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Isolated-atom reference at the model level on the dpmodel routes. + +With ``vacuum_ref`` in the fitting, an isolated atom contributes exactly the +bias of its type, and the energy of a cluster equals the energy of the same +model without the reference minus, for every atom, the network output of an +isolated atom of its type. The dense neighbor-list route evaluates the vacuum +descriptor on single-atom frames; the graph route carries one reference node +per type through the same descriptor call. +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, +) + +TYPE_MAP = ["O", "H"] +BIAS = np.array([[-3.0], [0.5]]) + + +def make_model(kind: str, vacuum_ref: bool, preset: bool = True) -> EnergyModel: + if kind == "se_e2_a": + ds = DescrptSeA( + rcut=4.0, rcut_smth=0.5, sel=[10, 10], neuron=[4, 8], axis_neuron=2, seed=1 + ) + else: + ds = DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[20], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + seed=1, + ) + ft = InvarFitting( + "energy", + 2, + ds.get_dim_out(), + 1, + neuron=[8, 8], + mixed_types=ds.mixed_types(), + vacuum_ref=vacuum_ref, + seed=1, + ) + ft["bias_atom_e"] = BIAS.copy() + preset_out_bias = {"energy": BIAS.tolist()} if preset else None + return EnergyModel(ds, ft, type_map=TYPE_MAP, preset_out_bias=preset_out_bias) + + +def atom_energies( + model: EnergyModel, coord: np.ndarray, atype: np.ndarray, method: str +) -> np.ndarray: + ret = model.call_common(coord, atype, None, neighbor_graph_method=method) + return ret["energy"][..., 0] + + +@pytest.mark.parametrize( + "kind, method", + [("se_e2_a", "legacy"), ("dpa1", "dense")], # dense nlist route / graph route +) +def test_isolated_atoms_and_reference_subtraction(kind: str, method: str) -> None: + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + vac = make_model(kind, True) + ref = make_model(kind, False) + + # an isolated atom of every type contributes exactly its bias + np.testing.assert_allclose( + atom_energies(vac, iso_coord, iso_atype, method), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the reference removes the isolated-atom network output of every atom + e_vac = atom_energies(vac, coord, atype, method) + e_ref = atom_energies(ref, coord, atype, method) + iso_ref = atom_energies(ref, iso_coord, iso_atype, method)[:, 0] + expected = e_ref - (iso_ref - BIAS[:, 0])[atype] + np.testing.assert_allclose(e_vac, expected, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize( + "kind, method", + [("se_e2_a", "legacy"), ("dpa1", "dense")], # dense nlist route / graph route +) +def test_fold_reproduces_the_referenced_model(kind: str, method: str) -> None: + """Folding the reference into the bias keeps every energy and drops the rows.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + model = make_model(kind, True) + e_cluster = atom_energies(model, coord, atype, method) + e_iso = atom_energies(model, iso_coord, iso_atype, method) + + model.fold_vacuum_reference() + assert not model.atomic_model.fitting_net.vacuum_ref + np.testing.assert_allclose( + atom_energies(model, coord, atype, method), e_cluster, rtol=1e-10, atol=1e-10 + ) + np.testing.assert_allclose( + atom_energies(model, iso_coord, iso_atype, method), + e_iso, + rtol=1e-10, + atol=1e-10, + ) + # the folded model is a plain model: its serialization carries no reference + assert model.serialize()["fitting"]["vacuum_ref"] is False + + +@pytest.mark.parametrize( + "kind, method", + [("se_e2_a", "legacy"), ("dpa1", "dense")], # dense nlist route / graph route +) +def test_without_preset_the_output_is_not_referenced(kind: str, method: str) -> None: + """A bias fitted from the data is no isolated-atom energy, so the output stays plain.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + plain = make_model(kind, False) + unreferenced = make_model(kind, True, preset=False) + fitting = unreferenced.atomic_model.fitting_net + assert not fitting.vacuum_ref + assert not fitting.needs_vacuum_descriptor() + np.testing.assert_allclose( + atom_energies(unreferenced, coord, atype, method), + atom_energies(plain, coord, atype, method), + rtol=1e-12, + atol=1e-12, + ) + unreferenced.fold_vacuum_reference() + np.testing.assert_array_equal(fitting["bias_atom_e"], BIAS) diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index e4379c1b12..7739cdd9f4 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -76,6 +76,37 @@ def test_builder_composes_linear_model(): assert float(dp_child.descriptor.inner_clamp.r_inner) == 0.8 +def test_bridged_preset_belongs_to_output_statistics() -> None: + cfg = copy.deepcopy(ZBL_CONFIG) + cfg["preset_out_bias"] = {"energy": {"Ni": 2.0, "O": 1.0}} + model = get_model(cfg) + am = model.atomic_model + assert am.preset_out_bias == {"energy": [[2.0], [1.0]]} + coord, atype, box = _close_pair_inputs() + before = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + sampled = [ + { + "coord": coord, + "atype": atype, + "box": box, + "natoms": np.array([[6, 6, 3, 3]]), + "energy": np.array([[9.0]]), + "find_energy": 1.0, + } + ] + am.change_out_bias(sampled, bias_adjust_mode="set-by-statistic") + np.testing.assert_allclose(np.asarray(am.out_bias).reshape(-1), [2.0, 1.0]) + after = model.call_common(coord, atype, box=box, neighbor_graph_method="dense")[ + "energy_redu" + ] + np.testing.assert_allclose(after - before, 9.0, atol=1e-10) + loaded = BaseModel.deserialize(model.serialize()) + assert loaded.atomic_model.preset_out_bias == am.preset_out_bias + np.testing.assert_allclose(loaded.atomic_model.out_bias, am.out_bias) + + def test_third_child_without_common_route_raises(): """[learned, inner_potential, pairtab] has no common execution route (pairtab is dense-only, the bridged pair is graph-only): the builder diff --git a/source/tests/common/test_argcheck_backend_docs.py b/source/tests/common/test_argcheck_backend_docs.py index ee5d40398a..fd42df7cfa 100644 --- a/source/tests/common/test_argcheck_backend_docs.py +++ b/source/tests/common/test_argcheck_backend_docs.py @@ -82,7 +82,7 @@ def test_representative_declared_support_labels(self) -> None: preset_out_bias = argcheck.model_args()["preset_out_bias"] self.assertTrue( preset_out_bias.doc.startswith( - "(Supported Backend: PyTorch, PaddlePaddle) " + "(Supported Backend: PyTorch, JAX, PaddlePaddle, PyTorch Exportable) " ) ) rglob_patterns = argcheck.training_data_args()["rglob_patterns"] diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index a86141f3fd..416ddb2b60 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -68,6 +68,7 @@ p_examples / "water" / "dpa3" / "input_torch_dynamic.json", p_examples / "water" / "dpa4" / "input.json", p_examples / "water" / "dpa4" / "input_preset.json", + p_examples / "water" / "dpa4" / "input_e0.json", p_examples / "water" / "dpa4" / "input-zbl.json", p_examples / "water" / "dpa4" / "lmp" / "input.json", p_examples / "water" / "dpa4c" / "input.json", @@ -85,6 +86,7 @@ p_examples / "hessian" / "multi_task" / "input.json", p_examples / "water" / "dpa4" / "input_multitask.json", p_examples / "water" / "dpa4" / "input_multitask_preset.json", + p_examples / "water" / "dpa4" / "input_multitask_e0.json", p_examples / "water_multi_task" / "pytorch_example" diff --git a/source/tests/common/test_preset_out_bias.py b/source/tests/common/test_preset_out_bias.py new file mode 100644 index 0000000000..02cadb51bc --- /dev/null +++ b/source/tests/common/test_preset_out_bias.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.utils.finetune import ( + get_index_between_two_maps, +) +from deepmd.utils.preset_out_bias import ( + bundled_preset_out_bias_tables, + load_preset_out_bias_table, + normalize_preset_out_bias, + override_assigned_bias, + preset_out_bias_rows, + remap_preset_out_bias, + resolve_preset_out_bias_tables, +) + + +class TestNormalizePresetOutBias(unittest.TestCase): + def setUp(self) -> None: + self.type_map = ["O", "H", "B"] + + def test_none(self) -> None: + self.assertIsNone(normalize_preset_out_bias(None, self.type_map)) + + def test_list_form(self) -> None: + preset = {"energy": [None, "1.", 3], "dipole": [None, [0, 1, 2], None]} + out = normalize_preset_out_bias(preset, self.type_map) + self.assertEqual(out["energy"], [None, [1.0], [3.0]]) + self.assertEqual(out["dipole"], [None, [0.0, 1.0, 2.0], None]) + # the configuration dict is left untouched + self.assertEqual(preset["energy"], [None, "1.", 3]) + + def test_element_dict_form(self) -> None: + out = normalize_preset_out_bias( + {"energy": {"B": 3.0, "H": [1.0]}, "dipole": {}}, self.type_map + ) + self.assertEqual(out["energy"], [None, [1.0], [3.0]]) + self.assertEqual(out["dipole"], [None, None, None]) + + def test_array_entries(self) -> None: + out = normalize_preset_out_bias( + { + "polar": [np.eye(2), None, np.array(4.0)], + "energy": np.array([[1.0], [2.0], [3.0]]), + }, + self.type_map, + ) + self.assertEqual(out["polar"], [[[1.0, 0.0], [0.0, 1.0]], None, [4.0]]) + self.assertEqual(out["energy"], [[1.0], [2.0], [3.0]]) + + def test_idempotent_and_json_round_trip(self) -> None: + out = normalize_preset_out_bias( + {"energy": {"H": -13.6}, "polar": {"O": [[1.0, 0.0], [0.0, 1.0]]}}, + self.type_map, + ) + self.assertEqual(normalize_preset_out_bias(out, self.type_map), out) + self.assertEqual( + normalize_preset_out_bias(json.loads(json.dumps(out)), self.type_map), out + ) + + def test_unknown_elements_ignored(self) -> None: + out = normalize_preset_out_bias( + {"energy": {"C": 3.0, "H": -13.6, "B": None}}, self.type_map + ) + self.assertEqual(out["energy"], [None, [-13.6], None]) + + def test_json_table(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + table = Path(tmp) / "energy_bias.json" + table.write_text(json.dumps({"H": -13.6, "B": None, "C": 3.0})) + out = normalize_preset_out_bias({"energy": str(table)}, self.type_map) + self.assertEqual(out["energy"], [None, [-13.6], None]) + # a relative path resolves against the working directory + cwd = os.getcwd() + os.chdir(tmp) + try: + out = normalize_preset_out_bias( + {"energy": "energy_bias.json"}, self.type_map + ) + finally: + os.chdir(cwd) + self.assertEqual(out["energy"], [None, [-13.6], None]) + (Path(tmp) / "list.json").write_text("[1, 2]") + with self.assertRaises(ValueError): + normalize_preset_out_bias( + {"energy": str(Path(tmp) / "list.json")}, self.type_map + ) + + def test_errors(self) -> None: + for bad in ( + {"energy": [None]}, + {"energy": 1.0}, + {"energy": [None, 1.0 + 2.0j, None]}, + {"energy": [None, "1.0 + 2.0j", None]}, + {"energy": [None, np.inf, None]}, + {"energy": [None, -np.inf, None]}, + {"dipole": [None, [0.0, None, 2.0], None]}, + ): + with self.assertRaises(ValueError): + normalize_preset_out_bias(bad, self.type_map) + + +class TestBundledTables(unittest.TestCase): + def test_name_before_path(self) -> None: + tables = bundled_preset_out_bias_tables() + self.assertEqual(load_preset_out_bias_table("omat24"), tables["omat24"]) + # an element of the type map outside the table stays unassigned + out = normalize_preset_out_bias({"energy": "omat24"}, ["O", "H", "Po"]) + self.assertEqual(out["energy"], [[-1.54797136], [-1.11700253], None]) + config = {"type_map": ["H"], "preset_out_bias": {"energy": "omol25"}} + resolved = resolve_preset_out_bias_tables(config)["preset_out_bias"] + self.assertEqual(resolved, {"energy": tables["omol25"]}) + with self.assertRaises(ValueError): + load_preset_out_bias_table("no_such_table.json") + + +class TestResolvePresetOutBiasTables(unittest.TestCase): + def test_single_and_multi_task(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + table = Path(tmp) / "energy_bias.json" + table.write_text(json.dumps({"H": -13.6})) + single = {"type_map": ["H"], "preset_out_bias": {"energy": str(table)}} + out = resolve_preset_out_bias_tables(single) + self.assertEqual(out["preset_out_bias"], {"energy": {"H": -13.6}}) + # the configuration passed in is left untouched + self.assertEqual(single["preset_out_bias"], {"energy": str(table)}) + multi = { + "preset_out_bias": {"energy": str(table)}, + "model_dict": { + "a": single, + "b": {"type_map": ["H"], "preset_out_bias": {"energy": {"H": 1.0}}}, + "c": {"type_map": ["H"]}, + }, + } + out = resolve_preset_out_bias_tables(multi) + # a preset next to model_dict is resolved as well + self.assertEqual(out["preset_out_bias"], {"energy": {"H": -13.6}}) + self.assertEqual( + out["model_dict"]["a"]["preset_out_bias"], {"energy": {"H": -13.6}} + ) + self.assertIs(out["model_dict"]["b"], multi["model_dict"]["b"]) + self.assertIs(out["model_dict"]["c"], multi["model_dict"]["c"]) + + def test_without_tables(self) -> None: + config = {"type_map": ["H"], "preset_out_bias": {"energy": [1.0]}} + self.assertIs(resolve_preset_out_bias_tables(config), config) + config = {"type_map": ["H"]} + self.assertIs(resolve_preset_out_bias_tables(config), config) + + +class TestPresetOutBiasRows(unittest.TestCase): + def setUp(self) -> None: + self.type_map = ["O", "H", "B"] + self.preset = normalize_preset_out_bias( + {"energy": {"H": -13.6}, "dipole": {"H": [1.0, 2.0, 3.0]}}, self.type_map + ) + self.stored = np.arange(2 * 3 * 3, dtype=np.float64).reshape(2, 3, 3) + self.keys = ["energy", "dipole"] + self.sizes = [1, 3] + + def rows(self, observed: list[str], keep_unassigned: bool) -> dict[str, np.ndarray]: + return preset_out_bias_rows( + self.preset, + self.type_map, + observed, + self.stored, + self.keys, + self.sizes, + keep_unassigned, + ) + + def test_zero_or_stored_for_unassigned_types(self) -> None: + rows = self.rows(["H"], keep_unassigned=False) + np.testing.assert_array_equal(rows["energy"], [[0.0], [-13.6], [0.0]]) + np.testing.assert_array_equal( + rows["dipole"], [[0.0, 0.0, 0.0], [1.0, 2.0, 3.0], [0.0, 0.0, 0.0]] + ) + rows = self.rows(["H"], keep_unassigned=True) + np.testing.assert_array_equal( + rows["energy"], [[self.stored[0, 0, 0]], [-13.6], [self.stored[0, 2, 0]]] + ) + np.testing.assert_array_equal( + rows["dipole"], [self.stored[1, 0], [1.0, 2.0, 3.0], self.stored[1, 2]] + ) + + def test_unassigned_output_is_not_fixed(self) -> None: + self.preset = normalize_preset_out_bias( + {"energy": [None, None, None], "dipole": {"H": [1.0, 2.0, 3.0]}}, + self.type_map, + ) + self.assertEqual(list(self.rows(["H"], keep_unassigned=False)), ["dipole"]) + + def test_observed_types_must_be_assigned(self) -> None: + with self.assertRaisesRegex(ValueError, r"energy.*\['O'\]"): + self.rows(["O", "H"], keep_unassigned=False) + + def test_unknown_and_excluded_observed_types(self) -> None: + # an observed name outside the type map is ignored and an excluded + # type needs no preset + rows = preset_out_bias_rows( + self.preset, + self.type_map, + ["H", "Fe", "B"], + self.stored, + self.keys, + self.sizes, + keep_unassigned=False, + excluded_types=[2], + ) + np.testing.assert_array_equal(rows["energy"], [[0.0], [-13.6], [0.0]]) + with self.assertRaisesRegex(ValueError, r"energy.*\['B'\]"): + self.rows(["H", "B"], keep_unassigned=False) + + +class TestRemapPresetOutBias(unittest.TestCase): + def test_follows_out_bias_remap(self) -> None: + old_map = ["O", "H", "B", "C"] + preset = normalize_preset_out_bias( + {"energy": {"O": 1.0, "B": 3.0, "C": 4.0}}, old_map + ) + # per-type reference values in the layout of a stored out_bias + out_bias = np.array([1.0, np.nan, 3.0, 4.0]) + for new_map in ( + ["C", "O", "H", "B"], + ["B", "N", "O"], + ["N", "F"], + ["H"], + ): + remap_index, _ = get_index_between_two_maps(old_map, new_map) + remapped = remap_preset_out_bias(preset, remap_index)["energy"] + reference = np.concatenate([out_bias, np.full(len(new_map), np.nan)])[ + remap_index + ] + self.assertEqual(len(remapped), len(new_map)) + for got, want in zip(remapped, reference, strict=True): + if np.isnan(want): + self.assertIsNone(got) + else: + self.assertEqual(got, [want]) + + def test_none(self) -> None: + self.assertIsNone(remap_preset_out_bias(None, [0, 1])) + + +class TestOverrideAssignedBias(unittest.TestCase): + def test_flattened_bias_against_shaped_preset(self) -> None: + bias = np.arange(12, dtype=np.float64).reshape(3, 4) + assigned = np.full((3, 2, 2), np.nan) + assigned[1] = [[10.0, 11.0], [12.0, 13.0]] + out = override_assigned_bias(bias, assigned) + expected = bias.copy() + expected[1] = [10.0, 11.0, 12.0, 13.0] + np.testing.assert_array_equal(out, expected) + self.assertEqual(out.shape, bias.shape) + # the input is left untouched + self.assertEqual(bias[1, 0], 4.0) + + def test_none(self) -> None: + bias = np.ones((2, 3)) + self.assertIs(override_assigned_bias(bias, None), bias) diff --git a/source/tests/common/test_vacuum_reference.py b/source/tests/common/test_vacuum_reference.py new file mode 100644 index 0000000000..b3ece81cb1 --- /dev/null +++ b/source/tests/common/test_vacuum_reference.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.utils.vacuum_reference import ( + reference_charge_spin, + reference_spin, + unpaired_electrons, +) + + +class TestVacuumReference(unittest.TestCase): + def test_unpaired_electrons(self) -> None: + # Hund's rule ground states: 1s1, 1s2, 2p3, 2p4, 3d6 4s2, 4f7 5d1 6s2 + type_map = ["H", "He", "N", "O", "Fe", "Gd"] + np.testing.assert_array_equal(unpaired_electrons(type_map), [1, 0, 3, 2, 4, 8]) + + def test_unknown_type_name(self) -> None: + with self.assertRaises(ValueError): + unpaired_electrons(["O", "H1"]) + + def test_reference_conditions(self) -> None: + type_map = ["O", "H"] + charge_spin = reference_charge_spin(type_map) + np.testing.assert_array_equal(charge_spin, [[0.0, 3.0], [0.0, 2.0]]) + self.assertEqual(charge_spin.dtype, np.float64) + spin = reference_spin(type_map) + np.testing.assert_array_equal(spin, [[0.0, 0.0, 2.0], [0.0, 0.0, 1.0]]) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/consistent/utils/test_stat.py b/source/tests/consistent/utils/test_stat.py index d2a7a6c1b6..9377cdb8dc 100644 --- a/source/tests/consistent/utils/test_stat.py +++ b/source/tests/consistent/utils/test_stat.py @@ -4,7 +4,14 @@ from collections import ( defaultdict, ) +from pathlib import ( + Path, +) +from typing import ( + Any, +) +import h5py import numpy as np import pytest @@ -15,6 +22,9 @@ _compute_output_stats_global as compute_output_stats_global_dp, ) from deepmd.dpmodel.utils.stat import compute_output_stats as compute_output_stats_dp +from deepmd.utils.path import ( + DPPath, +) from ..common import ( INSTALLED_PD, @@ -49,6 +59,71 @@ NLOC = 4 +@pytest.mark.parametrize("backend", ["dp", "pt", "pd"]) # statistics implementation +def test_output_stats_cache_depends_on_model_predictions( + tmp_path: Path, backend: str +) -> None: + """Absolute constrained statistics are reusable; model-dependent shifts are not.""" + if backend == "pt" and not INSTALLED_PT: + pytest.skip("PyTorch is not installed") + if backend == "pd" and not INSTALLED_PD: + pytest.skip("PaddlePaddle is not installed") + sampled, _, _ = _make_data(True, False, True, []) + sampled[0]["coord"] = np.zeros((NFRAMES, NLOC, 3)) + sampled[0]["box"] = None + if backend == "pt": + sampled = _np_to_torch(sampled) + compute = compute_output_stats_pt + as_numpy = to_numpy_array_pt + elif backend == "pd": + sampled = _np_to_paddle(sampled) + compute = compute_output_stats_pd + as_numpy = to_numpy_array_pd + else: + compute = compute_output_stats_dp + as_numpy = np.asarray + filename = str(tmp_path / "stat.h5") + with h5py.File(filename, "w"): + pass + path = DPPath(filename, "a") + preset_bias = {"energy": [None, [0.25]]} + original, _ = compute( + sampled, NTYPES, ["energy"], stat_file_path=path, preset_bias=preset_bias + ) + + def no_sampling() -> list[dict]: + pytest.fail("Absolute output statistics must remain reusable") + + restored, _ = compute( + no_sampling, NTYPES, ["energy"], stat_file_path=path, preset_bias=preset_bias + ) + np.testing.assert_array_equal( + as_numpy(restored["energy"]), as_numpy(original["energy"]) + ) + + for value in (0.25, 1.5): + + def predict(coord: Any, atype: Any, box: Any, **kwargs: Any) -> dict[str, Any]: + return {"energy": atype[..., None] * 0 + value, "mask": atype * 0 + 1} + + kwargs = {"model_forward": predict, "preset_bias": preset_bias} + expected, _ = compute(sampled, NTYPES, ["energy"], **kwargs) + actual, _ = compute(sampled, NTYPES, ["energy"], stat_file_path=path, **kwargs) + np.testing.assert_allclose( + as_numpy(actual["energy"]), as_numpy(expected["energy"]) + ) + fresh_dir = tmp_path / f"delta-{value}" + fresh_dir.mkdir() + fresh_path = DPPath(str(fresh_dir), "a") + compute(sampled, NTYPES, ["energy"], stat_file_path=fresh_path, **kwargs) + assert not list(fresh_dir.iterdir()) + + restored, _ = compute(no_sampling, NTYPES, ["energy"], stat_file_path=path) + np.testing.assert_array_equal( + as_numpy(restored["energy"]), as_numpy(original["energy"]) + ) + + def _make_data( has_global: bool, has_atomic: bool, diff --git a/source/tests/jax/test_preset_out_bias.py b/source/tests/jax/test_preset_out_bias.py new file mode 100644 index 0000000000..f0006d96bb --- /dev/null +++ b/source/tests/jax/test_preset_out_bias.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Preset output bias through the JAX model factory and bias change.""" + +import unittest + +import numpy as np + +from deepmd.jax.env import ( + jnp, +) +from deepmd.jax.model.ener_model import ( + EnergyModel, +) +from deepmd.jax.model.model import ( + get_model, +) + +from ..common.stat_file import ( + energy_model_params, + energy_stat_sample, +) + + +class TestPresetOutBias(unittest.TestCase): + def setUp(self) -> None: + params = energy_model_params() + params["preset_out_bias"] = {"energy": {"O": -10.0, "H": 5.0}} + self.model = get_model(params) + self.sampled = [ + { + key: jnp.asarray(value) if isinstance(value, np.ndarray) else value + for key, value in sample.items() + } + for sample in energy_stat_sample() + ] + + def out_bias(self) -> np.ndarray: + return np.asarray(self.model.get_out_bias()).reshape(-1) + + def test_preset_pinned_in_both_modes(self) -> None: + self.assertEqual( + self.model.atomic_model.preset_out_bias, {"energy": [[-10.0], [5.0]]} + ) + # the data would fit O to 1 and H to 2; the preset pins both in both modes + for mode in ("set-by-statistic", "change-by-statistic"): + self.model.change_out_bias(self.sampled, bias_adjust_mode=mode) + np.testing.assert_allclose(self.out_bias(), [-10.0, 5.0]) + + def test_serialize_round_trip(self) -> None: + self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") + loaded = EnergyModel.deserialize(self.model.serialize()) + self.assertEqual( + loaded.atomic_model.preset_out_bias, {"energy": [[-10.0], [5.0]]} + ) + np.testing.assert_allclose( + np.asarray(loaded.get_out_bias()), np.asarray(self.model.get_out_bias()) + ) diff --git a/source/tests/pd/model/test_atomic_model_global_stat.py b/source/tests/pd/model/test_atomic_model_global_stat.py index ec72973719..bacec7e682 100644 --- a/source/tests/pd/model/test_atomic_model_global_stat.py +++ b/source/tests/pd/model/test_atomic_model_global_stat.py @@ -289,9 +289,9 @@ def test_preset_bias(self): ).to(env.DEVICE) ft = FooFitting().to(env.DEVICE) type_map = ["foo", "bar"] + # both types occur in the data, so an assigned output assigns both preset_out_bias = { - # "foo": np.array(3.0, 2.0]).reshape(2, 1), - "foo": [None, 2], + "foo": [3, 2], "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), } md0 = DPAtomicModel( @@ -335,16 +335,13 @@ def cvt_ret(x): for kk in ["foo", "pix", "bar"]: np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - # 2. test bias is applied + # 2. the preset fixes foo and bar; pix has no label and keeps a zero bias md0.compute_or_load_out_stat( self.merged_output_stat, stat_file_path=self.stat_file_path ) ret1 = md0.forward_common_atomic(*args) ret1 = cvt_ret(ret1) - # foo sums: [5, 7], - # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] - # the solution of type 0 is 1.8 - foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) + foo_bias = np.array(preset_out_bias["foo"], dtype=np.float64).reshape(2, 1) bar_bias = preset_out_bias["bar"] expected_ret1 = {} expected_ret1["foo"] = ret0["foo"] + foo_bias[at] @@ -352,44 +349,26 @@ def cvt_ret(x): expected_ret1["bar"] = ret0["bar"] + bar_bias[at] for kk in ["foo", "pix", "bar"]: np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error(): - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args) - ret2 = cvt_ret(ret2) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - - # 4. test change bias + # 3. change-by-statistic keeps every assigned type at its preset BaseAtomicModel.change_out_bias( md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" ) - args = [ - to_paddle_tensor(ii) - for ii in [ - self.coord_ext, - to_numpy_array(self.merged_output_stat[0]["atype_ext"]), - self.nlist, - ] - ] ret3 = md0.forward_common_atomic(*args) ret3 = cvt_ret(ret3) - ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] - ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] - ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) - ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] - ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( - 2, 3, 1 - ) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret3[kk], ret1[kk]) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) + np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) + # 4. a preset that leaves an observed type unassigned is rejected + md1 = DPAtomicModel( + ds, + FooFitting().to(env.DEVICE), + type_map=type_map, + preset_out_bias={"foo": [None, 2]}, + ).to(env.DEVICE) + with self.assertRaisesRegex(ValueError, "foo"): + md1.compute_or_load_out_stat(self.merged_output_stat) def test_preset_bias_all_none(self): nf, nloc, nnei = self.nlist.shape diff --git a/source/tests/pd/model/test_ener_fitting.py b/source/tests/pd/model/test_ener_fitting.py index dd13f139dc..014296eda0 100644 --- a/source/tests/pd/model/test_ener_fitting.py +++ b/source/tests/pd/model/test_ener_fitting.py @@ -148,3 +148,14 @@ def test_get_set(self): np.testing.assert_allclose( foo, np.reshape(ifn0[ii].detach().cpu().numpy(), foo.shape) ) + + def test_vacuum_ref_is_rejected(self): + """A serialized ``vacuum_ref`` is rejected; a version-4 dictionary loads.""" + data = InvarFitting("energy", self.nt, 3, 1, seed=GLOBAL_SEED).serialize() + self.assertFalse(data["vacuum_ref"]) + with self.assertRaises(NotImplementedError): + InvarFitting.deserialize({**data, "vacuum_ref": True}) + older = {k: v for k, v in data.items() if k != "vacuum_ref"} + self.assertIsInstance( + InvarFitting.deserialize({**older, "@version": 4}), InvarFitting + ) diff --git a/source/tests/pd/model/test_get_model.py b/source/tests/pd/model/test_get_model.py index 7ace7c4e43..1cf6fa3332 100644 --- a/source/tests/pd/model/test_get_model.py +++ b/source/tests/pd/model/test_get_model.py @@ -2,7 +2,6 @@ import copy import unittest -import numpy as np import paddle from deepmd.pd.model.model import ( @@ -55,8 +54,8 @@ def test_model_attr(self): { "energy": [ None, - np.array([1.0]), - np.array([3.0]), + [1.0], + [3.0], ] }, ) @@ -73,8 +72,8 @@ def test_model_attr_energy_float(self): atomic_model.preset_out_bias, { "energy": [ - np.array([1.0]), - np.array([3.0]), + [1.0], + [3.0], None, ] }, @@ -82,6 +81,31 @@ def test_model_attr_energy_float(self): self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) + def test_model_attr_energy_element_dict(self): + model_params = copy.deepcopy(model_se_e2_a) + model_params["preset_out_bias"] = {"energy": {"B": 3.0, "H": [1.0]}} + self.model = get_model(model_params).to(env.DEVICE) + atomic_model = self.model.atomic_model + self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) + self.assertEqual( + atomic_model.preset_out_bias, + { + "energy": [ + None, + [1.0], + [3.0], + ] + }, + ) + + def test_model_attr_energy_unknown_element_ignored(self): + model_params = copy.deepcopy(model_se_e2_a) + model_params["preset_out_bias"] = {"energy": {"C": 3.0, "H": 1.0}} + self.model = get_model(model_params).to(env.DEVICE) + self.assertEqual( + self.model.atomic_model.preset_out_bias, {"energy": [None, [1.0], None]} + ) + def test_model_attr_energy_unsupported_type(self): model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": [1.0 + 2.0j, 3, None]} diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py index 95b94f0fe6..a4bca815b9 100644 --- a/source/tests/pt/model/test_atomic_model_global_stat.py +++ b/source/tests/pt/model/test_atomic_model_global_stat.py @@ -13,10 +13,16 @@ import torch from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel +from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel +from deepmd.dpmodel.model.model import get_model as get_dp_model from deepmd.dpmodel.output_def import ( FittingOutputDef, OutputVariableDef, ) +from deepmd.dpmodel.utils.serialization import ( + load_dp_model, + save_dp_model, +) from deepmd.pt.model.atomic_model import ( BaseAtomicModel, DPAtomicModel, @@ -25,6 +31,10 @@ DescrptDPA1, DescrptSeA, ) +from deepmd.pt.model.model import ( + EnergyModel, + get_model, +) from deepmd.pt.model.task.base_fitting import ( BaseFitting, ) @@ -306,9 +316,9 @@ def test_preset_bias(self) -> None: ).to(env.DEVICE) ft = FooFitting().to(env.DEVICE) type_map = ["foo", "bar"] + # both types occur in the data, so an assigned output assigns both preset_out_bias = { - # "foo": np.array(3.0, 2.0]).reshape(2, 1), - "foo": [None, 2], + "foo": [3, 2], "bar": np.array([7.0, 5.0, 13.0, 11.0]).reshape(2, 1, 2), } md0 = DPAtomicModel( @@ -351,17 +361,13 @@ def cvt_ret(x): ).reshape([nf, nloc] + md0.fitting_output_def()["bar"].shape) # noqa: RUF005 for kk in ["foo", "pix", "bar"]: np.testing.assert_almost_equal(ret0[kk], expected_ret0[kk]) - - # 2. test bias is applied + # 2. the preset fixes foo and bar; pix has no label and keeps a zero bias md0.compute_or_load_out_stat( self.merged_output_stat, stat_file_path=self.stat_file_path ) ret1 = md0.forward_common_atomic(*args, fparam=self.fparam) ret1 = cvt_ret(ret1) - # foo sums: [5, 7], - # given bias of type 1 being 2, the bias left for type 0 is [5-2*1, 7-2*2] = [3,3] - # the solution of type 0 is 1.8 - foo_bias = np.array([1.8, preset_out_bias["foo"][1]]).reshape(2, 1) + foo_bias = np.array(preset_out_bias["foo"], dtype=np.float64).reshape(2, 1) bar_bias = preset_out_bias["bar"] expected_ret1 = {} expected_ret1["foo"] = ret0["foo"] + foo_bias[at] @@ -369,44 +375,26 @@ def cvt_ret(x): expected_ret1["bar"] = ret0["bar"] + bar_bias[at] for kk in ["foo", "pix", "bar"]: np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - - # 3. test bias load from file - def raise_error() -> NoReturn: - raise RuntimeError - - md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) - ret2 = md0.forward_common_atomic(*args, fparam=self.fparam) - ret2 = cvt_ret(ret2) - for kk in ["foo", "pix", "bar"]: - np.testing.assert_almost_equal(ret1[kk], ret2[kk]) - - # 4. test change bias + # 3. change-by-statistic keeps every assigned type at its preset BaseAtomicModel.change_out_bias( md0, self.merged_output_stat, bias_adjust_mode="change-by-statistic" ) - args = [ - to_torch_tensor(ii) - for ii in [ - self.coord_ext, - to_numpy_array(self.merged_output_stat[0]["atype_ext"]), - self.nlist, - ] - ] ret3 = md0.forward_common_atomic(*args, fparam=self.fparam) ret3 = cvt_ret(ret3) - ## model output on foo: [[2.8, 3.8, 5], [5.8, 7., 8.]] given bias [1.8, 2] - ## foo sumed: [11.6, 20.8] compared with [5, 7], fit target is [-6.6, -13.8] - ## fit bias is [-7, 2] (2 is assigned. -7 is fit to [-8.6, -17.8]) - ## old bias[1.8,2] + fit bias[-7, 2] = [-5.2, 4] - ## new model output is [[-4.2, -3.2, 7], [-1.2, 9, 10]] - expected_ret3 = {} - expected_ret3["foo"] = np.array([[-4.2, -3.2, 7.0], [-1.2, 9.0, 10.0]]).reshape( - 2, 3, 1 - ) - expected_ret3["pix"] = ret0["pix"] - for kk in ["foo", "pix"]: - np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) - # bar is too complicated to be manually computed. + for kk in ["foo", "pix", "bar"]: + np.testing.assert_almost_equal(ret3[kk], ret1[kk]) + out_bias, _ = md0._fetch_out_stat(["foo", "bar"]) + np.testing.assert_almost_equal(to_numpy_array(out_bias["foo"]), foo_bias) + np.testing.assert_almost_equal(to_numpy_array(out_bias["bar"]), bar_bias) + # 4. a preset that leaves an observed type unassigned is rejected + md1 = DPAtomicModel( + ds, + FooFitting().to(env.DEVICE), + type_map=type_map, + preset_out_bias={"foo": [None, 2]}, + ).to(env.DEVICE) + with self.assertRaisesRegex(ValueError, "foo"): + md1.compute_or_load_out_stat(self.merged_output_stat) def test_preset_bias_all_none(self) -> None: nf, nloc, nnei = self.nlist.shape @@ -524,3 +512,227 @@ def cvt_ret(x): ret2 = md2.forward_common_atomic(*args) for kk in ["foo"]: np.testing.assert_almost_equal(ret0[kk], ret2[kk]) + + +class TestModelPresetBias(unittest.TestCase): + """Preset bias through the model-level bias change used by fine-tuning.""" + + def setUp(self) -> None: + # a mixed-types descriptor supports change_type_map + self.model = get_model( + { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "dpa1", + "sel": 20, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 8], + "axis_neuron": 2, + "attn_layer": 0, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, + "preset_out_bias": {"energy": {"O": -10.0, "H": -13.6}}, + } + ).to(env.DEVICE) + rng = np.random.default_rng(20260913) + # O and H occur in the data and are assigned; B is neither + self.sampled = [ + { + "coord": to_torch_tensor(rng.uniform(-1.5, 1.5, size=(2, 3, 3))), + "atype": to_torch_tensor( + np.array([[0, 0, 1], [0, 1, 1]], dtype=np.int32) + ), + "natoms": to_torch_tensor( + np.array([[3, 3, 2, 1, 0], [3, 3, 1, 2, 0]], dtype=np.int32) + ), + "energy": to_torch_tensor(np.array([-30.0, -20.0]).reshape(2, 1)), + "find_energy": np.float32(1.0), + } + ] + + def out_bias(self) -> np.ndarray: + return to_numpy_array(self.model.get_out_bias()).reshape(-1) + + def test_preset_pinned_in_both_modes(self) -> None: + preset = np.array([-10.0, -13.6, 0.0]) + self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") + np.testing.assert_allclose(self.out_bias(), preset) + # the output std of a fully preset output keeps its stored value + out_std = to_numpy_array(self.model.atomic_model.out_std).reshape(-1) + np.testing.assert_allclose(out_std, 1.0) + self.model.change_out_bias(self.sampled, bias_adjust_mode="change-by-statistic") + np.testing.assert_allclose(self.out_bias(), preset) + + def test_spin_model(self) -> None: + # the virtual spin types are excluded from the output and need no preset + model = get_model( + { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "dpa1", + "sel": 20, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 8], + "axis_neuron": 2, + "attn_layer": 0, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, + "spin": {"use_spin": [True, False, False], "virtual_scale": [0.3]}, + "preset_out_bias": {"energy": {"O": -10.0, "H": -13.6}}, + } + ).to(env.DEVICE) + rng = np.random.default_rng(1) + sampled = [ + {**self.sampled[0], "spin": to_torch_tensor(rng.normal(size=(2, 3, 3)))} + ] + for mode in ("set-by-statistic", "change-by-statistic"): + model.change_out_bias(sampled, bias_adjust_mode=mode) + out_bias = to_numpy_array(model.get_out_bias()).reshape(-1) + np.testing.assert_allclose(out_bias[:3], [-10.0, -13.6, 0.0]) + + def test_preset_needs_no_data(self) -> None: + def raise_error() -> NoReturn: + raise RuntimeError + + self.model.atomic_model.change_out_bias( + raise_error, bias_adjust_mode="set-by-statistic", observed_type=["O", "H"] + ) + np.testing.assert_allclose(self.out_bias(), [-10.0, -13.6, 0.0]) + + def test_missing_element_rejected(self) -> None: + model = get_model( + { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "dpa1", + "sel": 20, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 8], + "axis_neuron": 2, + "attn_layer": 0, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, + "preset_out_bias": {"energy": {"H": -13.6, "B": 3.0}}, + } + ).to(env.DEVICE) + with self.assertRaisesRegex(ValueError, "O"): + model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") + + def test_dp_file_round_trip(self) -> None: + self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") + with tempfile.TemporaryDirectory() as tmp: + filename = str(Path(tmp) / "model.dp") + save_dp_model(filename, {"model": self.model.serialize()}) + data = load_dp_model(filename)["model"] + # the file is read back by the pt backend and by the dpmodel backend + for loaded in (EnergyModel.deserialize(data), DPEnergyModel.deserialize(data)): + np.testing.assert_allclose( + to_numpy_array(loaded.get_out_bias()), + to_numpy_array(self.model.get_out_bias()), + ) + self.assertEqual( + loaded.atomic_model.preset_out_bias, + {"energy": [[-10.0], [-13.6], None]}, + ) + + def test_nonfinite_assigned_bias_from_checkpoint(self) -> None: + # the preset replaces a non-finite stored bias of an assigned type + for value in (np.nan, np.inf, -np.inf): + with self.subTest(value=value): + self.model.atomic_model.out_bias[0, :2, 0] = value + loaded = EnergyModel.deserialize(self.model.serialize()) + loaded.change_out_bias( + self.sampled, bias_adjust_mode="change-by-statistic" + ) + bias = to_numpy_array(loaded.get_out_bias()).reshape(-1) + np.testing.assert_allclose(bias, [-10.0, -13.6, 0.0]) + loaded.change_out_bias( + self.sampled, bias_adjust_mode="change-by-statistic" + ) + np.testing.assert_allclose( + to_numpy_array(loaded.get_out_bias()).reshape(-1), bias, atol=1e-10 + ) + + def test_change_type_map_remaps_preset(self) -> None: + self.model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") + # H is dropped, O keeps its preset, B and C are unassigned + self.model.change_type_map(["B", "O", "C"]) + self.assertEqual( + self.model.atomic_model.preset_out_bias, {"energy": [None, [-10.0], None]} + ) + sampled = [ + { + **self.sampled[0], + "atype": to_torch_tensor( + np.array([[1, 1, 2], [1, 2, 2]], dtype=np.int32) + ), + "natoms": to_torch_tensor( + np.array([[3, 3, 0, 2, 1], [3, 3, 0, 1, 2]], dtype=np.int32) + ), + } + ] + # C occurs in the new data without a preset + with self.assertRaisesRegex(ValueError, "C"): + self.model.change_out_bias(sampled, bias_adjust_mode="change-by-statistic") + + def test_dipole_preset_rejected(self) -> None: + params = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "neuron": [4, 8], + "axis_neuron": 2, + "rcut": 3.0, + "rcut_smth": 2.5, + }, + "fitting_net": {"type": "dipole", "neuron": [8]}, + "preset_out_bias": {"dipole": {"H": [0.0, 1.0, 2.0]}}, + } + for builder in (get_model, get_dp_model): + with self.subTest(builder=builder.__module__): + with self.assertRaisesRegex(ValueError, "do not apply an output bias"): + builder(params) + + def test_unknown_output_rejected(self) -> None: + params = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "se_e2_a", + "sel": [8, 8, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + }, + "fitting_net": {"neuron": [8]}, + "preset_out_bias": {"enrgy": {"H": -13.6}}, + } + with self.assertRaisesRegex(ValueError, "enrgy"): + get_model(params) + + def test_no_distinguish_rejected(self) -> None: + params = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "se_e2_a", + "sel": [8, 8, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + }, + "fitting_net": { + "type": "property", + "property_name": "band_prop", + "task_dim": 1, + "neuron": [8], + "distinguish_types": False, + }, + "preset_out_bias": {"band_prop": {"H": 1.0}}, + } + model = get_model(params).to(env.DEVICE) + with self.assertRaisesRegex(ValueError, "distinguish"): + model.change_out_bias(self.sampled, bias_adjust_mode="set-by-statistic") diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index 4d069cd75f..4f33c74e1c 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -543,7 +543,7 @@ def test_zero_block_descriptor(self) -> None: ) edge_vec = flat[edge_index[0]] - flat[edge_index[1]] edge_mask = torch.ones(2, dtype=torch.bool, device=self.device) - desc_e, latent = model.forward_with_edges( + desc_e, latent, _ = model.forward_with_edges( extended_coord=coord.reshape(1, -1), extended_atype=atype, edge_index=edge_index, @@ -978,7 +978,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: nlist, charge_spin=torch.tensor([[0.0, 1.0]], device=self.device), ) - desc_ref, _ = model.forward_with_edges( + desc_ref, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -986,7 +986,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: edge_mask=edge_mask, charge_spin=torch.tensor([[0.0, 1.0]], device=self.device), ) - desc_shifted, _ = model.forward_with_edges( + desc_shifted, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -995,7 +995,7 @@ def test_charge_spin_sparse_edge_conditioning(self) -> None: charge_spin=torch.tensor([[1.0, 1.0]], device=self.device), ) restored = DescrptSeZM.deserialize(model.serialize()) - desc_restored, _ = restored.forward_with_edges( + desc_restored, _, _ = restored.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1195,7 +1195,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: p.copy_(torch.randn_like(p) * 0.1) model.eval() - desc, _ = model.forward_with_edges( + desc, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1203,7 +1203,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: edge_mask=edge_mask, spin=spin, ) - desc_rot, _ = model.forward_with_edges( + desc_rot, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1214,7 +1214,7 @@ def test_descriptor_spin_joint_rotation_invariance(self) -> None: torch.testing.assert_close(desc, desc_rot, atol=1e-9, rtol=1e-9) # Spin actually changes the descriptor (injection is not a no-op). - desc_zero, _ = model.forward_with_edges( + desc_zero, _, _ = model.forward_with_edges( extended_coord=coord, extended_atype=atype, edge_index=edge_index, @@ -1287,7 +1287,7 @@ def test_zero_gate_receives_a_gradient(self) -> None: """ model = self._descriptor() kwargs, spin = self._inputs() - desc, _ = model.forward_with_edges(**kwargs, spin=spin) + desc, _, _ = model.forward_with_edges(**kwargs, spin=spin) desc.sum().backward() gate_grad = model.env_seed_embedding.spin_scale.grad self.assertIsNotNone(gate_grad) @@ -1306,10 +1306,10 @@ def test_gate_placement_differs_from_the_legacy_one_by_a_square(self) -> None: amplitude = 2.0 with torch.no_grad(): model.env_seed_embedding.spin_scale.fill_(amplitude**2) - migrated, _ = model.forward_with_edges(**kwargs, spin=spin) + migrated, _, _ = model.forward_with_edges(**kwargs, spin=spin) with torch.no_grad(): model.env_seed_embedding.spin_scale.fill_(1.0) - legacy, _ = model.forward_with_edges(**kwargs, spin=amplitude * spin) + legacy, _, _ = model.forward_with_edges(**kwargs, spin=amplitude * spin) torch.testing.assert_close(migrated, legacy, atol=1e-12, rtol=1e-12) def test_loading_a_legacy_state_squares_the_gate(self) -> None: diff --git a/source/tests/pt/model/test_fitting_vacuum_ref.py b/source/tests/pt/model/test_fitting_vacuum_ref.py new file mode 100644 index 0000000000..7443dcdc22 --- /dev/null +++ b/source/tests/pt/model/test_fitting_vacuum_ref.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fitting-side vacuum reference. + +With ``vacuum_ref`` the output of an atom is +``bias(t_i) + f(x_i; c_i) - f(x_vac(t_i); c_i)``, where ``x_vac`` is the +descriptor of an isolated atom of every type and ``c_i`` the conditioning of +the atom itself. An atom whose descriptor equals its vacuum descriptor thus +contributes exactly the bias of its type. +""" + +import itertools +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.pt.model.descriptor.sezm_nn.dens import ( + SeZMDeNSFittingNet, +) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.model.task.invar_fitting import ( + InvarFitting, +) +from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +from ...seed import ( + GLOBAL_SEED, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +NTYPES, ND, NF, NLOC = 3, 8, 2, 5 +CONDITIONING = [(0, 0), (2, 0), (0, 1), (2, 1)] + + +class TestVacuumRefModelSupport(unittest.TestCase): + def test_standard_model_rejects_vacuum_ref(self) -> None: + """Unsupported atomic models reject the option before their first forward.""" + config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "rcut": 4.0, + "rcut_smth": 3.5, + "sel": [4, 4], + "neuron": [4, 8], + "axis_neuron": 2, + "seed": GLOBAL_SEED, + }, + "fitting_net": {"type": "ener", "neuron": [8], "seed": GLOBAL_SEED}, + } + serialized = get_model(config).serialize() + config["fitting_net"]["vacuum_ref"] = True + serialized["fitting"]["vacuum_ref"] = True + for preset in (None, {"energy": [[-3.0], [0.5]]}): + config["preset_out_bias"] = preset + serialized["preset_out_bias"] = preset + with ( + self.subTest(preset=preset, entry="constructor"), + self.assertRaisesRegex(NotImplementedError, "DPA4/SeZM"), + ): + get_model(config) + with ( + self.subTest(preset=preset, entry="deserialize"), + self.assertRaisesRegex(NotImplementedError, "DPA4/SeZM"), + ): + BaseModel.deserialize(serialized) + + +class VacuumRefInputs(unittest.TestCase): + """Random descriptors, a random vacuum table and random conditioning.""" + + def setUp(self) -> None: + self.rng = np.random.default_rng(GLOBAL_SEED) + self.descriptor = self.tensor(self.rng.normal(size=(NF, NLOC, ND))) + self.vacuum = self.tensor(self.rng.normal(size=(NTYPES, ND))) + atype = self.rng.integers(0, NTYPES, size=(NF, NLOC)) + atype[0, :NTYPES] = np.arange(NTYPES) + self.atype = torch.tensor(atype, dtype=torch.long, device=env.DEVICE) + self.bias = self.rng.normal(size=(NTYPES, 1)) + + def tensor(self, array: np.ndarray) -> torch.Tensor: + return torch.tensor(array, dtype=dtype, device=env.DEVICE) + + def conditioning( + self, nfp: int, nap: int + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + fparam = self.tensor(self.rng.normal(size=(NF, nfp))) if nfp else None + aparam = self.tensor(self.rng.normal(size=(NF, NLOC, nap))) if nap else None + return fparam, aparam + + def expected_bias(self) -> np.ndarray: + return self.bias[to_numpy_array(self.atype)] + + def assert_reference_subtraction( + self, + ft_ref: torch.nn.Module, + ft_vac: torch.nn.Module, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + ) -> None: + """``ft_vac`` equals ``ft_ref`` minus ``ft_ref`` on the vacuum rows plus the bias.""" + out = ft_vac( + self.descriptor, + self.atype, + fparam=fparam, + aparam=aparam, + vacuum_descriptor=self.vacuum, + )["energy"] + ref = ft_ref(self.descriptor, self.atype, fparam=fparam, aparam=aparam) + ref_vac = ft_ref( + self.vacuum[self.atype], self.atype, fparam=fparam, aparam=aparam + ) + expected = ( + to_numpy_array(ref["energy"]) + - to_numpy_array(ref_vac["energy"]) + + self.expected_bias() + ) + np.testing.assert_allclose( + to_numpy_array(out), expected, rtol=1e-10, atol=1e-10 + ) + + +class TestInvarFittingVacuumRef(VacuumRefInputs): + def build(self, vacuum_ref: bool, **kwargs) -> InvarFitting: + ft = InvarFitting( + "energy", + NTYPES, + ND, + 1, + neuron=[6, 6], + bias_atom_e=self.bias, + vacuum_ref=vacuum_ref, + seed=GLOBAL_SEED, + **kwargs, + ).to(env.DEVICE) + if ft.dim_case_embd > 0: + ft.set_case_embd(1) + return ft + + def test_isolated_atom_gives_bias(self) -> None: + for mixed_types, (nfp, nap), ncase, mask in itertools.product( + [True, False], CONDITIONING, [0, 2], [False, True] + ): + ft = self.build( + True, + mixed_types=mixed_types, + numb_fparam=nfp, + numb_aparam=nap, + dim_case_embd=ncase, + use_aparam_as_mask=mask, + ) + fparam, aparam = self.conditioning(nfp, nap) + out = ft( + self.vacuum[self.atype], + self.atype, + fparam=fparam, + aparam=aparam, + vacuum_descriptor=self.vacuum, + )["energy"] + np.testing.assert_allclose( + to_numpy_array(out), self.expected_bias(), rtol=1e-10, atol=1e-10 + ) + + def test_matches_reference_subtraction(self) -> None: + for mixed_types, (nfp, nap), ncase in itertools.product( + [True, False], CONDITIONING, [0, 2] + ): + ft_ref = self.build( + False, + mixed_types=mixed_types, + numb_fparam=nfp, + numb_aparam=nap, + dim_case_embd=ncase, + ) + ft_vac = InvarFitting.deserialize( + {**ft_ref.serialize(), "vacuum_ref": True} + ).to(env.DEVICE) + self.assertTrue(ft_vac.vacuum_ref) + self.assert_reference_subtraction( + ft_ref, ft_vac, *self.conditioning(nfp, nap) + ) + + def test_aparam_layout_and_atom_count(self) -> None: + """Flat and per-atom layouts agree; atomic parameters match descriptor rows.""" + _, aparam = self.conditioning(0, 2) + for vacuum_ref in (False, True): + with self.subTest(vacuum_ref=vacuum_ref): + fitting = self.build(vacuum_ref, numb_aparam=2) + kwargs = {"vacuum_descriptor": self.vacuum} + expected = fitting(self.descriptor, self.atype, aparam=aparam, **kwargs) + actual = fitting( + self.descriptor, + self.atype, + aparam=aparam.reshape(NF, NLOC * 2), + **kwargs, + ) + torch.testing.assert_close(actual["energy"], expected["energy"]) + with self.assertRaisesRegex(ValueError, "input aparam"): + fitting( + self.descriptor, self.atype, aparam=aparam[:, :-1], **kwargs + ) + + def test_dpmodel_consistency(self) -> None: + for mixed_types, (nfp, nap), ncase in itertools.product( + [True, False], CONDITIONING, [0, 2] + ): + ft = self.build( + True, + mixed_types=mixed_types, + numb_fparam=nfp, + numb_aparam=nap, + dim_case_embd=ncase, + ) + ft_dp = DPInvarFitting.deserialize(ft.serialize()) + fparam, aparam = self.conditioning(nfp, nap) + out = ft( + self.descriptor, + self.atype, + fparam=fparam, + aparam=aparam, + vacuum_descriptor=self.vacuum, + )["energy"] + out_dp = ft_dp( + to_numpy_array(self.descriptor), + to_numpy_array(self.atype), + fparam=to_numpy_array(fparam), + aparam=to_numpy_array(aparam), + vacuum_descriptor=to_numpy_array(self.vacuum), + )["energy"] + np.testing.assert_allclose( + to_numpy_array(out), out_dp, rtol=1e-12, atol=1e-12 + ) + + def test_fold_vacuum_reference(self) -> None: + for mixed_types, ncase in itertools.product([True, False], [0, 2]): + ft = self.build(True, mixed_types=mixed_types, dim_case_embd=ncase) + expected = ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum) + ft.fold_vacuum_reference(self.vacuum) + self.assertFalse(ft.vacuum_ref) + out = ft(self.descriptor, self.atype)["energy"] + np.testing.assert_allclose( + to_numpy_array(out), + to_numpy_array(expected["energy"]), + rtol=1e-10, + atol=1e-10, + ) + # a conditioned fitting stores the table and references from it; the + # table is a deployment constant that checkpoints leave out and a + # type-map change drops + type_map = ["O", "H", "B"] + ft = self.build(True, numb_fparam=2, type_map=type_map) + fparam = self.tensor(self.rng.normal(size=(NF, 2))) + expected = ft( + self.descriptor, self.atype, fparam=fparam, vacuum_descriptor=self.vacuum + ) + ft.fold_vacuum_reference(self.vacuum) + self.assertTrue(ft.vacuum_ref) + self.assertFalse(ft.needs_vacuum_descriptor()) + out = ft(self.descriptor, self.atype, fparam=fparam)["energy"] + np.testing.assert_allclose( + to_numpy_array(out), + to_numpy_array(expected["energy"]), + rtol=1e-10, + atol=1e-10, + ) + state = ft.state_dict() + self.assertNotIn("vacuum_table", state) + fresh = self.build(True, numb_fparam=2, type_map=type_map) + fresh.load_state_dict(state) + self.assertTrue(fresh.needs_vacuum_descriptor()) + ft.change_type_map(["B", "O", "H"]) + self.assertTrue(ft.needs_vacuum_descriptor()) + + def test_default_fparam(self) -> None: + ft = self.build(True, numb_fparam=1, default_fparam=[0.3]) + explicit = ft( + self.descriptor, + self.atype, + fparam=self.tensor(np.full((NF, 1), 0.3)), + vacuum_descriptor=self.vacuum, + ) + default = ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum) + np.testing.assert_allclose( + to_numpy_array(default["energy"]), + to_numpy_array(explicit["energy"]), + rtol=1e-10, + atol=1e-10, + ) + + def test_jit(self) -> None: + for mixed_types, (nfp, nap) in itertools.product([True, False], CONDITIONING): + ft = self.build( + True, mixed_types=mixed_types, numb_fparam=nfp, numb_aparam=nap + ) + torch.jit.script(ft) + + def test_vacuum_descriptor_required(self) -> None: + ft = self.build(True) + with self.assertRaises(ValueError): + ft(self.descriptor, self.atype) + with self.assertRaises(ValueError): + ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum[:, :-1]) + # the table is ignored without the option + ft_ref = self.build(False) + out = ft_ref(self.descriptor, self.atype, vacuum_descriptor=self.vacuum) + np.testing.assert_allclose( + to_numpy_array(out["energy"]), + to_numpy_array(ft_ref(self.descriptor, self.atype)["energy"]), + ) + + def test_serialization(self) -> None: + data = self.build(True).serialize() + self.assertTrue(data["vacuum_ref"]) + self.assertTrue(InvarFitting.deserialize(data).vacuum_ref) + self.assertFalse( + InvarFitting.deserialize({**data, "vacuum_ref": False}).vacuum_ref + ) + # a dictionary of the previous version carries no key + older = {k: v for k, v in data.items() if k != "vacuum_ref"} + self.assertFalse(InvarFitting.deserialize({**older, "@version": 4}).vacuum_ref) + + def test_atom_ener_is_exclusive(self) -> None: + self.assertTrue(self.build(True, atom_ener=[None] * NTYPES).vacuum_ref) + with self.assertRaises(ValueError): + self.build(True, atom_ener=[1.0] + [None] * (NTYPES - 1)) + + +class TestSeZMFittingVacuumRef(VacuumRefInputs): + def build(self, vacuum_ref: bool, case_film_embd: bool) -> SeZMEnergyFittingNet: + ft = SeZMEnergyFittingNet( + NTYPES, + ND, + neuron=[16], + bias_atom_e=self.bias, + numb_fparam=2, + dim_case_embd=2, + case_film_embd=case_film_embd, + precision="float64", + vacuum_ref=vacuum_ref, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + ft.set_case_embd(1) + return ft + + def test_isolated_atom_gives_bias(self) -> None: + for case_film_embd in [False, True]: + ft = self.build(True, case_film_embd) + fparam, _ = self.conditioning(2, 0) + out = ft( + self.vacuum[self.atype], + self.atype, + fparam=fparam, + vacuum_descriptor=self.vacuum, + )["energy"] + np.testing.assert_allclose( + to_numpy_array(out), self.expected_bias(), rtol=1e-10, atol=1e-10 + ) + + def test_matches_reference_subtraction(self) -> None: + for case_film_embd in [False, True]: + ft_ref = self.build(False, case_film_embd) + ft_vac = SeZMEnergyFittingNet.deserialize( + {**ft_ref.serialize(), "vacuum_ref": True} + ).to(env.DEVICE) + self.assertTrue(ft_vac.vacuum_ref) + self.assert_reference_subtraction(ft_ref, ft_vac, *self.conditioning(2, 0)) + + def test_atom_ener_subtracts_the_zero_descriptor_output(self) -> None: + """``atom_ener`` removes the output of a zero descriptor, also under case FiLM.""" + for case_film_embd in [False, True]: + ft_ref = self.build(False, case_film_embd) + ft = SeZMEnergyFittingNet.deserialize( + {**ft_ref.serialize(), "atom_ener": [1.0] * NTYPES} + ).to(env.DEVICE) + fparam, _ = self.conditioning(2, 0) + zeros = torch.zeros_like(self.descriptor) + # an atom with a zero descriptor contributes exactly its bias + out_zero = ft(zeros, self.atype, fparam=fparam)["energy"] + np.testing.assert_allclose( + to_numpy_array(out_zero), self.expected_bias(), rtol=1e-10, atol=1e-10 + ) + out = ft(self.descriptor, self.atype, fparam=fparam)["energy"] + ref = ft_ref(self.descriptor, self.atype, fparam=fparam)["energy"] + ref_zero = ft_ref(zeros, self.atype, fparam=fparam)["energy"] + expected = ( + to_numpy_array(ref) - to_numpy_array(ref_zero) + self.expected_bias() + ) + np.testing.assert_allclose( + to_numpy_array(out), expected, rtol=1e-10, atol=1e-10 + ) + + def test_fold_vacuum_reference(self) -> None: + for case_film_embd in [False, True]: + ft = SeZMEnergyFittingNet( + NTYPES, + ND, + neuron=[16], + bias_atom_e=self.bias, + dim_case_embd=2, + case_film_embd=case_film_embd, + precision="float64", + vacuum_ref=True, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + ft.set_case_embd(1) + expected = ft(self.descriptor, self.atype, vacuum_descriptor=self.vacuum) + ft.fold_vacuum_reference(self.vacuum) + self.assertFalse(ft.vacuum_ref) + out = ft(self.descriptor, self.atype)["energy"] + np.testing.assert_allclose( + to_numpy_array(out), + to_numpy_array(expected["energy"]), + rtol=1e-10, + atol=1e-10, + ) + + +class TestDeNSFittingVacuumRef(VacuumRefInputs): + def test_energy_head_is_referenced(self) -> None: + channels, lmax = 4, 1 + ft = SeZMDeNSFittingNet( + ntypes=NTYPES, + dim_descrpt=ND, + condition_lmax=lmax, + latent_lmax=lmax, + channels=channels, + neuron=[8], + bias_atom_e=self.bias, + precision="float64", + vacuum_ref=True, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + self.assertTrue(ft.energy_head.vacuum_ref) + latent = self.tensor( + self.rng.normal(size=(NF * NLOC, (lmax + 1) ** 2, 1, channels)) + ) + out = ft( + self.vacuum[self.atype], + latent, + self.atype, + vacuum_descriptor=self.vacuum, + )["energy"] + np.testing.assert_allclose( + to_numpy_array(out), self.expected_bias(), rtol=1e-10, atol=1e-10 + ) + data = ft.serialize() + self.assertTrue(data["config"]["vacuum_ref"]) + self.assertTrue(SeZMDeNSFittingNet.deserialize(data).vacuum_ref) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py index 2db6370059..7c8a0b1a98 100644 --- a/source/tests/pt/model/test_get_model.py +++ b/source/tests/pt/model/test_get_model.py @@ -55,8 +55,8 @@ def test_model_attr(self) -> None: { "energy": [ None, - np.array([1.0]), - np.array([3.0]), + [1.0], + [3.0], ] }, ) @@ -79,8 +79,8 @@ def test_model_attr_energy_float(self) -> None: atomic_model.preset_out_bias, { "energy": [ - np.array([1.0]), - np.array([3.0]), + [1.0], + [3.0], None, ] }, @@ -88,6 +88,31 @@ def test_model_attr_energy_float(self) -> None: self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) + def test_model_attr_energy_element_dict(self) -> None: + model_params = copy.deepcopy(model_se_e2_a) + model_params["preset_out_bias"] = {"energy": {"B": 3.0, "H": [1.0]}} + self.model = get_model(model_params).to(env.DEVICE) + atomic_model = self.model.atomic_model + self.assertEqual(atomic_model.type_map, ["O", "H", "B"]) + self.assertEqual( + atomic_model.preset_out_bias, + { + "energy": [ + None, + [1.0], + [3.0], + ] + }, + ) + + def test_model_attr_energy_unknown_element_ignored(self) -> None: + model_params = copy.deepcopy(model_se_e2_a) + model_params["preset_out_bias"] = {"energy": {"C": 3.0, "H": 1.0}} + self.model = get_model(model_params).to(env.DEVICE) + self.assertEqual( + self.model.atomic_model.preset_out_bias, {"energy": [None, [1.0], None]} + ) + def test_model_attr_energy_unsupported_type(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": [1.0 + 2.0j, 3, None]} diff --git a/source/tests/pt/model/test_sezm_parallel.py b/source/tests/pt/model/test_sezm_parallel.py index be3b093fc1..bceec317d8 100644 --- a/source/tests/pt/model/test_sezm_parallel.py +++ b/source/tests/pt/model/test_sezm_parallel.py @@ -336,14 +336,15 @@ def test_descriptor_parity_cpu(self) -> None: sysm = _build_extended_system(model, device) comm = _self_comm_dict(sysm["mapping"], sysm["nloc"], sysm["nall"]) - ref, _ = descriptor.forward_with_edges( + ref, _, ref_vacuum = descriptor.forward_with_edges( extended_coord=sysm["coord"][:, : sysm["nloc"], :], extended_atype=sysm["atype"], edge_index=sysm["edge_index"], edge_vec=sysm["edge_vec"], edge_mask=sysm["edge_mask"], + vacuum_conditions={}, ) - par, _ = descriptor.forward_with_edges( + par, _, par_vacuum = descriptor.forward_with_edges( extended_coord=sysm["coord"], extended_atype=sysm["extended_atype"], edge_index=sysm["edge_scatter_index"], @@ -351,8 +352,15 @@ def test_descriptor_parity_cpu(self) -> None: edge_mask=sysm["edge_mask"], comm_dict=comm, nloc=sysm["nloc"], + vacuum_conditions={}, ) torch.testing.assert_close(par, ref, rtol=1e-8, atol=1e-9) + # the vacuum reference rows trail the ghost rows and stay untouched + # by the border exchange + self.assertEqual( + tuple(par_vacuum.shape), (descriptor.ntypes, descriptor.channels) + ) + torch.testing.assert_close(par_vacuum, ref_vacuum, rtol=1e-8, atol=1e-9) class TestSeZMNativeSpinParallelParity(unittest.TestCase): diff --git a/source/tests/pt/model/test_sezm_vacuum_freeze.py b/source/tests/pt/model/test_sezm_vacuum_freeze.py new file mode 100644 index 0000000000..4a2990aa9a --- /dev/null +++ b/source/tests/pt/model/test_sezm_vacuum_freeze.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Freezing a SeZM model with the isolated-atom reference. + +The ``.pt2`` freeze folds the vacuum reference into the fitting bias, so the +frozen model carries no reference atoms yet reproduces the referenced eager +model: an isolated atom gives exactly its bias and a cluster gives the eager +energies. +""" + +import copy +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np +import torch + +from deepmd.pt.entrypoints.freeze_pt2 import ( + freeze_sezm_to_pt2, +) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.train.wrapper import ( + ModelWrapper, +) + +from .test_sezm_export import ( + _CPU, + _SKIP_OFF_COMPILE_TORCH, + _SKIP_OFF_COMPILE_TORCH_REASON, + _clear_default_device, + _tiny_sezm_model_params, +) + +BIAS = np.array([[-3.0], [0.5]]) + + +@unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) +class TestSeZMVacuumFreeze(unittest.TestCase): + def test_frozen_model_folds_the_reference(self) -> None: + self.check_frozen_model(numb_fparam=0) + + def test_frozen_model_keeps_the_reference_with_fparam(self) -> None: + """With frame parameters the archive references from the stored vacuum table.""" + self.check_frozen_model(numb_fparam=1) + + @unittest.skipIf(not torch.cuda.is_available(), "the CUDA target needs a GPU") + def test_frozen_model_on_a_cuda_target(self) -> None: + """The fold and the compiled archive share the CUDA target.""" + self.check_frozen_model(numb_fparam=1, device=torch.device("cuda")) + + def test_dens_checkpoint_is_rejected(self) -> None: + """The DeNS head serves training alone, so a ``dens`` checkpoint is not frozen.""" + params = _tiny_sezm_model_params() + # the DeNS vector heads need an l=1 latent + params["descriptor"]["l_schedule"] = [1, 1] + params["fitting_net"]["vacuum_ref"] = True + model = get_model(params) + model.set_active_mode("dens") + with tempfile.TemporaryDirectory() as tmp, _clear_default_device(): + wrapper = ModelWrapper(model, model_params=copy.deepcopy(params)) + ckpt = Path(tmp) / "dens.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + with self.assertRaisesRegex(ValueError, "`ener` mode"): + freeze_sezm_to_pt2(str(ckpt), str(Path(tmp) / "dens.pt2")) + + def check_frozen_model(self, numb_fparam: int, device: torch.device = _CPU) -> None: + params = _tiny_sezm_model_params() + params["preset_out_bias"] = {"energy": {"A": BIAS[0, 0], "B": BIAS[1, 0]}} + params["fitting_net"]["vacuum_ref"] = True + params["fitting_net"]["numb_fparam"] = numb_fparam + fparam = None if numb_fparam == 0 else np.array([[0.7]]) + model = get_model(params) + model.eval() + model.to(device) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor(BIAS, dtype=fitting.bias_atom_e.dtype, device=device) + ) + self.assertTrue(fitting.vacuum_ref) + + box_edge = params["descriptor"]["rcut"] * 3.0 + cell = (np.eye(3) * box_edge).reshape(1, 9) + rng = np.random.default_rng(2026) + natoms = 5 + atype = np.array([0, 1, 0, 1, 0], dtype=np.int32) + coord = rng.random((1, natoms, 3)) * box_edge * 0.4 + box_edge * 0.3 + eager = ( + model.forward( + torch.tensor(coord, dtype=torch.float64, device=device), + torch.tensor(atype, dtype=torch.int64, device=device).unsqueeze(0), + torch.tensor(cell, dtype=torch.float64, device=device), + fparam=None + if fparam is None + else torch.tensor(fparam, dtype=torch.float64, device=device), + )["atom_energy"] + .detach() + .cpu() + .numpy() + ) + + import deepmd.pt_expt.utils.env as pt_expt_env + from deepmd.infer import ( + DeepPot, + ) + + with tempfile.TemporaryDirectory() as tmp, _clear_default_device(): + wrapper = ModelWrapper(model, model_params=copy.deepcopy(params)) + ckpt = Path(tmp) / "vacuum.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + out = Path(tmp) / "vacuum.pt2" + freeze_sezm_to_pt2(str(ckpt), str(out), device=device) + # the evaluator runs on the device the archive is compiled for + saved_device = pt_expt_env.DEVICE + pt_expt_env.DEVICE = device + try: + dp = DeepPot(str(out)) + for itype in range(2): + _, _, _, atom_energy, _ = dp.eval( + np.full((1, 1, 3), box_edge / 2.0), + cell, + np.array([itype], dtype=np.int32), + atomic=True, + fparam=fparam, + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8 + ) + _, _, _, atom_energy, _ = dp.eval( + coord, cell, atype, atomic=True, fparam=fparam + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), eager.reshape(-1), rtol=1e-8, atol=1e-8 + ) + finally: + pt_expt_env.DEVICE = saved_device + # the training model itself keeps its reference + self.assertTrue(fitting.vacuum_ref) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_sezm_vacuum_ref.py b/source/tests/pt/model/test_sezm_vacuum_ref.py new file mode 100644 index 0000000000..82e9781d46 --- /dev/null +++ b/source/tests/pt/model/test_sezm_vacuum_ref.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Isolated-atom reference of a SeZM model on the PyTorch edge route. + +The SeZM descriptor carries one reference atom per type through the same +forward as the real atoms, so an isolated atom contributes exactly its bias +and any cluster differs from the unreferenced model by the isolated-atom +network output of its atoms. +""" + +import copy +import unittest + +import numpy as np +import torch + +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils import ( + env, +) + +TYPE_MAP = ["O", "H"] +BIAS = np.array([[-3.0], [0.5]]) +PRESET = {"energy": {"O": float(BIAS[0, 0]), "H": float(BIAS[1, 0])}} +MODEL_PARAMS = { + "type": "SeZM", + "type_map": TYPE_MAP, + "preset_out_bias": PRESET, + "descriptor": { + "type": "SeZM", + "sel": [4, 4], + "rcut": 4.0, + "channels": 8, + "n_focus": 1, + "n_radial": 4, + "radial_mlp": [8], + "use_env_seed": True, + "l_schedule": [1, 0], + "mmax": 1, + "so2_norm": False, + "so2_layers": 1, + "n_atten_head": 1, + "sandwich_norm": [True, False, True, False], + "ffn_neurons": 8, + "ffn_blocks": 1, + "s2_activation": [False, True], + "mlp_bias": False, + "layer_scale": False, + "use_amp": False, + "activation_function": "silu", + "glu_activation": True, + "precision": "float64", + "seed": 7, + }, + "fitting_net": { + "neuron": [8], + "activation_function": "silu", + "precision": "float64", + "seed": 7, + }, +} + + +class TestSeZMVacuumRef(unittest.TestCase): + def make_model(self, vacuum_ref: bool, preset: bool = True) -> torch.nn.Module: + params = copy.deepcopy(MODEL_PARAMS) + if not preset: + params.pop("preset_out_bias") + params["fitting_net"]["vacuum_ref"] = vacuum_ref + model = get_model(params).to(env.DEVICE) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor( + BIAS, + dtype=fitting.bias_atom_e.dtype, + device=fitting.bias_atom_e.device, + ) + ) + return model.eval() + + def predict( + self, model: torch.nn.Module, coord: np.ndarray, atype: np.ndarray + ) -> dict[str, torch.Tensor]: + return model( + torch.as_tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.as_tensor(atype, dtype=torch.long, device=env.DEVICE), + None, + ) + + def atom_energies( + self, model: torch.nn.Module, coord: np.ndarray, atype: np.ndarray + ) -> np.ndarray: + return ( + self.predict(model, coord, atype)["atom_energy"][..., 0] + .detach() + .cpu() + .numpy() + ) + + def test_isolated_atoms_and_reference_subtraction(self) -> None: + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + vac = self.make_model(True) + ref = self.make_model(False) + + # an isolated atom of every type contributes exactly its bias + np.testing.assert_allclose( + self.atom_energies(vac, iso_coord, iso_atype), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the reference removes the isolated-atom network output of every atom + e_vac = self.atom_energies(vac, coord, atype) + e_ref = self.atom_energies(ref, coord, atype) + iso_ref = self.atom_energies(ref, iso_coord, iso_atype)[:, 0] + expected = e_ref - (iso_ref - BIAS[:, 0])[atype] + np.testing.assert_allclose(e_vac, expected, rtol=1e-10, atol=1e-10) + + def test_without_preset_the_output_is_not_referenced(self) -> None: + """A bias fitted from the data is no isolated-atom energy, so the output stays plain.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + plain = self.make_model(False) + unreferenced = self.make_model(True, preset=False) + fitting = unreferenced.atomic_model.fitting_net + self.assertFalse(fitting.vacuum_ref) + self.assertFalse(fitting.needs_vacuum_descriptor()) + np.testing.assert_allclose( + self.atom_energies(unreferenced, coord, atype), + self.atom_energies(plain, coord, atype), + rtol=1e-12, + atol=1e-12, + ) + + def test_shared_fitting_references_only_the_branch_with_a_preset(self) -> None: + """Branches sharing one fitting reference their output only where a preset fixes the bias.""" + from deepmd.pt.train.training import ( + get_model_for_wrapper, + prepare_model_for_loss, + ) + from deepmd.pt.train.wrapper import ( + ModelWrapper, + ) + from deepmd.pt.utils.multi_task import ( + preprocess_shared_params, + ) + + branch = { + "type": "SeZM", + "type_map": "type_map", + "descriptor": "descriptor", + "fitting_net": "fitting", + } + config = { + "shared_dict": { + "type_map": TYPE_MAP, + "descriptor": copy.deepcopy(MODEL_PARAMS["descriptor"]), + "fitting": { + **MODEL_PARAMS["fitting_net"], + "vacuum_ref": True, + "dim_case_embd": 2, + }, + }, + "model_dict": { + "with_table": {**branch, "preset_out_bias": PRESET}, + "without_table": dict(branch), + }, + } + config, shared_links = preprocess_shared_params(config) + models = get_model_for_wrapper(config) + prepare_model_for_loss(models, {key: {"type": "ener"} for key in models}) + wrapper = ModelWrapper(models) + wrapper.share_params(shared_links, dict.fromkeys(models, 0.5)) + referenced = wrapper.model["with_table"].to(env.DEVICE).eval() + unreferenced = wrapper.model["without_table"].to(env.DEVICE).eval() + fit_a = referenced.atomic_model.fitting_net + fit_b = unreferenced.atomic_model.fitting_net + # one network, one decision per branch + self.assertIs(fit_a.filter_layers, fit_b.filter_layers) + self.assertTrue(fit_a.vacuum_ref) + self.assertFalse(fit_b.vacuum_ref) + for fitting in (fit_a, fit_b): + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor( + BIAS, + dtype=fitting.bias_atom_e.dtype, + device=fitting.bias_atom_e.device, + ) + ) + + rng = np.random.default_rng(5) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + # the branch with a table pins its isolated atoms to the preset + np.testing.assert_allclose( + self.atom_energies(referenced, iso_coord, iso_atype), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the branch without a table is the plain model on the shared weights + params = copy.deepcopy(MODEL_PARAMS) + params.pop("preset_out_bias") + params["fitting_net"]["dim_case_embd"] = 2 + plain = get_model(params).to(env.DEVICE) + plain.load_state_dict(unreferenced.state_dict()) + plain.eval() + np.testing.assert_allclose( + self.atom_energies(unreferenced, coord, atype), + self.atom_energies(plain, coord, atype), + rtol=1e-12, + atol=1e-12, + ) + iso_energy = self.atom_energies(unreferenced, iso_coord, iso_atype) + self.assertGreater(np.abs(iso_energy - BIAS[:, 0][:, None]).max(), 1e-6) + + def test_forces_are_unchanged(self) -> None: + """The reference is independent of the coordinates, so forces do not move.""" + rng = np.random.default_rng(3) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + force_vac = self.predict(self.make_model(True), coord, atype)["force"] + force_ref = self.predict(self.make_model(False), coord, atype)["force"] + np.testing.assert_allclose( + force_vac.detach().cpu().numpy(), + force_ref.detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + ) + + def test_fold_reproduces_the_referenced_model(self) -> None: + """Folding the reference into the bias keeps every energy and drops the rows.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + model = self.make_model(True) + e_cluster = self.atom_energies(model, coord, atype) + e_iso = self.atom_energies(model, iso_coord, iso_atype) + + model.fold_vacuum_reference() + self.assertFalse(model.atomic_model.fitting_net.vacuum_ref) + np.testing.assert_allclose( + self.atom_energies(model, coord, atype), e_cluster, rtol=1e-10, atol=1e-10 + ) + np.testing.assert_allclose( + self.atom_energies(model, iso_coord, iso_atype), + e_iso, + rtol=1e-10, + atol=1e-10, + ) + + def test_fold_drops_the_compiled_graphs(self) -> None: + from deepmd.pt.model.model.sezm_model import ( + _sezm_structure_key, + ) + + model = self.make_model(True) + key_live = _sezm_structure_key(model) + model.compiled_core_compute_cache[(False, False)] = object() + model.fold_vacuum_reference() + self.assertEqual(model.compiled_core_compute_cache, {}) + # a folded model traces a graph without reference nodes + self.assertNotEqual(_sezm_structure_key(model), key_live) + + +class TestSeZMNativeSpinVacuumRef(unittest.TestCase): + """The reference atom carries the ground-state charge/spin condition and spin.""" + + def make_model(self, vacuum_ref: bool) -> torch.nn.Module: + params = copy.deepcopy(MODEL_PARAMS) + params["descriptor"]["add_chg_spin_ebd"] = True + params["descriptor"]["default_chg_spin"] = [0, 1] + params["fitting_net"]["vacuum_ref"] = vacuum_ref + params["spin"] = {"use_spin": [True, False], "scheme": "native"} + model = get_model(params).to(env.DEVICE) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor( + BIAS, + dtype=fitting.bias_atom_e.dtype, + device=fitting.bias_atom_e.device, + ) + ) + return model.eval() + + def atom_energies( + self, + model: torch.nn.Module, + coord: np.ndarray, + atype: np.ndarray, + spin: np.ndarray, + charge_spin: np.ndarray, + ) -> np.ndarray: + def tensor( + array: np.ndarray, dtype: torch.dtype = torch.float64 + ) -> torch.Tensor: + return torch.as_tensor(array, dtype=dtype, device=env.DEVICE) + + ret = model( + tensor(coord), + tensor(atype, torch.long), + tensor(spin), + None, + charge_spin=tensor(charge_spin), + ) + return ret["atom_energy"][..., 0].detach().cpu().numpy() + + def test_isolated_atoms_and_reference_subtraction(self) -> None: + from deepmd.utils.vacuum_reference import ( + reference_charge_spin, + reference_spin, + ) + + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + spin = rng.normal(size=(2, 6, 3)) * (atype == 0)[..., None] + charge_spin = np.array([[0.0, 1.0], [1.0, 2.0]]) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + iso_spin = reference_spin(TYPE_MAP)[:, None, :] + iso_charge_spin = reference_charge_spin(TYPE_MAP) + vac = self.make_model(True) + ref = self.make_model(False) + + # an isolated neutral ground-state atom contributes exactly its bias + np.testing.assert_allclose( + self.atom_energies(vac, iso_coord, iso_atype, iso_spin, iso_charge_spin), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the reference removes the isolated-atom network output of every atom + e_vac = self.atom_energies(vac, coord, atype, spin, charge_spin) + e_ref = self.atom_energies(ref, coord, atype, spin, charge_spin) + iso_ref = self.atom_energies( + ref, iso_coord, iso_atype, iso_spin, iso_charge_spin + )[:, 0] + expected = e_ref - (iso_ref - BIAS[:, 0])[atype] + np.testing.assert_allclose(e_vac, expected, rtol=1e-10, atol=1e-10) + + # the fold evaluates the reference under the ground-state conditions + vac.fold_vacuum_reference() + self.assertFalse(vac.atomic_model.fitting_net.vacuum_ref) + np.testing.assert_allclose( + self.atom_energies(vac, coord, atype, spin, charge_spin), + e_vac, + rtol=1e-10, + atol=1e-10, + ) + + +class TestSeZMDeNSVacuumRef(unittest.TestCase): + """The DeNS energy head references every atom to the isolated atom without force input.""" + + def make_model(self, vacuum_ref: bool) -> torch.nn.Module: + params = copy.deepcopy(MODEL_PARAMS) + # the DeNS vector heads need an l=1 latent + params["descriptor"]["l_schedule"] = [1, 1] + params["fitting_net"]["vacuum_ref"] = vacuum_ref + model = get_model(params).to(env.DEVICE) + model.set_active_mode("dens") + head = model.atomic_model.get_dens_fitting_net().energy_head + self.assertEqual(head.vacuum_ref, vacuum_ref) + with torch.no_grad(): + head.bias_atom_e.copy_( + torch.as_tensor( + BIAS, dtype=head.bias_atom_e.dtype, device=head.bias_atom_e.device + ) + ) + return model.eval() + + def test_without_preset_the_head_is_not_referenced(self) -> None: + """The DeNS energy head follows the preset of the branch like the energy fitting.""" + params = copy.deepcopy(MODEL_PARAMS) + params.pop("preset_out_bias") + params["descriptor"]["l_schedule"] = [1, 1] + params["fitting_net"]["vacuum_ref"] = True + model = get_model(params) + model.set_active_mode("dens") + dens = model.atomic_model.get_dens_fitting_net() + self.assertFalse(dens.vacuum_ref) + self.assertFalse(dens.energy_head.vacuum_ref) + self.assertFalse(dens.needs_vacuum_descriptor()) + self.assertFalse(model.atomic_model.fitting_net.vacuum_ref) + + def atom_energies( + self, + model: torch.nn.Module, + coord: np.ndarray, + atype: np.ndarray, + force: np.ndarray, + noise_mask: np.ndarray, + ) -> np.ndarray: + def tensor( + array: np.ndarray, dtype: torch.dtype = torch.float64 + ) -> torch.Tensor: + return torch.as_tensor(array, dtype=dtype, device=env.DEVICE) + + nf = coord.shape[0] + box = np.tile(np.eye(3).reshape(1, 9) * 20.0, (nf, 1)) + ret = model( + tensor(coord), + tensor(atype, torch.long), + box=tensor(box), + force_input=tensor(force), + noise_mask=tensor(noise_mask, torch.bool), + ) + return ret["atom_energy"][..., 0].detach().cpu().numpy() + + def test_isolated_atoms_and_reference_subtraction(self) -> None: + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + force = rng.normal(size=(2, 6, 3)) * 0.3 + noise_mask = rng.integers(0, 2, size=(2, 6)).astype(bool) + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + iso_force = np.zeros((2, 1, 3)) + iso_mask = np.zeros((2, 1), dtype=bool) + vac = self.make_model(True) + ref = self.make_model(False) + + # an isolated atom without force input contributes exactly its bias + np.testing.assert_allclose( + self.atom_energies(vac, iso_coord, iso_atype, iso_force, iso_mask), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the reference removes the isolated-atom network output of every atom + e_vac = self.atom_energies(vac, coord, atype, force, noise_mask) + e_ref = self.atom_energies(ref, coord, atype, force, noise_mask) + iso_ref = self.atom_energies(ref, iso_coord, iso_atype, iso_force, iso_mask)[ + :, 0 + ] + expected = e_ref - (iso_ref - BIAS[:, 0])[atype] + np.testing.assert_allclose(e_vac, expected, rtol=1e-10, atol=1e-10) + + +def test_isolated_atom_under_amp_and_fused_training_kernels(monkeypatch) -> None: + """The identity holds in the mixed-precision fused training forward. + + The reference atom of every type is carried through the same fused + kernels as the real atoms, so an isolated atom and its reference receive + the same descriptor up to the round-off of the float32 fitting. + """ + import pytest + + from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import ( + op_available as cuda_value_available, + ) + + if not torch.cuda.is_available() or not cuda_value_available(): + pytest.skip("the DPA4 CUDA training operators are unavailable") + monkeypatch.setenv("DP_CUDA_TRAIN", "1") + monkeypatch.setenv("DP_TRITON_TRAIN", "1") + params = { + "type": "dpa4", + "type_map": TYPE_MAP, + "preset_out_bias": PRESET, + "descriptor": { + "type": "dpa4", + "sel": 20, + "rcut": 4.0, + "channels": 32, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "use_amp": True, + "precision": "float32", + "seed": 7, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [16], + "precision": "float32", + "vacuum_ref": True, + "seed": 7, + }, + } + model = get_model(params).to(env.DEVICE) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor( + BIAS, dtype=fitting.bias_atom_e.dtype, device=fitting.bias_atom_e.device + ) + ) + model.train() + coord = torch.zeros((2, 1, 3), dtype=torch.float64, device=env.DEVICE) + atype = torch.tensor([[0], [1]], dtype=torch.long, device=env.DEVICE) + energy = model(coord, atype, None)["atom_energy"][..., 0].detach().cpu().numpy() + np.testing.assert_allclose(energy, BIAS[:, 0][:, None], rtol=0.0, atol=1e-5) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/fitting/test_dpa4_ener.py b/source/tests/pt_expt/fitting/test_dpa4_ener.py index 0c7469611e..bd1f46701a 100644 --- a/source/tests/pt_expt/fitting/test_dpa4_ener.py +++ b/source/tests/pt_expt/fitting/test_dpa4_ener.py @@ -107,3 +107,62 @@ def fn(descriptor, atype): np.testing.assert_allclose( grad_t.detach().cpu().numpy(), grad_e.detach().cpu().numpy() ) + + +class TestSeZMEnergyFittingNetVacuumRef(TestCaseSingleFrameWithNlist): + """``vacuum_ref`` on the torch-backed SeZM fitting matches dpmodel.""" + + def setup_method(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + self.device = env.DEVICE + rng = np.random.default_rng(GLOBAL_SEED) + self.descriptor = rng.normal(size=(self.nf, self.nloc, DIM_DESCRPT)) + self.vacuum = rng.normal(size=(self.nt, DIM_DESCRPT)) + self.atype = self.atype_ext[:, : self.nloc] + self.bias = rng.normal(size=(self.nt, 1)) + + @pytest.mark.parametrize("numb_fparam", [0, 2]) # per-type / per-atom references + @pytest.mark.parametrize("mixed_types", [True, False]) + def test_isolated_atom_gives_bias_and_matches_dpmodel( + self, numb_fparam, mixed_types + ) -> None: + ft = SeZMEnergyFittingNet( + self.nt, + DIM_DESCRPT, + neuron=[16], + bias_atom_e=self.bias, + numb_fparam=numb_fparam, + mixed_types=mixed_types, + precision="float64", + vacuum_ref=True, + seed=GLOBAL_SEED, + ).to(self.device) + ft_dp = SeZMEnergyFittingNetDP.deserialize(ft.serialize()) + assert ft_dp.vacuum_ref + rng = np.random.default_rng(GLOBAL_SEED + 1) + fparam = rng.normal(size=(self.nf, numb_fparam)) if numb_fparam else None + fparam_t = None if fparam is None else torch.from_numpy(fparam).to(self.device) + vacuum_t = torch.from_numpy(self.vacuum).to(self.device) + atype_t = torch.from_numpy(self.atype).to(self.device) + + isolated = ft( + vacuum_t[atype_t], atype_t, fparam=fparam_t, vacuum_descriptor=vacuum_t + )["energy"] + np.testing.assert_allclose( + isolated.detach().cpu().numpy(), + self.bias[self.atype], + rtol=1e-12, + atol=1e-12, + ) + out = ft( + torch.from_numpy(self.descriptor).to(self.device), + atype_t, + fparam=fparam_t, + vacuum_descriptor=vacuum_t, + )["energy"] + out_dp = ft_dp.call( + self.descriptor, self.atype, fparam=fparam, vacuum_descriptor=self.vacuum + )["energy"] + np.testing.assert_allclose( + out.detach().cpu().numpy(), out_dp, rtol=1e-12, atol=1e-14 + ) diff --git a/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py new file mode 100644 index 0000000000..e5524f7ec2 --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa4_vacuum_ref.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Isolated-atom reference of a DPA4 model on the pt_expt graph route. + +The reference node of every type is conditioned as the neutral ground-state +atom: zero charge with the ground-state multiplicity for the charge/spin +conditioning and a spin vector of one Bohr magneton per unpaired electron for +the native spin. An isolated atom under these conditions contributes exactly +its bias, and any cluster differs from the unreferenced model by the +isolated-atom network output of its atoms. +""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, +) +from deepmd.pt_expt.model import ( + EnergyModel, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.utils.vacuum_reference import ( + reference_charge_spin, + reference_spin, +) + +from ...dpa4_fixtures import ( + jitter_zero_arrays, +) + +TYPE_MAP = ["O", "H"] +BIAS = np.array([[-3.0], [0.5]]) +CONFIG = { + "type": "dpa4", + "type_map": TYPE_MAP, + "descriptor": { + "type": "dpa4", + "sel": 20, + "rcut": 4.0, + "channels": 16, + "n_radial": 8, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "precision": "float64", + "seed": 1, + "use_spin": [True, False], + "add_chg_spin_ebd": True, + "default_chg_spin": [0, 1], + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [16], + "precision": "float64", + "seed": 1, + }, +} + + +PRESET = {"energy": {"O": float(BIAS[0, 0]), "H": float(BIAS[1, 0])}} + + +def make_model(vacuum_ref: bool, preset: bool = True) -> EnergyModel: + config = { + **CONFIG, + "fitting_net": {**CONFIG["fitting_net"], "vacuum_ref": vacuum_ref}, + } + if preset: + config["preset_out_bias"] = PRESET + model = get_model(config) + data = jitter_zero_arrays( + model.atomic_model.descriptor.serialize(), np.random.default_rng(3) + ) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor(BIAS, dtype=fitting.bias_atom_e.dtype) + ) + return model.to(env.DEVICE).eval() + + +def atom_energies( + model: EnergyModel, + coord: np.ndarray, + atype: np.ndarray, + charge_spin: np.ndarray, + spin: np.ndarray, +) -> np.ndarray: + def tensor(array: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.as_tensor(array, dtype=dtype, device=env.DEVICE) + + ret = model.call_common( + tensor(coord), + tensor(atype, torch.long), + None, + charge_spin=tensor(charge_spin), + spin=tensor(spin), + ) + return ret["energy"][..., 0].detach().cpu().numpy() + + +def test_isolated_atoms_and_reference_subtraction() -> None: + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + charge_spin = np.array([[0.0, 1.0], [1.0, 2.0]]) + spin = rng.normal(size=(2, 6, 3)) * (atype == 0)[..., None] + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + iso_charge_spin = reference_charge_spin(TYPE_MAP) + iso_spin = reference_spin(TYPE_MAP)[:, None, :] + vac = make_model(True) + ref = make_model(False) + + # an isolated neutral ground-state atom contributes exactly its bias + np.testing.assert_allclose( + atom_energies(vac, iso_coord, iso_atype, iso_charge_spin, iso_spin), + BIAS[:, 0][:, None], + rtol=1e-10, + atol=1e-10, + ) + # the reference removes the isolated-atom network output of every atom, + # whatever the charge, spin and geometry of the cluster + e_vac = atom_energies(vac, coord, atype, charge_spin, spin) + e_ref = atom_energies(ref, coord, atype, charge_spin, spin) + iso_ref = atom_energies(ref, iso_coord, iso_atype, iso_charge_spin, iso_spin)[:, 0] + expected = e_ref - (iso_ref - BIAS[:, 0])[atype] + np.testing.assert_allclose(e_vac, expected, rtol=1e-10, atol=1e-10) + # the conditioning is essential: a differently conditioned isolated atom + # carries a learned deviation from the bias + other = atom_energies( + vac, + iso_coord, + iso_atype, + np.array([[0.0, 1.0], [0.0, 1.0]]), + np.zeros((2, 1, 3)), + ) + assert np.abs(other[0, 0] - BIAS[0, 0]) > 1e-8 + + +def test_fold_reproduces_the_referenced_model() -> None: + """Folding the reference into the bias keeps every energy and drops the rows.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + charge_spin = np.array([[0.0, 1.0], [1.0, 2.0]]) + spin = rng.normal(size=(2, 6, 3)) * (atype == 0)[..., None] + iso_coord = np.zeros((2, 1, 3)) + iso_atype = np.array([[0], [1]]) + iso_charge_spin = reference_charge_spin(TYPE_MAP) + iso_spin = reference_spin(TYPE_MAP)[:, None, :] + model = make_model(True) + e_cluster = atom_energies(model, coord, atype, charge_spin, spin) + e_iso = atom_energies(model, iso_coord, iso_atype, iso_charge_spin, iso_spin) + + model.fold_vacuum_reference() + assert not model.atomic_model.fitting_net.vacuum_ref + np.testing.assert_allclose( + atom_energies(model, coord, atype, charge_spin, spin), + e_cluster, + rtol=1e-10, + atol=1e-10, + ) + np.testing.assert_allclose( + atom_energies(model, iso_coord, iso_atype, iso_charge_spin, iso_spin), + e_iso, + rtol=1e-10, + atol=1e-10, + ) + + +def test_dense_and_graph_routes_share_the_vacuum_descriptor() -> None: + """The reference rows appended to a graph equal the single-atom-frame descriptor.""" + model = make_model(vacuum_ref=True) + am = model.atomic_model + rng = np.random.default_rng(5) + coord = torch.as_tensor( + rng.normal(size=(1, 4, 3)) * 1.2, dtype=torch.float64, device=env.DEVICE + ) + atype = torch.as_tensor([[0, 1, 1, 0]], dtype=torch.long, device=env.DEVICE) + graph = build_neighbor_graph( + coord, atype, None, CONFIG["descriptor"]["rcut"], with_csr=True + ) + charge_spin = torch.as_tensor([[0.0, 1.0]], dtype=torch.float64, device=env.DEVICE) + spin = torch.as_tensor( + rng.normal(size=(4, 3)), dtype=torch.float64, device=env.DEVICE + ) + graph, atype_all, charge_spin, spin = am.append_vacuum_frames( + graph, atype.reshape(-1), charge_spin, spin + ) + gg, _ = am.descriptor.call_graph( + graph, + atype_all, + type_embedding=am.descriptor.graph_type_embedding_table(), + spin=spin, + charge_spin=charge_spin, + ) + np.testing.assert_allclose( + gg[4:].detach().cpu().numpy(), + am.vacuum_descriptor().detach().cpu().numpy(), + rtol=1e-12, + atol=1e-12, + ) + + +def freeze_model( + tmp_path, numb_fparam: int, suffix: str = ".pt2" +) -> tuple[EnergyModel, DeepPot]: + """Freeze a vacuum-referenced model without charge/spin conditioning.""" + import copy + + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, + ) + + config = copy.deepcopy(CONFIG) + for key in ("use_spin", "add_chg_spin_ebd", "default_chg_spin"): + config["descriptor"].pop(key) + config["fitting_net"].update({"vacuum_ref": True, "numb_fparam": numb_fparam}) + config["preset_out_bias"] = PRESET + model = get_model(config) + data = jitter_zero_arrays( + model.atomic_model.descriptor.serialize(), np.random.default_rng(3) + ) + model.atomic_model.descriptor = DescrptDPA4.deserialize(data) + fitting = model.atomic_model.fitting_net + with torch.no_grad(): + fitting.bias_atom_e.copy_( + torch.as_tensor(BIAS, dtype=fitting.bias_atom_e.dtype) + ) + model = model.to(env.DEVICE).eval() + wrapper = ModelWrapper(model, model_params=copy.deepcopy(config)) + ckpt = tmp_path / "vacuum.pt" + torch.save({"model": wrapper.state_dict()}, ckpt) + frozen = tmp_path / f"vacuum_frozen{suffix}" + freeze(model=str(ckpt), output=str(frozen)) + return model, DeepPot(str(frozen)) + + +def eager_energies( + model: EnergyModel, coord: np.ndarray, atype: np.ndarray, fparam=None +) -> np.ndarray: + kwargs = {} + if fparam is not None: + kwargs["fparam"] = torch.as_tensor( + fparam, dtype=torch.float64, device=env.DEVICE + ) + ret = model.call_common( + torch.as_tensor(coord, dtype=torch.float64, device=env.DEVICE), + torch.as_tensor(atype[None], dtype=torch.long, device=env.DEVICE), + None, + **kwargs, + ) + return ret["energy"][..., 0].detach().cpu().numpy().reshape(-1) + + +@pytest.mark.parametrize("suffix", [".pt2", ".pte"]) +def test_freeze_folds_the_reference(tmp_path, suffix) -> None: + """A frozen model reproduces the referenced model without reference atoms.""" + model, dp = freeze_model(tmp_path, numb_fparam=0, suffix=suffix) + assert dp.deep_eval.get_ntypes() == 2 + rng = np.random.default_rng(2) + coord = rng.normal(size=(1, 6, 3)) * 1.2 + atype = np.array([0, 1, 1, 0, 1, 0]) + for itype in range(2): + _, _, _, atom_energy, _ = dp.eval( + np.zeros((1, 1, 3)), None, np.array([itype]), atomic=True + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8 + ) + _, _, _, atom_energy, _ = dp.eval(coord, None, atype, atomic=True) + np.testing.assert_allclose( + atom_energy.reshape(-1), + eager_energies(model, coord, atype), + rtol=1e-8, + atol=1e-8, + ) + + +def test_freeze_keeps_the_reference_with_fparam(tmp_path) -> None: + """With frame parameters the frozen model references from the stored vacuum table.""" + model, dp = freeze_model(tmp_path, numb_fparam=1) + rng = np.random.default_rng(2) + coord = rng.normal(size=(1, 6, 3)) * 1.2 + atype = np.array([0, 1, 1, 0, 1, 0]) + for fparam in (np.array([[0.3]]), np.array([[-1.1]])): + for itype in range(2): + _, _, _, atom_energy, _ = dp.eval( + np.zeros((1, 1, 3)), None, np.array([itype]), atomic=True, fparam=fparam + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8 + ) + _, _, _, atom_energy, _ = dp.eval( + coord, None, atype, atomic=True, fparam=fparam + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), + eager_energies(model, coord, atype, fparam), + rtol=1e-8, + atol=1e-8, + ) + + +def test_reference_rows_follow_a_padded_node_axis() -> None: + """Padding nodes after the real frames leave the reference conditioning intact.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + def tensor(array: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.as_tensor(array, dtype=dtype, device=env.DEVICE) + + rng = np.random.default_rng(4) + coord = tensor(rng.normal(size=(2, 6, 3)) * 1.2) + atype = tensor(np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]), torch.long) + charge_spin = tensor(np.array([[0.0, 1.0], [1.0, 2.0]])) + spin = tensor(rng.normal(size=(2, 6, 3)) * (atype.cpu().numpy() == 0)[..., None]) + model = make_model(True) + atomic_model = model.atomic_model + graph = build_neighbor_graph(coord, atype, None, CONFIG["descriptor"]["rcut"]) + atype_flat = atype.reshape(-1) + spin_flat = spin.reshape(-1, 3) + reference = atomic_model.forward_common_atomic_graph( + graph, atype_flat, charge_spin=charge_spin, spin=spin_flat + )["energy"] + + n_pad = 3 + atype_padded = torch.cat( + [atype_flat, torch.full((n_pad,), -1, dtype=torch.long, device=env.DEVICE)] + ) + spin_padded = torch.cat( + [spin_flat, torch.zeros((n_pad, 3), dtype=torch.float64, device=env.DEVICE)] + ) + padded = atomic_model.forward_common_atomic_graph( + graph, atype_padded, charge_spin=charge_spin, spin=spin_padded + )["energy"] + np.testing.assert_allclose( + padded[: atype_flat.shape[0]].detach().cpu().numpy(), + reference.detach().cpu().numpy(), + rtol=1e-12, + atol=1e-12, + ) + assert torch.all(padded[atype_flat.shape[0] :] == 0.0) + + +def test_archive_carries_the_live_model(tmp_path) -> None: + """The archive keeps the unfolded model, so a re-export resolves the reference once.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + serialize_from_file, + ) + + _, dp = freeze_model(tmp_path, numb_fparam=0) + data = serialize_from_file(str(tmp_path / "vacuum_frozen.pt2")) + model = BaseModel.deserialize(data["model"]) + fitting = model.atomic_model.fitting_net + assert fitting.needs_vacuum_descriptor() + np.testing.assert_allclose(fitting.bias_atom_e.detach().cpu().numpy(), BIAS) + again = tmp_path / "vacuum_again.pt2" + deserialize_to_file(str(again), data) + dp_again = DeepPot(str(again)) + for itype in range(2): + for evaluator in (dp, dp_again): + _, _, _, atom_energy, _ = evaluator.eval( + np.zeros((1, 1, 3)), None, np.array([itype]), atomic=True + ) + np.testing.assert_allclose( + atom_energy.reshape(-1), BIAS[itype], rtol=1e-8, atol=1e-8 + ) + + +def test_without_preset_the_output_is_not_referenced() -> None: + """A bias fitted from the data is no isolated-atom energy, so the output stays plain.""" + rng = np.random.default_rng(1) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + charge_spin = np.array([[0.0, 1.0], [1.0, 2.0]]) + spin = rng.normal(size=(2, 6, 3)) * (atype == 0)[..., None] + plain = make_model(False) + unreferenced = make_model(True, preset=False) + fitting = unreferenced.atomic_model.fitting_net + assert not fitting.vacuum_ref + assert not fitting.needs_vacuum_descriptor() + np.testing.assert_allclose( + atom_energies(unreferenced, coord, atype, charge_spin, spin), + atom_energies(plain, coord, atype, charge_spin, spin), + rtol=1e-12, + atol=1e-12, + ) diff --git a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py index d8f9ed31eb..f67a1f85b1 100644 --- a/source/tests/pt_expt/model/test_dpa4c_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa4c_graph_lower.py @@ -239,14 +239,16 @@ def spy(atype, *args, **kwargs): assert out["force"].shape == (len(nlocs), pad_nloc, 3) -def test_compiled_lower_accepts_a_compacted_node_axis() -> None: +@pytest.mark.parametrize("vacuum_ref", [False, True]) +def test_compiled_lower_accepts_a_compacted_node_axis(vacuum_ref: bool) -> None: """The compiled artifact must not carry ``N == nframes * nloc`` as a guard. Its trace is taken on a uniform system, where the flat node axis happens to be the product of the frame count and the atom count. Dropping the padding breaks that relation, so this exercises the compiled lower on a batch where it no longer holds, and holds the result against the eager graph path, - which takes the same compaction. + which takes the same compaction. Both public layouts preserve force-loss + gradients, including the isolated-atom reference when it is enabled. """ from deepmd.pt_expt.train.training import ( _CompiledModel, @@ -255,9 +257,15 @@ def test_compiled_lower_accepts_a_compacted_node_axis() -> None: torch.manual_seed(0) config = _config() + config["type_map"] = ["O", "H"] + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} config["descriptor"]["channels"] = 8 + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [0.0, 1.0] config["fitting_net"]["neuron"] = [8, 8] + config["fitting_net"]["vacuum_ref"] = vacuum_ref model = get_model(config).to(env.DEVICE).train() + assert model.get_fitting_net().needs_vacuum_descriptor() == vacuum_ref compiled = _CompiledModel(model, _get_model_structure_key(model)) rng = np.random.default_rng(0) @@ -282,10 +290,91 @@ def test_compiled_lower_accepts_a_compacted_node_axis() -> None: expected = model(*args) assert got["force"].shape == (len(nlocs), pad_nloc, 3) - torch.testing.assert_close(got["energy"], expected["energy"]) - torch.testing.assert_close(got["force"], expected["force"]) phantom = args[1] < 0 assert bool(torch.all(got["force"][phantom] == 0.0)) + ragged = compiled.forward_ragged( + args[0][~phantom], + args[1][~phantom], + torch.tensor(nlocs, dtype=torch.int64, device=env.DEVICE), + box=args[2], + ) + parameters = tuple(p for p in model.parameters() if p.requires_grad) + expected_gradients = torch.autograd.grad( + expected["energy"].sum() + expected["force"].square().sum(), + parameters, + allow_unused=True, + ) + for result, force in ( + (got, expected["force"]), + (ragged, expected["force"][~phantom]), + ): + torch.testing.assert_close(result["energy"], expected["energy"]) + torch.testing.assert_close(result["virial"], expected["virial"]) + torch.testing.assert_close(result["force"], force) + gradients = torch.autograd.grad( + result["energy"].sum() + result["force"].square().sum(), + parameters, + allow_unused=True, + ) + for actual, reference in zip(gradients, expected_gradients, strict=True): + if reference is None: + assert actual is None + else: + torch.testing.assert_close(actual, reference) + + +def test_shared_weights_preserve_each_tasks_vacuum_reference() -> None: + """A referenced task cannot supply the compiled graph of a plain task.""" + from deepmd.pt_expt.train.training import ( + _CompiledModel, + _detect_task_buffers, + _get_model_structure_key, + ) + + models = [] + for has_preset in (True, False): + config = _config() + config["type_map"] = ["O", "H"] + config["descriptor"].update({"channels": 8, "lmax": 2}) + config["fitting_net"].update( + {"neuron": [8, 8], "dim_case_embd": 2, "vacuum_ref": True} + ) + if has_preset: + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} + model = get_model(config).to(env.DEVICE).train() + assert model.get_fitting_net().vacuum_ref == has_preset + if models: + model.get_descriptor().share_params( + models[0].get_descriptor(), 0, resume=True + ) + model.get_fitting_net().share_params( + models[0].get_fitting_net(), 0, resume=True + ) + models.append(model) + + coord = torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=torch.float64, + device=env.DEVICE, + ) + atype = torch.tensor([0, 1, 1], device=env.DEVICE) + n_node = torch.tensor([3], device=env.DEVICE) + shared_graphs = {} + for model in models: + key = _get_model_structure_key(model) + group = [other for other in models if _get_model_structure_key(other) == key] + buffers = _detect_task_buffers(model, group) + compiled = _CompiledModel( + model, + key, + task_buf_order=tuple(buffers), + task_buffers=buffers, + compiled_by_structure=shared_graphs, + ) + expected = model.forward_ragged(coord, atype, n_node) + actual = compiled.forward_ragged(coord, atype, n_node) + for name in ("energy", "force", "virial"): + torch.testing.assert_close(actual[name], expected[name]) def test_ragged_and_padded_batches_agree() -> None: diff --git a/source/tests/pt_expt/model/test_fused_vacuum_ref.py b/source/tests/pt_expt/model/test_fused_vacuum_ref.py new file mode 100644 index 0000000000..311912651f --- /dev/null +++ b/source/tests/pt_expt/model/test_fused_vacuum_ref.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The fused inference operators serve the isolated-atom reference through the bias. + +On a fitting without conditioning the reference output of every type is a +constant, so the model-level fused energy/force route hands the operators the +bias minus that constant and reproduces the autograd route exactly: an +isolated atom gives its bias and a cluster gives the referenced energies. +""" + +import numpy as np +import pytest +import torch + +import deepmd.pt_expt.model.make_model as make_model +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) +from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.pt_expt.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.pt_expt.fitting.ener_fitting import ( + EnergyFittingNet, +) +from deepmd.pt_expt.kernels.cuda.dpa1.graph_energy_force import ( + op_available as dpa1_op_available, +) +from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ef_op_available as dpa4c_op_available, +) +from deepmd.pt_expt.model import ( + EnergyModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +RCUT = 4.0 +BIAS = np.array([[-3.0], [0.5]]) + + +def make_model_with( + descriptor: DescrptDPA1 | DescrptDPA4C, numb_fparam: int = 0 +) -> EnergyModel: + fitting = EnergyFittingNet( + 2, + descriptor.get_dim_out(), + neuron=[32, 32], + mixed_types=True, + precision="float32", + resnet_dt=False, + activation_function="tanh", + vacuum_ref=True, + numb_fparam=numb_fparam, + seed=1, + ) + fitting["bias_atom_e"] = BIAS.copy() + return ( + EnergyModel( + descriptor, + fitting, + type_map=["O", "H"], + preset_out_bias={"energy": BIAS.tolist()}, + ) + .to(env.DEVICE) + .eval() + ) + + +def make_dpa1(numb_fparam: int = 0) -> EnergyModel: + return make_model_with( + DescrptDPA1( + rcut=RCUT, + rcut_smth=0.5, + sel=[20], + ntypes=2, + attn_layer=0, + axis_neuron=4, + neuron=[8, 16, 32], + tebd_input_mode="concat", + precision="float32", + seed=1, + ), + numb_fparam, + ) + + +def make_dpa4c() -> EnergyModel: + model = make_model_with( + DescrptDPA4C( + rcut=RCUT, + ntypes=2, + channels=32, + lmax=2, + n_radial=8, + precision="float32", + seed=17, + ) + ) + model.atomic_model.descriptor.enable_compression(0.5) + return model + + +def fused_and_autograd( + model: EnergyModel, coord: np.ndarray, atype: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + coord_t = torch.as_tensor(coord, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.as_tensor(atype, dtype=torch.long, device=env.DEVICE) + graph = build_neighbor_graph( + coord_t, atype_t, None, RCUT, with_csr=True, canonicalize=True + ) + fused = make_model._fused_energy_force_graph( + model, graph, atype_t.reshape(-1), False + ) + assert fused is not None, "the fused energy/force route was not taken" + autograd = model.atomic_model.forward_common_atomic_graph( + graph, atype_t.reshape(-1) + ) + return ( + fused["energy"][:, 0].detach().cpu().numpy(), + autograd["energy"][:, 0].detach().cpu().numpy(), + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="the fused operators need CUDA" +) +@pytest.mark.parametrize( + "build, available", + [(make_dpa1, dpa1_op_available), (make_dpa4c, dpa4c_op_available)], + ids=["dpa1", "dpa4c_compressed"], +) +def test_fused_route_matches_autograd_with_vacuum_ref( + monkeypatch, build, available +) -> None: + if not available(): + pytest.skip("the fused operator library is unavailable") + monkeypatch.setenv("DP_CUDA_INFER", "2") + model = build() + rng = np.random.default_rng(0) + coord = rng.normal(size=(2, 6, 3)) * 1.2 + atype = np.array([[0, 1, 1, 0, 1, 0], [1, 1, 0, 0, 1, 0]]) + fused, autograd = fused_and_autograd(model, coord, atype) + np.testing.assert_allclose(fused, autograd, rtol=1e-6, atol=1e-6) + fused_iso, _ = fused_and_autograd(model, np.zeros((2, 1, 3)), np.array([[0], [1]])) + np.testing.assert_allclose(fused_iso, BIAS[:, 0], rtol=1e-6, atol=1e-6) + + +def test_fused_route_yields_to_autograd_with_fparam(monkeypatch) -> None: + """A reference that varies between atoms is served by the autograd lower.""" + monkeypatch.setenv("DP_CUDA_INFER", "2") + model = make_dpa1(numb_fparam=1) + coord_t = torch.zeros((1, 1, 3), dtype=torch.float64, device=env.DEVICE) + atype_t = torch.zeros((1, 1), dtype=torch.long, device=env.DEVICE) + graph = build_neighbor_graph( + coord_t, atype_t, None, RCUT, with_csr=True, canonicalize=True + ) + assert ( + make_model._fused_energy_force_graph(model, graph, atype_t.reshape(-1), False) + is None + ) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 6a174753b5..e4b9874519 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -88,6 +88,14 @@ def test_sezm_builder_rejects_bridging() -> None: get_sezm_model(data) +def test_bridged_dpa4_forwards_preset() -> None: + """The composition built from a bridged DPA4 config carries the preset.""" + data = _bridged(_dpa4_standard_config()) + data["preset_out_bias"] = {"energy": {"Ni": 2.0}} + model = get_model(data) + assert model.atomic_model.preset_out_bias == {"energy": [[2.0], None]} + + def test_standard_builder_without_bridging_is_unaffected() -> None: """The rejection keys on the flag, not on the DPA4 components: a plain DPA4 standard model still builds and carries no bridging switch. diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py index 86e180b378..9e2779bacd 100644 --- a/source/tests/pt_expt/model/test_get_model_dpa4.py +++ b/source/tests/pt_expt/model/test_get_model_dpa4.py @@ -238,10 +238,6 @@ def test_unsupported_keys_raise(self) -> None: cases = { "spin": ({"use_spin": [True, False], "virtual_scale": [0.3]}, "Spin DPA4"), "lora": ({"rank": 4}, "`lora` is not supported"), - "preset_out_bias": ( - {"energy": [None, 1.0]}, - "`preset_out_bias` is not supported", - ), } for key, (value, msg_regex) in cases.items(): raw = _make_raw_model_config() @@ -249,6 +245,13 @@ def test_unsupported_keys_raise(self) -> None: with self.assertRaisesRegex(NotImplementedError, msg_regex): get_model(raw) + def test_preset_out_bias_accepted(self) -> None: + """A model-level ``preset_out_bias`` reaches the atomic model.""" + raw = _make_raw_model_config() + raw["preset_out_bias"] = {"energy": [None, 1.0]} + model = get_model(raw) + self.assertEqual(model.atomic_model.preset_out_bias, {"energy": [None, [1.0]]}) + def test_native_spin_capability_gate_standard_config(self) -> None: """The generic ``supports_native_spin()`` gate rejects a dense descriptor. diff --git a/source/tests/tf/test_fitting_vacuum_ref.py b/source/tests/tf/test_fitting_vacuum_ref.py new file mode 100644 index 0000000000..9a0ed481c3 --- /dev/null +++ b/source/tests/tf/test_fitting_vacuum_ref.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""The TensorFlow fittings reject a serialized ``vacuum_ref`` and load older dictionaries.""" + +import unittest + +from deepmd.dpmodel.fitting import ( + DipoleFitting, + DOSFittingNet, + EnergyFittingNet, + PolarFitting, +) +from deepmd.tf.fit import ( + DipoleFittingSeA, + DOSFitting, + EnerFitting, + PolarFittingSeA, +) + +NTYPES, DIM_DESCRPT, EMBEDDING_WIDTH = 2, 6, 4 + + +class TestFittingVacuumRef(unittest.TestCase): + def cases(self) -> list[tuple[dict, type, int]]: + """The serialized dictionary, TensorFlow class and version of every fitting.""" + return [ + ( + EnergyFittingNet(NTYPES, DIM_DESCRPT, mixed_types=False).serialize(), + EnerFitting, + 5, + ), + ( + DOSFittingNet( + NTYPES, DIM_DESCRPT, numb_dos=3, mixed_types=False + ).serialize(), + DOSFitting, + 5, + ), + ( + DipoleFitting( + NTYPES, DIM_DESCRPT, EMBEDDING_WIDTH, mixed_types=False + ).serialize(), + DipoleFittingSeA, + 5, + ), + ( + PolarFitting( + NTYPES, DIM_DESCRPT, EMBEDDING_WIDTH, mixed_types=False + ).serialize(), + PolarFittingSeA, + 6, + ), + ] + + def test_vacuum_ref_is_rejected(self) -> None: + for data, tf_class, version in self.cases(): + self.assertEqual(data["@version"], version) + self.assertFalse(data["vacuum_ref"]) + with self.assertRaises(NotImplementedError): + tf_class.deserialize({**data, "vacuum_ref": True}, suffix="") + + def test_previous_version_loads(self) -> None: + for data, tf_class, version in self.cases(): + older = {k: v for k, v in data.items() if k != "vacuum_ref"} + older["@version"] = version - 1 + self.assertIsInstance(tf_class.deserialize(older, suffix=""), tf_class) + + +if __name__ == "__main__": + unittest.main()