feat(stat): isolated-atom energy reference through preset_out_bias and vacuum_ref

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,15 +77,15 @@ 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__()
         self.type_map = type_map
         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
 

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,