diff --git a/deepmd/infer/deep_density.py b/deepmd/infer/deep_density.py new file mode 100644 index 0000000000..b1655c45b6 --- /dev/null +++ b/deepmd/infer/deep_density.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + +from .deep_eval import ( + DeepEval, +) + + +class DeepDensity(DeepEval): + """Charge density evaluated on grid points. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model. + + The output definition is identical to that of the density + fitting net of the model: the density is predicted on grid + points, and is neither reducible nor differentiable. + """ + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "density", + [1], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + ) + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, + grid: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + mixed_type: bool = False, + grid_batch_size: int | None = 4096, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate the density on grid points. + + Parameters + ---------- + coords : np.ndarray + The coordinates of the atoms, in shape (nframes, natoms, 3). + cells : np.ndarray + The cell vectors of the system, in shape (nframes, 9). If the system + is not periodic, set it to None. + atom_types : list[int] or np.ndarray + The types of the atoms. If mixed_type is False, the shape is (natoms,); + otherwise, the shape is (nframes, natoms). + grid : np.ndarray + The coordinates of the grid points, in shape (nframes, ngrid, 3). + fparam : np.ndarray, optional + The frame parameters, by default None. + aparam : np.ndarray, optional + The atomic parameters, by default None. + mixed_type : bool, optional + Whether the atom_types is mixed type, by default False. + grid_batch_size : int or None, optional + Maximum number of grid points evaluated in one model call. Grid + chunks are concatenated in their original order. The default is + 4096; set to None to evaluate the full grid in one call. + **kwargs : dict[str, Any] + Keyword arguments. ``charge_spin`` (nframes, 2) is forwarded to + models with charge spin embedding; when omitted, such models use + their ``default_chg_spin`` if set. + + Returns + ------- + density + The density on the grid points, in shape (nframes, ngrid). + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + _, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + grid = self._standard_grid(grid, nframes) + if grid_batch_size is not None and grid_batch_size <= 0: + raise ValueError("grid_batch_size must be positive or None") + if grid_batch_size is None or grid.shape[1] <= grid_batch_size: + grid_slices = (grid,) + else: + grid_slices = tuple( + grid[:, begin : begin + grid_batch_size] + for begin in range(0, grid.shape[1], grid_batch_size) + ) + density_chunks = [] + for grid_chunk in grid_slices: + results = self.deep_eval.eval( + coords, + cells, + atom_types, + False, + fparam=fparam, + aparam=aparam, + grid=grid_chunk, + **kwargs, + ) + density_chunks.append(results["density"].reshape(nframes, -1)) + return np.concatenate(density_chunks, axis=1) + + @staticmethod + def _standard_grid(grid: np.ndarray, nframes: int) -> np.ndarray: + """Normalise the grid coordinates to ``(nframes, ngrid, 3)``. + + The auto batcher slices every argument with ``ndim > 1`` along axis + 0, so a natural single-frame ``(ngrid, 3)`` input would silently be + truncated to one grid point; carry the frame dimension explicitly. + """ + arr = np.asarray(grid) + if arr.ndim == 2: + if nframes != 1: + raise ValueError( + f"grid of shape {arr.shape} is ambiguous for {nframes} " + "frames; pass grid with shape (nframes, ngrid, 3)" + ) + arr = arr[None, ...] + if arr.ndim != 3 or arr.shape[-1] != 3: + raise ValueError( + f"grid must have shape (nframes, ngrid, 3), got {arr.shape}" + ) + if arr.shape[0] != nframes: + raise ValueError(f"grid has {arr.shape[0]} frames but coord has {nframes}") + if arr.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + return arr + + +__all__ = ["DeepDensity"] diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index ce23c5981b..96ef64a139 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -212,6 +212,7 @@ def eval( aparam=aparam, **kwargs, ) + energy = results["energy_redu"].reshape(nframes, 1) force = results["energy_derv_r"].reshape(nframes, natoms, 3) virial = results["energy_derv_c_redu"].reshape(nframes, 9) diff --git a/deepmd/infer/model_test/__init__.py b/deepmd/infer/model_test/__init__.py index 4afd626d0e..1152c3141d 100644 --- a/deepmd/infer/model_test/__init__.py +++ b/deepmd/infer/model_test/__init__.py @@ -14,6 +14,9 @@ Any, ) +from deepmd.infer.deep_density import ( + DeepDensity, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -36,6 +39,9 @@ save_txt_file, test_chunk_atoms, ) +from deepmd.infer.model_test.density import ( + DensityTester, +) from deepmd.infer.model_test.dos import ( DosTester, ) @@ -56,6 +62,7 @@ __all__ = [ "ChunkContext", + "DensityTester", "DipoleTester", "DosTester", "EnerTester", @@ -96,6 +103,8 @@ def build_tester(dp: Any, *, atomic: bool) -> ModelTester: return tester(dp, atomic=atomic) if isinstance(dp, DeepDOS): return DosTester(dp, atomic=atomic) + if isinstance(dp, DeepDensity): + return DensityTester(dp, atomic=atomic) if isinstance(dp, DeepProperty): return PropertyTester(dp, atomic=atomic) if isinstance(dp, DeepGlobalPolar): diff --git a/deepmd/infer/model_test/density.py b/deepmd/infer/model_test/density.py new file mode 100644 index 0000000000..0ec2821b8a --- /dev/null +++ b/deepmd/infer/model_test/density.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of models predicting charge density on grid points.""" + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + _write_per_frame_details, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + mae, + rmse, +) + +__all__ = ["DensityTester"] + + +class DensityTester(ModelTester): + """Test a model of charge density on grid points.""" + + report = ( + ("mae_density", "DENSITY MAE : {} units"), + ("rmse_density", "DENSITY RMSE : {} units"), + ) + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels a density test consumes.""" + dp = self.dp + # The grid and the density are defined on grid points rather than on + # atoms, and their extent (ngrid) is not known until the data is + # loaded. They are declared frame-major so the loader keeps the + # (nframes, ngrid, ndof) layout without reshaping to natoms. + data.add( + "grid", + 3, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + data.add( + "density", + 1, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", + dp.get_dim_fparam(), + atomic=False, + must=not dp.has_default_fparam(), + high_prec=False, + ) + if dp.get_dim_aparam() > 0: + data.add( + "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False + ) + if dp.has_chg_spin_ebd(): + data.add( + "charge_spin", + 2, + atomic=False, + must=not dp.has_default_chg_spin(), + high_prec=False, + ) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of a density test.""" + dp = self.dp + mixed_type = data.mixed_type + nframes = test_data["box"].shape[0] + + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + fparam = ( + test_data["fparam"] + if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0 + else None + ) + aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None + charge_spin = ( + test_data["charge_spin"] + if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0 + else None + ) + # the loader accepts both (nframes, ngrid, 3) and the flattened + # (nframes, ngrid*3) layout; DeepDensity.eval only accepts the + # former, so normalize here and keep one contract at the eval call + grid = test_data["grid"].reshape(nframes, -1, 3) + + prediction = dp.eval( + coord, + box, + atype, + grid, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + mixed_type=mixed_type, + ).reshape(nframes, -1) + label = test_data["density"].reshape(nframes, -1) + + diff = prediction - label + errors: dict[str, tuple[float, float]] = { + "mae_density": (mae(diff), diff.size), + "rmse_density": (rmse(diff), diff.size), + } + + if context.detail_path is not None: + _write_per_frame_details( + context, + suffix="density", + reference=label, + prediction=prediction, + ) + + return errors diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index bd8a0bc91e..90a87edcf6 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -57,4 +57,5 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N default=data_item["default"], dtype=data_item["dtype"], output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"], + special_shape=getattr(data_item, "special_shape", None), ) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..5c88ba0049 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -20,6 +20,9 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.infer.deep_density import ( + DeepDensity, +) from deepmd.infer.deep_dipole import ( DeepDipole, ) @@ -440,11 +443,26 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepWFC elif "population" in model_output_type: return DeepPopulation + elif "density" in model_output_type and self._model_has_grid( + self.dp.model["Default"] + ): + # key on the grid capability rather than the output name, so a + # property fitting with property_name "density" still dispatches + # to DeepProperty below + return DeepDensity elif self.get_var_name() in model_output_type: return DeepProperty else: raise RuntimeError("Unknown model type") + @staticmethod + def _model_has_grid(model: Any) -> bool: + has_grid = getattr(model, "has_grid", None) + try: + return bool(has_grid()) if callable(has_grid) else False + except AttributeError: + return False + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -551,6 +569,41 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + if "grid" in kwargs and kwargs["grid"] is not None: + grid_input = np.array(kwargs["grid"]) + # the directional neighbor list is dense in ngrid x nall, so the + # batching size proxy must account for the grid extent, not just + # the atom count + ngrid = grid_input.size // (numb_test * 3) + out = self._eval_func( + self._eval_model_density, numb_test, max(natoms, ngrid) + )( + coords, + cells, + atom_types, + grid_input, + fparam, + aparam, + self._get_request_defs(atomic), + charge_spin, + ) + # _eval_model_density returns a 1-element tuple; execute_all unwraps + # it when auto batching is enabled, but with auto_batch_size=False + # the inner function is called directly and the tuple survives. + # request_defs is ignored on this path: the density output is a + # single fixed variable, so atomic=True returns the same dict as + # atomic=False. Kept for signature parity with the other + # _eval_model_* paths. + if isinstance(out, tuple): + (out,) = out + return {"density": out} + if "density" in self.output_def.var_defs and self._model_has_grid( + self.dp.model["Default"] + ): + raise ValueError( + "grid is required to evaluate a density model; " + "pass grid=... with shape (nframes, ngrid, 3)" + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( @@ -916,6 +969,86 @@ def _eval_model_spin( ) # this is kinda hacky return tuple(results) + def _eval_model_density( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + grid: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None, + ) -> tuple[np.ndarray, ...]: + model = self.dp.to(DEVICE) + + nframes = coords.shape[0] + if len(atom_types.shape) == 1: + natoms = len(atom_types) + atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + else: + natoms = len(atom_types[0]) + + coord_input = torch.tensor( + coords.reshape([nframes, natoms, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE) + grid_input = torch.tensor( + grid.reshape([nframes, -1, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + ngrid = grid_input.shape[1] + if cells is not None: + box_input = torch.tensor( + cells.reshape([nframes, 3, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) + else: + box_input = None + if fparam is not None: + fparam_input = to_torch_tensor( + fparam.reshape(nframes, self.get_dim_fparam()) + ) + else: + fparam_input = None + if aparam is not None: + aparam_input = to_torch_tensor( + aparam.reshape(nframes, natoms, self.get_dim_aparam()) + ) + else: + aparam_input = None + if charge_spin is not None: + charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2)) + else: + charge_spin_input = None + + do_atomic_virial = any( + x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs + ) + batch_output = model( + coord_input, + type_input, + grid=grid_input, + box=box_input, + do_atomic_virial=do_atomic_virial, + fparam=fparam_input, + aparam=aparam_input, + charge_spin=charge_spin_input, + ) + if isinstance(batch_output, tuple): + batch_output = batch_output[0] + + results = [] + pt_name = "density" + density_shape = [nframes, ngrid] + out = batch_output[pt_name].reshape(density_shape).detach().cpu().numpy() + results.append(out) + return tuple(results) + def _get_output_shape( self, odef: OutputVariableDef, nframes: int, natoms: int ) -> list[int]: diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 7a47d13ae2..544320faf5 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from .charge import ( + GridDensityLoss, +) from .dens import ( DeNSLoss, ) @@ -31,6 +34,7 @@ "EnergyHessianStdLoss", "EnergySpinLoss", "EnergyStdLoss", + "GridDensityLoss", "PopulationLoss", "PropertyLoss", "TaskLoss", diff --git a/deepmd/pt/loss/charge.py b/deepmd/pt/loss/charge.py new file mode 100644 index 0000000000..53b2250434 --- /dev/null +++ b/deepmd/pt/loss/charge.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + GLOBAL_PT_FLOAT_PRECISION, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class GridDensityLoss(TaskLoss): + def __init__( + self, + starter_learning_rate: float = 1.0, + start_pref_d: float = 0.0, + limit_pref_d: float = 0.0, + inference: bool = False, + **kwargs: Any, + ) -> None: + r"""Construct a layer to compute loss on grid density. + + The residual is masked by the grid-point mask (excluded points do + not contribute) and reduced per frame before averaging over the + batch, consistent with the other losses in the package. + + Parameters + ---------- + starter_learning_rate : float + The learning rate at the start of the training. + start_pref_d : float + The prefactor of charge density loss at the start of the training. + limit_pref_d : float + The prefactor of charge density loss at the end of the training. + inference : bool + If true, it will output all losses found in output, ignoring the pre-factors. + **kwargs + Other keyword arguments. + """ + super().__init__() + self.starter_learning_rate = starter_learning_rate + self.has_d = (start_pref_d != 0.0 or limit_pref_d != 0.0) or inference + + self.start_pref_d = start_pref_d + self.limit_pref_d = limit_pref_d + self.inference = inference + + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: + """Return loss on grid density. + + Parameters + ---------- + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. + + Returns + ------- + model_pred: dict[str, torch.Tensor] + Model predictions. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + Other losses for display. + """ + model_pred = model(**input_dict) + coef = learning_rate / self.starter_learning_rate + pref_d = self.limit_pref_d + (self.start_pref_d - self.limit_pref_d) * coef + + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] + more_loss = {} + # more_loss['log_keys'] = [] # showed when validation on the fly + # more_loss['test_keys'] = [] # showed when doing dp test + if self.has_d and "density" in model_pred and "density" in label: + find_density = label.get("find_density", 1.0) + density_pred = model_pred["density"] + density_label = label["density"] + pref_d = pref_d * find_density + # mask out excluded grid points (the atomic model zeroes them + # and returns the mask) and reduce per frame, consistent with + # the other losses in the package, so that frames with larger + # grids do not dominate the batch + mask = model_pred.get("mask", None) + nframes = density_pred.shape[0] + if mask is None: + mask_f = torch.ones( + density_pred.shape[:2], + dtype=density_pred.dtype, + device=density_pred.device, + ) + else: + mask_f = mask.reshape(density_pred.shape[:2]).to(density_pred.dtype) + mask_sum = mask_f.sum(dim=-1).clamp(min=1.0) + residual = ( + density_label.reshape(nframes, -1) - density_pred.reshape(nframes, -1) + ) * mask_f + l2_density_loss = torch.square(residual).sum(dim=-1).div(mask_sum).mean() + rmse_d = l2_density_loss.sqrt() + more_loss["rmse_d"] = self.display_if_exist(rmse_d.detach(), find_density) + l1_density_loss = residual.abs().sum(dim=-1).div(mask_sum).mean() + mae_d = l1_density_loss + # minimise the squared error, consistent with every other loss + # in the package; the absolute error is only for display + loss += (pref_d * l2_density_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_d"] = self.display_if_exist(mae_d.detach(), find_density) + elif not self.inference and "density" in model_pred: + # the density term is disabled (zero prefactors), but the graph + # must stay connected so that backward() does not fail + loss = loss + (model_pred["density"].sum() * 0.0).to( + GLOBAL_PT_FLOAT_PRECISION + ) + return model_pred, loss, more_loss + + def serialize(self) -> dict: + """Serialize the loss module. + + Returns + ------- + dict + The serialized loss module + """ + return { + "@class": "GridDensityLoss", + "@version": 1, + "starter_learning_rate": self.starter_learning_rate, + "start_pref_d": self.start_pref_d, + "limit_pref_d": self.limit_pref_d, + } + + @classmethod + def deserialize(cls, data: dict) -> "TaskLoss": + """Deserialize the loss module. + + Parameters + ---------- + data : dict + The serialized loss module + + Returns + ------- + TaskLoss + The deserialized loss module + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + return cls(**data) + + @property + def training_metric_names(self) -> tuple[str, ...]: + """Return density metrics displayed by averaged training output.""" + return ("rmse_d", "mae_d") if self.has_d else () + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Return data label requirements needed for this loss calculation. + + Only the density label is declared here; the grid is a model input + and is declared via ``get_additional_data_requirement``. + """ + label_requirement = [] + if self.has_d: + # the density label is the only supervision signal of this model: + # a missing file must abort training, not silently optimise nothing + label_requirement.append( + DataRequirementItem( + "density", + ndof=1, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + ) + return label_requirement diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index e343287cad..f1ec92eb84 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -17,6 +17,9 @@ from .base_atomic_model import ( BaseAtomicModel, ) +from .density_atomic_model import ( + DPDensityAtomicModel, +) from .dipole_atomic_model import ( DPDipoleAtomicModel, ) @@ -53,6 +56,7 @@ "BaseAtomicModel", "DPAtomicModel", "DPDOSAtomicModel", + "DPDensityAtomicModel", "DPDipoleAtomicModel", "DPEnergyAtomicModel", "DPPolarAtomicModel", diff --git a/deepmd/pt/model/atomic_model/density_atomic_model.py b/deepmd/pt/model/atomic_model/density_atomic_model.py new file mode 100644 index 0000000000..3c79b714d3 --- /dev/null +++ b/deepmd/pt/model/atomic_model/density_atomic_model.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np +import torch + +from deepmd.dpmodel.utils.env_mat_stat import ( + EnvMatStatSe, +) +from deepmd.pt.model.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.pt.model.descriptor.descriptor import ( + DescriptorBlock, +) +from deepmd.pt.model.task.density import ( + DensityFittingNet, +) +from deepmd.pt.utils.nlist import ( + build_directional_neighbor_list, + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.region import ( + normalize_coord, +) +from deepmd.utils.model_preset import ( + PERIODIC_TABLE, +) +from deepmd.utils.path import ( + DPPath, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + +log = logging.getLogger(__name__) + + +class DPDensityAtomicModel(DPAtomicModel): + def __init__( + self, + descriptor: BaseDescriptor, + fitting: DensityFittingNet, + type_map: list[str], + **kwargs: Any, + ) -> None: + assert isinstance(fitting, DensityFittingNet) + super().__init__(descriptor, fitting, type_map, **kwargs) + self.rcut = self.descriptor.get_rcut() + self.rcut_smth = self.descriptor.get_rcut_smth() + if type_map and type_map[-1] in PERIODIC_TABLE: + raise ValueError( + "the last entry of type_map must be a reserved grid type " + "for a density model, but got the real element " + f"{type_map[-1]!r}; append a non-element placeholder " + '(e.g. "X") at the end of type_map' + ) + if self.descriptor.get_env_protection() == 0.0: + log.warning( + "The descriptor env_protection is 0.0; grid points coincident " + "with atoms would produce NaN densities. Setting it to 1e-6." + ) + self._set_descriptor_env_protection(1e-6) + self.sel = self.descriptor.get_sel() + self.nnei = self.descriptor.get_nsel() + + def forward_atomic( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + return_atomic_feature: bool = False, + ) -> dict[str, torch.Tensor]: + """Return atomic prediction. + + Parameters + ---------- + extended_coord + coodinates in extended region + extended_atype + atomic type in extended region + nlist + neighbor list. nf x nloc x nsel + mapping + mapps the extended indices to local indices + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + comm_dict + Data needed for communication during parallel inference. + grid + Grid coordinates. nf x ngrid x 3. + grid_type + Reserved grid-point types. nf x ngrid. + grid_nlist + Directional grid-to-atom neighbor list. nf x ngrid x nsel. + charge_spin + Frame charge and spin conditions. nf x 2. + return_atomic_feature + Whether atomic features are requested. Density models do not + currently return them. + + Returns + ------- + result_dict + the result dict, defined by the `FittingOutputDef`. + + """ + del return_atomic_feature + nframes, _, _ = nlist.shape + if self.do_grad_r() or self.do_grad_c(): + extended_coord.requires_grad_(True) + assert mapping is not None + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + # Handle default chg_spin if the descriptor supports it, the same + # way DPAtomicModel.forward_atomic does + if self.add_chg_spin_ebd and charge_spin is None: + default_cs_tensor = self.descriptor.get_default_chg_spin() + if default_cs_tensor is not None: + default_cs_tensor = default_cs_tensor.to(device=extended_coord.device) + charge_spin = torch.tile(default_cs_tensor.unsqueeze(0), [nframes, 1]) + else: + raise ValueError( + "charge_spin is required when the descriptor enables " + "charge/spin embedding without default_chg_spin" + ) + _, ngrid, _ = grid_nlist.shape + # nb x (ngrid+nall) x 3 + merged_coord = torch.cat([grid, extended_coord], dim=1) + + grid_atype = torch.ones( + [nframes, ngrid], device=extended_atype.device, dtype=extended_atype.dtype + ) * (self.descriptor.get_ntypes() - 1) + # nb x (ngrid+nall) + merged_atype = torch.cat([grid_atype, extended_atype], dim=1) + + # nb x ngrid + grid_nlist_mask = grid_nlist >= 0 + shifted_grid_nlist = torch.where(grid_nlist_mask, grid_nlist + ngrid, -1) + # nb x all + nlist_mask = nlist >= 0 + shifted_nlist = torch.where(nlist_mask, nlist + ngrid, -1) + # nb x (ngrid+nall) + merged_nlist = torch.cat([shifted_grid_nlist, shifted_nlist], dim=1) + + # nb x ngrid, row i is the grid index i + grid_mapping = torch.arange( + ngrid, device=mapping.device, dtype=mapping.dtype + ).expand(nframes, -1) + # nb x (ngrid+nall) + merged_mapping = torch.cat([grid_mapping, mapping + ngrid], dim=1) + + # the descriptor evaluates the environment of every merged point; + # the atom rows are computed and discarded (only the grid rows feed + # the fitting net), which is inherent to evaluating the merged + # grid+atom system in one call + descriptor, rot_mat, g2, h2, _sw = self.descriptor( + merged_coord, + merged_atype, + merged_nlist, + mapping=merged_mapping, + comm_dict=comm_dict, + charge_spin=charge_spin if self.add_chg_spin_ebd else None, + ) + assert descriptor is not None + + # the fitting net must see the same grid type as the descriptor: + # the reserved last entry of the type map + grid_ftype = torch.ones( + [nframes, ngrid], device=grid_type.device, dtype=grid_type.dtype + ) * (self.descriptor.get_ntypes() - 1) + ret = self.fitting_net( + descriptor[:, :ngrid, :], + grid_ftype, + gr=rot_mat, + g2=g2, + h2=h2, + fparam=fparam, + ) + return ret + + def forward_common_atomic( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Common interface for atomic inference. + + This method accept extended coordinates, extended atom typs, neighbor list, + and predict the atomic contribution of the fit property. + + Parameters + ---------- + extended_coord + extended coodinates, shape: nf x (nall x 3) + extended_atype + extended atom typs, shape: nf x nall + for a type < 0 indicating the atomic is virtual. + nlist + neighbor list, shape: nf x nloc x nsel + mapping + extended to local index mapping, shape: nf x nall + fparam + frame parameters, shape: nf x dim_fparam + aparam + atomic parameter, shape: nf x nloc x dim_aparam + comm_dict + The data needed for communication for parallel inference. + grid + Grid coordinates, shape: nf x ngrid x 3. + grid_type + Reserved grid-point types, shape: nf x ngrid. + grid_nlist + Directional grid-to-atom neighbor list, shape: nf x ngrid x nsel. + charge_spin + Frame charge and spin conditions, shape: nf x 2. + + Returns + ------- + ret_dict + dict of output atomic properties. + should implement the definition of `fitting_output_def`. + ret_dict["mask"] of shape nf x ngrid will be provided. + ret_dict["mask"][ff,ii] == 1 indicating the ii-th grid point of the ff-th frame is included. + ret_dict["mask"][ff,ii] == 0 indicating the ii-th grid point of the ff-th frame is excluded. + + """ + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + nframes, _, _ = nlist.shape + _, ngrid, _ = grid_nlist.shape + + if self.pair_excl is not None: + pair_mask = self.pair_excl(nlist, extended_atype) + # exclude neighbors in the nlist + nlist = torch.where(pair_mask == 1, nlist, -1) + # the directional grid-to-atom list: every center is the grid + # type, so exclusions involving the grid type drop the neighbor + reserved_type = self.descriptor.get_ntypes() - 1 + if any( + reserved_type in pair for pair in self.pair_excl.get_exclude_types() + ): + nall = extended_atype.shape[1] + nsel = grid_nlist.shape[-1] + virtual_type = self.pair_excl.ntypes * torch.ones( + [nframes, 1], + dtype=extended_atype.dtype, + device=extended_atype.device, + ) + ae = torch.cat([extended_atype, virtual_type], dim=-1) + index = torch.where(grid_nlist == -1, nall, grid_nlist).view( + nframes, ngrid * nsel + ) + type_j = torch.gather(ae, 1, index) + type_ij = reserved_type * (self.pair_excl.ntypes + 1) + type_j + grid_mask = self.pair_excl.type_mask[type_ij].view(nframes, ngrid, nsel) + grid_nlist = torch.where(grid_mask == 1, grid_nlist, -1) + + ext_atom_mask = self.make_atom_mask(extended_atype) + ret_dict = self.forward_atomic( + extended_coord, + torch.where(ext_atom_mask, extended_atype, 0), + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + comm_dict=comm_dict, + grid=grid, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + ret_dict = self.apply_out_stat(ret_dict, grid_type) + + # nf x ngrid + grid_mask = torch.ones( + [nframes, ngrid], dtype=torch.int32, device=ext_atom_mask.device + ) + if self.atom_excl is not None: + grid_mask *= self.atom_excl(grid_type) + + for kk in ret_dict.keys(): + out_shape = ret_dict[kk].shape + out_shape2 = 1 + for ss in out_shape[2:]: + out_shape2 *= ss + ret_dict[kk] = ( + ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2]) + * grid_mask[:, :, None] + ).view(out_shape) + ret_dict["mask"] = grid_mask + + return ret_dict + + def forward( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + raise NotImplementedError( + "DPDensityAtomicModel.forward requires grid, grid_type and grid_nlist; " + "use forward_common_atomic with explicit grid inputs instead." + ) + + def _get_forward_wrapper_func( + self, + ) -> Callable[..., dict[str, torch.Tensor]]: + """Get a forward wrapper of the atomic model for output bias calculation.""" + + def model_forward( + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + del spin + with torch.no_grad(): + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + coord, + atype, + self.get_rcut(), + self.get_sel(), + mixed_types=self.mixed_types(), + box=box, + ) + assert grid is not None + if box is not None: + # same wrapping convention as forward_common: grid points + # outside the primary cell are periodically equivalent + grid = normalize_coord( + grid, box.to(grid.device).reshape(box.shape[0], 3, 3) + ) + grid_type = torch.full( + grid.shape[:-1], + self.descriptor.get_ntypes() - 1, + device=grid.device, + dtype=atype.dtype, + ) + grid_nlist = build_directional_neighbor_list( + grid, + grid_type, + extended_coord, + extended_atype, + self.get_rcut(), + self.get_sel(), + distinguish_types=(not self.mixed_types()), + ) + atomic_ret = self.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + grid=grid, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + return {kk: vv.detach() for kk, vv in atomic_ret.items()} + + return model_forward + + def change_out_bias( + self, + 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. + + For density models, this operation is skipped because the output is + grid-based rather than atomic-based, and the standard atomic bias + adjustment (change-by-statistic / set-by-statistic) does not apply. + The fitting net will adapt to the target dataset through normal + gradient descent during training. + + This guard stays until output-stat support is designed for grid + models; tracked in https://github.com/deepmodeling/deepmd-kit/issues/6029. + """ + del ( + sample_merged, + stat_file_path, + bias_adjust_mode, + model_forward, + observed_type, + ) + log.warning("change_out_bias is not supported for density models; skipping.") + + def compute_or_load_stat( + self, + sampled_func: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + preset_observed_type: list[str] | None = None, + ) -> None: + """Compute or load statistics, with real statistics for the grid type. + + The reserved grid type (``ntypes - 1``) has no real atoms in any + training system, so the standard pass yields zero samples for it + and falls back to placeholder statistics (mean 0, stddev 0.1), + which are then applied to exactly the rows that produce the + density output. When the sampled data provides grids, a second + pass is run with the grid points injected as pseudo-atoms of the + grid type, and only the grid-type row of the statistics is taken + from it; the real types keep their clean statistics from the + standard pass. + + The injected pass only runs when the grid-type statistics are + missing (no cache, or a cache whose grid-type row has zero + samples), so the stat-file fast path still applies; when a stat + file is used, the patched grid-type row is written back to disk so + that the file and the in-memory statistics agree. + """ + sampled = None + grid_stats = None + if self._grid_stat_missing(stat_file_path): + sampled = sampled_func() if callable(sampled_func) else sampled_func + if any("grid" in sample for sample in sampled): + try: + # run the wrapped sampler first so the injected pass sees + # the same pair_exclude_types as the standard pass + wrapped = self._make_wrapped_sampler(lambda: sampled) + self.descriptor.compute_input_stats( + self._inject_grid_samples(wrapped()) + ) + grid_stats = [ + ( + block["davg"][-1].detach().clone(), + block["dstd"][-1].detach().clone(), + ) + for block in self._descriptor_stat_blocks() + ] + except (TypeError, KeyError) as err: + log.warning( + "Cannot compute input statistics for the grid type (%s); " + "falling back to the descriptor defaults.", + err, + ) + else: + log.warning( + "No grid data in the sampled frames; the grid type gets the " + "descriptor's default input statistics." + ) + super().compute_or_load_stat( + (lambda: sampled) if sampled is not None else sampled_func, + stat_file_path, + compute_or_load_out_stat=compute_or_load_out_stat, + preset_observed_type=preset_observed_type, + ) + if grid_stats is not None: + for block, (grid_davg, grid_dstd) in zip( + self._descriptor_stat_blocks(), grid_stats, strict=True + ): + davg = block["davg"] + dstd = block["dstd"] + davg[-1] = grid_davg.to(device=davg.device, dtype=davg.dtype) + dstd[-1] = grid_dstd.to(device=dstd.device, dtype=dstd.dtype) + if ( + stat_file_path is not None + and getattr(stat_file_path, "mode", None) != "r" + ): + self._write_back_grid_stat(stat_file_path, grid_stats) + elif stat_file_path is not None: + log.warning( + "Grid-type statistics were patched in memory but were not " + "written back because the statistics cache is read-only." + ) + + def _descriptor_stat_blocks(self) -> list: + """Collect every descriptor block that carries davg/dstd statistics. + + Descriptors such as DPA-2 and hybrids carry several blocks, each with + its own statistics; the grid-type row must be patched in all of them. + The blocks are discovered by walking the module tree (via + ``named_modules``) and probing the statistics item protocol, so no + per-descriptor attribute list is needed. + """ + blocks = [] + for _name, module in self.descriptor.named_modules(): + # filter networks etc. also implement __getitem__ with unrelated + # semantics, so only descriptor-level modules are probed for the + # statistics item protocol (some descriptors, e.g. se_r, carry + # the statistics on the descriptor itself) + if not isinstance(module, (BaseDescriptor, DescriptorBlock)): + continue + try: + module["davg"] + module["dstd"] + except (TypeError, KeyError): + continue + blocks.append(module) + if blocks: + return blocks + try: + self.descriptor["davg"] + self.descriptor["dstd"] + return [self.descriptor] + except (TypeError, KeyError): + pass + raise KeyError("davg/dstd not accessible on this descriptor") + + def _set_descriptor_env_protection(self, value: float) -> None: + """Set env_protection on every descriptor block, hybrids included. + + Wrapper-level fields, if any, are derived from the blocks; the whole + module tree is walked so no per-descriptor attribute list is needed. + """ + for _name, module in self.descriptor.named_modules(): + if hasattr(module, "env_protection"): + module.env_protection = value + + def _stat_cache_root(self, stat_file_path: DPPath) -> DPPath: + """Apply the same type_map subdirectory as the parent's stat path.""" + if self.type_map is not None: + stat_file_path = stat_file_path / " ".join(self.type_map) + return stat_file_path + + def _grid_stat_missing(self, stat_file_path: DPPath | None) -> bool: + """Tell whether the grid-type input statistics still need computing. + + True when there is no complete stat cache, or when the cache is + complete but its grid-type row has zero samples (placeholder + statistics saved by an older run). + """ + if stat_file_path is None: + return True + try: + blocks = self._descriptor_stat_blocks() + except KeyError: + # no accessible statistics: let the injected pass soft-fail below + return True + stat_file_path = self._stat_cache_root(stat_file_path) + grid_type = self.descriptor.get_ntypes() - 1 + for block in blocks: + env_stat = EnvMatStatSe(block) + cache = stat_file_path / env_stat.get_hash() + keys = env_stat.get_stat_keys() + if not cache.is_dir() or any(not (cache / kk).is_file() for kk in keys): + return True + # Every component of the grid-type row must carry samples. In + # particular, a valid radial row must not hide a placeholder + # angular row left by an interrupted or legacy cache update. + suffix = f"_{grid_type}" + for key in keys: + if key.endswith(suffix) and float((cache / key).load_numpy()[0]) == 0.0: + return True + return False + + def _write_back_grid_stat( + self, + stat_file_path: DPPath, + grid_stats: list[tuple[torch.Tensor, torch.Tensor]], + ) -> None: + """Write the patched grid-type row back to the stat cache files. + + The cache stores raw sums (number, sum, squared_sum) per stat key; + a single sample reproduces the patched mean and stddev exactly. + """ + grid_type = self.descriptor.get_ntypes() - 1 + stat_file_path = self._stat_cache_root(stat_file_path) + for block, (grid_davg, grid_dstd) in zip( + self._descriptor_stat_blocks(), grid_stats, strict=True + ): + env_stat = EnvMatStatSe(block) + cache = stat_file_path / env_stat.get_hash() + if not cache.is_dir(): + continue + davg_np = grid_davg.detach().cpu().numpy() + dstd_np = grid_dstd.detach().cpu().numpy() + mean_r = float(davg_np[0, 0]) + std_r = float(dstd_np[0, 0]) + (cache / f"r_{grid_type}").save_numpy( + np.array([1.0, mean_r, std_r**2 + mean_r**2]) + ) + if dstd_np.shape[-1] == 4: + std_a = float(dstd_np[0, 1]) + (cache / f"a_{grid_type}").save_numpy(np.array([1.0, 0.0, std_a**2])) + + def _inject_grid_samples(self, sampled: list[dict]) -> list[dict]: + """Append the grid points of each sample as pseudo-atoms of the + reserved grid type, so the descriptor input statistics get real + samples for it. Neighbor lists are rebuilt by the stat machinery, + so only ``coord`` and ``atype`` need to be extended. + + Grid-grid pairs are excluded from the injected neighbor lists so + that the grid-type statistics are computed over grid-to-atom + pairs, matching the directional neighbor list of the forward pass + (a no-op for per-type-sel descriptors with ``sel[-1] == 0``, but + required for mixed-type descriptors with a scalar sel). + """ + ntypes = self.descriptor.get_ntypes() + grid_type = ntypes - 1 + injected = [] + for sample in sampled: + grid = sample.get("grid") + if grid is None: + injected.append(sample) + continue + sample = dict(sample) + coord = sample["coord"] + atype = sample["atype"] + nframes = atype.shape[0] + gg = grid.reshape(nframes, -1, 3).to(coord.dtype) + ngrid = gg.shape[1] + sample["coord"] = torch.cat( + [coord.reshape(nframes, -1, 3), gg], dim=1 + ).reshape(nframes, -1) + sample["atype"] = torch.cat( + [ + atype, + torch.full( + (nframes, ngrid), + grid_type, + dtype=atype.dtype, + device=atype.device, + ), + ], + dim=1, + ) + pair_exclude_types = [ + tuple(pair) for pair in sample.get("pair_exclude_types", []) + ] + if (grid_type, grid_type) not in pair_exclude_types: + pair_exclude_types.append((grid_type, grid_type)) + sample["pair_exclude_types"] = pair_exclude_types + injected.append(sample) + return injected + + def compute_or_load_out_stat( + self, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the output statistics for the fitting net. + + Not implemented for density models: the output lives on grid points, + so per-type atomic output statistics do not apply. This guard stays + until output-stat support is designed for grid models; tracked in + https://github.com/deepmodeling/deepmd-kit/issues/6029. + """ + log.warning("Not implemented yet for density out stat!") diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d2e0a09e96..c3d6221b14 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -49,6 +49,9 @@ env_protection_descriptors, ) +from .density_model import ( + GridDensityModel, +) from .dipole_model import ( DipoleModel, ) @@ -433,6 +436,8 @@ def get_standard_model(model_params: dict) -> BaseModel: modelcls = DOSModel elif fitting_net_type in ["ener", "direct_force_ener"]: modelcls = EnergyModel + elif fitting_net_type == "density": + modelcls = GridDensityModel elif fitting_net_type == "property": modelcls = PropertyModel elif fitting_net_type == "population": diff --git a/deepmd/pt/model/model/density_model.py b/deepmd/pt/model/model/density_model.py new file mode 100644 index 0000000000..a319bf6cf5 --- /dev/null +++ b/deepmd/pt/model/model/density_model.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) +from typing import ( + Any, +) + +import torch + +from deepmd.pt.model.atomic_model import ( + DPDensityAtomicModel, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_density_model import ( + make_density_model, +) + +DPDensityModel_ = make_density_model(DPDensityAtomicModel) + + +@BaseModel.register("density") +@BaseModel.register("grid_density") +class GridDensityModel(DPModelCommon, DPDensityModel_): + model_type = "grid_density" + + def __init__( + self, + *args: Any, + **kwargs: Any, + ) -> None: + DPModelCommon.__init__(self) + DPDensityModel_.__init__(self, *args, **kwargs) + + def translated_output_def(self) -> dict[str, Any]: + out_def_data = self.model_output_def().get_data() + output_def = { + "density": deepcopy(out_def_data["density"]), + } + if "mask" in out_def_data: + output_def["mask"] = deepcopy(out_def_data["mask"]) + return output_def + + @torch.jit.export + def has_grid(self) -> bool: + """Returns whether it has grid input and output.""" + return True + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + assert grid is not None + model_ret = self.forward_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + grid=grid, + ) + model_predict = {} + model_predict["density"] = model_ret["density"] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + comm_dict: dict[str, torch.Tensor] | None = None, + charge_spin: torch.Tensor | None = None, + ) -> None: + raise NotImplementedError("forward_lower is not supported for density models.") diff --git a/deepmd/pt/model/model/make_density_model.py b/deepmd/pt/model/model/make_density_model.py new file mode 100644 index 0000000000..b41daf0180 --- /dev/null +++ b/deepmd/pt/model/model/make_density_model.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + Any, +) + +import torch + +from deepmd.pt.model.atomic_model.base_atomic_model import ( + BaseAtomicModel, +) +from deepmd.pt.model.model.make_model import ( + make_model, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.model.model.transform_output import ( + communicate_extended_output, + fit_output_to_model_output, +) +from deepmd.pt.utils.nlist import ( + build_directional_neighbor_list, + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.region import ( + normalize_coord, +) + +log = logging.getLogger(__name__) + + +def make_density_model(T_AtomicModel: type[BaseAtomicModel]) -> type[BaseModel]: + """Make a density model as a derived class of an atomic model. + + The model predicts a scalar density on grid points instead of atomic + energies. It is built by subclassing the standard model made by + `make_model` and overriding only the interfaces that differ: + + 1. the `forward_common_lower`, that takes extended coordinates, atypes, + neighbor list and the grid (with its type and directional neighbor list), + and outputs the density on the grid points; + + 2. the `forward_common`, that takes coordinates, atypes, cell and grid, + and predicts the density on the grid points. + + Parameters + ---------- + T_AtomicModel + The atomic model. + + Returns + ------- + CM + The model. + + """ + + class CM(make_model(T_AtomicModel)): + def forward_common( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + coord_corr_for_virial: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return model prediction. + + The base parameters keep the base positional order (including + ``coord_corr_for_virial``, which is accepted and ignored since + density has no virial correction); the density-specific + ``grid`` is appended at the end and is required for this model. + + Parameters + ---------- + coord + The coordinates of the atoms. + shape: nf x (nloc x 3) + atype + The type of atoms. shape: nf x nloc + box + The simulation box. shape: nf x 9 + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + If calculate the atomic virial. + grid + The coordinates of the grids. + shape: nf x (ngrid x 3) + + Returns + ------- + ret_dict + The result dict of type dict[str,torch.Tensor]. + The keys are defined by the `ModelOutputDef`. + + """ + del coord_corr_for_virial + assert grid is not None + cc, gg, bb, fp, ap, input_prec = self._input_type_cast( + coord, box=box, fparam=fparam, aparam=aparam, grid=grid + ) + del coord, grid, box, fparam, aparam + gg = gg.view(gg.shape[0], -1, 3) + if gg.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + if bb is not None: + # wrap grid points into the primary cell: the atomic + # coordinates are normalized and extended with ghosts, so + # periodically equivalent grids outside the cell would + # otherwise lose their neighbors + gg = normalize_coord(gg, bb.to(gg.device).view(bb.shape[0], 3, 3)) + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + cc, + atype, + self.get_rcut(), + self.get_sel(), + mixed_types=self.mixed_types(), + box=bb, + ) + grid_type = torch.full( + (gg.shape[0], gg.shape[1]), + self.atomic_model.descriptor.get_ntypes() - 1, + device=gg.device, + dtype=atype.dtype, + ) + grid_nlist = build_directional_neighbor_list( + gg, + grid_type, + extended_coord, + extended_atype, + self.get_rcut(), + self.get_sel(), + distinguish_types=(not self.mixed_types()), + ) + model_predict_lower = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + grid=gg, + grid_type=grid_type, + grid_nlist=grid_nlist, + mapping=mapping, + do_atomic_virial=do_atomic_virial, + fparam=fp, + aparam=ap, + charge_spin=charge_spin, + ) + model_predict = communicate_extended_output( + model_predict_lower, + self.model_output_def(), + mapping, + do_atomic_virial=do_atomic_virial, + ) + model_predict = self._output_type_cast(model_predict, input_prec) + return model_predict + + def forward_common_lower( + self, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + comm_dict: dict[str, torch.Tensor] | None = None, + extra_nlist_sort: bool = False, + extended_coord_corr: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + grid_type: torch.Tensor | None = None, + grid_nlist: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return model prediction. Lower interface that takes + extended atomic coordinates and types, nlist, mapping, and the + grid inputs as input, and returns the predictions on the grid + points. The predictions are not reduced. + + The base parameters keep the base positional order (including + ``extended_coord_corr``, which is accepted and ignored); the + density-specific grid parameters are appended at the end and + are required for this model. + + Parameters + ---------- + extended_coord + coodinates in extended region. nf x (nall x 3) + extended_atype + atomic type in extended region. nf x nall + nlist + neighbor list. nf x nloc x nsel. + grid + grid coordinates. nf x ngrid x 3 + grid_type + type of the grid points. nf x ngrid + grid_nlist + directional neighbor list from grid points to atoms. + nf x ngrid x nsel + mapping + mapps the extended indices to local indices. nf x nall. + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + whether calculate atomic virial. + comm_dict + The data needed for communication for parallel inference. + extra_nlist_sort + whether to forcibly sort the nlist. + + Returns + ------- + result_dict + the result dict, defined by the `FittingOutputDef`. + + """ + del extended_coord_corr + nframes, nall = extended_atype.shape[:2] + extended_coord = extended_coord.view(nframes, -1, 3) + nlist = self.format_nlist( + extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort + ) + assert grid is not None + assert grid_type is not None + assert grid_nlist is not None + if grid.shape[1] == 0: + raise ValueError("grid must contain at least one point per frame") + cc_ext, gg, _, fp, ap, input_prec = self._input_type_cast( + extended_coord, grid=grid, fparam=fparam, aparam=aparam + ) + del extended_coord, grid, fparam, aparam + atomic_ret = self.atomic_model.forward_common_atomic( + cc_ext, + extended_atype, + nlist, + mapping=mapping, + fparam=fp, + aparam=ap, + comm_dict=comm_dict, + grid=gg, + grid_type=grid_type, + grid_nlist=grid_nlist, + charge_spin=charge_spin, + ) + model_predict = fit_output_to_model_output( + atomic_ret, + self.atomic_output_def(), + cc_ext, + do_atomic_virial=do_atomic_virial, + create_graph=self.training, + ) + model_predict = self._output_type_cast(model_predict, input_prec) + return model_predict + + def _input_type_cast( + self, + coord: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + str, + ]: + """Cast the input data to global float type.""" + input_prec = self.reverse_precision_dict[coord.dtype] + # dtype mismatch warnings are not emitted here: type checking + # would not pass jit, so inputs are converted to coord prec anyway. + _lst: list[torch.Tensor | None] = [ + vv.to(coord.dtype) if vv is not None else None + for vv in [grid, box, fparam, aparam] + ] + grid, box, fparam, aparam = _lst + assert grid is not None + if ( + input_prec + == self.reverse_precision_dict[self.global_pt_float_precision] + ): + return coord, grid, box, fparam, aparam, input_prec + else: + pp = self.global_pt_float_precision + return ( + coord.to(pp), + grid.to(pp), + box.to(pp) if box is not None else None, + fparam.to(pp) if fparam is not None else None, + aparam.to(pp) if aparam is not None else None, + input_prec, + ) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + # directly call the forward_common method when no specific transform rule + assert grid is not None + return self.forward_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + grid=grid, + ) + + @torch.jit.export + def forward_embedding( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + # the base implementation would call the grid-aware caster without + # a grid; density models do not support embedding extraction + raise NotImplementedError( + "forward_embedding is not supported for density models." + ) + + def change_out_bias( + self, + merged: Any, + bias_adjust_mode: str = "change-by-statistic", + preset_out_bias: dict | None = None, + ) -> None: + """Change the output bias according to the input data. + + Not supported for density models: the output is defined on grid + points rather than atoms, so the standard bias adjustment (which + would run a grid-less forward through the stat wrapper) does not + apply. Overridden at the model level because the default + ``change-by-statistic`` mode never reaches the atomic-level + no-op. + """ + del merged, bias_adjust_mode, preset_out_bias + log.warning( + "change_out_bias is not supported for density models; skipping." + ) + + return CM diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index f9f8e6561d..c6e7df38ed 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -2,6 +2,9 @@ from .base_fitting import ( BaseFitting, ) +from .density import ( + DensityFittingNet, +) from .dipole import ( DipoleFittingNet, ) @@ -31,6 +34,7 @@ __all__ = [ "BaseFitting", "DOSFittingNet", + "DensityFittingNet", "DipoleFittingNet", "EnergyFittingNet", "EnergyFittingNetDirect", diff --git a/deepmd/pt/model/task/density.py b/deepmd/pt/model/task/density.py new file mode 100644 index 0000000000..5f386ef540 --- /dev/null +++ b/deepmd/pt/model/task/density.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.task.fitting import ( + Fitting, + GeneralFitting, +) +from deepmd.pt.model.task.invar_fitting import ( + InvarFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("density") +class DensityFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + dim_descrpt: int, + neuron: list[int] = [128, 128, 128], + bias_atom_e: torch.Tensor | None = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + mixed_types: bool = True, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + "density", + ntypes, + dim_descrpt, + 1, + neuron=neuron, + bias_atom_e=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + seed=seed, + type_map=type_map, + **kwargs, + ) + if numb_aparam > 0: + raise ValueError( + "density fitting does not support atomic parameters (aparam): " + "the fitting net consumes the grid-point descriptor rows, " + "which have no per-atom parameters" + ) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 5, 1) + data.pop("var_name") + data.pop("dim_out") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + **super().serialize(), + "type": "density", + } + + # make jit happy with torch 2.0.0 + exclude_types: list[int] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7efc76e4d9..5a1b904d40 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -51,6 +51,7 @@ DOSLoss, EnergySpinLoss, EnergyStdLoss, + GridDensityLoss, PopulationLoss, PropertyLoss, TaskLoss, @@ -2231,6 +2232,7 @@ def get_data( "coord", "atype", "spin", + "grid", "box", "fparam", "aparam", @@ -2255,6 +2257,10 @@ def get_data( if "fid" in batch_data: log_dict["fid"] = batch_data["fid"] log_dict["sid"] = batch_data["sid"] + # models without grid support do not accept a grid keyword in forward; + # drop it when the data does not provide grid + if input_dict.get("grid") is None: + input_dict.pop("grid", None) return input_dict, label_dict, log_dict def print_header( @@ -2423,6 +2429,22 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ] additional_data_requirement += spin_requirement_items + has_grid = getattr(_model, "has_grid", False) + if callable(has_grid): + has_grid = has_grid() + if has_grid: + # the grid is a model input (like fparam/aparam), not a label; it is + # frame-major (nframes, ngrid, 3) and not per-atom + additional_data_requirement.append( + DataRequirementItem( + "grid", + ndof=3, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + ) if _model.has_chg_spin_ebd(): has_default_cs = _model.has_default_chg_spin() cs_default = ( @@ -2526,6 +2548,9 @@ def get_loss( tensor_name = "polar" loss_params["tensor_name"] = tensor_name return TensorLoss(**loss_params) + elif loss_type == "grid_density": + loss_params["starter_learning_rate"] = start_lr + return GridDensityLoss(**loss_params) elif loss_type == "property": task_dim = _model.get_task_dim() var_name = _model.get_var_name() diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index b2e4d4f229..7d7308b0e9 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -172,6 +172,7 @@ def forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, + grid: torch.Tensor | None = None, ) -> tuple[Any, Any, Any]: if not self.multi_task: task_key = "Default" @@ -194,6 +195,12 @@ def forward( if has_spin: input_dict["spin"] = spin + has_grid = getattr(self.model[task_key], "has_grid", False) + if callable(has_grid): + has_grid = has_grid() + if has_grid: + input_dict["grid"] = grid + # A loss-free wrapper is a pure inference object, so parameters can be # treated as constants while coordinate gradients remain enabled. if self.inference_only: diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index d3a3583938..616d9cae9a 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -74,6 +74,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N default=data_item["default"], dtype=data_item["dtype"], output_natoms_for_type_sel=data_item["output_natoms_for_type_sel"], + special_shape=getattr(data_item, "special_shape", None), length_scale=data_item["length_scale"], ) diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 49ffd995cf..373f3bfad0 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -400,11 +400,22 @@ def _compute_model_predict( charge_spin = system.get("charge_spin", None) spin = system.get("model_spin", system.get("spin", None)) + size_proxy = system["atype"].shape[-1] + if "grid" in system: + # the directional neighbor list is dense in ngrid x nall, so the + # batching size proxy must account for the grid extent. Note: no + # density model reaches this out-stat path today (out-stat is a + # documented no-op for grid models, see issue #6029), so this is + # for a future out-stat design + grid = system["grid"] + ngrid = grid.shape[-2] if grid.ndim >= 3 else grid.shape[-1] // 3 + size_proxy = max(size_proxy, ngrid) + def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: return auto_batch_size.execute_all( model_forward, nframes, - system["atype"].shape[-1], + size_proxy, *args, **kwargs, ) @@ -414,6 +425,8 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: "aparam": aparam, "charge_spin": charge_spin, } + if "grid" in system: + model_kwargs["grid"] = system["grid"] if spin is not None: model_kwargs["spin"] = spin sample_predict = model_forward_auto_batch_size( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..560c1f7d7e 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2928,6 +2928,68 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: fitting_args_plugin = ArgsPlugin() +@fitting_args_plugin.register("density", doc=supported_backends("pt")) +def fitting_density() -> list[Argument]: + doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." + doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." + doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." + doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' + doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_trainable = "Whether the parameters in the fitting net are trainable. This option can be\n\n\ +- bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\ +- list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of this list should be equal to len(`neuron`)+1." + doc_seed = "Random seed for parameter initialization of the fitting net" + + # numb_aparam is deliberately not advertised: the fitting consumes + # grid-point descriptor rows, which have no per-atom parameters. + # rcond is deliberately not advertised: the per-type output shift it + # would control is disabled for density models. + return [ + Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_default_fparam, + ), + Argument( + "dim_case_embd", + int, + optional=True, + default=0, + doc=doc_dim_case_embd, + ), + Argument( + "neuron", + list[int], + optional=True, + default=[120, 120, 120], + alias=["n_neuron"], + doc=doc_neuron, + ), + Argument( + "activation_function", + str, + optional=True, + default="tanh", + doc=doc_activation_function, + ), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + @fitting_args_plugin.register( "ener", doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + doc_ener, @@ -5183,6 +5245,28 @@ def loss_ener() -> list[Argument]: ] +@loss_args_plugin.register("grid_density", doc=supported_backends("pt")) +def loss_grid_density() -> list[Argument]: + doc_start_pref_d = start_pref("density", abbr="d") + doc_limit_pref_d = limit_pref("density") + return [ + Argument( + "start_pref_d", + [float, int], + optional=True, + default=1.00, + doc=doc_start_pref_d, + ), + Argument( + "limit_pref_d", + [float, int], + optional=True, + default=1.00, + doc=doc_limit_pref_d, + ), + ] + + @loss_args_plugin.register("dens", doc=supported_backends("pt")) def loss_dens() -> list[Argument]: doc_start_pref_e = start_pref("energy", abbr="e") @@ -6967,6 +7051,57 @@ def validate_no_multitask_lora(data: dict[str, Any], multi_task: bool = False) - ) +def _apply_density_env_protection_default(data: dict[str, Any]) -> None: + """Default env_protection to 1e-6 for density models. + + Applied at normalization time so that the recorded model_def_script and + the built model agree on this field. Grid points may legitimately + coincide with atoms, and the default 0.0 would let the 1/r terms in the + environment matrix produce NaN densities. 1e-6 is chosen as the minimal + perturbation that keeps those terms finite (order 1e6 at exact + coincidence) without shifting normal environments measurably; users can + still set any positive value (e.g. 0.1) explicitly. An explicit + ``env_protection: 0.0`` cannot be distinguished from the schema default + (which normalizes to 0.0) and receives the same override with a warning. + """ + + def _fix(model: dict[str, Any]) -> None: + if model.get("fitting_net", {}).get("type") != "density": + return + descriptor = model.get("descriptor", {}) + fill = 1e-6 + # a hybrid descriptor has no top-level env_protection; each + # sub-descriptor carries its own, and the model build requires all + # of them to agree. Filling each omitted sub-descriptor with 1e-6 + # while a sibling carries an explicit value (e.g. 0.1) would produce + # a mixed list that the build then rejects, so omitted entries take + # the siblings' common explicit value instead. + if descriptor.get("type") == "hybrid": + sub_descriptors = descriptor.get("list", []) + explicit = {sub.get("env_protection", 0.0) for sub in sub_descriptors} + explicit.discard(0.0) + if len(explicit) == 1: + fill = explicit.pop() + else: + sub_descriptors = [descriptor] + for sub in sub_descriptors: + if sub.get("env_protection", 0.0) == 0.0: + log.warning( + "env_protection is 0.0 for a density model; grid points " + "coincident with atoms would produce NaN densities. " + "Setting env_protection to %s.", + fill, + ) + sub["env_protection"] = fill + + model = data.get("model", {}) + if "model_dict" in model: + for sub_model in model["model_dict"].values(): + _fix(sub_model) + else: + _fix(model) + + def normalize( data: dict[str, Any], multi_task: bool = False, *, check: bool = True ) -> dict[str, Any]: @@ -6978,6 +7113,7 @@ def normalize( validate_full_validation_config(data, multi_task=multi_task) _check_dpa3_chg_spin_migration(data) validate_no_multitask_lora(data, multi_task=multi_task) + _apply_density_env_protection_default(data) return data diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 4d73e58d87..86414b90ce 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -208,6 +208,9 @@ def add( special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + ``"frame_major"`` stores one variable-length ``(npoints, ndof)`` + array per frame (e.g. grid/density), kept as a frame-major + ``(nframes, npoints, ndof)`` tensor without any natoms reshaping. length_scale : str, optional Per-element length scale a derived geometric item measures its default against: ``"absolute"`` gives every element the unit radius, @@ -493,16 +496,19 @@ def iter_test( chunk_atoms: int, numb_test: float = float("inf"), ) -> Iterator[dict]: - """Yield the test data in chunks of at most ``chunk_atoms`` atoms. + """Yield test data in chunks bounded by ``chunk_atoms`` work units. The complete test system is materialized first, so chunking here bounds the evaluation batch and its predictions rather than source-data memory. + For ordinary systems one frame costs ``natoms`` work units. For + frame-major data such as density grids, the largest point extent also + contributes so grids with many more points than atoms remain bounded. Parameters ---------- chunk_atoms : int - Upper bound on the number of atoms per chunk. A chunk always - carries at least one frame, however many atoms it has. + Upper bound on the per-chunk work proxy. A chunk always carries at + least one frame, however many atoms or frame-major points it has. numb_test : float, optional Upper bound on the number of frames served. A non-finite bound serves the whole test set. @@ -517,7 +523,17 @@ def iter_test( total = int(self.test_set["type"].shape[0]) if np.isfinite(numb_test): total = min(total, int(numb_test)) - step = max(1, int(chunk_atoms) // max(1, self.natoms)) + frame_size = self.natoms + for key, requirement in self.data_dict.items(): + if requirement.get("special_shape") != "frame_major": + continue + value = self.test_set.get(key) + if not isinstance(value, np.ndarray) or value.ndim < 2: + continue + ndof = int(requirement["ndof"]) + extent = int(np.prod(value.shape[1:], dtype=np.int64)) // ndof + frame_size = max(frame_size, extent) + step = max(1, int(chunk_atoms) // max(1, frame_size)) for begin in range(0, total, step): idx = np.arange(begin, min(begin + step, total), dtype=np.int64) chunk = self._get_subdata(self.test_set, idx=idx) @@ -674,12 +690,36 @@ def get_single_frame(self, index: int, num_worker: int) -> dict: # 5. Standardize keys frame_data = {kk.replace("atomic", "atom"): vv for kk, vv in frame_data.items()} - # 6. Reshape atomic data to match expected format [natoms, ndof] + # frame-major keys of one frame must agree on the point extent: a + # grid/density mismatch would otherwise broadcast silently (or fail + # with a raw size error) far downstream in the loss. The per-frame + # payload is (npoints, ndof) for a 3-D source array and (npoints*ndof,) + # for a flattened one. + frame_major_extents = {} + for kk, vv in self.data_dict.items(): + if ( + vv.get("special_shape") == "frame_major" + and kk in frame_data + and float(frame_data.get("find_" + kk, 0.0)) == 1.0 + ): + arr = frame_data[kk] + extent = arr.shape[-2] if arr.ndim >= 2 else arr.shape[-1] // vv["ndof"] + frame_major_extents[kk] = extent + if len(set(frame_major_extents.values())) > 1: + raise ValueError( + f"The frame-major data in {set_dir} have inconsistent " + f"extents {frame_major_extents}; grid and density labels " + "must share the same number of points" + ) + + # 6. Reshape non-atomic data to match expected format [ndof]; + # frame-major arrays keep their per-frame (npoints, ndof) layout for kk in self.data_dict.keys(): if ( "find_" not in kk and kk in frame_data and not self.data_dict[kk]["atomic"] + and self.data_dict[kk].get("special_shape") != "frame_major" ): frame_data[kk] = frame_data[kk].reshape(-1) frame_data["atype"] = frame_data["type"] @@ -812,6 +852,26 @@ def _load_test_set(self, shuffle_test: bool) -> None: if "find_" in kk: self.test_set[kk] = test_sets[0][kk] else: + vv = self.data_dict.get(kk, {}) + if vv.get("special_shape") == "frame_major": + # sets of one system must agree on the frame-major extent: + # numpy would otherwise fail with a raw shape error + extents = { + ( + ss[kk].shape[-2] + if ss[kk].ndim >= 3 + else ss[kk].shape[-1] // vv["ndof"] + ) + for ss in test_sets + if float(ss.get("find_" + kk, 0.0)) == 1.0 + } + if len(extents) > 1: + raise ValueError( + f"The frame-major data {kk} has inconsistent " + f"extents {sorted(extents)} across the sets of " + "this system; every set must carry the same " + "number of points" + ) self.test_set[kk] = np.concatenate( [test_set[kk] for test_set in test_sets], axis=0 ) @@ -828,9 +888,14 @@ def _shuffle_data(self, data: dict[str, Any]) -> dict[str, Any]: for kk in data: if ( isinstance(data[kk], np.ndarray) - and len(data[kk].shape) == 2 and data[kk].shape[0] == nframes and "find_" not in kk + # frame-major arrays (e.g. grid/density) keep their 3D layout + # (nframes, ngrid, ...) and shuffle along the frame axis + and ( + len(data[kk].shape) == 2 + or self.data_dict.get(kk, {}).get("special_shape") == "frame_major" + ) ): ret[kk] = data[kk][idx] else: @@ -923,6 +988,28 @@ def _load_set(self, set_name: DPPath) -> dict[str, Any]: np.reshape(tmp_in, [nframes, self.natoms, ndof]), axis=1 ) + # frame-major keys that share the same extent must agree on it: + # a grid/density pair with different npoints would only fail far + # downstream when the residual is formed + frame_major_extents = {} + for kk, vv in self.data_dict.items(): + if ( + vv.get("special_shape") == "frame_major" + and kk in data + and float(data.get("find_" + kk, 0.0)) == 1.0 + ): + arr = data[kk] + extent = arr.shape[-2] if arr.ndim >= 3 else arr.shape[-1] // vv["ndof"] + frame_major_extents[kk] = extent + if frame_major_extents: + extents = set(frame_major_extents.values()) + if len(extents) > 1: + raise ValueError( + f"The frame-major data in {set_name} have inconsistent " + f"extents {frame_major_extents}; grid and density labels " + "must share the same number of points" + ) + if self.mixed_type: # nframes x natoms atom_type_mix = self._load_type_mix(set_name) @@ -991,6 +1078,50 @@ def _load_data( special_shape: str | None = None, ) -> np.ndarray: is_hessian = special_shape == "hessian" or key == "hessian" + if dtype is not None: + pass + elif high_prec: + dtype = GLOBAL_ENER_FLOAT_PRECISION + else: + dtype = GLOBAL_NP_FLOAT_PRECISION + path = self._get_data_path(set_name, key) + # Frame-major data does not have an atom axis, so validate and return + # it before preparing atom-selection bookkeeping. + if special_shape == "frame_major" and path.is_file(): + data = path.load_numpy().astype(dtype) + if data.ndim not in (2, 3): + raise ValueError( + f"The frame-major data {key} in {set_name} must have shape " + f"(nframes, npoints * ndof) or (nframes, npoints, ndof), " + f"but got {data.shape}" + ) + if data.shape[0] != nframes: + raise ValueError( + f"The frame count of data {key} in {set_name} is " + f"{data.shape[0]}, which doesn't match the set's nframes {nframes}" + ) + if data.ndim >= 3: + # (nframes, npoints, ndof): trailing dim must be exactly ndof + if data.shape[-1] != ndof_: + raise ValueError( + f"The data {key} in {set_name} has trailing dimension " + f"{data.shape[-1]}, which doesn't match the declared " + f"ndof {ndof_}" + ) + elif data.shape[-1] % ndof_ != 0: + # (nframes, npoints*ndof) flattened: width must be a multiple + raise ValueError( + f"The data {key} in {set_name} has trailing dimension " + f"{data.shape[-1]}, which is not a multiple of the " + f"declared ndof {ndof_}" + ) + npoints = data.shape[1] if data.ndim == 3 else data.shape[1] // ndof_ + if npoints == 0: + raise ValueError( + f"The frame-major data {key} in {set_name} must contain " + "at least one point per frame" + ) + return np.float32(1.0), data if atomic: natoms = self.natoms idx_map = self.idx_map @@ -1008,13 +1139,6 @@ def _load_data( ndof = ndof_ natoms_sel = 0 idx_map_sel = None - if dtype is not None: - pass - elif high_prec: - dtype = GLOBAL_ENER_FLOAT_PRECISION - else: - dtype = GLOBAL_NP_FLOAT_PRECISION - path = self._get_data_path(set_name, key) if path.is_file(): data = path.load_numpy().astype(dtype) try: # YWolfeee: deal with data shape error @@ -1081,7 +1205,12 @@ def _load_data( ndof = 9 * self.natoms * self.natoms elif atomic and type_sel is not None and not output_natoms_for_type_sel: ndof = ndof_ * natoms_sel - data = np.full([nframes, ndof], default, dtype=dtype) + if special_shape == "frame_major": + # frame extent is unknown without the file; keep an empty + # per-frame payload and rely on find=0 downstream + data = np.full([nframes, 0, ndof], default, dtype=dtype) + else: + data = np.full([nframes, ndof], default, dtype=dtype) if repeat != 1: data = np.repeat(data, repeat).reshape([nframes, -1]) return np.float32(0.0), data @@ -1146,6 +1275,10 @@ def _load_single_data( # Create a default array based on requirements if is_hessian: data = np.full([9 * natoms * natoms], vv["default"], dtype=dtype) + elif vv.get("special_shape") == "frame_major": + # frame extent is unknown without the file; an empty frame + # dimension is the honest shape and find=0 marks it absent + data = np.full([0, ndof], vv["default"], dtype=dtype) elif vv["atomic"]: if vv["type_sel"] is not None and not vv["output_natoms_for_type_sel"]: natoms = natoms_sel @@ -1163,13 +1296,53 @@ def _load_single_data( # For filesystem paths, use memmap for better performance mmap_obj = self._get_memmap(path) - # corner case: single frame - if set_nframes == 1: + # Validate frame-major layout before indexing, so that shuffling + # cannot pair a structure with another frame's grid or density label. + if vv.get("special_shape") == "frame_major": + if mmap_obj.ndim not in (2, 3): + raise ValueError( + f"The frame-major data {key} in {set_dir} must have shape " + f"(nframes, npoints * ndof) or (nframes, npoints, ndof), " + f"but got {mmap_obj.shape}" + ) + if mmap_obj.shape[0] != set_nframes: + raise ValueError( + f"The frame count of data {key} in {set_dir} is " + f"{mmap_obj.shape[0]}, which doesn't match the set's " + f"nframes {set_nframes}" + ) + if mmap_obj.ndim >= 3: + if mmap_obj.shape[-1] != ndof: + raise ValueError( + f"The data {key} in {set_dir} has trailing dimension " + f"{mmap_obj.shape[-1]}, which doesn't match the declared " + f"ndof {ndof}" + ) + elif mmap_obj.shape[-1] % ndof != 0: + raise ValueError( + f"The data {key} in {set_dir} has trailing dimension " + f"{mmap_obj.shape[-1]}, which is not a multiple of the " + f"declared ndof {ndof}" + ) + npoints = ( + mmap_obj.shape[1] if mmap_obj.ndim == 3 else mmap_obj.shape[1] // ndof + ) + if npoints == 0: + raise ValueError( + f"The frame-major data {key} in {set_dir} must contain " + "at least one point per frame" + ) + + # corner case: single frame. frame-major arrays always carry a + # leading frame dimension (validated above), so no expansion is needed. + if set_nframes == 1 and vv.get("special_shape") != "frame_major": mmap_obj = mmap_obj[None, ...] # Slice the single frame and make an in-memory copy for modification data = mmap_obj[frame_idx].copy().astype(dtype, copy=False) try: + if vv.get("special_shape") == "frame_major": + return np.float32(1.0), data if is_hessian: data = data.reshape(3 * natoms, 3 * natoms) num_chunks, chunk_size = len(idx_map), 3 @@ -1323,6 +1496,9 @@ class DataRequirementItem: special_shape : str, optional Name of a loader-defined non-standard shape contract. ``"hessian"`` stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. + ``"frame_major"`` stores one variable-length ``(npoints, ndof)`` array + per frame (e.g. grid/density), kept frame-major without natoms + reshaping. source_policy : {"tracked", "default", "derived"}, optional How source availability affects the consumer. ``"tracked"`` keeps source presence in the ``find_*`` contract. ``"default"`` treats the diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 890224e063..766efd0239 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -585,6 +585,15 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: ) b_data[kk] = merged.reshape(len(batch_data), -1) elif not vv["atomic"]: + if vv.get("special_shape") == "frame_major": + extents = {bb[kk].shape[1:] for bb in batch_data} + if len(extents) > 1: + raise ValueError( + f"The frame-major data {kk} has inconsistent " + f"extents {sorted(extents)} across systems; " + "mixed batching requires the same extent in " + "every system" + ) b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0) else: b_data[kk] = np.zeros( diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..854c4b11a4 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -22,6 +22,7 @@ Model train-fitting-tensor train-fitting-dos train-fitting-property + train-fitting-density train-se-e2-a-tebd train-se-a-mask train-se-e3-tebd diff --git a/doc/model/train-fitting-density.md b/doc/model/train-fitting-density.md new file mode 100644 index 0000000000..10fd86a2f1 --- /dev/null +++ b/doc/model/train-fitting-density.md @@ -0,0 +1,92 @@ +# Fit charge density on grid points {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch-TorchScript {{ pytorch_icon }} +> +> Training and evaluation are supported in Python (`dp train`, `dp test`, +> and the `DeepEval`/`DeepDensity` interfaces). C++/LAMMPS inference is +> **not** supported yet: the model's `forward_lower` interface is not +> implemented for density models. See the multi-backend tracking issue for +> the planned dpmodel implementation and backend wrappers. + +Here we present an API to the grid density model, which predicts the charge density on a set of grid points for a given atomic configuration. + +A complete example can be found in + +```bash +$deepmd_source_dir/examples/density/ +``` + +with DPA-2 and DPA-3 configurations, a small QM9 dataset, and a README describing the data format in detail. Note that **the data provided along with the examples are of limited amount, and should not be used to train a production model.** + +## Data format + +The training/validation data follows the standard `deepmd/npy` format, with two additional files in each `set.000/` directory: + +- `grid.npy` of shape `[nframes, ngrid, 3]`: the coordinates of the grid points; +- `density.npy` of shape `[nframes, ngrid, 1]`: the charge density labels on the grid points. + +The number of grid points `ngrid` may differ from the number of atoms. The last entry of `type_map` is reserved as a virtual "grid point type": internally, grid points are assigned this type when building the grid-to-atom neighbor list, so `type_map` must contain one more entry than the real element types. For descriptors with a per-type `sel` list (e.g. `se_e2_a`), set its last entry to 0, since grid points are only centers, never neighbors. + +Every frame in one system must use the same positive `ngrid`; empty grids are +rejected. Python inference evaluates large grids in chunks of 4096 points by +default, which bounds the temporary directional-neighbor matrix. Pass +`grid_batch_size=None` to `DeepDensity.eval` only when evaluating the complete +grid in one call is known to fit in memory. + +It is recommended to set a positive `env_protection` in the descriptor configuration (e.g. `"env_protection": 0.1`). Grid points may legitimately coincide with atoms, and the default `env_protection = 0.0` would let the `1/r` terms in the environment matrix produce `NaN` densities. If `env_protection` is not set for a density model, it is automatically set to `1e-6` with a warning. + +## The fitting network + +The {ref}`fitting_net ` section tells DP which fitting net to use. + +The JSON of `density` type should be provided like + +```json + "fitting_net" : { + "type": "density", + "neuron": [120,120,120], + "resnet_dt": true, + "seed": 1 + }, +``` + +- `type` specifies which type of fitting net should be used. It should be `density`. +- The rest arguments have the same meaning as they do in `ener` mode. `numb_aparam` is not supported, because the fitting net consumes the descriptor rows of the grid points, which have no per-atom parameters. + +## Loss + +The grid density model is trained with the `grid_density` loss, which minimises the squared error between the predicted and the reference densities on all grid points, consistent with the other losses in the package: + +```json + "loss" : { + "type": "grid_density", + "start_pref_d": 1.0, + "limit_pref_d": 1.0 + }, +``` + +- {ref}`type ` should be written as `grid_density`. +- `start_pref_d` and `limit_pref_d` specify the weight of the density loss at the start and at the end of the training. If both are set to 0, the density term is skipped (and the graph is kept connected). When either prefactor is nonzero, `density.npy` is **required** in every training and validation system; when both are zero, the file is optional. + +The loss gives every grid point equal weight. Non-uniform quadrature weights, +integrated-charge constraints, and output non-negativity are not imposed by +this fitting type; datasets that need those properties must account for them +in their sampling or preprocessing. + +## Evaluation + +`dp test` evaluates the model on a system containing `grid.npy` and `density.npy` and reports the MAE and RMSE of the density. For programmatic access, `DeepEval` dispatches density models to `DeepDensity`, whose `eval` takes the grid coordinates and returns the density on the grid points: + +```python +from deepmd.infer import DeepEval + +dp = DeepEval("model.pth") +density = dp.eval( + coord, + box, + atype, + grid=grid, + grid_batch_size=4096, +) # (nframes, ngrid) +``` diff --git a/examples/density/README.md b/examples/density/README.md new file mode 100644 index 0000000000..3e24d9705e --- /dev/null +++ b/examples/density/README.md @@ -0,0 +1,218 @@ +# Charge Density Prediction Example + +This example demonstrates how to train and evaluate a **charge density** model using DeePMD-kit with the PyTorch backend. + +The model predicts the charge density on a set of grid points (`grid`) for a given atomic configuration (`coord`, `atype`, `box`). + +______________________________________________________________________ + +## Directory Structure + +``` +. +├── dpa2/ # DPA-2 descriptor example +│ └── input.json +├── dpa3/ # DPA-3 descriptor example (recommended) +│ └── input.json +├── dataset/ +│ └── qm9/ +│ ├── C7H15NO_train/ # Training data (deepmd/npy format) +│ └── C7H15NO_val/ # Validation data (deepmd/npy format) +└── dptest_density_script.py # Standalone evaluation script (alternative to dp test) +``` + +______________________________________________________________________ + +## Data Format + +The training/validation data follows the standard **`deepmd/npy`** format, with two additional files in each `set.000/` directory: + +| File | Shape | Description | +| ----------------- | ----------------------- | ---------------------------------- | +| `coord.npy` | `[nframes, natoms * 3]` | Atomic coordinates | +| `box.npy` | `[nframes, 9]` | Simulation cell vectors | +| `type.raw` | `[natoms]` | Atom type indices | +| `type_map.raw` | `[ntypes]` | Real element types in this dataset | +| **`grid.npy`** | `[nframes, ngrid, 3]` | **Grid point coordinates** | +| **`density.npy`** | `[nframes, ngrid, 1]` | **Charge density labels on grids** | + +> **Notes:** +> +> - `grid.npy` and `density.npy` are required for the density model. The number of grid points (`ngrid`) must match between `grid.npy` and `density.npy`, and is allowed to differ from `natoms`. +> - The **last entry of the model configuration's `type_map` is reserved as a virtual "grid point type"** (e.g. `X` in the example): internally, grid points are assigned this type when building the grid-to-atom neighbor list. Make sure the model `type_map` contains one more entry than the real element types. `type_map.raw` may list only the real types actually present in that dataset, as the bundled example does. For descriptors with a per-type `sel` list (e.g. `se_e2_a`), set its **last entry to 0** (grid points are only centers, never neighbors); scalar `sel` descriptors such as DPA-2/DPA-3 need no change. The descriptor input statistics for the grid type are computed from the actual `grid.npy` data during training initialization. +> - Every frame in one system must contain the same positive number of grid points. Python inference chunks large grids into groups of 4096 points by default to limit temporary neighbor-list memory. +> - **Always set a positive `env_protection`** in the descriptor configuration (e.g. `"env_protection": 0.1`, as in `dpa3/input.json`). Grid points may legitimately coincide with atoms, and the default `env_protection = 0.0` would let the `1/r` terms in the environment matrix produce `NaN` densities and poison the training gradients at such points. + +______________________________________________________________________ + +## Training + +### 1. Choose a Configuration + +Two example configurations are provided: + +- **`dpa2/input.json`** — Uses the DPA-2 descriptor. +- **`dpa3/input.json`** — Uses the DPA-3 descriptor (recommended). + +Key settings in `input.json`: + +```json +{ + "model": { + "type_map": [ + "Li", + "Ni", + "Co", + "Mn", + "O", + "C", + "H", + "N", + "F", + "X" + ], + "descriptor": { + "type": "dpa3" + }, + "fitting_net": { + "type": "density", + "neuron": [ + 240, + 240, + 240 + ] + } + }, + "loss": { + "type": "grid_density", + "start_pref_d": 1, + "limit_pref_d": 1 + }, + "training": { + "training_data": { + "systems": [ + "../dataset/qm9/C7H15NO_train" + ], + "batch_size": "auto:128" + }, + "validation_data": { + "systems": [ + "../dataset/qm9/C7H15NO_val" + ], + "batch_size": 1, + "numb_btch": 3 + } + } +} +``` + +- The fitting net type is `"density"` and the loss type is `"grid_density"`. +- The model type `"grid_density"` is selected automatically from the fitting net type. + +### 2. Run Training + +```bash +cd dpa3 # or cd dpa2 +dp --pt train input.json +``` + +The training will output: + +- `model.ckpt-*.pt` — Model checkpoints +- `lcurve.out` — Training/validation loss curves (`mae_d` / `rmse_d` for density) +- `out.json` — Final training parameters + +### 3. Finetune from a Pretrained Model + +```bash +cd dpa3 +dp --pt train input.json --finetune model.ckpt-*.pt +``` + +> **Note:** For density models, `change_out_bias` (the energy-bias adjustment used in standard finetuning) is **automatically skipped** because density outputs are grid-based, not atomic-based. The descriptor weights are inherited, and the fitting net adapts via normal gradient descent. + +### 4. Freeze the Model + +Export a trained checkpoint into a frozen model for inference: + +```bash +cd dpa3 +dp --pt freeze +``` + +This reads `model.ckpt.pt` in the current directory and generates `frozen_model.pth`. + +______________________________________________________________________ + +## Testing / Evaluation + +### Option 1: `dp test` (recommended) + +`dp test` supports density models natively: + +```bash +cd dpa3 +dp --pt test -m frozen_model.pth -s ../dataset/qm9/C7H15NO_val -n 0 +``` + +- `-n 0` tests **all** frames; use a positive number to test only the first N frames. +- Add `--detail-file detail` to write per-frame `[label, prediction]` pairs to `detail.density.out.` (one file per frame, one line per grid point). +- Add `--shuffle-test --rand-seed 42` to test on a shuffled subset. + +Example output: + +``` +# testing system : ../dataset/qm9/C7H15NO_val +# number of test data : 30 +DENSITY MAE : x.xxxe+00 units +DENSITY RMSE : x.xxxe+00 units +``` + +### Option 2: Standalone script + +`dptest_density_script.py` evaluates a model and writes **one combined** `[prediction, label]` file (`result.d.out`), which is convenient for plotting: + +```bash +python dptest_density_script.py \ + dpa3/frozen_model.pth \ + dataset/qm9/C7H15NO_val \ + --ratio 1.0 \ + --output val_result.txt +``` + +| Argument | Description | +| ------------- | ------------------------------------------------------------------------- | +| `model` | Path to the frozen model file (`.pth`) | +| `data_dir` | Root directory of deepmd/npy datasets (searched recursively) | +| `--ratio` | Fraction of frames to randomly sample (default: `0.1`) | +| `--output` | If provided, save screen output to this file as well | +| `--pred-file` | File to save paired `[prediction, label]` array (default: `result.d.out`) | + +> **Note:** Each individual system must use one fixed `ngrid`, but different +> systems discovered by the script may use different grid extents. Large grids +> are evaluated in chunks by the Python inference interface. + +### Option 3: Python interface + +```python +import numpy as np +from deepmd.infer import DeepEval + +dp = DeepEval("frozen_model.pth") # dispatched to DeepDensity automatically +density = dp.eval( + coord, + box, + atype, + grid=grid, + grid_batch_size=4096, +) # (nframes, ngrid) +``` + +______________________________________________________________________ + +## Notes + +- **Backend:** PyTorch only (`--pt`). Make sure DeePMD-kit is installed with PyTorch support. +- **Deployment scope:** Python inference only (`dp test`, `DeepEval`/`DeepDensity`, and the standalone script). C++/LAMMPS inference is not supported yet (`forward_lower` is not implemented for density models). +- **Stat file:** `input.json` specifies `"stat_file": "./qm9_charge_density.hdf5"` for caching descriptor statistics; it is generated automatically on the first run. +- **Checkpoint vs. frozen model:** evaluation requires a frozen model (`.pth`). Freeze a checkpoint first with `dp --pt freeze`. diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy new file mode 100644 index 0000000000..37946ac645 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/box.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy new file mode 100644 index 0000000000..7927d65163 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/coord.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy new file mode 100644 index 0000000000..24289af963 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/density.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy b/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy new file mode 100644 index 0000000000..fb95783aa9 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_train/set.000/grid.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_train/type.raw b/examples/density/dataset/qm9/C7H15NO_train/type.raw new file mode 100644 index 0000000000..c46e48ed0a --- /dev/null +++ b/examples/density/dataset/qm9/C7H15NO_train/type.raw @@ -0,0 +1,24 @@ +5 +5 +5 +5 +5 +5 +5 +4 +7 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 diff --git a/examples/density/dataset/qm9/C7H15NO_train/type_map.raw b/examples/density/dataset/qm9/C7H15NO_train/type_map.raw new file mode 100644 index 0000000000..36c071bf5b --- /dev/null +++ b/examples/density/dataset/qm9/C7H15NO_train/type_map.raw @@ -0,0 +1,9 @@ +Li +Ni +Co +Mn +O +C +H +N +F diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy new file mode 100644 index 0000000000..05e520e22e Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/box.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy new file mode 100644 index 0000000000..f6cea16e30 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/coord.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy new file mode 100644 index 0000000000..e0c673e4e2 Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/density.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy b/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy new file mode 100644 index 0000000000..13a126f2bf Binary files /dev/null and b/examples/density/dataset/qm9/C7H15NO_val/set.000/grid.npy differ diff --git a/examples/density/dataset/qm9/C7H15NO_val/type.raw b/examples/density/dataset/qm9/C7H15NO_val/type.raw new file mode 100644 index 0000000000..c46e48ed0a --- /dev/null +++ b/examples/density/dataset/qm9/C7H15NO_val/type.raw @@ -0,0 +1,24 @@ +5 +5 +5 +5 +5 +5 +5 +4 +7 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 +6 diff --git a/examples/density/dataset/qm9/C7H15NO_val/type_map.raw b/examples/density/dataset/qm9/C7H15NO_val/type_map.raw new file mode 100644 index 0000000000..36c071bf5b --- /dev/null +++ b/examples/density/dataset/qm9/C7H15NO_val/type_map.raw @@ -0,0 +1,9 @@ +Li +Ni +Co +Mn +O +C +H +N +F diff --git a/examples/density/dpa2/input.json b/examples/density/dpa2/input.json new file mode 100644 index 0000000000..c25fa42d2c --- /dev/null +++ b/examples/density/dpa2/input.json @@ -0,0 +1,117 @@ +{ + "_comment": "that's all", + "model": { + "type_map": [ + "Li", + "Ni", + "Co", + "Mn", + "O", + "C", + "H", + "N", + "F", + "X" + ], + "descriptor": { + "type": "dpa2", + "repinit": { + "tebd_dim": 8, + "rcut": 6.0, + "rcut_smth": 0.5, + "nsel": 120, + "neuron": [ + 25, + 50, + 100 + ], + "axis_neuron": 12, + "activation_function": "tanh", + "three_body_sel": 48, + "three_body_rcut": 4.0, + "three_body_rcut_smth": 3.5, + "use_three_body": true + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 3.5, + "nsel": 48, + "nlayers": 12, + "g1_dim": 128, + "g2_dim": 32, + "attn2_hidden": 32, + "attn2_nhead": 4, + "attn1_hidden": 128, + "attn1_nhead": 4, + "axis_neuron": 4, + "update_h2": false, + "update_g1_has_conv": true, + "update_g1_has_grrg": true, + "update_g1_has_drrd": true, + "update_g1_has_attn": false, + "update_g2_has_g1g1": false, + "update_g2_has_attn": true, + "update_style": "res_residual", + "update_residual": 0.01, + "update_residual_init": "norm", + "attn2_has_gate": true, + "use_sqrt_nnei": true, + "g1_out_conv": true, + "g1_out_mlp": true + }, + "precision": "float32", + "add_tebd_to_repinit_out": false, + "env_protection": 0.1 + }, + "fitting_net": { + "type": "density", + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "precision": "float32", + "seed": 1, + "_comment": " that's all" + }, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment": "that's all" + }, + "loss": { + "type": "grid_density", + "start_pref_d": 1, + "limit_pref_d": 1, + "_comment": " that's all" + }, + "training": { + "stat_file": "./dpa2.hdf5", + "training_data": { + "systems": [ + "../dataset/qm9/C7H15NO_train" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../dataset/qm9/C7H15NO_val" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "numb_steps": 200, + "gradient_max_norm": 5.0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 2000, + "_comment": "that's all" + } +} diff --git a/examples/density/dpa3/input.json b/examples/density/dpa3/input.json new file mode 100644 index 0000000000..f23f5c6978 --- /dev/null +++ b/examples/density/dpa3/input.json @@ -0,0 +1,108 @@ +{ + "_comment": "that's all", + "model": { + "type_map": [ + "Li", + "Ni", + "Co", + "Mn", + "O", + "C", + "H", + "N", + "F", + "X" + ], + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 128, + "e_dim": 64, + "a_dim": 32, + "nlayers": 12, + "e_rcut": 6.0, + "e_rcut_smth": 5.3, + "e_sel": 120, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 30, + "axis_neuron": 4, + "skip_stat": true, + "a_compress_rate": 1, + "a_compress_e_rate": 2, + "a_compress_use_split": true, + "update_angle": true, + "smooth_edge_update": true, + "use_dynamic_sel": true, + "sel_reduce_factor": 10.0, + "edge_use_dist": true, + "use_env_envelope": true, + "optim_update": true, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const" + }, + "activation_function": "custom_silu:10.0", + "use_tebd_bias": false, + "env_protection": 0.1, + "precision": "float32", + "concat_output_tebd": false + }, + "fitting_net": { + "type": "density", + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "precision": "float32", + "activation_function": "silu", + "_comment": " that's all" + }, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 1e-05, + "_comment": "that's all" + }, + "loss": { + "type": "grid_density", + "start_pref_d": 1, + "limit_pref_d": 1, + "_comment": " that's all" + }, + "training": { + "stat_file": "./qm9_charge_density.hdf5", + "training_data": { + "systems": [ + "../dataset/qm9/C7H15NO_train" + ], + "batch_size": "auto:128", + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../dataset/qm9/C7H15NO_val" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" + }, + "numb_steps": 100, + "gradient_max_norm": 5.0, + "seed": 10, + "max_ckpt_keep": 100, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 10000, + "_comment": "that's all" + }, + "optimizer": { + "type": "AdamW" + } +} diff --git a/examples/density/dptest_density_script.py b/examples/density/dptest_density_script.py new file mode 100644 index 0000000000..9dfc2d7351 --- /dev/null +++ b/examples/density/dptest_density_script.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Evaluate a charge-density model on validation/test datasets. + +Usage +----- +python dptest_density_script.py model.pt /path/to/data --ratio 0.1 --output result.txt +""" + +import argparse +import glob +import math +import os +import random +import sys + +import numpy as np + +from deepmd.infer import ( + DeepEval, +) +from deepmd.utils.data import ( + DeepmdData, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Evaluate a charge-density model on deepmd/npy datasets." + ) + parser.add_argument("model", type=str, help="Path to the model file (.pt or .pth).") + parser.add_argument( + "data_dir", type=str, help="Root directory of deepmd/npy datasets." + ) + parser.add_argument( + "--ratio", + type=float, + default=0.1, + help="Fraction of frames to randomly sample from each system (default: 0.1).", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="If provided, write screen output to this file as well.", + ) + parser.add_argument( + "--pred-file", + type=str, + default="result.d.out", + help="File to save paired [prediction, label] array (default: result.d.out).", + ) + return parser.parse_args() + + +class TeeLogger: + """Redirect stdout to both the terminal and a file.""" + + def __init__(self, filepath: str) -> None: + self.terminal = sys.stdout + self.log = open(filepath, "w") + + def write(self, message: str) -> None: + self.terminal.write(message) + self.log.write(message) + + def flush(self) -> None: + self.terminal.flush() + self.log.flush() + + def close(self) -> None: + self.log.close() + + +def evaluate( + model_path: str, data_dir: str, ratio: float +) -> tuple[np.ndarray, np.ndarray]: + """Run inference and return (predictions, labels).""" + if not 0.0 < ratio <= 1.0: + raise ValueError(f"ratio must be in (0, 1], got {ratio}") + # DeepEval dispatches density models to DeepDensity automatically. + dm = DeepEval(model_path) + type_map = dm.get_type_map() + + pred_list, label_list = [], [] + pattern = os.path.join(data_dir, "**/type.raw") + systems = glob.glob(pattern, recursive=True) + + if not systems: + raise RuntimeError(f"No deepmd/npy systems found under '{data_dir}'") + + for f in systems: + sys_name = os.path.dirname(f) + data = DeepmdData(sys_name, type_map=type_map) + data.add( + "grid", + 3, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + data.add( + "density", + 1, + atomic=False, + must=True, + high_prec=True, + special_shape="frame_major", + ) + test_data = data.get_test() + + n_frames = test_data["coord"].shape[0] + n_sample = max(1, math.floor(n_frames * ratio)) + indices = random.sample(range(n_frames), n_sample) + + coord = test_data["coord"][indices].reshape(n_sample, -1) + if data.mixed_type: + atype = test_data["type"][indices].reshape(n_sample, -1) + else: + atype = test_data["type"][0] + cell = test_data["box"][indices] if data.pbc else None + grid = test_data["grid"][indices].reshape(n_sample, -1, 3) + + density_pred = dm.eval( + coord, + cell, + atype, + grid=grid, + mixed_type=data.mixed_type, + ) + density_label = test_data["density"][indices].reshape(n_sample, -1) + + # Flatten each system before concatenating so systems may carry + # different grid extents. + pred_list.append(density_pred.reshape(-1)) + label_list.append(density_label.reshape(-1)) + sys.stdout.write(f" {sys_name:60s} frames={n_sample}/{n_frames}" + "\n") + + predictions = np.concatenate(pred_list) + labels = np.concatenate(label_list) + return predictions, labels + + +def print_summary(pred: np.ndarray, label: np.ndarray) -> None: + diff = pred - label + mae = np.mean(np.abs(diff)) + rmse = np.sqrt(np.mean(diff**2)) + label_mean_abs = np.mean(np.abs(label)) + label_std = np.std(label) + + sys.stdout.write("\n" + "=" * 60 + "\n") + sys.stdout.write("Summary" + "\n") + sys.stdout.write("=" * 60 + "\n") + sys.stdout.write(f" Number of grid points : {label.size}" + "\n") + sys.stdout.write(f" Label std : {label_std:.6e}" + "\n") + sys.stdout.write(f" RMSE : {rmse:.6e}" + "\n") + sys.stdout.write(f" MAE : {mae:.6e}" + "\n") + sys.stdout.write(f" Mean |label| : {label_mean_abs:.6e}" + "\n") + sys.stdout.write( + f" epsilon_MAE (MAE/Mean|label|) : {mae / label_mean_abs:.6e}" + "\n" + ) + sys.stdout.write("=" * 60 + "\n") + + +def main() -> None: + args = parse_args() + + if args.output: + tee = TeeLogger(args.output) + sys.stdout = tee + sys.stdout.write(f"[INFO] Screen output will also be saved to: {args.output}\n") + + sys.stdout.write(f"[INFO] Model : {args.model}" + "\n") + sys.stdout.write(f"[INFO] Data : {args.data_dir}" + "\n") + sys.stdout.write(f"[INFO] Ratio : {args.ratio}\n") + + pred, label = evaluate(args.model, args.data_dir, args.ratio) + + # Save paired predictions & labels + out_array = np.stack([pred.reshape(-1), label.reshape(-1)], axis=1) + np.savetxt(args.pred_file, out_array) + sys.stdout.write(f"\n[INFO] Paired [pred, label] saved to: {args.pred_file}" + "\n") + + print_summary(pred, label) + + if args.output: + tee.close() + sys.stdout = tee.terminal + + +if __name__ == "__main__": + main() diff --git a/source/tests/common/test_deepmd_data_grid_density.py b/source/tests/common/test_deepmd_data_grid_density.py new file mode 100644 index 0000000000..2af6404cab --- /dev/null +++ b/source/tests/common/test_deepmd_data_grid_density.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the frame-major grid/density branches in DeepmdData. + +These keys are loaded verbatim (no reshaping to natoms) and must carry a +leading frame dimension matching the set's frame count, so that shuffling +never pairs a structure with another frame's grid or density label. +""" + +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.utils.data import ( + DeepmdData, +) + +NATOMS = 3 +NGRID = 5 +NFRAMES = 4 + + +class TestDeepmdDataGridDensity(unittest.TestCase): + def setUp(self) -> None: + self.tmpdir = tempfile.TemporaryDirectory() + self.root = Path(self.tmpdir.name) + self.set_dir = self.root / "set.000" + self.set_dir.mkdir() + rng = np.random.default_rng(42) + + np.savetxt(self.root / "type.raw", np.zeros(NATOMS, dtype=np.int32), fmt="%d") + np.save( + self.set_dir / "coord.npy", + rng.random((NFRAMES, NATOMS * 3), dtype=np.float32), + ) + np.save( + self.set_dir / "box.npy", + np.eye(3, dtype=np.float32).reshape(1, 9).repeat(NFRAMES, axis=0), + ) + np.save( + self.set_dir / "type.npy", + np.zeros((NFRAMES, NATOMS), dtype=np.int32), + ) + self.set_data("grid", rng.random((NFRAMES, NGRID, 3), dtype=np.float32)) + self.set_data("density", rng.random((NFRAMES, NGRID, 1), dtype=np.float32)) + + def tearDown(self) -> None: + self.tmpdir.cleanup() + + def set_data(self, key: str, arr: np.ndarray | None) -> None: + path = self.set_dir / f"{key}.npy" + if arr is None: + path.unlink(missing_ok=True) + else: + np.save(path, arr) + + def build_data(self, must: bool = True) -> DeepmdData: + data = DeepmdData(str(self.root)) + data.add("coord", 3, atomic=True, must=True) + data.add("box", 9, atomic=False, must=True) + data.add("type", 1, atomic=True, must=True) + data.add( + "grid", + 3, + atomic=False, + must=must, + high_prec=True, + special_shape="frame_major", + ) + data.add( + "density", + 1, + atomic=False, + must=must, + high_prec=True, + special_shape="frame_major", + ) + return data + + def test_load_set_frame_major(self) -> None: + data = self.build_data() + loaded = data._load_set(self.set_dir) + self.assertEqual(loaded["find_grid"], 1.0) + self.assertEqual(loaded["find_density"], 1.0) + self.assertEqual(loaded["grid"].shape, (NFRAMES, NGRID, 3)) + self.assertEqual(loaded["density"].shape, (NFRAMES, NGRID, 1)) + + def test_iter_test_uses_frame_major_extent_for_chunking(self) -> None: + data = self.build_data() + chunks = list(data.iter_test(chunk_atoms=NGRID * 2)) + self.assertEqual([chunk["box"].shape[0] for chunk in chunks], [2, 2]) + + def test_single_frame(self) -> None: + data = self.build_data() + frame = data.get_single_frame(0, 0) + self.assertEqual(frame["find_grid"], 1.0) + self.assertEqual(frame["find_density"], 1.0) + self.assertEqual(frame["grid"].shape, (NGRID, 3)) + self.assertEqual(frame["density"].shape, (NGRID, 1)) + + def test_frame_count_mismatch(self) -> None: + self.set_data("grid", np.zeros((NFRAMES - 1, NGRID, 3), dtype=np.float32)) + loaders = { + "set": lambda data: data._load_set(self.set_dir), + "single frame": lambda data: data.get_single_frame(0, 0), + } + for name, load in loaders.items(): + with self.subTest(loader=name): + with self.assertRaisesRegex(ValueError, "frame count"): + load(self.build_data()) + + def test_missing_frame_dimension(self) -> None: + # stored without a leading frame dimension: rejected rather than + # silently misaligned with the frames + self.set_data("grid", np.zeros((NGRID, 3), dtype=np.float32)) + loaders = { + "set": lambda data: data._load_set(self.set_dir), + "single frame": lambda data: data.get_single_frame(0, 0), + } + for name, load in loaders.items(): + with self.subTest(loader=name): + with self.assertRaisesRegex(ValueError, r"frame count|leading frame"): + load(self.build_data()) + + def test_missing_file_optional(self) -> None: + # optional labels absent: the loader must not fail and must signal + # find=0 so that downstream code can skip the residual + self.set_data("density", None) + data = self.build_data(must=False) + loaded = data._load_set(self.set_dir) + self.assertEqual(loaded["find_density"], 0.0) + frame = data.get_single_frame(0, 0) + self.assertEqual(frame["find_density"], 0.0) + + def test_invalid_frame_major_layout(self) -> None: + invalid = ( + ("wrong trailing width", (NFRAMES, NGRID, 4), "ndof"), + ("divisible wrong width", (NFRAMES, NGRID, 6), "ndof"), + ("invalid flattened width", (NFRAMES, NGRID * 3 + 1), "ndof"), + ("empty", (NFRAMES, 0, 3), "at least one point"), + ("rank four", (NFRAMES, 1, NGRID, 3), "must have shape"), + ) + loaders = { + "set": lambda data: data._load_set(self.set_dir), + "single frame": lambda data: data.get_single_frame(0, 0), + } + for case, shape, message in invalid: + self.set_data("grid", np.zeros(shape, dtype=np.float32)) + for loader_name, load in loaders.items(): + with self.subTest(case=case, loader=loader_name): + with self.assertRaisesRegex(ValueError, message): + load(self.build_data()) + + # A correctly flattened frame-major array remains accepted. + self.set_data("grid", np.zeros((NFRAMES, NGRID * 3), dtype=np.float32)) + loaded = self.build_data()._load_set(self.set_dir) + self.assertEqual(loaded["find_grid"], 1.0) + + def test_ragged_mixed_batch_clear_error(self) -> None: + # mixed batching across systems with different ngrid must fail with + # a clear message instead of a raw numpy shape error + from deepmd.utils.data_system import ( + DeepmdDataSystem, + ) + + class StubDS(DeepmdDataSystem): + """Test double: instantiated via __new__ to skip the heavy init.""" + + def get_ntypes(self) -> int: + return 2 + + def get_data_dict(self, idx: int) -> dict: + return { + "density": { + "atomic": False, + "special_shape": "frame_major", + "ndof": 1, + } + } + + batch_data = [ + { + "type": np.zeros((1, 3), dtype=np.int32), + "natoms_vec": np.array([3, 3, 1, 1], dtype=np.int32), + "default_mesh": np.zeros(9), + "density": np.zeros((1, 5), dtype=np.float32), + "find_density": 1.0, + }, + { + "type": np.zeros((1, 3), dtype=np.int32), + "natoms_vec": np.array([3, 3, 1, 1], dtype=np.int32), + "default_mesh": np.zeros(9), + "density": np.zeros((1, 6), dtype=np.float32), + "find_density": 1.0, + }, + ] + with self.assertRaisesRegex(ValueError, "extent"): + StubDS.__new__(StubDS)._merge_batch_data(batch_data) + + def test_grid_density_extent_mismatch(self) -> None: + # grid and density must carry the same number of points per frame + self.set_data("density", np.zeros((NFRAMES, NGRID + 2, 1), dtype=np.float32)) + with self.assertRaisesRegex(ValueError, "extents"): + self.build_data()._load_set(self.set_dir) + + def test_extent_mismatch_single_frame(self) -> None: + # the PT training path (get_item_torch -> get_single_frame) assembles + # each frame key by key; the extent cross-check must run there too, + # or a 1-point density broadcasts silently against the ngrid-point + # prediction in the loss + self.set_data("density", np.zeros((NFRAMES, 1, 1), dtype=np.float32)) + data = self.build_data() + with self.assertRaisesRegex(ValueError, "extents"): + data.get_item_torch(0) + + def test_cross_set_extent_mismatch(self) -> None: + # two sets of one system must agree on the frame-major extent; + # numpy would otherwise fail with a raw shape error in get_test + set1 = self.root / "set.001" + set1.mkdir() + rng = np.random.default_rng(1) + np.save(set1 / "coord.npy", rng.random((NFRAMES, NATOMS * 3), dtype=np.float32)) + np.save( + set1 / "box.npy", + np.eye(3, dtype=np.float32).reshape(1, 9).repeat(NFRAMES, axis=0), + ) + np.save(set1 / "type.npy", np.zeros((NFRAMES, NATOMS), dtype=np.int32)) + np.save( + set1 / "grid.npy", rng.random((NFRAMES, NGRID + 3, 3), dtype=np.float32) + ) + np.save( + set1 / "density.npy", rng.random((NFRAMES, NGRID + 3, 1), dtype=np.float32) + ) + with self.assertRaisesRegex(ValueError, "extents"): + self.build_data().get_test() + + def test_property_named_density_not_hijacked(self) -> None: + # a user property named "density" without a special_shape declaration + # must load through the ordinary path, not the frame-major branch + self.set_data("density", np.zeros((NFRAMES, 2), dtype=np.float32)) + data = DeepmdData(str(self.root)) + data.add("coord", 3, atomic=True, must=True) + data.add("box", 9, atomic=False, must=True) + data.add("type", 1, atomic=True, must=True) + data.add("density", 2, atomic=False, must=True) + loaded = data._load_set(self.set_dir) + self.assertEqual(loaded["find_density"], 1.0) + self.assertEqual(loaded["density"].shape, (NFRAMES, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index ff11aff68d..ba94b139d7 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -78,6 +78,8 @@ p_examples / "hessian" / "single_task" / "input.json", p_examples / "water" / "se_e2_a" / "input_torch_num_epoch.json", p_examples / "lmdb_downsample_data" / "input_lmdb.json", + p_examples / "density" / "dpa2" / "input.json", + p_examples / "density" / "dpa3" / "input.json", ) input_files_multi = ( diff --git a/source/tests/pt/test_dp_test_density.py b/source/tests/pt/test_dp_test_density.py new file mode 100644 index 0000000000..690867c1fa --- /dev/null +++ b/source/tests/pt/test_dp_test_density.py @@ -0,0 +1,1262 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import tempfile +import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from unittest.mock import ( + patch, +) + +import numpy as np +import torch + +from deepmd.entrypoints.test import test as dp_test +from deepmd.infer.deep_density import ( + DeepDensity, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) +from deepmd.infer.model_test import ( + build_tester, +) +from deepmd.pt.entrypoints.main import ( + get_trainer, +) +from deepmd.pt.model.atomic_model.density_atomic_model import ( + DPDensityAtomicModel, +) +from deepmd.pt.model.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.pt.model.descriptor.hybrid import ( + DescrptHybrid, +) +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.descriptor.se_r import ( + DescrptSeR, +) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.model.model.density_model import ( + GridDensityModel, +) +from deepmd.pt.model.task.density import ( + DensityFittingNet, +) +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.stat_file import ( + StatFileSpec, + open_stat_file, +) + +model_density = { + # the last entry is the reserved grid point type; its sel is 0 because + # grid points are never neighbors, only centers + "type_map": ["O", "H", "X"], + "descriptor": { + "type": "se_e2_a", + "sel": [8, 8, 0], + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [8, 16], + "resnet_dt": False, + "axis_neuron": 4, + # grid points may coincide with atoms; without protection the + # env matrix would divide by zero and produce NaN densities + "env_protection": 1e-6, + "seed": 1, + }, + "fitting_net": { + "type": "density", + "neuron": [8, 8], + "seed": 1, + }, +} + + +class TestDPTestDensity(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.tmpdir = tempfile.TemporaryDirectory() + root = Path(cls.tmpdir.name) + + # write a synthetic density system + rng = np.random.default_rng(42) + cls.nframes, natoms, cls.ngrid = 4, 5, 8 + cls.system = root / "system" + set_dir = cls.system / "set.000" + set_dir.mkdir(parents=True) + np.save(set_dir / "coord.npy", rng.random((cls.nframes, natoms * 3)) * 8.0) + np.save( + set_dir / "box.npy", + np.tile(np.eye(3).reshape(-1) * 10.0, (cls.nframes, 1)), + ) + np.save(set_dir / "grid.npy", rng.random((cls.nframes, cls.ngrid, 3)) * 8.0) + np.save(set_dir / "density.npy", rng.random((cls.nframes, cls.ngrid, 1))) + np.savetxt(cls.system / "type.raw", [0, 0, 1, 1, 1], fmt="%d") + np.savetxt(cls.system / "type_map.raw", ["O", "H", "X"], fmt="%s") + + cls.config = { + "model": deepcopy(model_density), + "learning_rate": { + "type": "exp", + "start_lr": 0.001, + "stop_lr": 1e-8, + "decay_steps": 10, + }, + "optimizer": {"type": "Adam"}, + "loss": { + "type": "grid_density", + "start_pref_d": 1.0, + "limit_pref_d": 1.0, + }, + "training": { + "training_data": { + "systems": [str(cls.system)], + "batch_size": 1, + }, + "validation_data": { + "systems": [str(cls.system)], + "batch_size": 1, + }, + "numb_steps": 1, + "seed": 1, + "disp_file": os.devnull, + "save_freq": 100, + }, + } + + # build and freeze a tiny density model through the trainer path + trainer = get_trainer(normalize(deepcopy(cls.config))) + with torch.device("cpu"): + input_dict, _, _ = trainer.get_data(is_train=False) + # the density model takes grid instead of spin as the extra input + input_dict.pop("spin", None) + trainer.model(**input_dict) + # keep the live torch model and a data sample for the gradient check + cls.torch_model = trainer.model + cls.input_dict = input_dict + model = torch.jit.script(trainer.model) + tmp_fd, cls.model_path = tempfile.mkstemp(suffix=".pth") + os.close(tmp_fd) + torch.jit.save(model, cls.model_path) + + @classmethod + def tearDownClass(cls) -> None: + os.unlink(cls.model_path) + cls.tmpdir.cleanup() + + def _stat_sample(self, module: torch.nn.Module) -> dict: + """Return a cloned sample on the device used by statistic collection.""" + device = next(module.parameters()).device + return { + key: value.detach().clone().to(device) + if isinstance(value, torch.Tensor) + else deepcopy(value) + for key, value in self.input_dict.items() + } + + def test_change_out_bias_noop(self) -> None: + # fine-tuning calls the model-level change_out_bias, whose default + # change-by-statistic mode would run the grid-less stat wrapper; + # for density models it must be a no-op instead + model = self.torch_model + bias_before = model.atomic_model.out_bias.detach().clone() + model.change_out_bias(self.input_dict) + model.change_out_bias(self.input_dict, bias_adjust_mode="set-by-statistic") + # Keep both no-op overrides compatible with their base APIs. + model.change_out_bias(self.input_dict, "set-by-statistic", {}) + model.atomic_model.change_out_bias( + self.input_dict, None, "set-by-statistic", None, None + ) + torch.testing.assert_close(model.atomic_model.out_bias, bias_before) + + def test_pair_excl_involving_grid_type_applies_in_forward(self) -> None: + # exclusions involving the grid type must drop the neighbor in the + # directional grid-to-atom list too, not just in the atom nlist + model = self.torch_model + grid_type = len(self.config["model"]["type_map"]) - 1 + input_dict = { + kk: vv.clone() if isinstance(vv, torch.Tensor) else vv + for kk, vv in self.input_dict.items() + } + model.atomic_model.reinit_pair_exclude([(0, grid_type)]) + try: + out_excl = model(**input_dict)["density"] + finally: + model.atomic_model.reinit_pair_exclude([]) + out_none = model(**input_dict)["density"] + self.assertFalse(torch.allclose(out_excl, out_none)) + + def test_atom_excl_does_not_mask_grid(self) -> None: + # grid points carry the reserved type X; excluding a real atom type + # must not zero the density predictions (it would if grid_type were 0) + model = self.torch_model + model.atomic_model.reinit_atom_exclude([0]) + try: + input_dict = { + kk: vv.clone() if isinstance(vv, torch.Tensor) else vv + for kk, vv in self.input_dict.items() + } + out = model(**input_dict) + self.assertFalse(bool((out["density"] == 0).all())) + finally: + model.atomic_model.reinit_atom_exclude([]) + + def test_property_named_density_dispatches_to_property(self) -> None: + # the dispatch keys on the grid capability, not the output name, so + # a property fitting with property_name "density" is not hijacked + from deepmd.infer.deep_property import ( + DeepProperty, + ) + from deepmd.pt.model.model.property_model import ( + PropertyModel, + ) + from deepmd.pt.model.task.property import ( + PropertyFittingNet, + ) + + descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8]) + fitting = PropertyFittingNet( + descriptor.get_ntypes(), + descriptor.get_dim_out(), + "density", + neuron=[8, 8], + ) + model = PropertyModel(descriptor, fitting, type_map=["O", "H"]) + fd, path = tempfile.mkstemp(suffix=".pth") + os.close(fd) + try: + with torch.device("cpu"): + torch.jit.save(torch.jit.script(model), path) + dp = DeepEval(path) + self.assertIsInstance(dp, DeepProperty) + # the property must stay evaluable: the eval guard keys on the + # grid capability, not the output-variable name + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + result = dp.eval(coord, box, atype) + self.assertEqual(result[0].shape[0], 2) + finally: + os.unlink(path) + + def _run_training(self, config: dict, workdir: str) -> None: + """Run a few real training steps (checkpoints land in workdir).""" + cwd = os.getcwd() + os.chdir(workdir) + try: + trainer = get_trainer(normalize(deepcopy(config))) + trainer.run() + finally: + os.chdir(cwd) + self.assertTrue(os.path.exists(os.path.join(workdir, "model.ckpt.pt"))) + + def _training_config( + self, system: Path, start: float = 1.0, limit: float = 1.0 + ) -> dict: + config = deepcopy(self.config) + config["training"]["training_data"]["systems"] = [str(system)] + config["training"]["validation_data"]["systems"] = [str(system)] + config["training"]["training_data"]["batch_size"] = 2 + config["training"]["validation_data"]["batch_size"] = 2 + config["training"]["numb_steps"] = 3 + config["loss"]["start_pref_d"] = start + config["loss"]["limit_pref_d"] = limit + return config + + def test_eval_single_frame_grid_2d(self) -> None: + # a natural single-frame grid of shape (ngrid, 3) must be carried + # through with its frame dimension, not sliced to one grid point + dp = DeepDensity(self.model_path) + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:1] + box = np.load(set_dir / "box.npy")[:1] + grid = np.load(set_dir / "grid.npy")[0] # (ngrid, 3): no frame dim + atype = np.loadtxt(self.system / "type.raw", dtype=int) + out = dp.eval(coord, box, atype, grid=grid) + self.assertEqual(out.shape, (1, self.ngrid)) + + def test_eval_rejects_invalid_grid_inputs(self) -> None: + dp = DeepDensity(self.model_path) + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + grid = np.load(set_dir / "grid.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + invalid = ( + ("frame mismatch", grid[:1], {}, "frames"), + ("ambiguous rank two", grid[0], {}, "ambiguous"), + ("wrong coordinate width", np.zeros((2, self.ngrid, 2)), {}, "shape"), + ("rank four", np.zeros((2, self.ngrid, 3, 1)), {}, "shape"), + ("empty", np.zeros((2, 0, 3)), {}, "at least one point"), + ("invalid batch size", grid, {"grid_batch_size": 0}, "grid_batch_size"), + ) + for case, invalid_grid, kwargs, message in invalid: + with self.subTest(case=case): + with self.assertRaisesRegex(ValueError, message): + dp.eval(coord, box, atype, grid=invalid_grid, **kwargs) + + def test_fitting_rejects_aparam(self) -> None: + # grid descriptor rows have no per-atom parameters + descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0]) + with self.assertRaisesRegex(ValueError, "aparam"): + DensityFittingNet( + descriptor.get_ntypes(), + descriptor.get_dim_out(), + numb_aparam=2, + ) + + def test_fitting_model_serialization_roundtrip(self) -> None: + # GeneralFitting.serialize() emits @version 5; the density fitting + # and the whole model must round-trip through serialize/deserialize, + # which is what model freezing relies on. The model must go through + # BaseModel.deserialize, which dispatches on the serialized fitting + # type ("density") and is what dp convert-backend uses. + from deepmd.pt.loss.charge import ( + GridDensityLoss, + ) + from deepmd.pt.model.model.model import ( + BaseModel, + ) + + fitting = self.torch_model.atomic_model.fitting_net + restored_fit = DensityFittingNet.deserialize(fitting.serialize()) + self._assert_state_dict_equal(fitting, restored_fit) + model = self.torch_model + restored_model = BaseModel.deserialize(model.serialize()) + self.assertIsInstance(restored_model, GridDensityModel) + self._assert_state_dict_equal(model, restored_model) + + loss_fn = GridDensityLoss( + starter_learning_rate=0.001, start_pref_d=2.0, limit_pref_d=0.5 + ) + data = loss_fn.serialize() + self.assertEqual(data["@class"], "GridDensityLoss") + restored_loss = GridDensityLoss.deserialize(data) + self.assertEqual(restored_loss.serialize(), data) + + def _assert_state_dict_equal( + self, module0: torch.nn.Module, module1: torch.nn.Module + ) -> None: + sd0 = module0.state_dict() + sd1 = module1.state_dict() + self.assertEqual(set(sd0), set(sd1)) + for key in sd0: + torch.testing.assert_close(sd0[key], sd1[key]) + + def test_loss_inference_mode(self) -> None: + # inference=True reports the metrics even with zero prefactors + from deepmd.pt.loss.charge import ( + GridDensityLoss, + ) + from deepmd.pt.utils import ( + env, + ) + + loss_fn = GridDensityLoss(inference=True) + + class FakeModel(torch.nn.Module): + def forward(self, **kwargs): + return { + "density": torch.tensor([[[2.0], [4.0]]], device=env.DEVICE), + "mask": torch.tensor( + [[1, 1]], dtype=torch.int32, device=env.DEVICE + ), + } + + label = { + "density": torch.tensor([[[1.0], [2.0]]], device=env.DEVICE), + "find_density": 1.0, + } + _, loss, more_loss = loss_fn({}, FakeModel(), label, 1, 1.0) + self.assertIn("rmse_d", more_loss) + self.assertIn("mae_d", more_loss) + self.assertAlmostEqual(more_loss["rmse_d"].item(), ((1.0 + 4.0) / 2) ** 0.5) + + def test_loss_masks_excluded_grid_points(self) -> None: + # excluded grid points (mask == 0) must not contribute to the + # residual, which is normalised by the per-frame mask sum + from deepmd.pt.loss.charge import ( + GridDensityLoss, + ) + from deepmd.pt.utils import ( + env, + ) + + loss_fn = GridDensityLoss(start_pref_d=1.0, limit_pref_d=1.0) + + class FakeModel(torch.nn.Module): + def forward(self, **kwargs): + return { + "density": torch.tensor( + [[[1.0], [2.0], [3.0], [4.0]]], device=env.DEVICE + ), + "mask": torch.tensor( + [[1, 1, 0, 0]], dtype=torch.int32, device=env.DEVICE + ), + } + + label = { + "density": torch.tensor( + [[[1.0], [1.0], [100.0], [100.0]]], device=env.DEVICE + ), + "find_density": 1.0, + } + _, loss, more_loss = loss_fn({}, FakeModel(), label, 1, 1.0) + # only the two unmasked points count: (0^2 + 1^2) / 2 + self.assertAlmostEqual(loss.item(), 0.5) + self.assertAlmostEqual(more_loss["rmse_d"].item(), 0.5**0.5) + + def test_training_steps(self) -> None: + # exercises GridDensityLoss.forward on the main path + workdir = tempfile.mkdtemp(dir=self.tmpdir.name) + self._run_training(self._training_config(self.system), workdir) + + def test_training_requires_density_label(self) -> None: + # density is the only supervision signal: a missing density.npy must + # abort training instead of silently optimising nothing + system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system" + shutil.copytree(self.system, system) + (system / "set.000" / "density.npy").unlink() + with self.assertRaisesRegex(RuntimeError, "not found"): + get_trainer(normalize(self._training_config(system))) + + def test_training_zero_prefactor(self) -> None: + # start_pref_d = limit_pref_d = 0 disables the density term; + # backward() must still work on the graph-connected zero loss + workdir = tempfile.mkdtemp(dir=self.tmpdir.name) + self._run_training( + self._training_config(self.system, start=0.0, limit=0.0), workdir + ) + + def test_training_with_disp_avg(self) -> None: + # TaskLoss.training_metric_names is consulted before training starts + # whenever averaged metrics are requested. + config = self._training_config(self.system) + config["training"]["disp_avg"] = True + self._run_training(config, tempfile.mkdtemp(dir=self.tmpdir.name)) + + def test_env_protection_default(self) -> None: + # a density model built without an explicit env_protection gets + # 1e-6 by default (with a warning), and the recorded def script + # agrees with the built model + config = deepcopy(self.config) + config["model"]["descriptor"].pop("env_protection") + trainer = get_trainer(normalize(config)) + descriptor = trainer.model.atomic_model.descriptor + self.assertEqual(descriptor.get_env_protection(), 1e-6) + recorded = json.loads(trainer.model.model_def_script) + self.assertEqual(recorded["descriptor"]["env_protection"], 1e-6) + + def test_hybrid_descriptor_stat_blocks(self) -> None: + # hybrid descriptors keep their blocks in descrpt_list: the block + # discovery must walk the module tree instead of a hard-coded + # attribute list, and a stat file must not abort the run + descriptor = DescrptHybrid( + [ + DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0]), + DescrptSeR(rcut=2.0, rcut_smth=1.0, sel=[4, 4, 0]), + ] + ) + fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out()) + model = DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H", "X"]) + blocks = model._descriptor_stat_blocks() + self.assertEqual(len(blocks), 2) + # a stat file must take the soft-failure path, not abort + stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w") + + def sampler() -> list: + return [self._stat_sample(model)] + + model.compute_or_load_stat(sampler, stat_file_path=stat_dir) + + def test_grid_type_requires_reserved_slot(self) -> None: + # a real element in the reserved slot must be rejected, not + # silently used as the grid type + descriptor = DescrptSeA(rcut=4.0, rcut_smth=0.5, sel=[8, 8]) + fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out()) + with self.assertRaisesRegex(ValueError, "reserved grid type"): + DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H"]) + # non-element reserved slot is accepted (all other tests use X) + DPDensityAtomicModel(descriptor, fitting, type_map=["O", "X"]) + + def test_env_protection_enforced_at_model_level(self) -> None: + # models constructed directly with env_protection == 0.0 get the + # guard set on the descriptor block, not just a warning + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=0.5, sel=[8, 8, 0], env_protection=0.0 + ) + fitting = DensityFittingNet(descriptor.get_ntypes(), descriptor.get_dim_out()) + model = DPDensityAtomicModel(descriptor, fitting, type_map=["O", "H", "X"]) + self.assertEqual(model.descriptor.get_env_protection(), 1e-6) + + def test_env_protection_hybrid_defaults(self) -> None: + cases = ( + ("all omitted", [None], [1e-6]), + ("shared explicit value", [0.1, None], [0.1, 0.1]), + ) + for case, values, expected in cases: + config = deepcopy(self.config) + base = config["model"].pop("descriptor") + descriptors = [] + for value in values: + descriptor = deepcopy(base) + descriptor.pop("env_protection", None) + if value is not None: + descriptor["env_protection"] = value + descriptors.append(descriptor) + config["model"]["descriptor"] = { + "type": "hybrid", + "list": descriptors, + } + with self.subTest(case=case): + normalized = normalize(config) + actual = [ + sub["env_protection"] + for sub in normalized["model"]["descriptor"]["list"] + ] + self.assertEqual(actual, expected) + + def test_injected_pass_honors_pair_exclude_types(self) -> None: + # the injected stat pass must see the model's pair_exclude_types, + # which are only written by the wrapped sampler + atomic_model = self.torch_model.atomic_model + atomic_model.reinit_pair_exclude([(0, 1)]) + seen: list = [] + original = atomic_model.descriptor.compute_input_stats + + def spy(merged, path=None): + samples = merged() if callable(merged) else merged + seen.extend(samples) + original(merged, path) + + atomic_model.descriptor.compute_input_stats = spy # type: ignore[method-assign] + try: + + def sampler() -> list: + return [self._stat_sample(atomic_model)] + + atomic_model.compute_or_load_stat(sampler, stat_file_path=None) + finally: + atomic_model.descriptor.compute_input_stats = original # type: ignore[method-assign] + atomic_model.reinit_pair_exclude([]) + # injected samples carry the grid pseudo-atoms (more rows than atoms) + natoms = len(np.loadtxt(self.system / "type.raw", dtype=int)) + injected = [s for s in seen if s["atype"].shape[1] > natoms and "grid" in s] + self.assertTrue(injected, "no injected samples reached the descriptor") + for sample in injected: + excluded = [tuple(pair) for pair in sample["pair_exclude_types"]] + self.assertIn((0, 1), excluded, "model exclusions not honoured") + self.assertIn((2, 2), excluded, "grid-grid pairs not excluded") + + def test_stat_file_grid_row_writeback(self) -> None: + # the patched grid-type row is written back to the stat cache, and a + # complete cache then takes the fast path without the injected pass + stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w") + atomic_model = self.torch_model.atomic_model + + def sampler() -> list: + return [self._stat_sample(atomic_model)] + + atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir) + grid_type = len(self.config["model"]["type_map"]) - 1 + r_x = list(stat_dir.rglob(f"r_{grid_type}")) + self.assertTrue(r_x, "no grid-type stat item written to the cache") + for path in r_x: + self.assertNotEqual( + float(path.load_numpy()[0]), + 0.0, + f"{path} still holds zero samples (placeholder row)", + ) + # second call with the complete cache: the injected pass must not run + original_inject = atomic_model._inject_grid_samples + + def boom(_sampled: list) -> list: + raise AssertionError("injected pass should not run on a complete cache") + + atomic_model._inject_grid_samples = boom # type: ignore[method-assign] + try: + atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir) + finally: + atomic_model._inject_grid_samples = original_inject # type: ignore[method-assign] + + def test_stat_cache_rejects_zero_sample_angular_grid_row(self) -> None: + stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w") + atomic_model = self.torch_model.atomic_model + + def sampler() -> list: + return [self._stat_sample(atomic_model)] + + atomic_model.compute_or_load_stat(sampler, stat_file_path=stat_dir) + grid_type = len(self.config["model"]["type_map"]) - 1 + angular_rows = list(stat_dir.rglob(f"a_{grid_type}")) + self.assertTrue(angular_rows, "no angular grid-type statistic in cache") + for path in angular_rows: + value = path.load_numpy() + value[0] = 0.0 + path.save_numpy(value) + + self.assertTrue(atomic_model._grid_stat_missing(stat_dir)) + + def test_stat_without_grid_warns_and_uses_defaults(self) -> None: + # Statistics can be requested before a sampler is augmented with the + # density input; this must take the documented fallback rather than + # assuming every sample contains grid coordinates. + atomic_model = get_model(deepcopy(model_density)).atomic_model + sample = self._stat_sample(atomic_model) + sample.pop("grid") + with self.assertLogs( + "deepmd.pt.model.atomic_model.density_atomic_model", level="WARNING" + ) as logs: + atomic_model.compute_or_load_stat(lambda: [sample]) + self.assertIn("No grid data", "\n".join(logs.output)) + + def test_stat_writeback_skips_missing_cache_directory(self) -> None: + atomic_model = get_model(deepcopy(model_density)).atomic_model + grid_stats = [ + (block["davg"][-1].detach().clone(), block["dstd"][-1].detach().clone()) + for block in atomic_model._descriptor_stat_blocks() + ] + stat_dir = DPPath(tempfile.mkdtemp(dir=self.tmpdir.name), "w") + # A cache can lack one descriptor block (for example, after an older + # partial run); that block is deliberately skipped during write-back. + atomic_model._write_back_grid_stat(stat_dir, grid_stats) + + def test_read_only_stat_cache_patches_grid_row_in_memory(self) -> None: + # Older complete caches can contain a zero-sample placeholder for the + # reserved grid type. Read mode may patch that row in memory, but must + # never attempt the write-back that update mode performs. + stat_file = Path(self.tmpdir.name) / "density-stat.hdf5" + model = get_model(deepcopy(model_density)) + atomic_model = model.atomic_model + + def sampler() -> list: + return [self._stat_sample(atomic_model)] + + with open_stat_file(StatFileSpec(str(stat_file))) as update_path: + self.assertIsNotNone(update_path) + # The parent pass writes a complete cache with the ordinary + # zero-sample placeholder for the reserved grid type. + DPAtomicModel.compute_or_load_stat(atomic_model, sampler, update_path) + original = stat_file.read_bytes() + + with open_stat_file(StatFileSpec(str(stat_file), "read")) as read_path: + self.assertIsNotNone(read_path) + with self.assertLogs( + "deepmd.pt.model.atomic_model.density_atomic_model", level="WARNING" + ) as logs: + atomic_model.compute_or_load_stat(sampler, read_path) + self.assertIn("read-only", "\n".join(logs.output)) + self.assertEqual(stat_file.read_bytes(), original) + + def test_grid_type_statistics(self) -> None: + # the reserved grid type X gets real input statistics from the + # injected grid samples, not the descriptor's placeholder defaults + descriptor = self.torch_model.atomic_model.descriptor + dstd = descriptor.sea["dstd"].detach().cpu().numpy() + self.assertEqual(dstd.shape[0], 3) + self.assertTrue(np.isfinite(dstd).all()) + # placeholder default is 0.1; real statistics differ from it + self.assertFalse( + np.allclose(dstd[-1], 0.1, atol=1e-3), + f"grid type still has placeholder statistics: {dstd[-1]}", + ) + + def test_grid_type_statistics_dpa2(self) -> None: + # DPA-2 carries several stat blocks (repinit, repformers, + # repinit_three_body); the grid-type row must be patched in all of + # them, not just the first one + config = deepcopy(self.config) + config["model"]["descriptor"] = { + "type": "dpa2", + "repinit": { + "tebd_dim": 4, + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 16, + "neuron": [8, 16], + "axis_neuron": 4, + "activation_function": "tanh", + "use_three_body": True, + "three_body_sel": 8, + "three_body_rcut": 2.0, + "three_body_rcut_smth": 1.0, + }, + "repformer": { + "rcut": 2.0, + "rcut_smth": 1.5, + "nsel": 8, + "nlayers": 2, + "g1_dim": 16, + "g2_dim": 8, + "attn2_hidden": 8, + "attn2_nhead": 2, + "attn1_hidden": 16, + "attn1_nhead": 2, + "axis_neuron": 4, + }, + "env_protection": 1e-6, + "seed": 1, + } + trainer = get_trainer(normalize(config)) + descriptor = trainer.model.atomic_model.descriptor + for name in ("repinit", "repformers", "repinit_three_body"): + block = getattr(descriptor, name) + dstd = block["dstd"].detach().cpu().numpy() + self.assertTrue(np.isfinite(dstd).all(), name) + self.assertFalse( + np.allclose(dstd[-1], 0.1, atol=1e-3), + f"{name} still has placeholder statistics for the grid type", + ) + # the reserved grid type X gets real input statistics from the + # injected grid samples, not the descriptor's placeholder defaults + descriptor = self.torch_model.atomic_model.descriptor + dstd = descriptor.sea["dstd"].detach().cpu().numpy() + self.assertEqual(dstd.shape[0], 3) + self.assertTrue(np.isfinite(dstd).all()) + # placeholder default is 0.1; real statistics differ from it + self.assertFalse( + np.allclose(dstd[-1], 0.1, atol=1e-3), + f"grid type still has placeholder statistics: {dstd[-1]}", + ) + + def test_eval_requires_grid(self) -> None: + # evaluating a density model without grid must fail with a clear + # error instead of a bare KeyError on the output name table + dp = DeepDensity(self.model_path) + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + with self.assertRaisesRegex(ValueError, "grid is required"): + dp.deep_eval.eval(coord, box, atype) + + def test_eval_no_auto_batch(self) -> None: + # auto_batch_size=False bypasses execute_all's single-tuple unwrapping; + # the density result must still be normalized to a bare ndarray + dp = DeepDensity(self.model_path, auto_batch_size=False) + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + grid = np.load(set_dir / "grid.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + out = dp.eval(coord, box, atype, grid=grid) + self.assertEqual(out.shape, (2, self.ngrid)) + # and match the default-batching result numerically + out_default = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid) + np.testing.assert_allclose(out, out_default) + + def test_eval_without_periodic_box(self) -> None: + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:1] + grid = np.load(set_dir / "grid.npy")[:1] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + out = DeepDensity(self.model_path).eval(coord, None, atype, grid=grid) + self.assertEqual(out.shape, (1, self.ngrid)) + self.assertTrue(np.isfinite(out).all()) + + def test_grid_at_atoms_finite(self) -> None: + # grid points coincident with atoms are legitimate inputs; with + # env_protection > 0 (set in the test config) the predictions and + # gradients must stay finite instead of turning NaN via 1/r terms + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + grid_at_atoms = coord.reshape(coord.shape[0], -1, 3) + out = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid_at_atoms) + self.assertTrue(np.isfinite(out).all()) + + # gradients wrt the grid input and all parameters stay finite + input_dict = { + kk: vv.clone() if isinstance(vv, torch.Tensor) else vv + for kk, vv in self.input_dict.items() + } + coord_t = input_dict["coord"] + grid_t = ( + coord_t.reshape(coord_t.shape[0], -1, 3).detach().clone().requires_grad_() + ) + input_dict["grid"] = grid_t + model_out = self.torch_model(**input_dict) + self.assertTrue(torch.isfinite(model_out["density"]).all()) + model_out["density"].sum().backward() + self.assertIsNotNone(grid_t.grad) + self.assertTrue(torch.isfinite(grid_t.grad).all()) + for param in self.torch_model.parameters(): + if param.grad is not None: + self.assertTrue(torch.isfinite(param.grad).all()) + + def test_grid_periodic_translation(self) -> None: + # grids shifted by integer cell vectors are periodically equivalent + # and must produce the same density as the wrapped originals + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy")[:2] + box = np.load(set_dir / "box.npy")[:2] + grid = np.load(set_dir / "grid.npy")[:2] + atype = np.loadtxt(self.system / "type.raw", dtype=int) + dp = DeepDensity(self.model_path) + out = dp.eval(coord, box, atype, grid=grid) + cell_vec = box.reshape(box.shape[0], 3, 3)[:, 0] + out_shift = dp.eval(coord, box, atype, grid=grid + 2.0 * cell_vec[:, None, :]) + np.testing.assert_allclose(out_shift, out, rtol=1e-5, atol=1e-5) + + def test_dp_test_without_detail_file(self) -> None: + dp = DeepDensity(self.model_path) + data = DeepmdData( + str(self.system), + set_prefix="set", + shuffle_test=False, + type_map=dp.get_type_map(), + sort_atoms=False, + ) + errors = build_tester(dp, atomic=False).run( + data, + str(self.system), + numb_test=self.nframes, + detail_file=None, + ) + self.assertEqual(set(errors), {"mae_density", "rmse_density"}) + + def test_dp_test_chunking_matches_whole_run(self) -> None: + def run(chunk_size: int, detail_file: str) -> dict: + dp = DeepDensity(self.model_path) + data = DeepmdData( + str(self.system), + set_prefix="set", + shuffle_test=False, + type_map=dp.get_type_map(), + sort_atoms=False, + ) + with patch.dict( + os.environ, {"DP_TEST_CHUNK_ATOMS": str(chunk_size)}, clear=False + ): + return build_tester(dp, atomic=False).run( + data, + str(self.system), + numb_test=self.nframes, + detail_file=detail_file, + ) + + whole_detail = os.path.join(self.tmpdir.name, "detail_whole") + chunk_detail = os.path.join(self.tmpdir.name, "detail_chunk") + whole = run(1_000_000, whole_detail) + chunked = run(self.ngrid, chunk_detail) + self.assertEqual(whole.keys(), chunked.keys()) + for key in whole: + np.testing.assert_allclose(whole[key], chunked[key]) + for frame in range(self.nframes): + np.testing.assert_allclose( + np.loadtxt(f"{whole_detail}.density.out.{frame}"), + np.loadtxt(f"{chunk_detail}.density.out.{frame}"), + ) + + def test_dp_test_mixed_type(self) -> None: + system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "mixed-system" + shutil.copytree(self.system, system) + atom_types = np.array( + [ + [0, 0, 1, 1, 1], + [0, 1, 0, 1, 1], + [1, 0, 0, 1, 1], + [1, 1, 0, 0, 1], + ], + dtype=np.int32, + ) + np.save(system / "set.000" / "real_atom_types.npy", atom_types) + detail_file = os.path.join(self.tmpdir.name, "detail_mixed") + dp_test( + model=self.model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + set_dir = system / "set.000" + ref = DeepDensity(self.model_path).eval( + np.load(set_dir / "coord.npy"), + np.load(set_dir / "box.npy"), + atom_types, + grid=np.load(set_dir / "grid.npy"), + mixed_type=True, + ) + for frame in range(self.nframes): + detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1) + np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6) + + def test_dp_test_forwards_fparam(self) -> None: + params = deepcopy(model_density) + params["fitting_net"]["numb_fparam"] = 1 + model = get_model(params) + fd, model_path = tempfile.mkstemp(suffix=".pth") + os.close(fd) + system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "fparam-system" + shutil.copytree(self.system, system) + fparam = np.arange(1, self.nframes + 1, dtype=np.float64).reshape(-1, 1) + np.save(system / "set.000" / "fparam.npy", fparam) + try: + torch.jit.save(torch.jit.script(model), model_path) + detail_file = os.path.join(self.tmpdir.name, "detail_fparam") + dp_test( + model=model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + set_dir = system / "set.000" + dp = DeepDensity(model_path) + args = ( + np.load(set_dir / "coord.npy"), + np.load(set_dir / "box.npy"), + np.loadtxt(system / "type.raw", dtype=int), + ) + grid = np.load(set_dir / "grid.npy") + ref = dp.eval(*args, grid=grid, fparam=fparam) + without_fparam_effect = dp.eval( + *args, grid=grid, fparam=np.zeros_like(fparam) + ) + self.assertFalse(np.allclose(ref, without_fparam_effect)) + for frame in range(self.nframes): + detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1) + np.testing.assert_allclose( + detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6 + ) + finally: + os.unlink(model_path) + + def test_dp_test_uses_default_fparam(self) -> None: + params = deepcopy(model_density) + params["fitting_net"].update(numb_fparam=1, default_fparam=[0.5]) + fd, model_path = tempfile.mkstemp(suffix=".pth") + os.close(fd) + default_system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "default" + explicit_system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "explicit" + shutil.copytree(self.system, default_system) + shutil.copytree(self.system, explicit_system) + np.save( + explicit_system / "set.000" / "fparam.npy", + np.full((self.nframes, 1), 0.5), + ) + try: + torch.jit.save(torch.jit.script(get_model(params)), model_path) + detail_paths = [] + for name, system in ( + ("default", default_system), + ("explicit", explicit_system), + ): + detail_file = os.path.join(self.tmpdir.name, f"detail_fparam_{name}") + dp_test( + model=model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + detail_paths.append(detail_file) + for frame in range(self.nframes): + np.testing.assert_allclose( + np.loadtxt(f"{detail_paths[0]}.density.out.{frame}"), + np.loadtxt(f"{detail_paths[1]}.density.out.{frame}"), + ) + finally: + os.unlink(model_path) + + def test_dp_test_without_periodic_box(self) -> None: + system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "nopbc-system" + shutil.copytree(self.system, system) + (system / "set.000" / "box.npy").unlink() + (system / "nopbc").touch() + detail_file = os.path.join(self.tmpdir.name, "detail_nopbc") + dp_test( + model=self.model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + set_dir = system / "set.000" + ref = DeepDensity(self.model_path).eval( + np.load(set_dir / "coord.npy"), + None, + np.loadtxt(system / "type.raw", dtype=int), + grid=np.load(set_dir / "grid.npy"), + ) + for frame in range(self.nframes): + detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1) + np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6) + + def test_dp_test_shuffle(self) -> None: + # grid/density must be shuffled together with the frames: every + # (label, pred) row written by the shuffled run must match a row + # from the unshuffled reference run + detail_ref = os.path.join(self.tmpdir.name, "detail_ref") + dp_test( + model=self.model_path, + system=str(self.system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_ref, + atomic=False, + ) + ref_rows = set() + for frame in range(self.nframes): + detail = np.loadtxt(f"{detail_ref}.density.out.{frame}", skiprows=1) + for row in detail: + ref_rows.add(tuple(np.round(row, decimals=5))) + + detail_shuf = os.path.join(self.tmpdir.name, "detail_shuf") + dp_test( + model=self.model_path, + system=str(self.system), + datafile=None, + numb_test=2, + rand_seed=42, + shuffle_test=True, + detail_file=detail_shuf, + atomic=False, + ) + shuf_rows = [] + for frame in range(2): + detail = np.loadtxt(f"{detail_shuf}.density.out.{frame}", skiprows=1) + shuf_rows.extend(tuple(np.round(row, decimals=5)) for row in detail) + self.assertEqual(len(shuf_rows), 2 * self.ngrid) + for row in shuf_rows: + self.assertIn(row, ref_rows) + + def test_dp_test_flattened_grid(self) -> None: + # the loader also accepts a flattened (nframes, ngrid*3) grid; the + # tester must normalize it to the (nframes, ngrid, 3) contract that + # DeepDensity.eval expects, so both layouts give the same details + system = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system" + shutil.copytree(self.system, system) + set_dir = system / "set.000" + grid3d = np.load(set_dir / "grid.npy") + np.save(set_dir / "grid.npy", grid3d.reshape(self.nframes, -1)) + detail_file = os.path.join(self.tmpdir.name, "detail_flat") + dp_test( + model=self.model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + coord = np.load(set_dir / "coord.npy") + box = np.load(set_dir / "box.npy") + atype = np.loadtxt(system / "type.raw", dtype=int) + ref = DeepDensity(self.model_path).eval(coord, box, atype, grid=grid3d) + for frame in range(self.nframes): + detail = np.loadtxt(f"{detail_file}.density.out.{frame}", skiprows=1) + self.assertEqual(detail.shape, (self.ngrid, 2)) + np.testing.assert_allclose(detail[:, 1], ref[frame], rtol=1e-5, atol=1e-6) + + +model_density_dpa3_chg_spin = { + "type_map": ["O", "H", "X"], + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 8, + "e_dim": 6, + "a_dim": 4, + "nlayers": 1, + "e_rcut": 4.0, + "e_rcut_smth": 0.5, + "e_sel": 12, + "a_rcut": 3.5, + "a_rcut_smth": 0.5, + "a_sel": 8, + "axis_neuron": 4, + "update_angle": False, + }, + "env_protection": 1e-6, + "add_chg_spin_ebd": True, + "default_chg_spin": [1.0, 0.0], + "seed": 1, + }, + "fitting_net": { + "type": "density", + "neuron": [8, 8], + "seed": 1, + }, +} + + +class TestDPTestDensityChgSpin(unittest.TestCase): + """charge_spin forwarding on the density inference and dp test paths.""" + + @classmethod + def setUpClass(cls) -> None: + cls.tmpdir = tempfile.TemporaryDirectory() + root = Path(cls.tmpdir.name) + rng = np.random.default_rng(7) + cls.nframes, natoms, cls.ngrid = 4, 5, 8 + cls.system = root / "system" + set_dir = cls.system / "set.000" + set_dir.mkdir(parents=True) + np.save(set_dir / "coord.npy", rng.random((cls.nframes, natoms * 3)) * 8.0) + np.save( + set_dir / "box.npy", + np.tile(np.eye(3).reshape(-1) * 10.0, (cls.nframes, 1)), + ) + np.save(set_dir / "grid.npy", rng.random((cls.nframes, cls.ngrid, 3)) * 8.0) + np.save(set_dir / "density.npy", rng.random((cls.nframes, cls.ngrid, 1))) + # non-default charge/spin conditions, so forwarding has a visible effect + np.save(set_dir / "charge_spin.npy", np.tile([2.0, 1.0], (cls.nframes, 1))) + np.savetxt(cls.system / "type.raw", [0, 0, 1, 1, 1], fmt="%d") + np.savetxt(cls.system / "type_map.raw", ["O", "H", "X"], fmt="%s") + + model = get_model(deepcopy(model_density_dpa3_chg_spin)) + scripted = torch.jit.script(model) + tmp_fd, cls.model_path = tempfile.mkstemp(suffix=".pth") + os.close(tmp_fd) + torch.jit.save(scripted, cls.model_path) + + @classmethod + def tearDownClass(cls) -> None: + os.unlink(cls.model_path) + cls.tmpdir.cleanup() + + def _inputs(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + set_dir = self.system / "set.000" + coord = np.load(set_dir / "coord.npy") + box = np.load(set_dir / "box.npy") + grid = np.load(set_dir / "grid.npy") + atype = np.loadtxt(self.system / "type.raw", dtype=int) + return coord, box, grid, atype + + def test_eval_forwards_charge_spin(self) -> None: + dp = DeepDensity(self.model_path) + self.assertTrue(dp.has_chg_spin_ebd()) + self.assertTrue(dp.has_default_chg_spin()) + coord, box, grid, atype = self._inputs() + # without explicit charge_spin the model falls back to its default + out_default = dp.eval(coord, box, atype, grid=grid) + out_explicit_default = dp.eval( + coord, + box, + atype, + grid=grid, + charge_spin=np.tile([1.0, 0.0], (self.nframes, 1)), + ) + np.testing.assert_allclose( + out_default, out_explicit_default, rtol=1e-6, atol=1e-8 + ) + # a non-default condition must change the prediction: the value has + # to reach the descriptor, not be dropped on the eval path + out_other = dp.eval( + coord, + box, + atype, + grid=grid, + charge_spin=np.tile([2.0, 1.0], (self.nframes, 1)), + ) + self.assertFalse(np.allclose(out_default, out_other)) + + def test_eval_without_charge_spin_or_default(self) -> None: + # A charge/spin-conditioned descriptor needs either an explicit + # condition or a configured default; fail clearly before the + # descriptor's internal assertion. + config = deepcopy(model_density_dpa3_chg_spin) + config["descriptor"].pop("default_chg_spin") + path = os.path.join(self.tmpdir.name, "no-default-charge-spin.pth") + with torch.device("cpu"): + torch.jit.save(torch.jit.script(get_model(config)), path) + coord, box, grid, atype = self._inputs() + with self.assertRaisesRegex(torch.jit.Error, "charge_spin is required"): + DeepDensity(path).eval(coord, box, atype, grid=grid) + + def test_eval_grid_chunking_dpa3_matches_full_grid(self) -> None: + coord, box, grid, atype = self._inputs() + dp = DeepDensity(self.model_path) + + full = dp.eval(coord, box, atype, grid=grid, grid_batch_size=None) + chunked = dp.eval(coord, box, atype, grid=grid, grid_batch_size=3) + + np.testing.assert_allclose(chunked, full, rtol=1e-6, atol=1e-8) + + def test_dp_test_charge_spin_data_and_default(self) -> None: + coord, box, grid, atype = self._inputs() + dp = DeepDensity(self.model_path) + explicit_ref = dp.eval( + coord, + box, + atype, + grid=grid, + charge_spin=np.tile([2.0, 1.0], (self.nframes, 1)), + ) + default_ref = dp.eval(coord, box, atype, grid=grid) + self.assertFalse( + np.allclose(explicit_ref, default_ref), + "the test is vacuous: charge_spin has no effect on this model", + ) + + system_without_data = Path(tempfile.mkdtemp(dir=self.tmpdir.name)) / "system" + shutil.copytree(self.system, system_without_data) + (system_without_data / "set.000" / "charge_spin.npy").unlink() + cases = ( + ("explicit data", self.system, explicit_ref), + ("model default", system_without_data, default_ref), + ) + for case, system, reference in cases: + detail_file = os.path.join(self.tmpdir.name, f"detail_cs_{case}") + with self.subTest(case=case): + dp_test( + model=self.model_path, + system=str(system), + datafile=None, + numb_test=0, + rand_seed=None, + shuffle_test=False, + detail_file=detail_file, + atomic=False, + ) + for frame in range(self.nframes): + detail = np.loadtxt( + f"{detail_file}.density.out.{frame}", skiprows=1 + ) + self.assertEqual(detail.shape, (self.ngrid, 2)) + np.testing.assert_allclose( + detail[:, 1], reference[frame], rtol=1e-5, atol=1e-6 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_wrapper.py b/source/tests/pt/test_wrapper.py index 9fd64bdb06..5d74ff7b1b 100644 --- a/source/tests/pt/test_wrapper.py +++ b/source/tests/pt/test_wrapper.py @@ -21,6 +21,7 @@ def __init__(self, *, fail_forward: bool = False) -> None: self.scale = torch.nn.Parameter(torch.ones((), device="cpu")) self.fail_forward = fail_forward self.last_requires_grad: tuple[bool, ...] | None = None + self.last_box: torch.Tensor | None = None def forward( self, @@ -32,7 +33,8 @@ def forward( aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: - del atype, box, do_atomic_virial, fparam, aparam, charge_spin + del atype, do_atomic_virial, fparam, aparam, charge_spin + self.last_box = box self.last_requires_grad = tuple( param.requires_grad for param in self.parameters() ) @@ -131,6 +133,15 @@ def test_skip_loss_keeps_training_gradients(self) -> None: self.assertIsNotNone(model.linear.weight.grad) self.assertIsNotNone(model.scale.grad) + def test_positional_box_keeps_legacy_argument_order(self) -> None: + model = _LinearToyModel() + wrapper = ModelWrapper(model) + box = torch.eye(3, device="cpu").view(1, 3, 3).expand(2, -1, -1) + + wrapper(self.coord, self.atype, None, box) + + self.assertIs(model.last_box, box) + if __name__ == "__main__": unittest.main()