diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 6e940e3eb9..024abda6be 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -15,6 +15,9 @@ from .ener_spin import ( EnergySpinLoss, ) +from .group_property import ( + GroupPropertyLoss, +) from .loss import ( TaskLoss, ) @@ -35,6 +38,7 @@ "EnergyHessianStdLoss", "EnergySpinLoss", "EnergyStdLoss", + "GroupPropertyLoss", "PopulationLoss", "PropertyLoss", "TaskLoss", diff --git a/deepmd/pt/loss/group_property.py b/deepmd/pt/loss/group_property.py new file mode 100644 index 0000000000..2c362592fe --- /dev/null +++ b/deepmd/pt/loss/group_property.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Loss for grouped frame-level properties.""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import torch +import torch.nn.functional as F + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils.grouped import ( + GROUP_ID_KEY, + GROUP_WEIGHT_KEY, + POOL_MASK_KEY, + group_data_requirements, + normalize_group_id_tensor, +) +from deepmd.utils.data import ( + DataRequirementItem, +) + + +class GroupPropertyLoss(TaskLoss): + """Compute property loss after frame embeddings are aggregated by group.""" + + def __init__( + self, + task_dim: int, + var_name: str, + loss_func: str = "mse", + metric: list[str] | None = None, + beta: float = 1.0, + label_tol: float = 1e-8, + **kwargs: Any, + ) -> None: + super().__init__() + self.task_dim = task_dim + self.var_name = var_name + self.loss_func = loss_func + self.metric = metric or ["mae"] + self.beta = beta + self.label_tol = label_tol + + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float = 0.0, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: + del natoms, learning_rate, mae + var_name = self.var_name + frame_label = label[var_name] + nframes = frame_label.shape[0] + if frame_label.shape != (nframes, self.task_dim): + raise ValueError( + f"{var_name} label must have shape (nframes, {self.task_dim}); " + f"got {frame_label.shape}." + ) + + group_id = None + if GROUP_ID_KEY in label and label.get(f"find_{GROUP_ID_KEY}", 1) is not None: + find_group = label.get(f"find_{GROUP_ID_KEY}", 1) + if not torch.as_tensor(find_group).eq(0).all(): + group_id = label[GROUP_ID_KEY] + if group_id is None: + group_id = torch.arange( + nframes, dtype=torch.long, device=frame_label.device + ) + else: + group_id = normalize_group_id_tensor(group_id, nframes).to( + frame_label.device + ) + + model_pred = model( + **input_dict, + group_id=group_id, + weight=label.get(GROUP_WEIGHT_KEY), + pool_mask=label.get(POOL_MASK_KEY), + ) + pred = model_pred[var_name] + # Reuse the model's group ordering (group_inverse) as the single source of + # truth -- do NOT call torch.unique here, or ordering could silently drift. + group_label = self._group_labels( + frame_label, + model_pred["group_inverse"], + model_pred["group_id"].shape[0], + ).to(device=pred.device, dtype=pred.dtype) + if pred.shape != group_label.shape: + raise ValueError( + f"Prediction shape {pred.shape} does not match grouped labels " + f"{group_label.shape}." + ) + + loss = self._loss(pred, group_label) + more_loss = {} + if "mae" in self.metric: + more_loss["mae"] = F.l1_loss(pred, group_label, reduction="mean").detach() + if "mse" in self.metric: + more_loss["mse"] = F.mse_loss(pred, group_label, reduction="mean").detach() + if "rmse" in self.metric: + more_loss["rmse"] = torch.sqrt( + F.mse_loss(pred, group_label, reduction="mean") + ).detach() + return model_pred, loss, more_loss + + def _loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + if self.loss_func == "smooth_mae": + return F.smooth_l1_loss(pred, target, reduction="sum", beta=self.beta) + if self.loss_func == "mae": + return F.l1_loss(pred, target, reduction="sum") + if self.loss_func == "mse": + return F.mse_loss(pred, target, reduction="sum") + if self.loss_func == "rmse": + return torch.sqrt(F.mse_loss(pred, target, reduction="mean")) + raise RuntimeError(f"Unknown loss function : {self.loss_func}") + + def _group_labels( + self, + frame_label: torch.Tensor, + group_inverse: torch.Tensor, + n_groups: int, + ) -> torch.Tensor: + # Aggregate frame labels into group labels using the model's group_inverse + # (shared ordering). Scatter each frame's label into its group row, then + # gather back to verify every frame in a group carried the same label. + group_label = frame_label.new_zeros((n_groups, self.task_dim)) + group_label[group_inverse] = frame_label + gathered = group_label[group_inverse] + if not torch.allclose(gathered, frame_label, atol=self.label_tol): + mismatch = (~torch.isclose(gathered, frame_label, atol=self.label_tol)).any( + dim=1 + ) + bad_frame = int(torch.nonzero(mismatch, as_tuple=False).flatten()[0]) + raise ValueError( + f"Inconsistent {self.var_name} labels within a group " + f"(frame {bad_frame}, group_inverse={int(group_inverse[bad_frame])})." + ) + return group_label + + @property + def label_requirement(self) -> list[DataRequirementItem]: + return [ + DataRequirementItem( + self.var_name, + ndof=self.task_dim, + atomic=False, + must=True, + high_prec=True, + ), + *group_data_requirements(), + ] diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 8f9b63dc80..19185145b0 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -63,6 +63,9 @@ from .frozen import ( FrozenModel, ) +from .group_property_model import ( + GroupPropertyModel, +) from .make_hessian_model import ( make_hessian_model, ) @@ -326,6 +329,8 @@ def get_standard_model(model_params: dict) -> BaseModel: modelcls = EnergyModel elif fitting_net_type == "property": modelcls = PropertyModel + elif fitting_net_type == "group_property": + modelcls = GroupPropertyModel elif fitting_net_type == "population": modelcls = PopulationModel else: diff --git a/deepmd/pt/model/model/group_property_model.py b/deepmd/pt/model/model/group_property_model.py new file mode 100644 index 0000000000..7f9474aaea --- /dev/null +++ b/deepmd/pt/model/model/group_property_model.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Group-level property model for end-to-end grouped embeddings.""" + +from __future__ import ( + annotations, +) + +from types import ( + SimpleNamespace, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +from deepmd.dpmodel import ( + ModelOutputDef, +) +from deepmd.pt.model.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.pt.model.model.dp_model import ( + DPModelCommon, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.model.task import ( + BaseFitting, +) +from deepmd.pt.model.task.group_property import ( + GroupPropertyFittingNet, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.grouped import ( + normalize_group_id_tensor, + normalize_pool_mask_tensor, + normalize_weight_tensor, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + from deepmd.dpmodel import ( + OutputVariableDef, + ) + from deepmd.utils.path import ( + DPPath, + ) + + +@BaseModel.register("group_property") +class GroupPropertyModel(DPModelCommon, BaseModel): + """Predict one property per weighted group of frames.""" + + model_type = "group_property" + + def __init__( + self, + descriptor: Any, + fitting: Any, + type_map: list[str] | None = None, + **kwargs: Any, + ) -> None: + BaseModel.__init__(self) + if not isinstance(fitting, GroupPropertyFittingNet): + raise TypeError( + "fitting must be a GroupPropertyFittingNet for GroupPropertyModel" + ) + self.descriptor = descriptor + self.fitting_net = fitting + self.type_map = list(type_map or []) + self.atomic_model = SimpleNamespace( + descriptor=self.descriptor, + fitting_net=self.fitting_net, + observed_type=self.type_map, + get_dim_fparam=self.fitting_net.get_dim_fparam, + has_default_fparam=self.fitting_net.has_default_fparam, + get_default_fparam=self.fitting_net.get_default_fparam, + get_dim_aparam=self.fitting_net.get_dim_aparam, + ) + + def model_output_def(self) -> ModelOutputDef: + return ModelOutputDef(self.fitting_net.output_def()) + + def translated_output_def(self) -> dict[str, OutputVariableDef]: + return self.model_output_def().get_data() + + def compute_or_load_stat( + self, + sampled_func: Callable[[], Any], + stat_file_path: DPPath | None = None, + preset_observed_type: list[str] | None = None, + ) -> None: + if stat_file_path is not None and self.type_map is not None: + stat_file_path /= " ".join(self.type_map) + self.descriptor.compute_input_stats(sampled_func, stat_file_path) + + @torch.jit.export + def get_observed_type_list(self) -> list[str]: + return self.type_map + + def get_descriptor(self): # noqa: ANN201 + return self.descriptor + + def get_fitting_net(self): # noqa: ANN201 + return self.fitting_net + + @torch.jit.export + def get_task_dim(self) -> int: + return self.fitting_net.task_dim + + @torch.jit.export + def get_var_name(self) -> str: + return self.fitting_net.var_name + + def get_rcut(self) -> float: + return self.descriptor.get_rcut() + + def get_sel(self) -> list[int]: + return self.descriptor.get_sel() + + @torch.jit.export + def get_sel_type(self) -> list[int]: + get_sel_type = getattr(self.descriptor, "get_sel_type", None) + return [] if get_sel_type is None else get_sel_type() + + @torch.jit.export + def is_aparam_nall(self) -> bool: + return False + + @torch.jit.export + def model_output_type(self) -> list[str]: + return [self.get_var_name()] + + @torch.jit.export + def get_nsel(self) -> int: + return int(sum(self.get_sel())) + + @torch.jit.export + def get_nnei(self) -> int: + return self.get_nsel() + + def mixed_types(self) -> bool: + mixed_types = getattr(self.descriptor, "mixed_types", None) + return True if mixed_types is None else mixed_types() + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_dim_fparam(self) -> int: + return self.fitting_net.get_dim_fparam() + + def has_default_fparam(self) -> bool: + return self.fitting_net.has_default_fparam() + + def get_default_fparam(self) -> torch.Tensor | None: + return self.fitting_net.get_default_fparam() + + def get_dim_aparam(self) -> int: + return self.fitting_net.get_dim_aparam() + + def get_out_bias(self) -> torch.Tensor: + # GroupPropertyFittingNet zero-inits its output bias by design and + # learns the offset during training (setting it from replicated group + # labels would bias toward large groups). Expose a zero bias so the + # finetune bias-adjust step is a well-defined no-op. + return torch.zeros( + self.get_task_dim(), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + + def set_out_bias(self, out_bias: torch.Tensor) -> None: + # No externally managed per-type/group output bias; see get_out_bias. + return None + + def change_out_bias( + self, + merged: Any, + bias_adjust_mode: str = "change-by-statistic", + ) -> None: + # The group offset is learned by the fitting net during training, so + # there is no statistical output bias to (re)adjust on finetune. + return None + + @torch.jit.export + def has_chg_spin_ebd(self) -> bool: + if bool(getattr(self.descriptor, "add_chg_spin_ebd", False)): + return True + return ( + getattr(self.descriptor, "chg_embedding", None) is not None + and getattr(self.descriptor, "spin_embedding", None) is not None + ) + + @torch.jit.export + def has_default_chg_spin(self) -> bool: + if not self.has_chg_spin_ebd(): + return False + has_default = getattr(self.descriptor, "has_default_chg_spin", None) + if has_default is not None: + return has_default() + return getattr(self.descriptor, "default_chg_spin", None) is not None + + @torch.jit.export + def get_default_chg_spin(self) -> torch.Tensor | None: + if not self.has_default_chg_spin(): + return None + get_default = getattr(self.descriptor, "get_default_chg_spin", None) + if get_default is not None: + return get_default() + default_chg_spin = getattr(self.descriptor, "default_chg_spin", None) + if default_chg_spin is None: + return None + if isinstance(default_chg_spin, torch.Tensor): + return default_chg_spin + return torch.tensor( + default_chg_spin, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + + @torch.jit.export + def has_message_passing(self) -> bool: + has_message_passing = getattr(self.descriptor, "has_message_passing", None) + return False if has_message_passing is None else has_message_passing() + + def need_sorted_nlist_for_lower(self) -> bool: + need_sorted = getattr(self.descriptor, "need_sorted_nlist_for_lower", None) + return False if need_sorted is None else need_sorted() + + 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, + group_id: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + pool_mask: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + del aparam, do_atomic_virial + coord = coord.to(env.GLOBAL_PT_FLOAT_PRECISION) + if coord.dim() == 2: + coord = coord.view(coord.shape[0], -1, 3) + atype = atype.long() + box = None if box is None else box.to(env.GLOBAL_PT_FLOAT_PRECISION) + nframes, natoms = atype.shape + + 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, + ) + ) + descriptor_atype = extended_atype.clamp_min(0) + descriptor_kwargs = {"mapping": mapping} + if self.has_chg_spin_ebd(): + if charge_spin is None: + default_chg_spin = self.get_default_chg_spin() + if default_chg_spin is None: + default_chg_spin = extended_coord.new_tensor([0.0, 1.0]) + charge_spin = default_chg_spin + if charge_spin is not None: + charge_spin = charge_spin.to( + device=extended_coord.device, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + ) + if charge_spin.dim() == 1: + if charge_spin.numel() != 2: + raise ValueError( + "charge_spin must have shape (2,) or (nframes, 2)." + ) + charge_spin = charge_spin.unsqueeze(0) + if charge_spin.dim() != 2 or charge_spin.shape[1] != 2: + raise ValueError( + "charge_spin must have shape (2,) or (nframes, 2)." + ) + if charge_spin.shape[0] == 1 and nframes != 1: + charge_spin = charge_spin.expand(nframes, -1) + elif charge_spin.shape[0] != nframes: + raise ValueError( + "charge_spin must have one row or match the number of frames." + ) + # DPA3 descriptors consume charge/spin through their fparam slot. + # The group-level fparam argument is reserved for the fitting net. + descriptor_kwargs["fparam"] = charge_spin + descriptor, _rot_mat, _g2, _h2, _sw = self.descriptor( + extended_coord, + descriptor_atype, + nlist, + **descriptor_kwargs, + ) + if descriptor is None: + raise RuntimeError("Descriptor returned None for group_property.") + descriptor = descriptor[:, :natoms, :] + + if pool_mask is None: + pool_mask = (atype[:, :natoms] >= 0).to( + device=descriptor.device, dtype=descriptor.dtype + ) + else: + pool_mask = normalize_pool_mask_tensor(pool_mask, nframes, natoms).to( + descriptor.device, descriptor.dtype + ) + mask_sum = pool_mask.sum(dim=1) + if bool((mask_sum == 0).any()): + raise ValueError("all-zero pool_mask is not allowed for any frame.") + denom = mask_sum.clamp_min(1.0) + # Zero out non-pooled atoms (padding/virtual atoms and excluded caps) + # before the weighted sum so a non-finite descriptor on those rows -- + # e.g. a virtual padding atom -- cannot poison the frame embedding via + # ``0 * NaN``. Kept atoms retain their (possibly fractional) weights. + keep = pool_mask[:, :, None] > 0 + descriptor = torch.where(keep, descriptor, torch.zeros_like(descriptor)) + frame_embedding = (descriptor * pool_mask[:, :, None]).sum(dim=1) / denom[ + :, None + ] + + if group_id is None: + group_id = torch.arange(nframes, dtype=torch.long, device=descriptor.device) + else: + group_id = normalize_group_id_tensor(group_id, nframes).to( + descriptor.device + ) + if weight is None: + weight = torch.ones( + nframes, dtype=descriptor.dtype, device=descriptor.device + ) + else: + weight = normalize_weight_tensor(weight, nframes).to( + descriptor.device, descriptor.dtype + ) + weight = weight.detach() + + group_order, inverse = torch.unique(group_id, sorted=True, return_inverse=True) + n_groups = group_order.shape[0] + group_embedding = torch.zeros( + (n_groups, frame_embedding.shape[1]), + dtype=frame_embedding.dtype, + device=frame_embedding.device, + ) + group_embedding.index_add_(0, inverse, frame_embedding * weight[:, None]) + if self.fitting_net.group_reduce == "mean": + # weighted mean over frames: divide by the per-group weight sum so the + # embedding scale is independent of group size (see group_reduce doc). + weight_sum = torch.zeros( + (n_groups, 1), + dtype=frame_embedding.dtype, + device=frame_embedding.device, + ) + weight_sum.index_add_(0, inverse, weight[:, None]) + group_embedding = group_embedding / weight_sum.clamp_min(1e-12) + if self.fitting_net.get_dim_fparam() > 0 and fparam is not None: + # fparam is a per-group side feature (constant within a group). Take + # each group's value and concat AFTER aggregation, so it never passes + # through the weighted sum over frames. + fparam = fparam.reshape(nframes, -1).to( + group_embedding.device, group_embedding.dtype + ) + grouped_fparam: list[torch.Tensor] = [] + for group_index, group_value in enumerate(group_order): + values = fparam[inverse == group_index] + first = values[0] + if not torch.allclose(values, first.expand_as(values), atol=1e-8): + raise ValueError( + "fparam must be constant within each assembly group; " + f"group_id {int(group_value)} has inconsistent rows." + ) + grouped_fparam.append(first) + group_fparam = torch.stack(grouped_fparam, dim=0) + group_embedding = torch.cat([group_embedding, group_fparam], dim=-1) + prediction = self.fitting_net(group_embedding) + return { + self.get_var_name(): prediction, + "group_id": group_order, + "unique_group_ids": group_order, + "frame_group_id": group_id, + "group_inverse": inverse, + "frame_embedding": frame_embedding, + "weight": weight, + "pool_mask": pool_mask, + } + + def serialize(self) -> dict[str, Any]: + serialize_descriptor = getattr(self.descriptor, "serialize", None) + return { + "type": self.model_type, + "descriptor": serialize_descriptor() if serialize_descriptor else {}, + "fitting": self.fitting_net.serialize(), + "type_map": self.type_map, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> GroupPropertyModel: + data = data.copy() + data.pop("type", None) + descriptor = BaseDescriptor.deserialize(data.pop("descriptor")) + fitting = BaseFitting.deserialize(data.pop("fitting")) + return cls( + descriptor=descriptor, + fitting=fitting, + type_map=data.pop("type_map", None), + **data, + ) diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index d5b10b6f08..f9140e7049 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -18,6 +18,9 @@ from .fitting import ( Fitting, ) +from .group_property import ( + GroupPropertyFittingNet, +) from .polarizability import ( PolarFittingNet, ) @@ -42,6 +45,7 @@ "EnergyFittingNet", "EnergyFittingNetDirect", "Fitting", + "GroupPropertyFittingNet", "PolarFittingNet", "PopulationFittingNet", "PropertyFittingNet", diff --git a/deepmd/pt/model/task/group_property.py b/deepmd/pt/model/task/group_property.py new file mode 100644 index 0000000000..62f04662d0 --- /dev/null +++ b/deepmd/pt/model/task/group_property.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Group-level property fitting network.""" + +from __future__ import ( + annotations, +) + +import logging +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.task.fitting import ( + Fitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) + +_log = logging.getLogger(__name__) + + +def _activation(name: str) -> torch.nn.Module: + if name == "tanh": + return torch.nn.Tanh() + if name == "relu": + return torch.nn.ReLU() + if name == "gelu": + return torch.nn.GELU() + if name in {"linear", "none"}: + return torch.nn.Identity() + raise ValueError(f"Unsupported activation_function: {name!r}") + + +@Fitting.register("group_property") +class GroupPropertyFittingNet(Fitting): + """MLP that maps aggregated group embeddings to group properties.""" + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + property_name: str, + task_dim: int = 1, + neuron: list[int] | None = None, + # GELU (not tanh): the group head consumes an un-normalized + # frame/group embedding concatenated with fparam. tanh saturates at + # that input scale, so the head collapses to a constant (bias-only) + # prediction; GELU keeps gradients flowing and learns per-group signal. + activation_function: str = "gelu", + precision: str = DEFAULT_PRECISION, + trainable: bool | list[bool] = True, + type_map: list[str] | None = None, + numb_fparam: int = 0, + group_reduce: str = "mean", + seed: int | None = None, + # Injected unconditionally by the generic model-building path + # (deepmd.pt.model.model._get_standard_model_components) for every + # fitting type; "type" selects this class via the Fitting registry + # and "mixed_types" is the descriptor's own property (read via + # GroupPropertyModel.mixed_types()), so neither is used here. Kept + # as explicit, named, no-op parameters -- not **kwargs -- so any + # other unrecognized field fails loudly instead of being silently + # accepted and ignored. + type: str = "group_property", + mixed_types: bool = True, + ) -> None: + del type, mixed_types + super().__init__() + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.var_name = property_name + self.task_dim = task_dim + self.dim_out = task_dim + self.neuron = list(neuron or [128, 128]) + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.type_map = list(type_map or []) + self.seed = seed + self.numb_fparam = int(numb_fparam) + if group_reduce not in ("mean", "sum"): + raise ValueError( + f"group_reduce must be 'mean' or 'sum'; got {group_reduce!r}." + ) + # frame->group reduction (executed in GroupPropertyModel.forward): + # "mean": group_emb = index_add(w_i*e_i) / index_add(w_i) -- weighted + # mean; embedding scale independent of group size. (default) + # "sum" : group_emb = index_add(w_i*e_i) -- norm grows with group size; + # for additive-property semantics only. + self.group_reduce = group_reduce + + # Per-group side features (fparam) are concatenated to the aggregated + # group embedding, so the fitting input widens by ``numb_fparam``. + dims = [dim_descrpt + self.numb_fparam, *self.neuron, task_dim] + + def _build_layers() -> list[torch.nn.Module]: + layers: list[torch.nn.Module] = [] + for ii in range(len(dims) - 1): + layers.append(torch.nn.Linear(dims[ii], dims[ii + 1], dtype=self.prec)) + if ii < len(dims) - 2: + layers.append(_activation(self.activation_function)) + return layers + + if seed is None: + # No seed requested: draw from (and advance) the caller's global + # RNG stream exactly as an unseeded torch.nn.Linear always does. + layers = _build_layers() + else: + # Seed only this net's own initialization, without disturbing the + # caller's global RNG state (torch.nn.Linear.reset_parameters() + # has no generator= argument, so fork-then-seed is the way to + # scope a seed to a plain nn.Linear-based net). + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + layers = _build_layers() + self.network = torch.nn.Sequential(*layers).to(env.DEVICE) + # Task E (route 3): group_property does NOT init the output bias from + # label statistics. The property path's frame-level stat would count each + # group's replicated label once per component (biasing toward large groups) + # and is computed at frame level while this net operates at group level. + # Zero-init the output bias explicitly; training learns the offset. + last = self.network[-1] + if isinstance(last, torch.nn.Linear) and last.bias is not None: + torch.nn.init.zeros_(last.bias) + _log.warning( + "group_property output bias is zero-initialized (label-stat bias " + "init is disabled for grouped labels); training may need more steps." + ) + + linear_layers = [m for m in self.network if isinstance(m, torch.nn.Linear)] + if isinstance(trainable, list): + if len(trainable) != len(linear_layers): + raise ValueError( + f"trainable has {len(trainable)} entries; expected " + f"{len(linear_layers)} (one per Linear layer, i.e. " + "len(neuron)+1)." + ) + self.trainable = trainable + for layer, layer_trainable in zip(linear_layers, trainable, strict=True): + for param in layer.parameters(): + param.requires_grad = bool(layer_trainable) + else: + self.trainable = bool(trainable) + for param in self.parameters(): + param.requires_grad = self.trainable + + def forward(self, group_embedding: torch.Tensor) -> torch.Tensor: + return self.network(group_embedding.to(self.prec)) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.task_dim], + reducible=False, + r_differentiable=False, + c_differentiable=False, + intensive=True, + ), + ] + ) + + def get_dim_fparam(self) -> int: + return self.numb_fparam + + def has_default_fparam(self) -> bool: + return False + + def get_default_fparam(self) -> torch.Tensor | None: + return None + + def get_dim_aparam(self) -> int: + return 0 + + def get_type_map(self) -> list[str]: + return self.type_map + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat: Any | None = None + ) -> None: + self.type_map = list(type_map) + + def compute_input_stats(self, *args: Any, **kwargs: Any) -> None: + return None + + def serialize(self) -> dict[str, Any]: + return { + "type": "group_property", + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "property_name": self.var_name, + "task_dim": self.task_dim, + "numb_fparam": self.numb_fparam, + "group_reduce": self.group_reduce, + "neuron": self.neuron, + "activation_function": self.activation_function, + "precision": self.precision, + "type_map": self.type_map, + } diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index df883c8ee6..ad8e3bdfc7 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -44,6 +44,7 @@ EnergyHessianStdLoss, EnergySpinLoss, EnergyStdLoss, + GroupPropertyLoss, PopulationLoss, PropertyLoss, TaskLoss, @@ -894,8 +895,23 @@ def collect_single_finetune_params( not use_random_initialization and new_key not in _origin_state_dict ): + # GroupPropertyModel attaches the descriptor + # directly (its ``atomic_model`` is a plain + # namespace, not an nn.Module), so its keys read + # ``model..descriptor.*``. A standard + # (energy/property) pretrained model stores the + # descriptor one level deeper, under + # ``model..atomic_model.descriptor.*``. + # Bridge the two by inserting ``atomic_model``. + atomic_key = new_key.replace( + f".{_model_key_from}.descriptor.", + f".{_model_key_from}.atomic_model.descriptor.", + 1, + ) # for ZBL models finetuning from standard models - if ".models.0." in new_key: + if atomic_key in _origin_state_dict: + new_key = atomic_key + elif ".models.0." in new_key: new_key = new_key.replace(".models.0.", ".") elif ".models.1." in new_key: use_random_initialization = True @@ -2534,6 +2550,12 @@ def get_loss( loss_params["var_name"] = var_name loss_params["intensive"] = intensive return PropertyLoss(**loss_params) + elif loss_type == "group_property": + task_dim = _model.get_task_dim() + var_name = _model.get_var_name() + loss_params["task_dim"] = task_dim + loss_params["var_name"] = var_name + return GroupPropertyLoss(**loss_params) elif loss_type == "population": loss_params["starter_learning_rate"] = start_lr return PopulationLoss(**loss_params) diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 807c1f5ba1..3a237e8322 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os +from collections.abc import ( + Iterator, +) from multiprocessing.dummy import ( Pool, ) @@ -16,6 +19,7 @@ from torch.utils.data import ( DataLoader, Dataset, + Sampler, WeightedRandomSampler, ) from torch.utils.data._utils.collate import ( @@ -35,6 +39,12 @@ from deepmd.pt.utils.dataset import ( DeepmdDataSetForLoader, ) +from deepmd.pt.utils.grouped import ( + distributed_grouped_frame_batches, + grouped_frame_batches, + has_group_requirement, + load_group_ids_for_system, +) from deepmd.pt.utils.utils import ( mix_entropy, ) @@ -162,23 +172,71 @@ def construct_dataset(system: str) -> DeepmdDataSetForLoader: else: self.batch_sizes = batch_size * np.ones(len(systems), dtype=int) assert len(self.systems) == len(self.batch_sizes) + self._group_complete_batches = False + self._shuffle = shuffle + self._seed = seed + self._build_dataloaders() + + def _build_dataloaders(self) -> None: + """Build per-system dataloaders.""" + self.sampler_list = [] + self.dataloaders = [] + self.index = [] + self.total_batch = 0 for system, batch_size in zip(self.systems, self.batch_sizes): - if dist.is_available() and dist.is_initialized(): + distributed = dist.is_available() and dist.is_initialized() + system_sampler = None + batch_sampler = None + if self._group_complete_batches: + group_ids = load_group_ids_for_system(system.data_system) + if group_ids is None: + # GroupPropertyLoss falls back to one group per frame + # (torch.arange(nframes)) when group_id is absent, i.e. + # "no explicit grouping" means every frame stands alone. + # Mirror that here: treating a whole system as a single + # group instead would silently merge unrelated frames + # into one oversized batch (ignoring batch_size) and can + # break DDP by handing an implicit single group to more + # ranks than it has frames for. + group_ids = np.arange(len(system), dtype=np.int64) + if distributed: + batch_sampler = GroupDistributedBatchSampler( + group_ids, + max_frames=int(batch_size), + num_replicas=dist.get_world_size(), + rank=dist.get_rank(), + shuffle=self._shuffle, + seed=self._seed, + ) + else: + batch_sampler = GroupCompleteBatchSampler( + group_ids, + max_frames=int(batch_size), + shuffle=self._shuffle, + seed=self._seed, + ) + elif distributed: system_sampler = DistributedSampler(system) self.sampler_list.append(system_sampler) + if batch_sampler is None: + system_dataloader = DataLoader( + dataset=system, + batch_size=int(batch_size), + num_workers=0, # Should be 0 to avoid too many threads forked + sampler=system_sampler, + collate_fn=collate_batch, + shuffle=( + not (dist.is_available() and dist.is_initialized()) + ) # distributed sampler will do the shuffling by default + and self._shuffle, + ) else: - system_sampler = None - system_dataloader = DataLoader( - dataset=system, - batch_size=int(batch_size), - num_workers=0, # Should be 0 to avoid too many threads forked - sampler=system_sampler, - collate_fn=collate_batch, - shuffle=( - not (dist.is_available() and dist.is_initialized()) - ) # distributed sampler will do the shuffling by default - and shuffle, - ) + system_dataloader = DataLoader( + dataset=system, + batch_sampler=batch_sampler, + num_workers=0, + collate_fn=collate_batch, + ) self.dataloaders.append(system_dataloader) self.index.append(len(system_dataloader)) self.total_batch += len(system_dataloader) @@ -216,6 +274,9 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N """Add data requirement for each system in multiple systems.""" for system in self.systems: system.add_data_requirement(data_requirement) + if has_group_requirement(data_requirement): + self._group_complete_batches = True + self._build_dataloaders() def print_summary( self, @@ -263,6 +324,115 @@ def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: return result +def _normalize_group_sampler_seed( + seed: int | list[int] | tuple[int, ...] | None, +) -> int | None: + if seed is None: + return None + if isinstance(seed, (list, tuple)): + if not seed: + return None + # DDP group batching first builds a global group order, then slices it + # by rank. Every rank must therefore use the same base seed. + return int(seed[0]) + return int(seed) + + +class GroupDistributedBatchSampler(Sampler[list[int]]): + """Yield group-complete frame batches for one distributed rank.""" + + def __init__( + self, + group_ids: np.ndarray, + max_frames: int, + num_replicas: int, + rank: int, + shuffle: bool = True, + seed: int | list[int] | tuple[int, ...] | None = None, + ) -> None: + self.group_ids = np.asarray(group_ids, dtype=np.int64) + self.max_frames = max(int(max_frames), 1) + self.num_replicas = int(num_replicas) + self.rank = int(rank) + self.shuffle = shuffle + self.seed = _normalize_group_sampler_seed(seed) + self._epoch = 0 + self._batches = self._make_batches() + + def _make_batches(self) -> list[list[int]]: + base_seed = 0 if self.seed is None else self.seed + # Partition groups across ranks with an epoch-INDEPENDENT seed so every + # rank always agrees on the split. If the group shuffle depended on a + # per-rank ``_epoch`` (which drifts when ranks restart ``cycle_iterator`` + # at different times), the ``group_items[rank::num_replicas]`` slices + # would come from different shuffles and duplicate/drop groups. + rng = np.random.default_rng(base_seed) + batches = distributed_grouped_frame_batches( + self.group_ids, + self.max_frames, + num_replicas=self.num_replicas, + rank=self.rank, + shuffle=self.shuffle, + rng=rng, + ) + # Vary only THIS rank's own batch order per epoch. Reordering a rank's + # batches never moves groups between ranks, so the split stays intact + # while training still sees a fresh batch order each epoch even if the + # per-rank epoch counters are not in lock-step. + if self.shuffle and self._epoch > 0: + local = np.random.default_rng( + base_seed + 1 + self.rank + self._epoch * (self.num_replicas + 1) + ) + order = local.permutation(len(batches)) + batches = [batches[int(ii)] for ii in order] + return batches + + def __iter__(self) -> Iterator[list[int]]: + self._batches = self._make_batches() + self._epoch += 1 + yield from self._batches + + def __len__(self) -> int: + return len(self._batches) + + +class GroupCompleteBatchSampler(Sampler[list[int]]): + """Yield frame batches that never split a group inside one system.""" + + def __init__( + self, + group_ids: np.ndarray, + max_frames: int, + shuffle: bool = True, + seed: int | list[int] | tuple[int, ...] | None = None, + ) -> None: + self.group_ids = np.asarray(group_ids, dtype=np.int64) + self.max_frames = max(int(max_frames), 1) + self.shuffle = shuffle + self.seed = _normalize_group_sampler_seed(seed) + self._epoch = 0 + self._batches = self._make_batches() + + def _make_batches(self) -> list[list[int]]: + rng = np.random.default_rng( + None if self.seed is None else self.seed + self._epoch + ) + return grouped_frame_batches( + self.group_ids, + self.max_frames, + shuffle=self.shuffle, + rng=rng, + ) + + def __iter__(self) -> Iterator[list[int]]: + self._batches = self._make_batches() + self._epoch += 1 + yield from self._batches + + def __len__(self) -> int: + return len(self._batches) + + def get_weighted_sampler( training_data: Any, prob_style: str, sys_prob: bool = False ) -> WeightedRandomSampler: diff --git a/deepmd/pt/utils/grouped.py b/deepmd/pt/utils/grouped.py new file mode 100644 index 0000000000..b174a61be9 --- /dev/null +++ b/deepmd/pt/utils/grouped.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Utilities for grouped frame-level property training.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import numpy as np + +from deepmd.utils.data import ( + DataRequirementItem, +) + +if TYPE_CHECKING: + import torch + + from deepmd.utils.data import ( + DeepmdData, + ) + +GROUP_ID_KEY = "group_id" +GROUP_WEIGHT_KEY = "weight" +POOL_MASK_KEY = "pool_mask" + + +def group_data_requirements() -> list[DataRequirementItem]: + """Return the auxiliary data fields consumed by ``group_property``.""" + return [ + DataRequirementItem( + GROUP_ID_KEY, + ndof=1, + atomic=False, + must=False, + high_prec=False, + default=0.0, + dtype=np.int64, + ), + DataRequirementItem( + GROUP_WEIGHT_KEY, + ndof=1, + atomic=False, + must=False, + high_prec=True, + default=1.0, + ), + DataRequirementItem( + POOL_MASK_KEY, + ndof=1, + atomic=True, + must=False, + high_prec=True, + default=1.0, + output_natoms_for_type_sel=True, + ), + ] + + +def has_group_requirement(data_requirement: list[DataRequirementItem]) -> bool: + """Return True when a data requirement asks for grouped auxiliaries.""" + keys = {item["key"] for item in data_requirement} + return GROUP_ID_KEY in keys and GROUP_WEIGHT_KEY in keys and POOL_MASK_KEY in keys + + +def normalize_group_id_tensor(group_id: torch.Tensor, nframes: int) -> torch.Tensor: + """Normalize collated group ids to shape ``(nframes,)``.""" + group_id = group_id.reshape(nframes, -1) + if group_id.shape[1] != 1: + raise ValueError("group_id must have one value per frame.") + return group_id[:, 0].long() + + +def normalize_weight_tensor(weight: torch.Tensor, nframes: int) -> torch.Tensor: + """Normalize collated group weights to shape ``(nframes,)``.""" + weight = weight.reshape(nframes, -1) + if weight.shape[1] != 1: + raise ValueError("weight must have one value per frame.") + return weight[:, 0] + + +def normalize_pool_mask_tensor( + pool_mask: torch.Tensor, nframes: int, natoms: int +) -> torch.Tensor: + """Normalize collated pool masks to shape ``(nframes, natoms)``.""" + pool_mask = pool_mask.reshape(nframes, natoms, -1) + if pool_mask.shape[2] != 1: + raise ValueError("pool_mask must have one value per atom.") + return pool_mask[:, :, 0] + + +def load_group_ids_for_system(data_system: DeepmdData) -> np.ndarray | None: + """Load frame-level group ids from a DeePMD system, if present. + + Reads ``set.*/group_id.npy`` through the system's own sorted ``dirs`` + (``DPPath`` set directories, resolved once by ``DeepmdData.__init__``) + instead of re-globbing a raw filesystem path, so HDF5-backed systems are + checked the same way as on-disk ones: a plain ``pathlib.Path.glob`` can't + see inside an HDF5 archive and would silently -- and incorrectly -- report + an in-data ``group_id`` as missing. The returned ids follow DeePMD frame + order across those sets. Missing data returns ``None`` so callers can + fall back to ordinary frame batching. + """ + set_dirs = data_system.dirs + if not set_dirs: + return None + + chunks: list[np.ndarray] = [] + for set_dir in set_dirs: + path = set_dir / f"{GROUP_ID_KEY}.npy" + if not path.is_file(): + return None + arr = np.asarray(path.load_numpy()).reshape(-1) + chunks.append(arr.astype(np.int64, copy=False)) + return np.concatenate(chunks) if chunks else None + + +def _group_frame_indices(group_ids: np.ndarray) -> list[list[int]]: + """Return frame indices grouped by first-seen group id.""" + if group_ids.ndim != 1: + raise ValueError(f"{GROUP_ID_KEY} must be 1D; got shape {group_ids.shape}.") + groups: dict[int, list[int]] = {} + for frame_idx, group_id in enumerate(group_ids.astype(np.int64, copy=False)): + groups.setdefault(int(group_id), []).append(frame_idx) + return list(groups.values()) + + +def _shuffle_group_items( + group_items: list[list[int]], + shuffle: bool, + rng: np.random.Generator | None, +) -> list[list[int]]: + if not shuffle: + return list(group_items) + rng = rng or np.random.default_rng() + order = rng.permutation(len(group_items)) + return [group_items[int(ii)] for ii in order] + + +def _pack_group_items( + group_items: list[list[int]], + max_frames: int, +) -> list[list[int]]: + batches: list[list[int]] = [] + current: list[int] = [] + limit = max(int(max_frames), 1) + for indices in group_items: + if current and len(current) + len(indices) > limit: + batches.append(current) + current = [] + current.extend(indices) + if len(current) >= limit: + batches.append(current) + current = [] + if current: + batches.append(current) + return batches + + +def grouped_frame_batches( + group_ids: np.ndarray, + max_frames: int, + shuffle: bool = True, + rng: np.random.Generator | None = None, +) -> list[list[int]]: + """Pack complete groups into batches without splitting a group.""" + group_items = _shuffle_group_items( + _group_frame_indices(group_ids), shuffle=shuffle, rng=rng + ) + return _pack_group_items(group_items, max_frames) + + +def distributed_grouped_frame_batches( + group_ids: np.ndarray, + max_frames: int, + num_replicas: int, + rank: int, + shuffle: bool = True, + rng: np.random.Generator | None = None, +) -> list[list[int]]: + """Pack complete groups for one distributed rank. + + Groups, not frames, are assigned to ranks. Therefore every frame with the + same ``group_id`` is consumed by exactly one rank/GPU for the epoch. + """ + num_replicas = int(num_replicas) + rank = int(rank) + if num_replicas < 1: + raise ValueError(f"num_replicas must be >= 1; got {num_replicas}.") + if rank < 0 or rank >= num_replicas: + raise ValueError(f"rank must be in [0, {num_replicas}); got rank={rank}.") + if shuffle and rng is None: + rng = np.random.default_rng(0) + group_items = _shuffle_group_items( + _group_frame_indices(group_ids), shuffle=shuffle, rng=rng + ) + if len(group_items) < num_replicas: + raise ValueError( + "distributed grouped batching requires at least one group per rank; " + f"got {len(group_items)} groups for {num_replicas} ranks." + ) + rank_items = group_items[rank::num_replicas] + return _pack_group_items(rank_items, max_frames) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 84b04e02ba..16ac9d0ad7 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2825,6 +2825,56 @@ def fitting_property() -> list[Argument]: ] +@fitting_args_plugin.register("group_property", doc=doc_only_pt_supported) +def fitting_group_property() -> list[Argument]: + """Grouped frame-property fitting reuses the property head schema, but the + head pools an un-normalized frame embedding + fparam where tanh saturates + (collapsing to a constant), so it defaults to GELU instead of tanh. + + It also adds ``group_reduce``, the frame-embedding -> group-embedding + reduction (mean or weighted sum), which GroupPropertyFittingNet supports + but the reused property schema has no field for. + + GroupPropertyFittingNet is a small, standalone MLP (it does not go + through the shared GeneralFitting base class other fitting nets use), so + several property-schema fields have no effect on it: ``numb_aparam`` + (there is no per-atom input at group level), ``default_fparam``, + ``dim_case_embd`` (no multitask case embedding), ``resnet_dt`` (plain + Linear layers, no residual connections), ``intensive`` and + ``distinguish_types`` (the group-level output bias is always + zero-initialized rather than set from per-type label statistics -- see + GroupPropertyFittingNet.get_out_bias). Removing them makes setting one an + immediate, clear validation error instead of a silently ignored no-op. + """ + doc_group_reduce = ( + "How per-frame embeddings are reduced to one per-group embedding: " + "`mean` (weight-normalized average) or `sum` (weighted sum, for " + "extensive group properties)." + ) + _unsupported = { + "numb_aparam", + "default_fparam", + "dim_case_embd", + "resnet_dt", + "intensive", + "distinguish_types", + } + args = [arg for arg in fitting_property() if arg.name not in _unsupported] + for arg in args: + if arg.name == "activation_function": + arg.default = "gelu" + args.append( + Argument( + "group_reduce", + str, + optional=True, + default="mean", + doc=doc_group_reduce, + ) + ) + return args + + @fitting_args_plugin.register("polar", doc=doc_polar) def fitting_polar() -> 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." @@ -4875,6 +4925,12 @@ def loss_property() -> list[Argument]: ] +@loss_args_plugin.register("group_property") +def loss_group_property() -> list[Argument]: + """Grouped property loss uses the property loss hyper-parameters.""" + return loss_property() + + # YWolfeee: Modified to support tensor type of loss args. @loss_args_plugin.register("tensor") def loss_tensor() -> list[Argument]: diff --git a/doc/dpa_adapt/assembly_scenarios.md b/doc/dpa_adapt/assembly_scenarios.md new file mode 100644 index 0000000000..ba9119e8f5 --- /dev/null +++ b/doc/dpa_adapt/assembly_scenarios.md @@ -0,0 +1,103 @@ +# Assembly Groups for Multi-Component Properties + +Grouped property training is activated by ordinary DeePMD data fields. If a +system contains `group_id.npy`, DPA-ADAPT automatically treats frames with the +same group id as one labeled training example. `weight.npy` and +`pool_mask.npy` are optional; missing weights default to `1.0`, and missing +pool masks default to all real atoms. + +Preferred dataset-level entry for existing `deepmd/npy` data: + +```python +from dpa_adapt import mark_groups + +mark_groups("oer/dpdata", target="overpotential", group_by="system") +``` + +`mark_groups()` only adds grouped markers. It does not re-declare labels, +`fparam.npy`, coordinates, boxes, or type maps; those remain standard DeePMD / +DPA-ADAPT data fields and are read automatically during `fit()`. + +Each group becomes one labeled training example. Each component is one DeePMD +frame. The model computes frame embeddings, applies `pool_mask`, aggregates +`weight * frame_embedding` within the group (`group_reduce="mean"` by default), +concatenates per-group `fparam` after group aggregation, and predicts the +target. + +The low-level `Assembly.group().add(coords, symbols, ...)` API is intended only +for programmatic converters that already have arrays in memory. Public data +preparation should use `convert()` for raw formats and `mark_groups()` for +existing DeePMD systems. + +## Electrocatalysis Reaction Assemblies + +Use this when the measured or computed label depends on several adsorbate states on the same catalyst surface. + +- `Assembly.target`: `overpotential`, `activity`, `selectivity`, or a reaction free-energy descriptor. +- `group`: one catalyst composition, slab, active site, or equation directory. +- `component`: clean slab and/or adsorbate states such as `*`, `H*`, `O*`, `OH*`, `OOH*`, `CO*`, `CHO*`. +- `weight`: stoichiometric coefficient or reaction-specific linear-combination coefficient; use `1.0` for pure learned state aggregation. +- `pool_mask`: exclude virtual adsorbate slots, cap atoms, or atoms intentionally not participating in the descriptor pool. +- `fparam`: pH, potential, electrolyte identity encoded upstream, surface coverage, temperature. + +Example roles: `clean`, `O*`, `OH*`, `OOH*`. The OER retrofit path is +`mark_groups(path, target="overpotential", group_by="system")`. + +## Solvent and Electrolyte Mixtures + +Use this when a formulation property belongs to a recipe rather than a single molecule. + +- `Assembly.target`: `conductivity`, `viscosity`, `flash_point`, `solubility`, `dielectric_constant`. +- `group`: one mixture recipe. +- `component`: solvent molecules, salt ion pairs, additives, cosolvents. +- `weight`: mole fraction, mass fraction, or normalized concentration contribution. +- `fparam`: temperature, total salt concentration, pH, pressure, measurement protocol flags. +- `pool_mask`: usually all atoms; exclude generated caps if fragments are represented by capped structures. + +Suggested component roles: `solvent`, `salt`, `additive`, `cosolvent`. + +## Polymer and Copolymer Properties + +Use this when a polymer is represented by repeat units and end groups. + +- `Assembly.target`: `cloud_point`, `glass_transition_temperature`, `lcst`, `solubility`, `modulus`. +- `group`: one polymer or copolymer record. +- `component`: repeat units, start group, end group, optional side-chain fragments. +- `weight`: repeat-unit mole fraction; end-group share computed from `Mn` or supplied explicitly. +- `fparam`: standardized `Mn`, concentration, pH, salt concentrations, solvent identity. +- `pool_mask`: exclude artificial caps or attachment placeholders. + +For polymer CSVs, keep conversion in the ordinary data-preparation layer. A +converter should write standard DeePMD fields plus `group_id.npy`, +`weight.npy`, and `pool_mask.npy`; training then uses the same automatic grouped +path as every other scenario. + +## Battery Cell Recipes + +Use this when a performance label belongs to a full cell or formulation. + +- `Assembly.target`: `cycle_life`, `capacity_retention`, `rate_capability`, `safety_window`, `formation_loss`. +- `group`: one cell recipe and protocol. +- `component`: cathode active material, anode active material, electrolyte solvents, salts, additives, coating or binder fragments. +- `weight`: mass fraction, capacity-normalized contribution, or molar recipe fraction. +- `fparam`: C-rate, temperature, voltage window, state of charge, formation protocol. +- `pool_mask`: exclude caps/placeholders in fragment representations. + +Suggested component roles: `cathode`, `anode`, `solvent`, `salt`, `additive`, `binder`. + +## MOF/COF and Porous Framework Assemblies + +Use this when a material property can be represented by structural building blocks plus optional guest molecules. + +- `Assembly.target`: `uptake`, `selectivity`, `band_gap`, `stability`, `diffusivity`. +- `group`: one framework topology or material entry. +- `component`: metal node, linker, functional group, counter-ion, guest molecule. +- `weight`: topology count, formula-normalized count, occupancy, or guest loading. +- `fparam`: temperature, pressure, guest identity encoding, humidity, activation protocol. +- `pool_mask`: exclude caps used to make linker/node fragments chemically valid. + +Suggested component roles: `node`, `linker`, `functional_group`, `guest`. + +## Naming Rationale + +`Grouped` describes the tensor operation. `Assembly` describes the scientific object users build: several components assembled into one labeled group. API names should therefore stay on the scientific layer (`Assembly`, `group`, `component`, `weight`, `fparam`) while tensor files keep the implementation layer (`group_id`, `pool_mask`, `fparam`). diff --git a/dpa_adapt/__init__.py b/dpa_adapt/__init__.py index 5af6edff26..0e26e60db8 100644 --- a/dpa_adapt/__init__.py +++ b/dpa_adapt/__init__.py @@ -10,6 +10,14 @@ __version__ = "0.1.0" _LAZY = { + "Assembly": (".grouped", "Assembly"), + "ComponentSpec": (".grouped", "ComponentSpec"), + "GroupSpec": (".grouped", "GroupSpec"), + "PoolMask": (".grouped", "PoolMask"), + "SiteSelector": (".grouped", "SiteSelector"), + "SubstitutionSpec": (".grouped", "SubstitutionSpec"), + "GroupMarkerResult": (".grouped", "GroupMarkerResult"), + "mark_groups": (".grouped", "mark_groups"), "ConditionManager": (".conditions", "ConditionManager"), "DPAConditionError": (".conditions", "DPAConditionError"), "cross_validate": (".cv", "cross_validate"), diff --git a/dpa_adapt/cli.py b/dpa_adapt/cli.py index b5a7b1f104..186dfbb1d4 100644 --- a/dpa_adapt/cli.py +++ b/dpa_adapt/cli.py @@ -353,6 +353,52 @@ def _cmd_data_attach_labels(args: argparse.Namespace) -> int: return 0 +def _cmd_data_mark_groups(args: argparse.Namespace) -> int: + from dpa_adapt import ( + mark_groups, + ) + + group_by: str | int = args.group_by + if isinstance(group_by, str) and group_by.isdigit(): + group_by = int(group_by) + + results = mark_groups( + args.input, + group_by=group_by, + target=args.target, + weight=args.weight, + overwrite=args.overwrite, + dry_run=args.dry_run, + ) + n_updated = sum( + 1 + for item in results + if item.wrote_group_id or item.wrote_pool_mask or item.wrote_weight + ) + _LOG.info( + "%s%s systems discovered; %s updated; %s groups total.", + "[dry-run] " if args.dry_run else "", + len(results), + n_updated, + sum(item.n_groups for item in results), + ) + for item in results[:10]: + suffix = f" skipped({item.reason})" if item.skipped else "" + _LOG.info( + "%s: frames=%s groups=%s group_id=%s pool_mask=%s weight=%s%s", + item.system, + item.n_frames, + item.n_groups, + item.wrote_group_id, + item.wrote_pool_mask, + item.wrote_weight, + suffix, + ) + if len(results) > 10: + _LOG.info("... (+%s more)", len(results) - 10) + return 0 + + # --------------------------------------------------------------------------- # Dispatch table # --------------------------------------------------------------------------- @@ -369,6 +415,7 @@ def _cmd_data_attach_labels(args: argparse.Namespace) -> int: "convert": _cmd_data_convert, "validate": _cmd_data_validate, "attach-labels": _cmd_data_attach_labels, + "mark-groups": _cmd_data_mark_groups, } @@ -675,6 +722,32 @@ def get_parser() -> argparse.ArgumentParser: parser_data_attach.add_argument("--head-json", action="store_true") parser_data_attach.add_argument("--values", required=True) + # data mark-groups + parser_data_mark = data_subparsers.add_parser( + "mark-groups", + help="Add grouped-training markers to existing deepmd/npy directories", + parents=[parser_log], + ) + parser_data_mark.add_argument("--input", required=True, nargs="+") + parser_data_mark.add_argument( + "--group-by", + default="system", + help="'system' (default), 'label', or an integer group size.", + ) + parser_data_mark.add_argument( + "--target", + default="property", + help="Label key used when --group-by label.", + ) + parser_data_mark.add_argument( + "--weight", + type=float, + default=None, + help="Optional constant weight.npy value. Missing weight defaults to 1.0.", + ) + parser_data_mark.add_argument("--overwrite", action="store_true") + parser_data_mark.add_argument("--dry-run", action="store_true") + return parser diff --git a/dpa_adapt/data/desc_cache.py b/dpa_adapt/data/desc_cache.py index 92bd240539..07949aa78f 100644 --- a/dpa_adapt/data/desc_cache.py +++ b/dpa_adapt/data/desc_cache.py @@ -31,6 +31,9 @@ from dpa_adapt._backend import ( resolve_pretrained_path, ) +from dpa_adapt.data.loader import ( + _get_source, +) if TYPE_CHECKING: import dpdata @@ -74,6 +77,13 @@ def _system_fingerprint(system: dpdata.System) -> str: element costs O(total array size), but that is negligible next to the descriptor extraction the cache guards, and it makes the key collision-safe for changed systems. + + Grouped (heterogeneous) systems also fold in ``set.*/pool_mask.npy`` and + ``set.*/real_atom_types.npy`` when present: their per-atom, per-frame + masking/typing is what the extractor actually consumes for those systems + (``atom_types``/``coords`` alone are the uniform placeholder + padded + geometry, which can be identical across two systems with different group + markers or real atom types). """ d = system.data @@ -87,6 +97,17 @@ def _system_fingerprint(system: dpdata.System) -> str: _hash_array(h, np.asarray(d["coords"])) if "cells" in d: _hash_array(h, np.asarray(d["cells"])) + # grouped-only per-frame arrays: real local atom types (with -1 padding) + # and the pooling mask, read straight from disk since dpdata does not + # surface either as a per-frame field. + source = _get_source(system) + if source is not None: + for set_dir in sorted(Path(source).glob("set.*")): + for name in ("real_atom_types.npy", "pool_mask.npy"): + path = set_dir / name + if path.is_file(): + h.update(name.encode()) + _hash_array(h, np.load(path)) return h.hexdigest()[:16] diff --git a/dpa_adapt/finetuner.py b/dpa_adapt/finetuner.py index 9c409a8eda..11bd95ecf8 100644 --- a/dpa_adapt/finetuner.py +++ b/dpa_adapt/finetuner.py @@ -58,6 +58,165 @@ # Module-level helpers # --------------------------------------------------------------------------- +# Canonical order of the composable descriptor-pooling primitives. Feature +# columns are concatenated in this order, so the tuple must stay fixed -- +# reordering it would silently shift the meaning of every pooled feature. +POOLING_PRIMITIVES: tuple[str, ...] = ("mean", "sum", "std", "max", "min") + + +def parse_pooling(spec: Any) -> tuple[str, ...]: + """Parse a pooling spec into a canonical, de-duplicated primitive tuple. + + Accepts a ``"+"``-joined string (e.g. ``"mean+std"``) or a sequence of + primitive names. The output is de-duplicated and ordered by + :data:`POOLING_PRIMITIVES`, so ``"std+mean"`` and ``"mean+std"`` both yield + ``("mean", "std")`` and the concatenated feature layout is input-order + independent. + """ + if isinstance(spec, str): + tokens = [tok for tok in spec.split("+") if tok] + else: + tokens = list(spec) + if not tokens: + raise ValueError("empty pooling spec") + seen: set[str] = set() + for tok in tokens: + if tok not in POOLING_PRIMITIVES: + raise ValueError( + f"unknown pooling primitive {tok!r}; " + f"valid primitives are {POOLING_PRIMITIVES}" + ) + seen.add(tok) + return tuple(prim for prim in POOLING_PRIMITIVES if prim in seen) + + +def _pool_descriptor(descrpt: Any, primitives: Any, mask: Any = None) -> Any: + """Pool a per-atom descriptor ``(nframes, natoms, ndim)`` over atoms. + + Each primitive in ``primitives`` reduces the atom axis; results are + concatenated along the feature axis in the given order. ``std`` uses + ``nan_to_num`` so a single-atom frame contributes zeros instead of NaN. + With ``mask`` ``None`` this matches the historical per-string pooling + byte-for-byte; a ``(nframes, natoms)`` ``mask`` (1 = real, 0 = virtual) + excludes padding/virtual atoms from every reduction. + """ + import torch + + if mask is not None: + m = mask.to(dtype=descrpt.dtype, device=descrpt.device).unsqueeze(-1) + count = m.sum(dim=1).clamp_min(1.0) + # Sanitize masked (virtual/padding) rows to a finite value before any + # reduction below. Multiplying a non-finite descriptor by a zero + # mask keeps 0 * NaN == NaN, which would poison the whole frame's + # pooled feature instead of just dropping the masked atom. + descrpt = torch.where(m > 0, descrpt, torch.zeros_like(descrpt)) + + parts = [] + for prim in primitives: + if prim == "mean": + if mask is None: + parts.append(descrpt.mean(dim=1)) + else: + parts.append((descrpt * m).sum(dim=1) / count) + elif prim == "sum": + if mask is None: + parts.append(descrpt.sum(dim=1)) + else: + parts.append((descrpt * m).sum(dim=1)) + elif prim == "std": + if mask is None: + parts.append(torch.nan_to_num(descrpt.std(dim=1), nan=0.0)) + else: + mean = (descrpt * m).sum(dim=1, keepdim=True) / count.unsqueeze(1) + var = (((descrpt - mean) ** 2) * m).sum(dim=1) / count + parts.append(torch.nan_to_num(torch.sqrt(var), nan=0.0)) + elif prim == "max": + if mask is None: + parts.append(descrpt.max(dim=1).values) + else: + parts.append( + descrpt.masked_fill(m == 0, float("-inf")).max(dim=1).values + ) + elif prim == "min": + if mask is None: + parts.append(descrpt.min(dim=1).values) + else: + parts.append( + descrpt.masked_fill(m == 0, float("inf")).min(dim=1).values + ) + else: + raise ValueError(f"unknown pooling primitive {prim!r}") + if len(parts) == 1: + return parts[0] + return torch.cat(parts, dim=-1) + + +def _pool_mask_for_system(system: Any, n_frames: int, n_atoms: int) -> Any: + """Per-frame ``(n_frames, n_atoms)`` pool mask for a grouped system. + + Reads ``set.*/pool_mask.npy`` (or derives it from ``real_atom_types >= 0``) + in dpdata's set order so padded/virtual atoms can be excluded from offline + descriptor pooling. Returns ``None`` for non-grouped systems or when a + complete, frame-aligned mask cannot be built, keeping the pooling + byte-identical to the unmasked path in that case. + """ + source = _get_source(system) + if source is None: + return None + masks: list[np.ndarray] = [] + for set_dir in sorted(Path(source).glob("set.*")): + pool_mask_path = set_dir / "pool_mask.npy" + real_types_path = set_dir / "real_atom_types.npy" + if pool_mask_path.is_file(): + arr = np.asarray(np.load(pool_mask_path), dtype=np.float64) + elif real_types_path.is_file(): + arr = (np.asarray(np.load(real_types_path)) >= 0).astype(np.float64) + else: + return None + arr = arr.reshape(arr.shape[0], -1) + if arr.shape[1] != n_atoms: + return None + masks.append(arr) + if not masks: + return None + mask = np.concatenate(masks, axis=0) + if mask.shape[0] != n_frames or bool(np.all(mask == 1.0)): + # frame-count mismatch, or no virtual atoms -> nothing to mask out. + return None + return mask + + +def _real_atom_types_for_system( + system: Any, n_frames: int, n_atoms: int +) -> np.ndarray | None: + """Per-frame ``(n_frames, n_atoms)`` local atom types for a grouped system. + + Grouped systems write a uniform ``type.raw`` placeholder and store the + real, per-frame local atom-type indices (including ``-1`` for virtual + padding atoms) in ``set.*/real_atom_types.npy``. Returns ``None`` for + non-grouped systems or when a complete, frame-aligned array cannot be + built, so callers fall back to the constant-per-system ``atom_types``. + """ + source = _get_source(system) + if source is None: + return None + chunks: list[np.ndarray] = [] + for set_dir in sorted(Path(source).glob("set.*")): + real_types_path = set_dir / "real_atom_types.npy" + if not real_types_path.is_file(): + return None + arr = np.asarray(np.load(real_types_path), dtype=np.int64) + arr = arr.reshape(arr.shape[0], -1) + if arr.shape[1] != n_atoms: + return None + chunks.append(arr) + if not chunks: + return None + real_types = np.concatenate(chunks, axis=0) + if real_types.shape[0] != n_frames: + return None + return real_types + def _load_labels( systems: list[dpdata.System], @@ -654,6 +813,19 @@ def remap_atom_types( return local_to_global[atom_types] + def remap_atom_types_preserving_padding( + self, atom_types: np.ndarray, system: dpdata.System + ) -> np.ndarray: + """Like :meth:`remap_atom_types`, but keeps ``-1`` (virtual/padding + atom) entries untouched instead of wrapping them to the last + checkpoint type via numpy's negative-index fancy indexing. + """ + real = atom_types >= 0 + remapped = np.full_like(atom_types, -1) + if real.any(): + remapped[real] = self.remap_atom_types(atom_types[real], system) + return remapped + # ------------------------------------------------------------------ # Feature extraction (extract_features_cached is on DPAFineTuner # so that patches on DPAFineTuner._extract_features are honoured) @@ -692,8 +864,22 @@ def extract_features(self, systems: list[dpdata.System]) -> np.ndarray: n_frames = coords.shape[0] n_atoms = len(atom_types) - # Remap local atom-type indices to checkpoint-global indices. - atom_types_global = self.remap_atom_types(atom_types, system) + # Grouped (heterogeneous) systems write a uniform type.raw + # placeholder and store the real, per-frame local atom types + # (with -1 padding) in real_atom_types.npy. Use those per-frame + # types when present instead of tiling the single, uniform + # atom_types array -- otherwise every atom in every frame would + # be described as if it had the placeholder's type. + real_atom_types = _real_atom_types_for_system(system, n_frames, n_atoms) + if real_atom_types is not None: + atom_types_global = self.remap_atom_types_preserving_padding( + real_atom_types, system + ) + else: + # Remap local atom-type indices to checkpoint-global indices. + atom_types_global = np.tile( + self.remap_atom_types(atom_types, system), (n_frames, 1) + ) # Non-periodic structures must NOT use all-zero box: # the descriptor produces NaN in that case. @@ -709,7 +895,7 @@ def extract_features(self, systems: list[dpdata.System]) -> np.ndarray: device=self._device, ).requires_grad_(True) atype_t = torch.tensor( - np.tile(atom_types_global, (n_frames, 1)), + atom_types_global, dtype=torch.long, device=self._device, ) @@ -717,26 +903,13 @@ def extract_features(self, systems: list[dpdata.System]) -> np.ndarray: # Shape: (n_frames, n_atoms, feat_dim) descrpt = extractor._run_forward(coord_t, atype_t, box_t) - if self.pooling == "mean": - feat = descrpt.mean(dim=1) - elif self.pooling == "sum": - feat = descrpt.sum(dim=1) - elif self.pooling == "mean+std": - mean = descrpt.mean(dim=1) - std = torch.nan_to_num(descrpt.std(dim=1), nan=0.0) - feat = torch.cat([mean, std], dim=-1) - elif self.pooling == "mean+std+max+min": - mean = descrpt.mean(dim=1) - std = torch.nan_to_num(descrpt.std(dim=1), nan=0.0) - feat = torch.cat( - [ - mean, - std, - descrpt.max(dim=1).values, - descrpt.min(dim=1).values, - ], - dim=-1, - ) + mask_np = _pool_mask_for_system(system, n_frames, n_atoms) + mask_t = ( + None + if mask_np is None + else torch.tensor(mask_np, dtype=torch.float64, device=self._device) + ) + feat = _pool_descriptor(descrpt, parse_pooling(self.pooling), mask=mask_t) feat = torch.nan_to_num(feat, nan=0.0, posinf=0.0, neginf=0.0) all_features.append(feat.detach().cpu().numpy()) @@ -865,6 +1038,7 @@ def __init__( # ---- training paradigms ---- strategy: str = "frozen_sklearn", property_name: str = "property", + target: str | None = None, task_dim: int = 1, intensive: bool = True, init_branch: str = "SPICE2", @@ -889,15 +1063,22 @@ def __init__( aux_batch_size: str | int | None = None, downstream_batch_size: str | int | None = None, ) -> None: - if pooling not in self._VALID_POOLING: - raise ValueError( - f"pooling must be one of {sorted(self._VALID_POOLING)}, got {pooling!r}" - ) + if target is not None: + # ``target`` is a user-facing alias for ``property_name``. + property_name = target if strategy not in self._VALID_STRATEGIES: raise ValueError( f"strategy must be one of {sorted(self._VALID_STRATEGIES)}; " f"got {strategy!r}" ) + pooling_primitives = parse_pooling(pooling) + if strategy != "frozen_sklearn" and pooling_primitives != ("mean",): + raise ValueError( + f"pooling {pooling!r} is not size-intensive; strategy " + f"{strategy!r} trains a fitting head and requires intensive " + "pooling. Use pooling='mean', or strategy='frozen_sklearn' for " + "composite pooling." + ) validate_fparam_dim(fparam_dim) self.strategy = strategy @@ -1331,13 +1512,16 @@ def _run_training_predict( def fit( self, - train_data: str | list[str], + train_data: str | list[str] | None = None, valid_data: str | list[str] | None = None, type_map: list[str] | None = None, target_key: str | list[str] | None = None, labels: np.ndarray | None = None, fmt: str | None = None, aux_data: str | list[str] | None = None, + *, + train: str | list[str] | None = None, + valid: str | list[str] | None = None, ) -> str | None: """Train the model. @@ -1365,6 +1549,15 @@ def fit( (mft only) Auxiliary training system directories. Required when ``strategy='mft'``; must be absent otherwise. """ + # ``train=`` / ``valid=`` are user-facing aliases for the positional + # ``train_data`` / ``valid_data`` arguments. + if train_data is None: + train_data = train + if valid_data is None: + valid_data = valid + if train_data is None: + raise ValueError("fit() requires train_data (or the train= alias).") + if self.strategy == "frozen_sklearn": return self._fit_sklearn(train_data, type_map, target_key, labels, fmt) @@ -1455,6 +1648,16 @@ def _fit_sklearn( Refactored: logic extracted to ``_FrozenSklearnPipeline``; this method now orchestrates the pipeline and mirrors its state for backward compat. """ + from dpa_adapt.grouped._offline import ( + has_grouped_markers, + ) + + if has_grouped_markers(data): + # Grouped input carries its own per-group labels (read from + # set.*/.npy), so target_key/labels are optional here. + self._fit_sklearn_grouped(data, type_map, target_key, fmt) + return + if target_key is not None and labels is not None: raise ValueError( "target_key and labels are mutually exclusive; provide only one." @@ -1521,6 +1724,68 @@ def _fit_sklearn( p._condition_manager = self._condition_manager p._fitted = True + def _fit_sklearn_grouped( + self, + data: str | list[str], + type_map: list[str] | None, + target_key: str | list[str] | None, + fmt: str | None, + ) -> None: + """Fit the frozen-sklearn head on one pooled row per assembly group. + + Descriptors are extracted per frame, weighted-pooled into one embedding + per group id, and regressed against the group's shared label. + """ + from sklearn.pipeline import ( + make_pipeline, + ) + from sklearn.preprocessing import ( + StandardScaler, + ) + + from dpa_adapt.grouped._offline import ( + GroupedDataset, + ) + from dpa_adapt.utils.sklearn_heads import ( + build_sklearn_head, + ) + + p = self._ensure_sklearn() + self.type_map = type_map or [] + self._target_key = target_key if target_key is not None else "property" + + dataset = GroupedDataset( + data, + pretrained=self.pretrained, + model_branch=self.model_branch, + type_map=type_map, + target_key=self._target_key, + fmt=fmt, + ) + features = dataset.get_embeddings() + y = dataset.get_labels() + self._task_dim = 1 if y.ndim == 1 else y.shape[-1] + y_flat = y.ravel() if self._task_dim == 1 else y + + head = build_sklearn_head( + self._predictor_type, + seed=self.seed, + n_outputs=self._task_dim, + ) + self.predictor = make_pipeline(StandardScaler(), head) + self.predictor.fit(features, y_flat) + self._fitted = True + self._grouped = True + self._condition_manager = None + + # Mirror pipeline state for backward compat. + p.predictor = self.predictor + p.type_map = self.type_map + p._target_key = self._target_key + p._task_dim = self._task_dim + p._condition_manager = None + p._fitted = True + def predict(self, data: str | list[str], fmt: str | None = None) -> DotDict: """ Predict with the adapted model. @@ -1556,6 +1821,23 @@ def predict(self, data: str | list[str], fmt: str | None = None) -> DotDict: "predict() was called before fit(). Train the model with fit() first." ) + if getattr(self, "_grouped", False): + from dpa_adapt.grouped._offline import ( + GroupedDataset, + ) + + dataset = GroupedDataset( + data, + pretrained=self.pretrained, + model_branch=self.model_branch, + type_map=self.type_map or None, + target_key=self._target_key, + fmt=fmt, + ) + raw = self.predictor.predict(dataset.get_embeddings()) + predictions = np.asarray(raw).reshape(-1, self._task_dim) + return DotDict({"predictions": predictions}) + systems = load_data(data, fmt=fmt) features = self._extract_features(systems) @@ -1625,9 +1907,26 @@ def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: result = self.predict(data, fmt=fmt) predictions = result.predictions - systems = load_data(data, fmt=fmt) - labels = _load_labels(systems, self._target_key) - labels = labels.reshape(predictions.shape) + if getattr(self, "_grouped", False): + # predict() returns one row per group; use matching group-level + # labels instead of frame-level ones (which would not reshape). + from dpa_adapt.grouped._offline import ( + GroupedDataset, + ) + + dataset = GroupedDataset( + data, + pretrained=self.pretrained, + model_branch=self.model_branch, + type_map=self.type_map or None, + target_key=self._target_key, + fmt=fmt, + ) + labels = np.asarray(dataset.get_labels()).reshape(predictions.shape) + else: + systems = load_data(data, fmt=fmt) + labels = _load_labels(systems, self._target_key) + labels = labels.reshape(predictions.shape) if predictions.shape != labels.shape: raise DPADataError( diff --git a/dpa_adapt/grouped/__init__.py b/dpa_adapt/grouped/__init__.py new file mode 100644 index 0000000000..dc684e3f85 --- /dev/null +++ b/dpa_adapt/grouped/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Grouped property data helpers. + +Most users should keep data in ordinary ``deepmd/npy`` form and call +``mark_groups`` only when existing systems need grouped markers: + + from dpa_adapt import mark_groups + + mark_groups("oer/dpdata", target="overpotential", group_by="system") + +``Assembly`` is the low-level writer for converter implementations or tests +that already have component arrays in memory: + + from dpa_adapt import Assembly + + a = Assembly(target="x") + a.group(label=..., fparam={...}).add(coords, symbols, weight=...) + a.write(PATH) + +The DeePMD tensor names still use ``group_id`` internally because that is the +training primitive, and users describe assemblies and groups. +""" + +_LAZY = { + "Assembly": ("._core", "Assembly"), + "ComponentSpec": ("._core", "ComponentSpec"), + "GroupSpec": ("._core", "GroupSpec"), + "PoolMask": ("._core", "PoolMask"), + "SiteSelector": ("._core", "SiteSelector"), + "SubstitutionSpec": ("._core", "SubstitutionSpec"), + "GroupMarkerResult": ("._convert", "GroupMarkerResult"), + "mark_groups": ("._convert", "mark_groups"), +} + +__all__ = list(_LAZY) + + +def __getattr__(name: str) -> object: + if name in _LAZY: + import importlib + + mod_name, attr_name = _LAZY[name] + module = importlib.import_module(mod_name, __package__) + value = getattr(module, attr_name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/dpa_adapt/grouped/_aggregation.py b/dpa_adapt/grouped/_aggregation.py new file mode 100644 index 0000000000..b4af1a2245 --- /dev/null +++ b/dpa_adapt/grouped/_aggregation.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generic weighted many-to-one aggregation utilities.""" + +from __future__ import ( + annotations, +) + +import numpy as np + +from dpa_adapt.data.errors import ( + DPADataError, +) + + +def aggregate_weighted_groups( + features: np.ndarray, + group_ids: np.ndarray, + weights: np.ndarray, + labels: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Aggregate row features into one weighted vector per group. + + Parameters + ---------- + features + Per-item feature matrix with shape ``(n_items, feature_dim)``. + group_ids + Integer group id for each feature row, shape ``(n_items,)``. + weights + Weight for each feature row within its group, shape ``(n_items,)``. + labels + Per-item labels, shape ``(n_items,)`` or ``(n_items, task_dim)``. + The first label encountered for each group is used. + + Returns + ------- + embeddings, labels, group_ids + Group-level embeddings, group labels, and sorted group ids. + """ + features = np.asarray(features) + group_ids = np.asarray(group_ids) + weights = np.asarray(weights, dtype=float) + labels = np.asarray(labels) + + if features.ndim != 2: + raise DPADataError( + f"features has shape {features.shape}; expected (n_items, feature_dim)." + ) + n_items = features.shape[0] + if group_ids.shape != (n_items,): + raise DPADataError( + f"group_ids has shape {group_ids.shape}; expected ({n_items},)." + ) + if weights.shape != (n_items,): + raise DPADataError(f"weights has shape {weights.shape}; expected ({n_items},).") + if labels.shape[0] != n_items: + raise DPADataError(f"labels has {labels.shape[0]} rows; expected {n_items}.") + if n_items == 0: + raise DPADataError("Cannot aggregate an empty feature matrix.") + + ordered_ids = np.array( + sorted(np.unique(group_ids.astype(np.int64))), dtype=np.int64 + ) + embeddings = [] + grouped_labels = [] + for group_id in ordered_ids: + mask = group_ids == group_id + embeddings.append(np.sum(features[mask] * weights[mask, None], axis=0)) + grouped_labels.append(labels[mask][0]) + + label_arr = np.asarray(grouped_labels) + if label_arr.ndim == 2 and label_arr.shape[1] == 1: + label_arr = label_arr.reshape(-1) + return np.vstack(embeddings), label_arr, ordered_ids diff --git a/dpa_adapt/grouped/_convert.py b/dpa_adapt/grouped/_convert.py new file mode 100644 index 0000000000..89b0d17c42 --- /dev/null +++ b/dpa_adapt/grouped/_convert.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Materialize grouped auxiliaries for existing deepmd/npy systems. + +Several upstream builders (e.g. the OER ``O*/OH*/OOH*`` adsorption datasets) +already write mixed-type ``real_atom_types.npy`` with ``-1`` marking masked / +virtual atoms and one shared label per frame group, but do **not** write the +``group_id`` / ``pool_mask`` files that the grouped training route +(:mod:`deepmd.pt.utils.grouped`, ``strategy="finetune"``) needs: + +* without ``group_id.npy`` the finetuner never routes into the grouped path and, + even if it did, the loss falls back to one group per frame -- so the shared + label is never aggregated across the group; +* without ``pool_mask.npy`` the model pools the ``-1`` virtual atoms in with the + real atoms (``pool_mask`` defaults to all-ones), so e.g. ``O*``/``OH*``/``OOH*`` + frames become nearly indistinguishable. + +This module fills that gap **in place**, deriving ``pool_mask`` from the +mixed-type ``real_atom_types`` and assigning ``group_id`` by a configurable +policy, without touching the existing coord/box/label data. The files it writes +match :func:`dpa_adapt.grouped._core._write_group_system` exactly (``group_id`` +shape ``(nframes,)`` int64, ``pool_mask`` shape ``(nframes, natoms)`` float64). +""" + +from __future__ import ( + annotations, +) + +import os +from collections.abc import ( + Iterable, +) +from dataclasses import ( + dataclass, + field, +) +from glob import glob as _glob +from glob import ( + has_magic, +) +from pathlib import ( + Path, +) + +import numpy as np + +from dpa_adapt.data.errors import ( + DPADataError, +) + +GROUP_ID_KEY = "group_id" +WEIGHT_KEY = "weight" +POOL_MASK_KEY = "pool_mask" +REAL_ATYPE_KEY = "real_atom_types" + + +@dataclass +class GroupMarkerResult: + """Per-system summary of what :func:`mark_groups` wrote.""" + + system: Path + n_frames: int = 0 + n_groups: int = 0 + wrote_group_id: bool = False + wrote_pool_mask: bool = False + wrote_weight: bool = False + set_dirs: list[Path] = field(default_factory=list) + skipped: bool = False + reason: str = "" + + +def mark_groups( + data: str | Path | Iterable[str | Path], + *, + group_by: str | int = "system", + target: str = "property", + weight: float | None = None, + overwrite: bool = False, + dry_run: bool = False, +) -> list[GroupMarkerResult]: + """Write ``group_id`` / ``pool_mask`` markers into deepmd/npy systems. + + Parameters + ---------- + data + A system directory (one that directly contains ``set.*/``), a parent + directory / glob / arbitrarily deep tree containing many such systems, + or an iterable mixing any of the above. The tree is searched + recursively for every directory that directly holds ``set.*`` subdirs. + group_by + How frames are grouped **within one system** (frame order follows the + DeePMD convention: sorted ``set.*`` directories concatenated): + + * ``"system"`` (default) -- every frame in the system is one group + (group id ``0``). Matches "one system directory == one group", the + layout emitted by the OER writer (one equation directory holds the + ``O*``/``OH*``/``OOH*`` triplet). + * ``"label"`` -- frames whose ``property_name`` label rows are equal + form a group; distinct label rows get ids ``0, 1, 2, ...`` in + first-appearance order. Use when several groups were merged into one + system. + * ``int`` -- fixed group size: every ``group_by`` consecutive frames + form a group. A trailing remainder becomes a smaller final group. + target + Label key read from ``set.*/{property_name}.npy`` when + ``group_by="label"``. Ignored otherwise. + weight + When not ``None`` a constant ``weight.npy`` of this value is written for + every frame. Rarely needed: the grouped model defaults missing weights + to ``1.0`` (an unweighted sum over the group). + overwrite + When ``False`` (default) an existing ``group_id.npy`` / ``pool_mask.npy`` + / ``weight.npy`` is left untouched; only missing files are written. + ``True`` regenerates them (the derivation is deterministic). + dry_run + Compute and report what would be written without touching disk. + + Returns + ------- + list[GroupMarkerResult] + One entry per discovered system. + """ + systems = _discover_systems(data) + if not systems: + raise DPADataError( + f"No deepmd system directories (containing set.*/) found under {data!r}." + ) + if isinstance(group_by, bool) or ( + not isinstance(group_by, int) and group_by not in {"system", "label"} + ): + raise DPADataError( + f"group_by must be 'system', 'label', or a positive int; got {group_by!r}." + ) + if isinstance(group_by, int) and group_by < 1: + raise DPADataError(f"group_by size must be >= 1; got {group_by}.") + + return [ + _process_system( + sysdir, + group_by=group_by, + property_name=target, + weight=weight, + overwrite=overwrite, + dry_run=dry_run, + ) + for sysdir in systems + ] + + +# --------------------------------------------------------------------------- +# system discovery +# --------------------------------------------------------------------------- + + +def _as_paths(data: str | Path | Iterable[str | Path]) -> list[Path]: + if isinstance(data, (str, Path)): + raw = str(data) + if has_magic(raw): + return [Path(p) for p in sorted(_glob(raw))] + return [Path(raw)] + if isinstance(data, Iterable): + out: list[Path] = [] + for item in data: + out.extend(_as_paths(item)) + return out + raise DPADataError(f"Unsupported data spec: {data!r}.") + + +def _is_system_dir(path: Path) -> bool: + return path.is_dir() and any(child.is_dir() for child in path.glob("set.*")) + + +def _discover_systems(data: str | Path | Iterable[str | Path]) -> list[Path]: + found: list[Path] = [] + seen: set[Path] = set() + for root in _as_paths(data): + for sysdir in _walk_systems(root): + resolved = sysdir.resolve() + if resolved not in seen: + seen.add(resolved) + found.append(sysdir) + return found + + +def _walk_systems(root: Path) -> Iterable[Path]: + if not root.exists(): + return + if _is_system_dir(root): + yield root + return + if not root.is_dir(): + return + for dirpath, dirnames, _ in os.walk(root): + path = Path(dirpath) + if _is_system_dir(path): + yield path + # Do not descend into this system's own set.* directories. + dirnames[:] = [d for d in dirnames if not d.startswith("set.")] + + +# --------------------------------------------------------------------------- +# per-system processing +# --------------------------------------------------------------------------- + + +def _process_system( + sysdir: Path, + *, + group_by: str | int, + property_name: str, + weight: float | None, + overwrite: bool, + dry_run: bool, +) -> GroupMarkerResult: + set_dirs = sorted(d for d in sysdir.glob("set.*") if d.is_dir()) + if not set_dirs: + return GroupMarkerResult(sysdir, skipped=True, reason="no set.* directories") + + per_set: list[tuple[Path, int, int, np.ndarray | None]] = [] + for set_dir in set_dirs: + nframes, natoms, real_types = _read_set_shape(set_dir) + if nframes is None: + return GroupMarkerResult( + sysdir, + skipped=True, + reason=f"{set_dir.name} has no coord/real_atom_types", + ) + per_set.append((set_dir, nframes, natoms, real_types)) + + total = sum(n for _, n, _, _ in per_set) + group_ids = _assign_group_ids(group_by, per_set, property_name) + + result = GroupMarkerResult( + sysdir, n_frames=total, n_groups=len(np.unique(group_ids)) + ) + offset = 0 + for set_dir, nframes, _natoms, real_types in per_set: + result.set_dirs.append(set_dir) + gid_slice = group_ids[offset : offset + nframes].astype(np.int64, copy=False) + if _should_write(set_dir / f"{GROUP_ID_KEY}.npy", overwrite): + if not dry_run: + np.save(set_dir / f"{GROUP_ID_KEY}.npy", gid_slice) + result.wrote_group_id = True + + pool_mask_path = set_dir / f"{POOL_MASK_KEY}.npy" + needs_pool_mask = real_types is not None and bool((real_types < 0).any()) + if needs_pool_mask: + if _should_write(pool_mask_path, overwrite): + pool_mask = (real_types >= 0).astype(np.float64) + if not dry_run: + np.save(pool_mask_path, pool_mask) + result.wrote_pool_mask = True + elif overwrite and pool_mask_path.is_file(): + # The deterministic derivation says "no virtual atoms -> no mask". + # Drop any stale pool_mask.npy from an earlier run so it cannot + # silently mask real atoms on the next load. + if not dry_run: + pool_mask_path.unlink() + result.wrote_pool_mask = True + + if weight is not None and _should_write( + set_dir / f"{WEIGHT_KEY}.npy", overwrite + ): + if not dry_run: + np.save( + set_dir / f"{WEIGHT_KEY}.npy", + np.full((nframes,), float(weight), dtype=np.float64), + ) + result.wrote_weight = True + offset += nframes + return result + + +def _should_write(path: Path, overwrite: bool) -> bool: + return overwrite or not path.is_file() + + +def _read_set_shape(set_dir: Path) -> tuple[int | None, int, np.ndarray | None]: + """Return ``(nframes, natoms, real_atom_types|None)`` for one set directory.""" + real_path = set_dir / f"{REAL_ATYPE_KEY}.npy" + if real_path.is_file(): + real_types = np.load(real_path) + if real_types.ndim != 2: + raise DPADataError( + f"{real_path} has shape {real_types.shape}; expected (nframes, natoms)." + ) + return int(real_types.shape[0]), int(real_types.shape[1]), real_types + + coord_path = set_dir / "coord.npy" + if coord_path.is_file(): + coord = np.load(coord_path, mmap_mode="r") + nframes = int(coord.shape[0]) + natoms = int(coord.shape[1] // 3) if coord.ndim == 2 else int(coord.shape[1]) + return nframes, natoms, None + return None, 0, None + + +def _assign_group_ids( + group_by: str | int, + per_set: list[tuple[Path, int, int, np.ndarray | None]], + property_name: str, +) -> np.ndarray: + total = sum(n for _, n, _, _ in per_set) + if group_by == "system": + return np.zeros(total, dtype=np.int64) + if isinstance(group_by, int): + return np.arange(total, dtype=np.int64) // group_by + # group_by == "label" + labels = _load_system_label(per_set, property_name) + if labels.shape[0] != total: + raise DPADataError( + f"label rows ({labels.shape[0]}) do not match frame count ({total})." + ) + ids = np.empty(total, dtype=np.int64) + seen: dict[tuple, int] = {} + nxt = 0 + for i, row in enumerate(labels): + key = tuple(np.round(np.asarray(row, dtype=float), 8).tolist()) + if key not in seen: + seen[key] = nxt + nxt += 1 + ids[i] = seen[key] + return ids + + +def _load_system_label( + per_set: list[tuple[Path, int, int, np.ndarray | None]], + property_name: str, +) -> np.ndarray: + chunks: list[np.ndarray] = [] + for set_dir, nframes, _, _ in per_set: + label_path = set_dir / f"{property_name}.npy" + if not label_path.is_file(): + raise DPADataError( + f"group_by='label' needs {label_path}, which is missing. " + f"Pass the correct property_name (got {property_name!r})." + ) + arr = np.load(label_path) + chunks.append(arr.reshape(nframes, -1)) + return np.concatenate(chunks, axis=0) + + +def _main() -> None: # pragma: no cover - thin CLI wrapper + import argparse + + parser = argparse.ArgumentParser( + description=( + "Add group_id / pool_mask markers to mixed-type deepmd/npy systems " + "so they can be trained via the grouped route (strategy='finetune')." + ) + ) + parser.add_argument( + "data", nargs="+", help="System dir(s), parent tree(s), or glob(s)." + ) + parser.add_argument( + "--group-by", + default="system", + help="'system' (default), 'label', or an integer group size.", + ) + parser.add_argument("--property-name", default="property") + parser.add_argument("--weight", type=float, default=None) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + group_by: str | int = args.group_by + if isinstance(group_by, str) and group_by.isdigit(): + group_by = int(group_by) + + results = mark_groups( + args.data, + group_by=group_by, + target=args.property_name, + weight=args.weight, + overwrite=args.overwrite, + dry_run=args.dry_run, + ) + n_written = sum(1 for r in results if r.wrote_group_id or r.wrote_pool_mask) + print( # noqa: T201 + f"{'[dry-run] ' if args.dry_run else ''}" + f"{len(results)} systems discovered; " + f"{n_written} updated; " + f"{sum(r.n_groups for r in results)} groups total." + ) + for r in results[:10]: + print( # noqa: T201 + f" {r.system}: frames={r.n_frames} groups={r.n_groups} " + f"group_id={r.wrote_group_id} pool_mask={r.wrote_pool_mask}" + + (f" SKIPPED({r.reason})" if r.skipped else "") + ) + if len(results) > 10: + print(f" ... (+{len(results) - 10} more)") # noqa: T201 + + +if __name__ == "__main__": # pragma: no cover + _main() diff --git a/dpa_adapt/grouped/_core.py b/dpa_adapt/grouped/_core.py new file mode 100644 index 0000000000..37eea0ba69 --- /dev/null +++ b/dpa_adapt/grouped/_core.py @@ -0,0 +1,571 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""High-level assembly dataset writer for grouped property training. + +The DeePMD system stores only tensors needed at train time. Scientific +semantics, generation provenance, source paths, roles, blocks, and fparam +schemas live in the adapt manifest so the user-facing API can evolve without +turning a DeePMD ``set.*`` directory into a metadata dump. +""" + +from __future__ import ( + annotations, +) + +import json +import re +import shutil +from dataclasses import ( + dataclass, + field, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import numpy as np + +from dpa_adapt.data.errors import ( + DPADataError, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Mapping, + Sequence, + ) + +GROUP_ID_KEY = "group_id" +WEIGHT_KEY = "weight" +POOL_MASK_KEY = "pool_mask" +MANIFEST_NAME = "manifest.json" + + +@dataclass(frozen=True) +class SiteSelector: + """Declarative site selection spec stored in manifest provenance.""" + + mode: str + value: Any + + @classmethod + def indices(cls, indices: Sequence[int]) -> SiteSelector: + return cls("indices", [int(i) for i in indices]) + + @classmethod + def element(cls, element: str) -> SiteSelector: + return cls("element", str(element)) + + @classmethod + def tag(cls, tag: str) -> SiteSelector: + return cls("tag", str(tag)) + + @classmethod + def top_layer( + cls, element: str | None = None, topk: int | None = None + ) -> SiteSelector: + value: dict[str, Any] = {} + if element is not None: + value["element"] = element + if topk is not None: + value["topk"] = int(topk) + return cls("top_layer", value) + + def to_dict(self) -> dict[str, Any]: + return {"mode": self.mode, "value": self.value} + + +@dataclass(frozen=True) +class SubstitutionSpec: + """Declarative substitution/doping provenance for manifest storage.""" + + sites: SiteSelector + composition: Mapping[str, float] + mode: str = "random_by_fraction" + seed: int | None = None + + def to_dict(self) -> dict[str, Any]: + data: dict[str, Any] = { + "mode": self.mode, + "sites": self.sites.to_dict(), + "composition": {str(k): float(v) for k, v in self.composition.items()}, + } + if self.seed is not None: + data["seed"] = int(self.seed) + return data + + +@dataclass(frozen=True) +class PoolMask: + """Helpers for constructing frame-level pooling masks.""" + + include: Sequence[int] | None = None + exclude: Sequence[int] | None = None + + @classmethod + def all(cls) -> PoolMask: + return cls() + + @classmethod + def exclude_indices(cls, indices: Sequence[int]) -> PoolMask: + return cls(exclude=[int(i) for i in indices]) + + @classmethod + def include_indices(cls, indices: Sequence[int]) -> PoolMask: + return cls(include=[int(i) for i in indices]) + + def as_array(self, natoms: int) -> np.ndarray: + mask = np.ones(natoms, dtype=np.float64) + if self.include is not None: + mask[:] = 0.0 + mask[list(self.include)] = 1.0 + if self.exclude is not None: + mask[list(self.exclude)] = 0.0 + return mask + + +@dataclass +class ComponentSpec: + """One structure/component that becomes one DeePMD frame.""" + + coords: np.ndarray + symbols: list[str] + box: np.ndarray | None = None + weight: float = 1.0 + pool_mask: np.ndarray | PoolMask | None = None + role: str | None = None + block: str | None = None + source: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_arrays( + cls, + coords: Sequence[Sequence[float]], + symbols: Sequence[str], + *, + box: Sequence[Sequence[float]] | Sequence[float] | None = None, + weight: float = 1.0, + pool_mask: Sequence[float] | PoolMask | None = None, + role: str | None = None, + block: str | None = None, + source: str | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> ComponentSpec: + return cls( + coords=np.asarray(coords, dtype=np.float64), + symbols=[str(s) for s in symbols], + box=None if box is None else np.asarray(box, dtype=np.float64), + weight=float(weight), + pool_mask=pool_mask + if isinstance(pool_mask, PoolMask) or pool_mask is None + else np.asarray(pool_mask, dtype=np.float64), + role=role, + block=block, + source=source, + metadata=dict(metadata or {}), + ) + + def normalized_box(self) -> np.ndarray: + if self.box is None: + return np.eye(3, dtype=np.float64) * 100.0 + box = np.asarray(self.box, dtype=np.float64) + if box.shape == (9,): + return box.reshape(3, 3) + if box.shape != (3, 3): + raise DPADataError(f"box has shape {box.shape}; expected (3,3) or (9,).") + return box + + def normalized_pool_mask(self) -> np.ndarray: + natoms = len(self.symbols) + if self.pool_mask is None: + return np.ones(natoms, dtype=np.float64) + if isinstance(self.pool_mask, PoolMask): + return self.pool_mask.as_array(natoms) + mask = np.asarray(self.pool_mask, dtype=np.float64).reshape(-1) + if mask.shape != (natoms,): + raise DPADataError( + f"pool_mask has shape {mask.shape}; expected ({natoms},)." + ) + return mask + + def validate(self) -> None: + if self.coords.shape != (len(self.symbols), 3): + raise DPADataError( + f"coords has shape {self.coords.shape}; expected " + f"({len(self.symbols)}, 3)." + ) + self.normalized_box() + mask = self.normalized_pool_mask() + if float(mask.sum()) == 0.0: + raise DPADataError( + "pool_mask is all-zero: every atom of this component is excluded " + "from pooling, giving an undefined frame embedding. An all-masked " + "frame is a data bug, not a numerical edge case." + ) + + +@dataclass +class GroupSpec: + """A labeled group made of one or more component frames.""" + + key: str + label: float | Sequence[float] + components: list[ComponentSpec] = field(default_factory=list) + fparam: dict[str, float] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + def add_component( + self, component: ComponentSpec, **overrides: Any + ) -> ComponentSpec: + for key, value in overrides.items(): + if not hasattr(component, key): + raise TypeError(f"Unknown ComponentSpec field: {key}") + setattr(component, key, value) + self.components.append(component) + return component + + def add( + self, + coords: Sequence[Sequence[float]], + symbols: Sequence[str], + *, + weight: float = 1.0, + pool_mask: Sequence[float] | PoolMask | None = None, + block: str | None = None, + box: Sequence[Sequence[float]] | Sequence[float] | None = None, + role: str | None = None, + ) -> ComponentSpec: + """Build a component from ``coords`` + ``symbols`` and append it. + + Sugar over ``add_component(ComponentSpec.from_arrays(...))``. + """ + component = ComponentSpec.from_arrays( + coords, + symbols, + box=box, + weight=weight, + pool_mask=pool_mask, + block=block, + role=role, + ) + self.components.append(component) + return component + + +class Assembly: + """Build assembly DeePMD data plus an adapt manifest. + + Each group is written as one DeePMD system. Components within a group may + differ in size and composition: every frame is padded up to the group's max + atom count with virtual atoms (real type -1) and stored in the DeePMD + ``mixed_type`` layout, with padding atoms masked out of pooling. Richer + assembly semantics live in ``manifest.json``. + """ + + def __init__( + self, + *, + target: str = "property", + type_map: Sequence[str] | None = None, + schema: str = "dpa_adapt.assembly.v1", + ) -> None: + self.property_name = str(target) + self.type_map = [str(t) for t in type_map] if type_map is not None else None + self.schema = schema + self.groups: list[GroupSpec] = [] + self.fparam_schema: list[dict[str, Any]] = [] + + def group( + self, + *, + key: str | None = None, + label: float | Sequence[float], + fparam: Mapping[str, float] | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> GroupSpec: + group = GroupSpec( + key=str(key if key is not None else f"group_{len(self.groups)}"), + label=label, + fparam={str(k): float(v) for k, v in (fparam or {}).items()}, + metadata=dict(metadata or {}), + ) + self.groups.append(group) + return group + + def set_fparam_schema(self, schema: Sequence[Mapping[str, Any]]) -> None: + self.fparam_schema = [dict(item) for item in schema] + + def write( + self, + out: str | Path, + *, + system_dir: str = "systems", + overwrite: bool = False, + ) -> dict[str, Any]: + out_path = Path(out) + if out_path.exists() and any(out_path.iterdir()) and not overwrite: + raise DPADataError(f"Output directory is not empty: {out_path}") + out_path.mkdir(parents=True, exist_ok=True) + systems_root = out_path / system_dir + if overwrite and systems_root.exists(): + # Regenerate from scratch: drop stale system dirs so groups removed + # since the last run are not rediscovered by grouped loaders. + shutil.rmtree(systems_root) + systems_root.mkdir(parents=True, exist_ok=True) + + resolved_type_map = self._resolved_type_map() + fparam_columns, fparam_defaults = self._fparam_columns() + + manifest_groups = [] + systems: list[str] = [] + used_system_names: set[str] = set() + for group_idx, group in enumerate(self.groups): + base_name = _safe_name(group.key, fallback=f"group_{group_idx}") + system_name = base_name + suffix = 1 + while system_name in used_system_names: + system_name = f"{base_name}_{suffix}" + suffix += 1 + used_system_names.add(system_name) + system_path = systems_root / system_name + _write_group_system( + group, + group_idx=group_idx, + system_path=system_path, + property_name=self.property_name, + type_map=resolved_type_map, + fparam_columns=fparam_columns, + fparam_defaults=fparam_defaults, + ) + systems.append(str(system_path.relative_to(out_path))) + manifest_groups.append( + _group_manifest( + group, group_idx, str(system_path.relative_to(out_path)) + ) + ) + + manifest = { + "schema": self.schema, + "property_name": self.property_name, + "system_dir": system_dir, + "tensor_fields": { + "group_id": GROUP_ID_KEY, + "weight": WEIGHT_KEY, + "pool_mask": POOL_MASK_KEY, + "label": self.property_name, + "fparam": "fparam" if self._has_fparam() else None, + }, + "type_map": resolved_type_map, + "fparam_schema": self.fparam_schema, + "groups": manifest_groups, + } + manifest_path = out_path / MANIFEST_NAME + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8" + ) + return { + "output_dir": str(out_path.resolve()), + "manifest": str(manifest_path.resolve()), + "systems": systems, + "n_groups": len(self.groups), + } + + def _resolved_type_map(self) -> list[str] | None: + if self.type_map is not None: + return list(self.type_map) + symbols = [ + sym + for group in self.groups + for component in group.components + for sym in component.symbols + ] + return _stable_unique(symbols) if symbols else None + + def _has_fparam(self) -> bool: + return any(group.fparam for group in self.groups) + + def _fparam_columns(self) -> tuple[list[str], dict[str, float]]: + """Canonical fparam columns + per-name defaults for missing values. + + Once any group declares ``fparam`` the manifest advertises a uniform + ``fparam`` field, so every system must write ``fparam.npy`` with the same + columns. Order follows ``fparam_schema`` when set (matching the + manifest), otherwise the sorted union of the groups' keys. Missing + per-group values fall back to the schema ``default`` (or ``0.0``). + """ + if not self._has_fparam(): + return [], {} + if self.fparam_schema: + names = [str(item["name"]) for item in self.fparam_schema] + defaults = { + str(item["name"]): float(item.get("default", 0.0)) + for item in self.fparam_schema + } + return names, defaults + names = sorted({key for group in self.groups for key in group.fparam}) + return names, {} + + +def write_grouped_deepmd( + groups: Iterable[GroupSpec], + out: str | Path, + *, + target: str = "property", + type_map: Sequence[str] | None = None, + fparam_schema: Sequence[Mapping[str, Any]] | None = None, + overwrite: bool = False, +) -> dict[str, Any]: + """Write assembly data from pre-built :class:`GroupSpec` objects.""" + builder = Assembly(target=target, type_map=type_map) + builder.groups.extend(groups) + if fparam_schema is not None: + builder.set_fparam_schema(fparam_schema) + return builder.write(out, overwrite=overwrite) + + +def _write_group_system( + group: GroupSpec, + *, + group_idx: int, + system_path: Path, + property_name: str, + type_map: Sequence[str] | None, + fparam_columns: Sequence[str] = (), + fparam_defaults: Mapping[str, float] | None = None, +) -> None: + if not group.components: + raise DPADataError(f"Group {group.key!r} has no components.") + for component in group.components: + component.validate() + + # Components in a group may differ in size and composition (e.g. OER + # O*/OH*/OOH*). Pad every frame up to the group's max atom count with + # virtual atoms (real type -1) and emit the DeePMD ``mixed_type`` layout so + # one shared descriptor can consume the whole group in a single system. + # Padding atoms are excluded from pooling (``pool_mask`` = 0) and ignored by + # the neighbor list (type < 0), so they never affect real-atom embeddings. + natoms = max(len(c.symbols) for c in group.components) + all_symbols = [sym for c in group.components for sym in c.symbols] + resolved_type_map = ( + list(type_map) if type_map is not None else _stable_unique(all_symbols) + ) + type_index = {el: ii for ii, el in enumerate(resolved_type_map)} + missing = sorted({sym for sym in all_symbols if sym not in type_index}) + if missing: + raise DPADataError( + f"type_map is missing symbols for group {group.key!r}: {missing}" + ) + + set_dir = system_path / "set.000" + set_dir.mkdir(parents=True, exist_ok=True) + (system_path / "type_map.raw").write_text( + "".join(f"{el}\n" for el in resolved_type_map), encoding="utf-8" + ) + # ``mixed_type`` placeholder: a uniform (all-zero) ``type.raw`` makes + # DeePMD's per-atom sort a no-op, keeping coord/pool_mask/real_atom_types + # aligned in written order. Real per-frame types live in real_atom_types. + (system_path / "type.raw").write_text("0\n" * natoms, encoding="utf-8") + + nframes = len(group.components) + coord = np.zeros((nframes, natoms * 3), dtype=np.float64) + real_atom_types = np.full((nframes, natoms), -1, dtype=np.int32) + pool_mask = np.zeros((nframes, natoms), dtype=np.float64) + for frame, component in enumerate(group.components): + n_i = len(component.symbols) + coord[frame, : n_i * 3] = component.coords.reshape(n_i * 3) + real_atom_types[frame, :n_i] = [type_index[sym] for sym in component.symbols] + pool_mask[frame, :n_i] = component.normalized_pool_mask() + # Task D: place padding (virtual, real_atom_types == -1) atoms at a large, + # spread-out non-physical offset so they lie outside every real atom's + # cutoff -- and each other's -- even on a backend whose neighbor list does + # NOT relocate atype<0. (The pt backend also masks atype<0 in nlist, so + # for the pt/PBC path this is defensive; see the data-format note below.) + box_diag = float(np.linalg.norm(component.normalized_box().sum(axis=0))) + for pad_index in range(natoms - n_i): + offset = box_diag + 100.0 * (pad_index + 1) + start = (n_i + pad_index) * 3 + coord[frame, start : start + 3] = offset + + box = np.stack([c.normalized_box().reshape(9) for c in group.components]) + group_id = np.full((nframes,), int(group_idx), dtype=np.int64) + weight = np.asarray([c.weight for c in group.components], dtype=np.float64) + label = np.asarray(group.label, dtype=np.float64).reshape(1, -1) + label = np.repeat(label, nframes, axis=0) + + # Data-format notes: + # - ``role`` (e.g. repeat_unit, solvent) is CONSTRUCTION-TIME metadata: it + # informs weight/pool_mask generation in Assembly readers and is NOT + # serialized to npy or consumed by the model (only kept in the manifest). + # - Padding atoms (real_atom_types == -1) carry non-physical coords by + # construction (large offset above). On the pt backend the nlist also + # relocates atype<0, so for periodic systems -- where the offset may wrap + # back into the cell -- that nlist masking is the guarantee; the offset is + # the fallback for PBC-off / other backends. + np.save(set_dir / "coord.npy", coord) + np.save(set_dir / "box.npy", box) + np.save(set_dir / "real_atom_types.npy", real_atom_types) + np.save(set_dir / f"{property_name}.npy", label) + np.save(set_dir / f"{GROUP_ID_KEY}.npy", group_id) + np.save(set_dir / f"{WEIGHT_KEY}.npy", weight) + np.save(set_dir / f"{POOL_MASK_KEY}.npy", pool_mask) + if fparam_columns: + # The manifest advertises a uniform fparam field, so write fparam.npy + # for every group -- filling any key this group omits with the schema + # default -- so downstream readers never hit a missing/ragged file. + defaults = fparam_defaults or {} + row = np.asarray( + [ + [ + float(group.fparam.get(name, defaults.get(name, 0.0))) + for name in fparam_columns + ] + ], + dtype=np.float64, + ) + np.save(set_dir / "fparam.npy", np.repeat(row, nframes, axis=0)) + + +def _group_manifest(group: GroupSpec, group_idx: int, system: str) -> dict[str, Any]: + components = [] + for frame, component in enumerate(group.components): + item = { + "frame": frame, + "weight": float(component.weight), + "role": component.role, + "block": component.block, + "source": component.source, + "pool_mask_excluded": np.where(component.normalized_pool_mask() == 0)[0] + .astype(int) + .tolist(), + "metadata": component.metadata, + } + components.append({k: v for k, v in item.items() if v is not None and v != {}}) + return { + "group_id": int(group_idx), + "key": group.key, + "label": np.asarray(group.label, dtype=float).reshape(-1).tolist(), + "system": system, + "fparam": group.fparam, + "metadata": group.metadata, + "components": components, + } + + +def _safe_name(raw: str, fallback: str) -> str: + name = re.sub(r"[^A-Za-z0-9_.+-]+", "_", raw).strip("._") + return name or fallback + + +def _stable_unique(values: Sequence[str]) -> list[str]: + seen: set[str] = set() + result: list[str] = [] + for value in values: + if value not in seen: + seen.add(value) + result.append(value) + return result diff --git a/dpa_adapt/grouped/_offline.py b/dpa_adapt/grouped/_offline.py new file mode 100644 index 0000000000..e3fc57f11c --- /dev/null +++ b/dpa_adapt/grouped/_offline.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Grouped descriptor dataset for shared-target structures.""" + +from __future__ import ( + annotations, +) + +import glob as _glob +import os +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, +) + +import dpdata +import numpy as np + +from dpa_adapt.data.errors import ( + DPADataError, +) +from dpa_adapt.data.loader import ( + _get_source, + _resolve_label_key, + load_data, +) +from dpa_adapt.grouped._aggregation import ( + aggregate_weighted_groups, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + ) + + +def load_or_extract(*args: object, **kwargs: object) -> np.ndarray: + """Delegate to finetuner.load_or_extract without a module import cycle.""" + import importlib + + finetuner = importlib.import_module("dpa_adapt.finetuner") + return finetuner.load_or_extract(*args, **kwargs) + + +class GroupedDataset: + """Aggregate per-frame descriptors into one row per group id.""" + + def __init__( + self, + data: str | Path | list[str | Path | dpdata.System], + pretrained: str, + model_branch: str | None = None, + type_map: list[str] | tuple[str, ...] | None = None, + target_key: str | list[str] = "property", + fmt: str | None = None, + cache: bool = True, + ) -> None: + self.data = data + self.pretrained = pretrained + self.model_branch = model_branch + self.type_map = list(type_map) if type_map else None + # target_key may be a single string or a list (multi-property), same + # CLI convention as the non-grouped _load_labels(): --target-key a,b + # is split into a list before reaching here. + keys = [target_key] if isinstance(target_key, str) else list(target_key) + self.target_keys = [_resolve_label_key(key) for key in keys] + self.fmt = fmt + self.cache = cache + + self.systems = _load_groupable_systems(data, fmt=fmt) + self._embeddings, self._labels = self._build() + + def get_embeddings(self) -> np.ndarray: + return self._embeddings + + def get_labels(self) -> np.ndarray: + return self._labels + + def _build(self) -> tuple[np.ndarray, np.ndarray]: + group_ids: list[int] = [] + weights: list[float] = [] + labels: list[np.ndarray] = [] + + next_group_id = 0 + for system in self.systems: + source = _get_source(system) + if source is None: + raise DPADataError( + "Assembly input must come from deepmd/npy directories so " + "set.*/group_id.npy can be read." + ) + source_path = Path(source) + system_frames = _read_system_group_rows(source_path, self.target_keys) + # group_id.npy is scoped to one DeePMD system. Many assembly writers + # naturally use group_id=0 in every system, so remap each system's + # local ids into a process-wide id space before offline aggregation. + local_to_global: dict[int, int] = {} + for local_group_id, weight, label in system_frames: + if local_group_id not in local_to_global: + local_to_global[local_group_id] = next_group_id + next_group_id += 1 + group_ids.append(local_to_global[local_group_id]) + weights.append(float(weight)) + labels.append(np.asarray(label)) + + if not group_ids: + raise DPADataError("Grouped input contains no frames to aggregate.") + + descriptors = load_or_extract( + self.systems, + pretrained=self.pretrained, + model_branch=self.model_branch, + pooling="mean", + cache=self.cache, + type_map=self.type_map, + ) + if descriptors.shape[0] != len(group_ids): + raise DPADataError( + f"Descriptor rows ({descriptors.shape[0]}) do not match grouped " + f"frame rows ({len(group_ids)})." + ) + + embeddings, label_arr, _ = aggregate_weighted_groups( + descriptors, + np.asarray(group_ids, dtype=np.int64), + np.asarray(weights, dtype=float), + np.asarray(labels), + ) + return embeddings, label_arr + + +def has_grouped_markers(data: object) -> bool: + """Return True when any resolved system directory has group_id.npy.""" + return any(_system_has_marker(path) for path in _candidate_paths(data)) + + +def _load_groupable_systems( + data: str | Path | list[str | Path | dpdata.System], + fmt: str | None = None, +) -> list[dpdata.System]: + if isinstance(data, list): + systems: list[dpdata.System] = [] + for item in data: + systems.extend(_load_groupable_systems(item, fmt=fmt)) + return systems + if isinstance(data, (dpdata.System, dpdata.LabeledSystem)): + return [data] + + paths = _candidate_paths(data) + if not paths: + raise DPADataError(f"No grouped system directories found under {data!r}.") + systems: list[dpdata.System] = [] + for path in paths: + systems.extend(load_data(str(path), fmt=fmt)) + return systems + + +def _candidate_paths(data: object) -> list[Path]: + if isinstance(data, list): + paths: list[Path] = [] + for item in data: + paths.extend(_candidate_paths(item)) + return _unique_paths(paths) + if isinstance(data, (dpdata.System, dpdata.LabeledSystem)): + source = _get_source(data) + return [Path(source)] if source is not None else [] + if not isinstance(data, (str, Path)): + return [] + + raw = str(data) + if _glob.has_magic(raw): + return _unique_paths( + path + for match in sorted(_glob.glob(raw)) + for path in _paths_from_directory(Path(match)) + ) + return _paths_from_directory(Path(raw)) + + +def _paths_from_directory(path: Path) -> list[Path]: + if not path.exists(): + return [] + if _is_system_dir(path): + return [path] + if not path.is_dir(): + return [] + # Recurse arbitrarily deep so discovery matches ``mark_groups`` / + # ``_walk_systems`` (which uses ``os.walk``); scanning only immediate + # children would silently drop systems nested more than one level down. + found: list[Path] = [] + for dirpath, dirnames, _ in os.walk(path): + candidate = Path(dirpath) + if _is_system_dir(candidate): + found.append(candidate) + # Do not descend into this system's own set.* directories. + dirnames[:] = [d for d in dirnames if not d.startswith("set.")] + return sorted(found) + + +def _is_system_dir(path: Path) -> bool: + return path.is_dir() and any(path.glob("set.*")) + + +def _system_has_marker(path: Path) -> bool: + return any( + set_dir.joinpath("group_id.npy").is_file() + for set_dir in sorted(path.glob("set.*")) + ) + + +def _unique_paths(paths: Iterable[Path]) -> list[Path]: + result: list[Path] = [] + seen: set[Path] = set() + for path in paths: + resolved = path.resolve() + if resolved not in seen: + seen.add(resolved) + result.append(path) + return result + + +def _read_system_group_rows( + source_path: Path, target_keys: list[str] +) -> list[tuple[int, float, np.ndarray]]: + """Read (group_id, weight, label) rows for one system. + + *target_keys* mirrors the non-grouped ``_load_labels()`` multi-property + convention: one label column is read per key from ``set.*/{key}.npy``. + A single key keeps that column's own shape (unchanged from before + multi-target support); several keys are stacked into one row per frame, + same as ``_load_labels()``'s ``np.column_stack``. + """ + rows: list[tuple[int, float, np.ndarray]] = [] + set_dirs = sorted(source_path.glob("set.*")) + if not set_dirs: + raise DPADataError(f"No set.* directories found in {source_path}.") + + for set_dir in set_dirs: + group_id_path = set_dir / "group_id.npy" + weight_path = set_dir / "weight.npy" + label_paths = [set_dir / f"{key}.npy" for key in target_keys] + missing = [str(p) for p in (group_id_path, *label_paths) if not p.is_file()] + if missing: + raise DPADataError(f"Grouped input is missing required files: {missing}.") + + group_ids = np.asarray(np.load(str(group_id_path)), dtype=np.int64).reshape(-1) + label_arrays = [np.asarray(np.load(str(path))) for path in label_paths] + n_frames = label_arrays[0].shape[0] + for key, arr, path in zip(target_keys, label_arrays, label_paths, strict=True): + if arr.shape[0] != n_frames: + raise DPADataError( + f"{path} has {arr.shape[0]} frames; expected {n_frames} " + f"(from {label_paths[0]}, key={target_keys[0]!r} vs {key!r})." + ) + if weight_path.is_file(): + weight = np.asarray(np.load(str(weight_path)), dtype=float).reshape(-1) + else: + weight = np.ones((n_frames,), dtype=float) + if group_ids.shape != (n_frames,): + raise DPADataError( + f"{group_id_path} has shape {group_ids.shape}; expected ({n_frames},)." + ) + if weight.shape != (n_frames,): + raise DPADataError( + f"{weight_path} has shape {weight.shape}; expected ({n_frames},)." + ) + for frame_idx in range(n_frames): + if len(label_arrays) == 1: + label = np.asarray(label_arrays[0][frame_idx]) + else: + label = np.concatenate( + [np.asarray(arr[frame_idx]).reshape(-1) for arr in label_arrays] + ) + rows.append( + ( + int(group_ids[frame_idx]), + float(weight[frame_idx]), + label, + ) + ) + return rows diff --git a/dpa_adapt/grouped/_polymer.py b/dpa_adapt/grouped/_polymer.py new file mode 100644 index 0000000000..7e16798b35 --- /dev/null +++ b/dpa_adapt/grouped/_polymer.py @@ -0,0 +1,420 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Build grouped DeePMD data for polymer property prediction. + +A polymer is one **group**: each repeating unit / end group becomes one +component (frame), weighted by its mole fraction (end groups by their computed +share). Per-polymer non-structural context (Mn, salts, pH, concentration) is +standardized and written as ``fparam.npy`` -- a per-group side-feature vector the +grouped ``group_property`` head concatenates after aggregation. + +Thin wrapper over :class:`dpa_adapt.grouped.Assembly`: + + PolymerBuilder.from_csv("cloud_points_data.csv", target="cloud_point").write(PATH) + +SMILES are embedded to 3D once each (RDKit ETKDG+MMFF, open valences capped), +the type_map is the element union across all monomers, and every auxiliary npy +(coord / real_atom_types / pool_mask / weight / group_id / fparam / label) is +written automatically. +""" + +from __future__ import ( + annotations, +) + +import json +import math +from collections.abc import ( + Mapping, + Sequence, +) +from dataclasses import ( + dataclass, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np + +from dpa_adapt.data.errors import ( + DPADataError, +) + +MN_KEY = "mn_log" +SALT_PREFIX = "salt:" + + +def _isnan(value: object) -> bool: + try: + return math.isnan(float(value)) + except (TypeError, ValueError): + return value is None + + +def _as_unit_pairs( + units: Mapping[str, float] | Sequence[tuple[str, float]], +) -> list[tuple[str, float]]: + """Normalize ``units`` to a list of ``(smiles, mole_fraction)`` pairs.""" + if isinstance(units, Mapping): + items = list(units.items()) + else: + items = [tuple(item) for item in units] + pairs: list[tuple[str, float]] = [] + for smiles, frac in items: + if smiles is None or _isnan(frac): + continue + pairs.append((str(smiles), float(frac))) + if not pairs: + raise DPADataError("A polymer needs at least one repeating unit.") + return pairs + + +@dataclass +class _PolymerRow: + units: list[tuple[str, float]] + ends: list[str] + mol_weight: float | None + scalars: dict[str, float] + salts: dict[str, float] + target: float + key: str + + +class PolymerBuilder: + """Accumulate polymers and write them as grouped DeePMD systems.""" + + def __init__( + self, + target: str = "cloud_point", + *, + type_map: Sequence[str] | None = None, + seed: int = 42, + ) -> None: + self.target = str(target) + self.type_map = list(type_map) if type_map is not None else None + self.seed = int(seed) + self._rows: list[_PolymerRow] = [] + + # ------------------------------------------------------------------ + # low-level primitive + # ------------------------------------------------------------------ + + def add( + self, + *, + units: Mapping[str, float] | Sequence[tuple[str, float]], + target: float, + ends: Sequence[str] | None = None, + mol_weight: float | None = None, + fparam: Mapping[str, Any] | None = None, + key: str | None = None, + ) -> PolymerBuilder: + """Add one polymer (becomes one group). + + Parameters + ---------- + units + Repeating units as ``{smiles: mole_fraction}`` or ``[(smiles, frac)]``. + target + Group-level label (e.g. cloud point). + ends + End-group SMILES; weighted by the computed end share (needs + *mol_weight*), else by the mean unit fraction. + mol_weight + Number-average molar mass (Mn): used for the end share and, as + ``log10(Mn)``, as a side feature. + fparam + Per-polymer context. Scalar entries become fparam columns; a nested + ``"salts": {name: molar_conc}`` becomes one-hot concentration columns. + key + Optional system name; defaults to ``polymer_{i}``. + """ + conds = dict(fparam or {}) + salts_raw = conds.pop("salts", None) or {} + scalars = { + str(k): float(v) + for k, v in conds.items() + if v is not None and not _isnan(v) + } + salts = { + str(name): float(conc) + for name, conc in salts_raw.items() + if name is not None and conc is not None and not _isnan(conc) + } + self._rows.append( + _PolymerRow( + units=_as_unit_pairs(units), + ends=[str(s) for s in (ends or []) if s is not None and not _isnan(s)], + mol_weight=None + if mol_weight is None or _isnan(mol_weight) + else float(mol_weight), + scalars=scalars, + salts=salts, + target=float(target), + key=str(key) if key is not None else f"polymer_{len(self._rows)}", + ) + ) + return self + + # ------------------------------------------------------------------ + # CSV ingestion (standard cloud-point schema) + # ------------------------------------------------------------------ + + @classmethod + def from_csv( + cls, + path: str | Path, + *, + target: str = "cloud_point", + sep: str = ";", + decimal: str = ",", + seed: int = 42, + ) -> PolymerBuilder: + """Ingest the standard cloud-point polymer CSV. + + Recognizes ``SMILES_repeating_unitA..E`` + ``molpercent_repeating_unitA..E``, + ``SMILES_start_group`` / ``SMILES_end_group``, ``Mn``, ``pH``, + ``polymer_concentration_wpercent`` (or ``_mass_conc``), ``additive1/2`` + + ``additive*_concentration_molar``, and ``{target}``. Rows missing Mn or + the target are skipped. + """ + import pandas as pd + + df = pd.read_csv(str(path), sep=sep, decimal=decimal, encoding="utf8") + builder = cls(target=target, seed=seed) + + def cell(row: Any, col: str) -> Any: + return row[col] if col in df.columns else None + + for idx, row in df.iterrows(): + if _isnan(cell(row, target)) or _isnan(cell(row, "Mn")): + continue + units: list[tuple[str, float]] = [] + for suffix in ("A", "B", "C", "D", "E"): + smiles = cell(row, f"SMILES_repeating_unit{suffix}") + frac = cell(row, f"molpercent_repeating_unit{suffix}") + if smiles is None or _isnan(smiles) or _isnan(frac): + continue + units.append((str(smiles), float(frac))) + if not units: + continue + + ends = [ + str(cell(row, c)) + for c in ("SMILES_start_group", "SMILES_end_group") + if cell(row, c) is not None and not _isnan(cell(row, c)) + ] + + conc = cell(row, "polymer_concentration_wpercent") + if _isnan(conc): + conc = cell(row, "polymer_concentration_mass_conc") + pH = cell(row, "pH") + fparam: dict[str, Any] = { + "pH": 7.0 if _isnan(pH) else float(pH), + } + if not _isnan(conc): + fparam["conc"] = float(conc) + + salts: dict[str, float] = {} + for name_col, conc_col in ( + ("additive1", "additive1_concentration_molar"), + ("additive2", "additive2_concentration_molar"), + ): + name = cell(row, name_col) + c = cell(row, conc_col) + if name is not None and not _isnan(name) and not _isnan(c): + salts[str(name)] = salts.get(str(name), 0.0) + float(c) + if salts: + fparam["salts"] = salts + + builder.add( + units=units, + ends=ends, + mol_weight=float(cell(row, "Mn")), + fparam=fparam, + target=float(cell(row, target)), + key=f"row_{idx}", + ) + if not builder._rows: + raise DPADataError(f"No usable polymer rows parsed from {path!r}.") + return builder + + # ------------------------------------------------------------------ + # write + # ------------------------------------------------------------------ + + def write( + self, + out: str | Path, + *, + overwrite: bool = False, + scaler: dict[str, Any] | str | Path | None = None, + ) -> dict[str, Any]: + """Embed, standardize, and write grouped DeePMD data. + + ``scaler`` reuses a previously written scaler (dict or path to the + ``polymer_scaler.json`` written by an earlier ``write``) so a validation + split is standardized with the training statistics. When ``None`` the + scaler is fit on these rows and saved next to the data. + """ + from dpa_adapt.grouped._core import ( + Assembly, + ComponentSpec, + ) + + if not self._rows: + raise DPADataError("PolymerBuilder is empty; call add()/from_csv() first.") + + coords_by_smiles = self._embed_all() + type_map = self.type_map or self._collect_type_map(coords_by_smiles) + + loaded_scaler = _load_scaler(scaler) + schema = loaded_scaler["columns"] if loaded_scaler else self._build_schema() + stats = loaded_scaler["stats"] if loaded_scaler else self._fit_stats(schema) + + builder = Assembly(target=self.target, type_map=type_map) + for row in self._rows: + fparam_vec = self._standardized_fparam(row, schema, stats) + group = builder.group(key=row.key, label=row.target, fparam=fparam_vec) + end_w = _end_share(row.ends, row.units, row.mol_weight, coords_by_smiles) + for smiles, frac in row.units: + symbols, xyz = coords_by_smiles[smiles] + group.add_component( + ComponentSpec.from_arrays(xyz, symbols, weight=frac, block="rep") + ) + for smiles in row.ends: + symbols, xyz = coords_by_smiles[smiles] + group.add_component( + ComponentSpec.from_arrays(xyz, symbols, weight=end_w, block="end") + ) + + result = builder.write(out, overwrite=overwrite) + + scaler_out = {"columns": schema, "stats": stats} + scaler_path = Path(result["output_dir"]) / "polymer_scaler.json" + scaler_path.write_text(json.dumps(scaler_out, indent=2), encoding="utf-8") + result["fparam_dim"] = len(schema) + result["scaler"] = scaler_out + result["type_map"] = type_map + # Systems live under ``/systems/*``; hand back a ready-to-train glob. + result["train_glob"] = str(Path(result["output_dir"]) / "systems" / "*") + return result + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _embed_all(self) -> dict[str, tuple[list[str], np.ndarray]]: + from dpa_adapt.data.smiles import ( + smiles_to_3d_coords, + ) + + smiles_set = {s for row in self._rows for s, _ in row.units} + smiles_set.update(s for row in self._rows for s in row.ends) + cache: dict[str, tuple[list[str], np.ndarray]] = {} + for smiles in sorted(smiles_set): + symbols, xyz = smiles_to_3d_coords(smiles, random_seed=self.seed) + cache[smiles] = (symbols, np.asarray(xyz, dtype=np.float64)) + return cache + + @staticmethod + def _collect_type_map( + coords_by_smiles: dict[str, tuple[list[str], np.ndarray]], + ) -> list[str]: + seen: dict[str, None] = {} + for symbols, _ in coords_by_smiles.values(): + for sym in symbols: + seen.setdefault(sym, None) + return list(seen) + + def _build_schema(self) -> list[str]: + """Ordered fparam column names: mn_log, scalar fparam fields, salt one-hots.""" + columns: list[str] = [MN_KEY] + scalar_keys: set[str] = set() + salt_names: set[str] = set() + for row in self._rows: + scalar_keys.update(row.scalars) + salt_names.update(row.salts) + columns.extend(sorted(scalar_keys)) + columns.extend(f"{SALT_PREFIX}{name}" for name in sorted(salt_names)) + return columns + + def _raw_vector(self, row: _PolymerRow, schema: Sequence[str]) -> np.ndarray: + vec = np.zeros(len(schema), dtype=np.float64) + for i, col in enumerate(schema): + if col == MN_KEY: + vec[i] = math.log10(row.mol_weight) if row.mol_weight else 0.0 + elif col.startswith(SALT_PREFIX): + vec[i] = row.salts.get(col[len(SALT_PREFIX) :], 0.0) + else: + vec[i] = row.scalars.get(col, 0.0) + return vec + + def _fit_stats(self, schema: Sequence[str]) -> dict[str, list[float]]: + matrix = np.vstack([self._raw_vector(row, schema) for row in self._rows]) + mean = matrix.mean(axis=0) + std = matrix.std(axis=0) + std[std == 0] = 1.0 + return {"mean": mean.tolist(), "std": std.tolist()} + + def _standardized_fparam( + self, + row: _PolymerRow, + schema: Sequence[str], + stats: Mapping[str, Sequence[float]], + ) -> dict[str, float]: + raw = self._raw_vector(row, schema) + mean = np.asarray(stats["mean"], dtype=np.float64) + std = np.asarray(stats["std"], dtype=np.float64) + z = (raw - mean) / std + return {col: float(z[i]) for i, col in enumerate(schema)} + + +def _load_scaler(scaler: dict[str, Any] | str | Path | None) -> dict[str, Any] | None: + if scaler is None: + return None + if isinstance(scaler, (str, Path)): + return json.loads(Path(scaler).read_text(encoding="utf-8")) + return dict(scaler) + + +def _molar_weight( + smiles: str, cache: Mapping[str, tuple[list[str], np.ndarray]] +) -> float: + """Approximate molar weight from the embedded atom symbols.""" + from rdkit import ( + Chem, + ) + + periodic_table = Chem.GetPeriodicTable() + symbols, _ = cache[smiles] + return float( + sum( + periodic_table.GetAtomicWeight(periodic_table.GetAtomicNumber(symbol)) + for symbol in symbols + ) + ) + + +def _end_share( + ends: Sequence[str], + units: Sequence[tuple[str, float]], + mol_weight: float | None, + cache: Mapping[str, tuple[list[str], np.ndarray]], +) -> float: + """Mole-fraction share carried by each end group (mirrors calc_ends_share). + + Falls back to the mean unit fraction when Mn is unavailable. + """ + if not ends: + return 0.0 + if not mol_weight: + return float(np.mean([frac for _, frac in units])) + total_minus = mol_weight - sum(_molar_weight(s, cache) for s in ends) + total_moles = 0.0 + for smiles, frac in units: + total_moles += (total_minus * frac) / _molar_weight(smiles, cache) + return 1.0 / total_moles if total_moles > 0 else 0.0 diff --git a/dpa_adapt/trainer.py b/dpa_adapt/trainer.py index 36ce953c1d..974595267a 100644 --- a/dpa_adapt/trainer.py +++ b/dpa_adapt/trainer.py @@ -101,6 +101,111 @@ _VALID_LOSSES = ("mse", "smooth_mae") +# Group markers written by the grouped data writer; their presence flips the +# fitting/loss to the ``group_property`` head. +_GROUP_MARKERS = ("group_id", "weight", "pool_mask") + + +def _first_set_dir(systems: list) -> str | None: + """Return the first ``set.*`` directory across the resolved systems.""" + for sysdir in systems: + sets = sorted(_glob.glob(os.path.join(sysdir, "set.*"))) + if sets: + return sets[0] + return None + + +def _all_set_dirs(systems: list) -> list[str]: + """Return every ``set.*`` directory across the resolved systems, in order.""" + dirs: list[str] = [] + for sysdir in systems: + dirs.extend(sorted(_glob.glob(os.path.join(sysdir, "set.*")))) + return dirs + + +def _set_marker_status(setdir: str) -> bool | None: + """Whether one ``set.*`` directory carries the group markers. + + Returns ``True`` when all of ``_GROUP_MARKERS`` are present, ``False`` + when none are, and ``None`` for a partial set (some but not all) -- a + partial set is always an error, independent of any other set. + """ + present = [ + os.path.isfile(os.path.join(setdir, f"{name}.npy")) for name in _GROUP_MARKERS + ] + if all(present): + return True + if not any(present): + return False + return None + + +def _systems_are_grouped(*labeled_systems: tuple[str, list]) -> bool: + """Whether the given (label, systems) groups are grouped training data. + + Scans *every* ``set.*`` directory of *every* given system list -- not + just the first set of the first system -- and requires all of them to + agree. Checking only the first set previously meant: a later system + that had markers when the first one didn't left grouped mode disabled + (and that system silently dropped its labels), while a later system + *missing* markers when the first one had them enabled grouped mode and + then failed, or trained inconsistently, once that system was reached. + Any partial marker set, or any mix of grouped and ungrouped sets across + the given system lists (e.g. grouped train_systems with ungrouped + valid_systems), raises a clear error instead of guessing. + """ + from dpa_adapt.data.errors import ( + DPADataError, + ) + + grouped_dirs: list[str] = [] + ungrouped_dirs: list[str] = [] + partial_dirs: list[str] = [] + for label, systems in labeled_systems: + for setdir in _all_set_dirs(systems): + status = _set_marker_status(setdir) + if status is True: + grouped_dirs.append(f"{label}:{setdir}") + elif status is False: + ungrouped_dirs.append(f"{label}:{setdir}") + else: + partial_dirs.append(f"{label}:{setdir}") + + if partial_dirs: + raise DPADataError( + "Inconsistent grouped markers: the following set.* directories " + f"have only some of {_GROUP_MARKERS} (all three, or none, are " + "required):\n " + "\n ".join(partial_dirs[:10]) + ) + if grouped_dirs and ungrouped_dirs: + raise DPADataError( + "Inconsistent grouped markers across systems: " + f"{len(grouped_dirs)} set.* director{'y' if len(grouped_dirs) == 1 else 'ies'} " + f"carry {_GROUP_MARKERS} and {len(ungrouped_dirs)} do not. Mixing " + "grouped and ungrouped systems in the same train/valid run is " + "not supported.\n grouped, e.g.:\n " + + "\n ".join(grouped_dirs[:5]) + + "\n ungrouped, e.g.:\n " + + "\n ".join(ungrouped_dirs[:5]) + ) + return bool(grouped_dirs) + + +def _detect_fparam_dim(systems: list) -> int: + """Per-frame side-feature width from ``set.*/fparam.npy`` (0 if absent).""" + setdir = _first_set_dir(systems) + if setdir is None: + return 0 + fpath = os.path.join(setdir, "fparam.npy") + if not os.path.isfile(fpath): + return 0 + import numpy as np + + arr = np.load(fpath) + if arr.ndim == 0: + return 0 + return int(arr.reshape(arr.shape[0], -1).shape[1]) + # --------------------------------------------------------------------------- # DPATrainer @@ -173,6 +278,7 @@ def __init__( # ---- model overrides ---- fitting_net_params: dict | None = None, fparam_dim: int = 0, + grouped: bool | None = None, # ---- training ---- learning_rate: float = 1e-3, stop_lr: float = 1e-5, @@ -235,6 +341,8 @@ def __init__( self.type_map = type_map self.fitting_net_params = fitting_net_params self.fparam_dim = fparam_dim + # None => auto-detect from the resolved training systems in _build_config. + self.grouped = grouped self.learning_rate = learning_rate self.stop_lr = stop_lr self.decay_steps = decay_steps @@ -333,6 +441,14 @@ def _build_fitting_net(self) -> dict: "seed": self.seed, } ) + if self.grouped: + # Grouped data pools frame embeddings per assembly, so the head and + # loss switch to the group_property variants (same property schema). + fn["type"] = "group_property" + # The grouped head consumes an un-normalized frame embedding + fparam; + # tanh saturates at that scale and the head collapses to a constant. + # Default to GELU (a user override via fitting_net_params still wins). + fn["activation_function"] = "gelu" # NB: dim_case_embd is intentionally NOT injected for FT/LP. The paper # qm9_gap input.json omits it: single-task `--finetune` (without # --model-branch) copies only the backbone and random-inits the @@ -344,6 +460,26 @@ def _build_fitting_net(self) -> dict: fn.update(self.fitting_net_params) return fn + def _infer_grouped(self) -> None: + """Resolve grouped mode and ``numb_fparam`` from the training systems. + + Idempotent, so ``fit()`` can call it *before* the fparam preflight (so + auto-detected grouped ``fparam.npy`` files are still validated) and + ``_build_config()`` can call it again cheaply. + """ + if self.grouped is False: + return + if self.grouped and self.fparam_dim: + return + train_sys = self._expand_systems(self.train_systems, "train_systems") + if self.grouped is None: + valid_sys = self._expand_systems(self.valid_systems, "valid_systems") + self.grouped = _systems_are_grouped( + ("train_systems", train_sys), ("valid_systems", valid_sys) + ) + if self.grouped and not self.fparam_dim: + self.fparam_dim = _detect_fparam_dim(train_sys) + def _build_config(self) -> dict: # Seed propagation in DeePMD-kit v3.1.3 (deepmd/utils/argcheck.py): # - model.descriptor.seed verified: descrpt_dpa3_args() L1428 @@ -357,6 +493,10 @@ def _build_config(self) -> dict: self._resolved_train_systems = train_sys self._resolved_valid_systems = valid_sys + # Grouped training is inferred from the resolved systems unless the + # caller forced it; a grouped set also auto-sizes numb_fparam. + self._infer_grouped() + descriptor = self._get_descriptor() descriptor["seed"] = self.seed # verified: descrpt_dpa3_args (deepmd v3.1.3) fitting_net = self._build_fitting_net() @@ -368,7 +508,7 @@ def _build_config(self) -> dict: "fitting_net": fitting_net, }, "loss": { - "type": "property", + "type": "group_property" if self.grouped else "property", "loss_func": self.loss_function, "metric": ["mae", "rmse"], }, @@ -546,6 +686,11 @@ def fit(self) -> str: ) return str(latest) + # Infer grouped mode / numb_fparam first so an auto-detected grouped + # fparam.npy is covered by the preflight below (otherwise fparam_dim is + # still 0 here and the shape/frame checks are skipped). + self._infer_grouped() + if self.fparam_dim > 0: self._validate_fparam(self.train_systems, self.fparam_dim) if self.valid_systems is not None: diff --git a/source/tests/common/test_argcheck_group_property.py b/source/tests/common/test_argcheck_group_property.py new file mode 100644 index 0000000000..b7bca050fc --- /dev/null +++ b/source/tests/common/test_argcheck_group_property.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""``group_reduce`` was implemented on GroupPropertyFittingNet but the +argcheck schema for the ``group_property`` fitting type reused the +``property`` schema verbatim, so dargs' strict-mode key check (the same +check ``dp --pt train`` runs on every ``input.json``) rejected +``group_reduce`` as an unknown key -- even though ``group_reduce="sum"`` is +a supported, tested code path once past validation. + +The same reused schema also let a config set ``numb_aparam``, +``default_fparam``, ``dim_case_embd``, ``resnet_dt``, ``intensive``, or +``distinguish_types`` -- fields GroupPropertyFittingNet has no wiring for -- +and pass strict validation while silently doing nothing. Those fields are +removed from the group_property schema so setting one is now a validation +error like any other typo, not a silent no-op. +""" + +from __future__ import ( + annotations, +) + +from dargs import ( + Argument, +) +from dargs.dargs import ( + ArgumentKeyError, +) + +from deepmd.utils.argcheck import ( + fitting_group_property, + fitting_variant_type_args, +) + + +def _fitting_net_schema() -> Argument: + # Mirrors how model_args() builds the real "fitting_net" field: a + # type-selected variant, not a bare argument list, so "type" is handled + # by the Variant dispatch rather than treated as an unknown key. + return Argument("fitting_net", dict, [], [fitting_variant_type_args()]) + + +def _normalize_and_check(value: dict) -> dict: + schema = _fitting_net_schema() + normalized = schema.normalize_value(value) + schema.check_value(normalized, strict=True) + return normalized + + +def test_group_reduce_argument_is_declared(): + names = [arg.name for arg in fitting_group_property()] + assert "group_reduce" in names + + +def test_group_reduce_sum_passes_strict_validation(): + out = _normalize_and_check( + {"type": "group_property", "property_name": "y", "group_reduce": "sum"} + ) + assert out["group_reduce"] == "sum" + + +def test_group_reduce_defaults_to_mean_when_omitted(): + out = _normalize_and_check({"type": "group_property", "property_name": "y"}) + assert out["group_reduce"] == "mean" + + +def test_group_property_still_defaults_activation_to_gelu(): + # guard against the group_reduce addition disturbing the existing + # gelu-default override. + out = _normalize_and_check({"type": "group_property", "property_name": "y"}) + assert out["activation_function"] == "gelu" + + +def test_unsupported_property_fields_are_removed_from_the_schema(): + names = {arg.name for arg in fitting_group_property()} + for unsupported in ( + "numb_aparam", + "default_fparam", + "dim_case_embd", + "resnet_dt", + "intensive", + "distinguish_types", + ): + assert unsupported not in names + + +def test_setting_an_unsupported_field_fails_strict_validation(): + for unsupported, value in ( + ("numb_aparam", 3), + ("resnet_dt", False), + ("intensive", True), + ("distinguish_types", False), + ): + payload = { + "type": "group_property", + "property_name": "y", + unsupported: value, + } + normalized = _fitting_net_schema().normalize_value(payload) + try: + _fitting_net_schema().check_value(normalized, strict=True) + except ArgumentKeyError: + pass + else: + raise AssertionError(f"expected {unsupported!r} to fail strict check") + + +def test_group_reduce_key_was_rejected_by_the_reused_property_schema(): + """Regression pin: the bug this fixes. Before adding ``group_reduce`` to + ``fitting_group_property()``, the same strict-mode check that accepts it + above raised ``ArgumentKeyError`` -- confirmed here against the plain + ``property`` schema, which is exactly what ``fitting_group_property()`` + used to return unmodified (plus the activation-default override). + """ + from deepmd.utils.argcheck import ( + fitting_property, + ) + + old_schema = Argument("fitting_net", dict, fitting_property()) + value = {"property_name": "y", "group_reduce": "sum"} + normalized = old_schema.normalize_value(value) + try: + old_schema.check_value(normalized, strict=True) + except ArgumentKeyError: + pass + else: + raise AssertionError( + "expected ArgumentKeyError: the plain property schema has no " + "group_reduce field" + ) diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index cce2ca1850..45e8341c64 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -272,3 +272,52 @@ def test_finetune_rule_builder_rejects_multitask_cli_branch(): assert "Multi-task fine-tuning" in str(exc) else: raise AssertionError("expected ValueError") + + +def test_finetune_rule_builder_single_task_target_auto_picks_multitask_branch(): + """A single-task target (e.g. GroupPropertyModel) finetuning directly off a + multi-task foundation-model checkpoint, with no CLI ``--model-branch`` and + no ``finetune_head`` set, must auto-select the checkpoint's only branch and + randomly re-initialize the fitting net rather than erroring out. + """ + pretrained = { + "model_dict": { + "Alex2D": _model_config(["O", "H"], descriptor_sel=[8, 16]), + } + } + target = _model_config(["O", "H"], descriptor_sel=[1, 1], fitting_neuron=[2]) + + updated, links = finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=True, + ).build() + + rule = links["Default"] + assert rule.get_model_branch() == "Alex2D" + assert rule.get_random_fitting() + assert updated["descriptor"]["sel"] == [8, 16] + # fitting net stays the target's own (freshly-initialized) spec + assert updated["fitting_net"]["neuron"] == [2] + + +def test_finetune_rule_builder_single_task_target_picks_first_of_several_branches(): + """Documents current (non-error) behavior: with several branches and no + explicit selection, the builder deterministically takes the first one in + insertion order rather than raising an ambiguity error. + """ + pretrained = { + "model_dict": { + "first": _model_config(["O", "H"], descriptor_sel=[8, 16]), + "second": _model_config(["O", "H"], descriptor_sel=[4, 4]), + } + } + target = _model_config(["O", "H"], descriptor_sel=[1, 1]) + + _, links = finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=True, + ).build() + + assert links["Default"].get_model_branch() == "first" diff --git a/source/tests/dpa_adapt/test_assemblies.py b/source/tests/dpa_adapt/test_assemblies.py new file mode 100644 index 0000000000..36a9328c30 --- /dev/null +++ b/source/tests/dpa_adapt/test_assemblies.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from __future__ import ( + annotations, +) + +import json + +import numpy as np +import pytest + +from dpa_adapt import ( + Assembly, + ComponentSpec, + PoolMask, + SiteSelector, + SubstitutionSpec, +) +from dpa_adapt.data.errors import ( + DPADataError, +) +from dpa_adapt.grouped._core import ( + GROUP_ID_KEY, + POOL_MASK_KEY, + WEIGHT_KEY, +) + + +def _component(offset: float = 0.0, *, role: str = "state") -> ComponentSpec: + return ComponentSpec.from_arrays( + coords=[[offset, 0.0, 0.0], [0.0, 1.0 + offset, 0.0], [0.0, 0.0, 1.0]], + symbols=["Ni", "O", "H"], + box=np.eye(3) * 12.0, + weight=0.5, + pool_mask=PoolMask.exclude_indices([2]), + role=role, + block="oer_adsorbates", + source=f"{role}.vasp", + metadata={"anchor_atom": 0}, + ) + + +def test_assembly_builder_writes_minimal_deepmd_tensors_and_manifest(tmp_path) -> None: + builder = Assembly(target="overpotential", type_map=["Ni", "O", "H"]) + sub = SubstitutionSpec( + sites=SiteSelector.element("Ni"), + composition={"Ni": 0.8, "Fe": 0.2}, + seed=123, + ) + group = builder.group( + key="Ni0.8Fe0.2O2H1", + label=291.9, + metadata={"substitution": sub.to_dict()}, + ) + group.add_component(_component(0.0, role="O*")) + group.add_component(_component(0.1, role="OH*")) + group.add_component(_component(0.2, role="OOH*")) + + result = builder.write(tmp_path) + system = tmp_path / result["systems"][0] + set_dir = system / "set.000" + + assert (tmp_path / "manifest.json").is_file() + assert sorted(p.name for p in set_dir.iterdir()) == sorted( + [ + "box.npy", + "coord.npy", + f"{GROUP_ID_KEY}.npy", + "overpotential.npy", + f"{POOL_MASK_KEY}.npy", + "real_atom_types.npy", + f"{WEIGHT_KEY}.npy", + ] + ) + assert np.load(set_dir / "coord.npy").shape == (3, 9) + assert np.load(set_dir / "box.npy").shape == (3, 9) + assert np.load(set_dir / "overpotential.npy").shape == (3, 1) + assert np.load(set_dir / f"{GROUP_ID_KEY}.npy").tolist() == [0, 0, 0] + assert np.load(set_dir / f"{WEIGHT_KEY}.npy").tolist() == [0.5, 0.5, 0.5] + assert np.load(set_dir / f"{POOL_MASK_KEY}.npy").tolist() == [ + [1.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + ] + # mixed_type layout: real per-frame types live in real_atom_types.npy and + # type.raw is a uniform all-zero placeholder (Ni/O/H -> 0/1/2 in type_map). + assert np.load(set_dir / "real_atom_types.npy").tolist() == [ + [0, 1, 2], + [0, 1, 2], + [0, 1, 2], + ] + assert (system / "type.raw").read_text().splitlines() == ["0", "0", "0"] + + manifest = json.loads((tmp_path / "manifest.json").read_text()) + assert manifest["schema"] == "dpa_adapt.assembly.v1" + assert manifest["tensor_fields"] == { + "fparam": None, + "group_id": "group_id", + "label": "overpotential", + "pool_mask": "pool_mask", + "weight": "weight", + } + g0 = manifest["groups"][0] + assert g0["key"] == "Ni0.8Fe0.2O2H1" + assert g0["metadata"]["substitution"]["sites"] == {"mode": "element", "value": "Ni"} + assert [c["role"] for c in g0["components"]] == ["O*", "OH*", "OOH*"] + assert g0["components"][0]["pool_mask_excluded"] == [2] + + +def test_fparam_write_fparam_but_schema_stays_in_manifest(tmp_path) -> None: + builder = Assembly(target="cloud_point", type_map=["C", "H"]) + builder.set_fparam_schema( + [ + {"name": "log_mn", "source": "Mn", "transform": "log10(x)/6"}, + {"name": "pH", "default": 7.0}, + ] + ) + group = builder.group( + key="polymer_0", label=32.1, fparam={"log_mn": 0.67, "pH": 7.0} + ) + group.add_component( + ComponentSpec.from_arrays( + coords=[[0, 0, 0], [0, 0, 1]], + symbols=["C", "H"], + weight=1.0, + role="repeat_unit_A", + block="repeat_units", + ) + ) + builder.write(tmp_path) + fparam = np.load(tmp_path / "systems" / "polymer_0" / "set.000" / "fparam.npy") + assert fparam.tolist() == [[0.67, 7.0]] + manifest = json.loads((tmp_path / "manifest.json").read_text()) + assert manifest["tensor_fields"]["fparam"] == "fparam" + assert [item["name"] for item in manifest["fparam_schema"]] == ["log_mn", "pH"] + assert manifest["groups"][0]["components"][0]["block"] == "repeat_units" + + +def test_writer_pads_heterogeneous_components_into_mixed_type(tmp_path) -> None: + builder = Assembly(target="property", type_map=["C", "O", "H"]) + group = builder.group(key="oer", label=1.0) + group.add_component(ComponentSpec.from_arrays([[0, 0, 0], [0, 0, 1]], ["C", "O"])) + group.add_component( + ComponentSpec.from_arrays( + [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + ["C", "O", "H"], + pool_mask=PoolMask.exclude_indices([2]), + ) + ) + + result = builder.write(tmp_path) + system = tmp_path / result["systems"][0] + set_dir = system / "set.000" + + # every frame is padded to the group's max atom count (3) + assert (system / "type.raw").read_text().splitlines() == ["0", "0", "0"] + # the smaller component gets a trailing virtual atom (real type -1) + assert np.load(set_dir / "real_atom_types.npy").tolist() == [ + [0, 1, -1], + [0, 1, 2], + ] + coord = np.load(set_dir / "coord.npy") + assert coord.shape == (2, 9) + assert (np.abs(coord[0, 6:]) > 20.0).all() # padding coords are far off-system + assert coord[1].tolist() == [0, 0, 0, 0, 0, 1, 0, 0, 2] + # padding atom and the excluded H cap are both masked out of pooling + assert np.load(set_dir / f"{POOL_MASK_KEY}.npy").tolist() == [ + [1.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + ] + assert np.load(set_dir / f"{GROUP_ID_KEY}.npy").tolist() == [0, 0] + assert np.load(set_dir / "property.npy").shape == (2, 1) + + +def test_writer_infers_global_type_map_across_groups(tmp_path) -> None: + builder = Assembly(target="property") + g0 = builder.group(key="only_c", label=1.0) + g0.add_component(ComponentSpec.from_arrays([[0, 0, 0]], ["C"])) + g1 = builder.group(key="o_h", label=2.0) + g1.add_component(ComponentSpec.from_arrays([[0, 0, 0], [0, 0, 1]], ["O", "H"])) + + result = builder.write(tmp_path) + expected = ["C", "O", "H"] + for system in result["systems"]: + assert (tmp_path / system / "type_map.raw").read_text().splitlines() == expected + + assert np.load( + tmp_path / "systems" / "only_c" / "set.000" / "real_atom_types.npy" + ).tolist() == [[0]] + assert np.load( + tmp_path / "systems" / "o_h" / "set.000" / "real_atom_types.npy" + ).tolist() == [[1, 2]] + + manifest = json.loads((tmp_path / "manifest.json").read_text()) + assert manifest["type_map"] == expected + + +def test_writer_rejects_symbols_missing_from_type_map(tmp_path) -> None: + builder = Assembly(target="property", type_map=["C"]) + group = builder.group(key="bad", label=1.0) + group.add_component(ComponentSpec.from_arrays([[0, 0, 0], [0, 0, 1]], ["C", "H"])) + with pytest.raises(DPADataError, match="type_map is missing symbols"): + builder.write(tmp_path) diff --git a/source/tests/dpa_adapt/test_cache.py b/source/tests/dpa_adapt/test_cache.py index e16be6199b..8627cd437f 100644 --- a/source/tests/dpa_adapt/test_cache.py +++ b/source/tests/dpa_adapt/test_cache.py @@ -52,6 +52,43 @@ def test_different_elements_different_fp(self, tmp_path): s2 = _make_system(tmp_path, "s2", elements=["Cu", "O"]) assert _system_fingerprint(s1) != _system_fingerprint(s2) + def test_different_real_atom_types_different_fp(self, tmp_path): + # Grouped systems: type.raw/atom_types is a uniform placeholder, so + # coords/atom_types/cells alone are identical across these two + # systems -- only set.000/real_atom_types.npy (per-frame, with -1 + # padding) differs. The fingerprint must still change. + s1 = _make_system(tmp_path, "s1", natoms=3, nframes=2) + s2 = _make_system(tmp_path, "s2", natoms=3, nframes=2) + # force identical coords/box so only real_atom_types.npy differs + coord = np.load(tmp_path / "s1" / "set.000" / "coord.npy") + np.save(tmp_path / "s2" / "set.000" / "coord.npy", coord) + s1, s2 = load_data(str(tmp_path / "s1"))[0], load_data(str(tmp_path / "s2"))[0] + assert _system_fingerprint(s1) == _system_fingerprint(s2) # sanity + + np.save( + tmp_path / "s1" / "set.000" / "real_atom_types.npy", + np.array([[0, 1, -1], [0, 1, 2]]), + ) + np.save( + tmp_path / "s2" / "set.000" / "real_atom_types.npy", + np.array([[0, 1, 2], [0, 1, 2]]), + ) + assert _system_fingerprint(s1) != _system_fingerprint(s2) + + def test_different_pool_mask_different_fp(self, tmp_path): + s1 = _make_system(tmp_path, "s1", natoms=3, nframes=2) + s2 = _make_system(tmp_path, "s2", natoms=3, nframes=2) + + np.save( + tmp_path / "s1" / "set.000" / "pool_mask.npy", + np.array([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]), + ) + np.save( + tmp_path / "s2" / "set.000" / "pool_mask.npy", + np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), + ) + assert _system_fingerprint(s1) != _system_fingerprint(s2) + class TestFingerprint: def test_identical_data_same_fp(self, tmp_path): diff --git a/source/tests/dpa_adapt/test_finetuner_strategies.py b/source/tests/dpa_adapt/test_finetuner_strategies.py index 6f280ffa90..ed9b157c09 100644 --- a/source/tests/dpa_adapt/test_finetuner_strategies.py +++ b/source/tests/dpa_adapt/test_finetuner_strategies.py @@ -501,3 +501,79 @@ def __init__(self): features = ft.extract_features([FakeSystem()]) np.testing.assert_allclose(features, np.array([[2.0, 4.0, 6.0]])) + + +def test_extract_features_uses_per_frame_real_atom_types_for_grouped_systems( + monkeypatch, +): + """Grouped systems store a uniform ``type.raw`` placeholder and the real, + per-frame local types (with ``-1`` padding) in ``real_atom_types.npy``. + ``extract_features`` must feed the model those per-frame types instead of + tiling the single, uniform ``atom_types`` array -- otherwise every atom in + every frame is described as if it had the placeholder's type. + """ + import numpy as np + import torch + + from dpa_adapt import finetuner as finetuner_mod + + seen_atype = [] + + class FakeExtractor: + def __init__(self, model): + self.model = model + + def _enable_hook(self): + pass + + def _disable_hook(self): + pass + + def _run_forward(self, coord_t, atype_t, box_t): + seen_atype.append(atype_t.clone()) + return torch.zeros( + atype_t.shape[0], atype_t.shape[1], 1, dtype=coord_t.dtype + ) + + class FakeSystem: + orig = "fake" + + def __init__(self): + self.data = {"atom_names": ["H", "O", "N"]} + + # frame 0: H, O, padding; frame 1: H, O, N -- deliberately NOT what the + # uniform placeholder (all zeros) would tile. + real_types = np.array([[0, 1, -1], [0, 1, 2]], dtype=np.int64) + + monkeypatch.setattr(finetuner_mod, "_DescriptorExtraction", FakeExtractor) + monkeypatch.setattr( + finetuner_mod, + "_load_npy_system", + lambda system: ( + np.zeros((2, 3, 3)), + np.tile(np.eye(3).ravel(), (2, 1)), + np.array([0, 0, 0], dtype=np.int64), + ), + ) + monkeypatch.setattr( + finetuner_mod, + "_real_atom_types_for_system", + lambda system, n_frames, n_atoms: real_types, + ) + + ft = finetuner_mod._FrozenSklearnPipeline( + pretrained="fake.pt", + model_branch=None, + predictor_type="linear", + pooling="mean", + seed=42, + ) + ft._model = object() + ft._device = torch.device("cpu") + ft.type_map = ["H", "O", "N"] + ft._checkpoint_type_map = ["H", "O", "N"] + + ft.extract_features([FakeSystem()]) + + assert len(seen_atype) == 1 + np.testing.assert_array_equal(seen_atype[0].numpy(), real_types) diff --git a/source/tests/dpa_adapt/test_grouped_convert.py b/source/tests/dpa_adapt/test_grouped_convert.py new file mode 100644 index 0000000000..fea531e4ba --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_convert.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the mixed-type -> grouped marker converter (mark_groups). + +Pure numpy / filesystem; no compiled deepmd backend needed, so these run +anywhere. +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + +from dpa_adapt import ( + mark_groups, +) + + +def _write_mixed_system( + set_dir, + *, + natoms: int = 5, + masked_per_frame=(2, 1, 0), + overpotential: float = 324.9, +) -> None: + """Write an OER-style mixed-type set.000 (real_atom_types with -1, one label).""" + set_dir.mkdir(parents=True, exist_ok=True) + nframes = len(masked_per_frame) + coord = np.random.default_rng(0).random((nframes, natoms * 3)) + box = np.tile((np.eye(3) * 20.0).reshape(9), (nframes, 1)) + real = np.zeros((nframes, natoms), dtype=np.int32) + real[:, 0] = 1 # a second element, just to be non-trivial + for frame, k in enumerate(masked_per_frame): + if k: + real[frame, natoms - k :] = -1 + label = np.full((nframes, 1), float(overpotential), dtype=np.float64) + np.save(set_dir / "coord.npy", coord) + np.save(set_dir / "box.npy", box) + np.save(set_dir / "real_atom_types.npy", real) + np.save(set_dir / "overpotential.npy", label) + (set_dir.parent / "type.raw").write_text("0\n" * natoms) + (set_dir.parent / "type_map.raw").write_text("Ni\nO\n") + + +def test_group_by_system_writes_group_id_and_pool_mask(tmp_path): + sysdir = tmp_path / "Fe0.2Ni0.8" / "05" + _write_mixed_system(sysdir / "set.000", natoms=5, masked_per_frame=(2, 1, 0)) + + results = mark_groups(tmp_path, target="overpotential") + assert len(results) == 1 + r = results[0] + assert r.n_frames == 3 + assert r.n_groups == 1 + assert r.wrote_group_id and r.wrote_pool_mask + + set_dir = sysdir / "set.000" + gid = np.load(set_dir / "group_id.npy") + assert gid.dtype == np.int64 + assert gid.shape == (3,) + assert gid.tolist() == [0, 0, 0] # one group per system + + pool_mask = np.load(set_dir / "pool_mask.npy") + assert pool_mask.dtype == np.float64 + assert pool_mask.shape == (3, 5) + real = np.load(set_dir / "real_atom_types.npy") + np.testing.assert_array_equal(pool_mask, (real >= 0).astype(np.float64)) + # O*/OH*/OOH*: 2 / 1 / 0 masked -> 3 / 4 / 5 pooled atoms + assert pool_mask.sum(axis=1).tolist() == [3.0, 4.0, 5.0] + + +def test_discovers_deeply_nested_systems(tmp_path): + # mimic dpdata/set_XX/{equation}/{natoms}/set.000 layout + for si in (1, 2): + for eq in ("Fe0.2Ni0.8", "Co1.0"): + _write_mixed_system( + tmp_path / "dpdata" / f"set_{si:02d}" / eq / "05" / "set.000" + ) + results = mark_groups(tmp_path, target="overpotential") + assert len(results) == 4 + assert all(r.wrote_group_id for r in results) + # every leaf now has both markers + for gid in tmp_path.rglob("group_id.npy"): + assert np.load(gid).tolist() == [0, 0, 0] + + +def test_group_by_label_splits_distinct_labels(tmp_path): + # one system holding two triplets with different overpotentials + set_dir = tmp_path / "merged" / "set.000" + set_dir.mkdir(parents=True) + real = np.zeros((6, 4), dtype=np.int32) + real[0, -2:] = -1 # a couple of masked atoms so pool_mask is exercised + np.save(set_dir / "coord.npy", np.zeros((6, 12))) + np.save(set_dir / "real_atom_types.npy", real) + np.save(set_dir / "overpotential.npy", np.array([[1.0]] * 3 + [[2.0]] * 3)) + + mark_groups(tmp_path, group_by="label", target="overpotential") + gid = np.load(set_dir / "group_id.npy") + assert gid.tolist() == [0, 0, 0, 1, 1, 1] + + +def test_group_by_int_chunks(tmp_path): + set_dir = tmp_path / "chunked" / "set.000" + set_dir.mkdir(parents=True) + np.save(set_dir / "coord.npy", np.zeros((7, 3))) # 7 frames, 1 atom + mark_groups(tmp_path, group_by=3) + gid = np.load(set_dir / "group_id.npy") + assert gid.tolist() == [0, 0, 0, 1, 1, 1, 2] # trailing remainder is its own group + + +def test_no_masked_atoms_skips_pool_mask(tmp_path): + set_dir = tmp_path / "homog" / "set.000" + set_dir.mkdir(parents=True) + np.save(set_dir / "coord.npy", np.zeros((3, 9))) + np.save(set_dir / "real_atom_types.npy", np.zeros((3, 3), dtype=np.int32)) + r = mark_groups(tmp_path)[0] + assert r.wrote_group_id + assert not r.wrote_pool_mask # pool_mask defaults to 1.0 at train time + assert not (set_dir / "pool_mask.npy").exists() + + +def test_overwrite_false_preserves_existing(tmp_path): + set_dir = tmp_path / "sys" / "set.000" + _write_mixed_system(set_dir, masked_per_frame=(1, 0, 0)) + np.save(set_dir / "group_id.npy", np.array([7, 7, 7], dtype=np.int64)) + + mark_groups(tmp_path, target="overpotential", overwrite=False) + assert np.load(set_dir / "group_id.npy").tolist() == [7, 7, 7] # untouched + + mark_groups(tmp_path, target="overpotential", overwrite=True) + assert np.load(set_dir / "group_id.npy").tolist() == [0, 0, 0] # regenerated + + +def test_dry_run_writes_nothing(tmp_path): + set_dir = tmp_path / "sys" / "set.000" + _write_mixed_system(set_dir, masked_per_frame=(2, 1, 0)) + results = mark_groups(tmp_path, target="overpotential", dry_run=True) + assert results[0].wrote_group_id and results[0].wrote_pool_mask + assert not (set_dir / "group_id.npy").exists() + assert not (set_dir / "pool_mask.npy").exists() + + +def test_no_systems_found_raises(tmp_path): + from dpa_adapt.data.errors import ( + DPADataError, + ) + + (tmp_path / "empty").mkdir() + with pytest.raises(DPADataError): + mark_groups(tmp_path / "empty") diff --git a/source/tests/dpa_adapt/test_grouped_dataset.py b/source/tests/dpa_adapt/test_grouped_dataset.py new file mode 100644 index 0000000000..8418dcdf2d --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_dataset.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for grouped descriptor aggregation.""" + +from __future__ import ( + annotations, +) + +import numpy as np + +from dpa_adapt.finetuner import ( + DPAFineTuner, +) +from dpa_adapt.grouped._aggregation import ( + aggregate_weighted_groups, +) +from dpa_adapt.grouped._offline import ( + GroupedDataset, +) + + +def _write_system( + root, + name, + label=2.5, + group_id=None, + weight=None, + n_atoms=2, + n_frames=1, +): + sys_dir = root / name + sys_dir.mkdir(parents=True) + type_rows = "\n".join(str(i % 2) for i in range(n_atoms)) + "\n" + (sys_dir / "type.raw").write_text(type_rows) + (sys_dir / "type_map.raw").write_text("H\nO\n") + set_dir = sys_dir / "set.000" + set_dir.mkdir() + np.save(set_dir / "coord.npy", np.zeros((n_frames, n_atoms * 3))) + np.save(set_dir / "box.npy", np.tile(np.eye(3).ravel(), (n_frames, 1))) + np.save(set_dir / "energy.npy", np.full((n_frames,), label, dtype=float)) + np.save(set_dir / "property.npy", np.full((n_frames,), label, dtype=float)) + if group_id is not None: + np.save( + set_dir / "group_id.npy", np.full((n_frames,), group_id, dtype=np.int64) + ) + if weight is not None: + np.save(set_dir / "weight.npy", np.asarray(weight, dtype=float)) + return sys_dir + + +def test_aggregate_weighted_groups_core(): + features = np.array([[1.0, 0.0], [3.0, 2.0], [10.0, 5.0]]) + group_ids = np.array([2, 1, 2]) + weights = np.array([0.25, 1.0, 0.75]) + labels = np.array([8.0, 4.0, 8.0]) + + embeddings, group_labels, ordered_ids = aggregate_weighted_groups( + features, group_ids, weights, labels + ) + + np.testing.assert_array_equal(ordered_ids, np.array([1, 2])) + np.testing.assert_allclose(embeddings, np.array([[3.0, 2.0], [7.75, 3.75]])) + np.testing.assert_allclose(group_labels, np.array([4.0, 8.0])) + + +def test_grouped_dataset_weighted_embedding(monkeypatch, tmp_path): + parent = tmp_path / "data" + sys_a = _write_system(parent, "a", group_id=0, weight=[0.7, 0.3], n_frames=2) + expected_rows = { + str(sys_a.resolve()): np.array([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]), + } + + def fake_load_or_extract(systems, **kwargs): + return np.vstack([expected_rows[system._dpa_source] for system in systems]) + + monkeypatch.setattr( + "dpa_adapt.grouped._offline.load_or_extract", + fake_load_or_extract, + ) + + dataset = GroupedDataset(str(parent), pretrained="fake.pt") + + np.testing.assert_allclose( + dataset.get_embeddings(), + np.array([[0.7, 1.4, 2.1]]) + np.array([[3.0, 6.0, 9.0]]), + ) + np.testing.assert_allclose(dataset.get_labels(), np.array([2.5])) + + +def test_grouped_dataset_target_key_list_stacks_multi_property_labels( + monkeypatch, tmp_path +): + """The CLI's --target-key a,b (comma-separated) reaches GroupedDataset as + target_key=["a", "b"]; each key must be read from its own set.*/{key}.npy + and stacked into one (n_groups, 2) label column per group, the same + convention the non-grouped _load_labels() multi-property path uses. + """ + parent = tmp_path / "data" + sys_a = _write_system(parent, "a", group_id=0, weight=[1.0], n_frames=1) + # two extra label columns beyond the default energy/property ones + np.save(sys_a / "set.000" / "gap.npy", np.array([1.5])) + np.save(sys_a / "set.000" / "homo.npy", np.array([-4.0])) + + def fake_load_or_extract(systems, **kwargs): + return np.ones((1, 2)) + + monkeypatch.setattr( + "dpa_adapt.grouped._offline.load_or_extract", + fake_load_or_extract, + ) + + dataset = GroupedDataset( + str(parent), pretrained="fake.pt", target_key=["gap", "homo"] + ) + + assert dataset.get_labels().shape == (1, 2) + np.testing.assert_allclose(dataset.get_labels(), np.array([[1.5, -4.0]])) + + +def test_grouped_dataset_group_ids_are_scoped_per_system(monkeypatch, tmp_path): + parent = tmp_path / "data" + sys_a = _write_system(parent, "a", label=1.0, group_id=0, weight=[0.7]) + sys_b = _write_system(parent, "b", label=2.0, group_id=0, weight=[0.3]) + rows = { + str(sys_a.resolve()): np.array([[1.0, 2.0]]), + str(sys_b.resolve()): np.array([[10.0, 20.0]]), + } + + def fake_load_or_extract(systems, **kwargs): + return np.vstack([rows[system._dpa_source] for system in systems]) + + monkeypatch.setattr( + "dpa_adapt.grouped._offline.load_or_extract", + fake_load_or_extract, + ) + + dataset = GroupedDataset(str(parent), pretrained="fake.pt") + + np.testing.assert_allclose(dataset.get_embeddings(), [[0.7, 1.4], [3.0, 6.0]]) + np.testing.assert_allclose(dataset.get_labels(), [1.0, 2.0]) + + +def test_grouped_dataset_missing_weight_defaults_to_one(monkeypatch, tmp_path): + parent = tmp_path / "data" + sys_a = _write_system(parent, "a", group_id=0, weight=None, n_frames=2) + rows = { + str(sys_a.resolve()): np.array([[1.0, 2.0], [10.0, 20.0]]), + } + + def fake_load_or_extract(systems, **kwargs): + return np.vstack([rows[system._dpa_source] for system in systems]) + + monkeypatch.setattr( + "dpa_adapt.grouped._offline.load_or_extract", + fake_load_or_extract, + ) + + dataset = GroupedDataset(str(parent), pretrained="fake.pt") + + np.testing.assert_allclose(dataset.get_embeddings(), [[11.0, 22.0]]) + np.testing.assert_allclose(dataset.get_labels(), [2.5]) + + +def test_grouped_fit_and_predict(monkeypatch, tmp_path): + parent = tmp_path / "data" + sys_a = _write_system( + parent, "a", label=4.0, group_id=0, weight=[0.7, 0.3], n_frames=2 + ) + rows = { + str(sys_a.resolve()): np.array([[2.0, 0.0], [0.0, 8.0]]), + } + calls = [] + + def fake_load_or_extract(systems, **kwargs): + calls.append(kwargs) + return np.vstack([rows[system._dpa_source] for system in systems]) + + monkeypatch.setattr( + "dpa_adapt.grouped._offline.load_or_extract", + fake_load_or_extract, + ) + + model = DPAFineTuner( + pretrained="fake.pt", strategy="frozen_sklearn", predictor="linear" + ) + model.fit(str(parent)) + result = model.predict(str(parent)) + + assert result.predictions.shape == (1, 1) + np.testing.assert_allclose( + model.predictor[:-1].transform([[1.4, 2.4]]), [[0.0, 0.0]] + ) + assert calls[0]["pooling"] == "mean" + + +def test_plain_input_keeps_existing_fit_path(monkeypatch, tmp_path): + sys_dir = _write_system(tmp_path, "plain", n_frames=2) + called = [] + + def fake_fit(self, data, type_map=None, target_key=None, labels=None, fmt=None): + called.append((data, target_key)) + self._fitted = True + + monkeypatch.setattr(DPAFineTuner, "_fit_sklearn", fake_fit) + + model = DPAFineTuner(pretrained="fake.pt", strategy="frozen_sklearn") + model.fit(str(sys_dir), target_key="energy") + + assert called == [(str(sys_dir), "energy")] diff --git a/source/tests/dpa_adapt/test_grouped_detection.py b/source/tests/dpa_adapt/test_grouped_detection.py new file mode 100644 index 0000000000..443a22223d --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_detection.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for grouped-mode auto-detection (dpa_adapt.trainer._systems_are_grouped). + +Detection must scan every ``set.*`` directory of every given system list, not +just the first set of the first system, and must reject any mix of grouped +and ungrouped sets (including a partially-marked set, or grouped +train_systems paired with ungrouped valid_systems) with a clear error instead +of silently under- or over-detecting grouped mode. +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + +from dpa_adapt.data.errors import ( + DPADataError, +) +from dpa_adapt.trainer import ( + _systems_are_grouped, +) + +_MARKERS = ("group_id", "weight", "pool_mask") + + +def _make_system(tmp_path, name: str, *, markers: tuple[str, ...] | None) -> str: + """Create a system dir with one set.000, optionally carrying markers. + + ``markers=None`` writes no marker files; an explicit tuple writes only + those (so a strict subset of _MARKERS produces a "partial" set). + """ + sysdir = tmp_path / name + set_dir = sysdir / "set.000" + set_dir.mkdir(parents=True) + for marker in markers or (): + np.save(set_dir / f"{marker}.npy", np.zeros(1)) + return str(sysdir) + + +def test_all_ungrouped_returns_false(tmp_path): + systems = [ + _make_system(tmp_path, "a", markers=None), + _make_system(tmp_path, "b", markers=None), + ] + assert _systems_are_grouped(("train_systems", systems)) is False + + +def test_all_grouped_returns_true(tmp_path): + systems = [ + _make_system(tmp_path, "a", markers=_MARKERS), + _make_system(tmp_path, "b", markers=_MARKERS), + ] + assert _systems_are_grouped(("train_systems", systems)) is True + + +def test_first_system_ungrouped_later_grouped_raises(tmp_path): + """The bug this fix targets: checking only the first system's first set + would leave grouped mode disabled here, silently dropping the second + system's group labels instead of erroring. + """ + systems = [ + _make_system(tmp_path, "a", markers=None), + _make_system(tmp_path, "b", markers=_MARKERS), + ] + with pytest.raises(DPADataError, match="Inconsistent grouped markers"): + _systems_are_grouped(("train_systems", systems)) + + +def test_first_system_grouped_later_ungrouped_raises(tmp_path): + """The mirror bug: checking only the first system's first set would + enable grouped mode here and fail (or train inconsistently) once the + second, marker-less system was reached. + """ + systems = [ + _make_system(tmp_path, "a", markers=_MARKERS), + _make_system(tmp_path, "b", markers=None), + ] + with pytest.raises(DPADataError, match="Inconsistent grouped markers"): + _systems_are_grouped(("train_systems", systems)) + + +def test_partial_markers_raises(tmp_path): + systems = [_make_system(tmp_path, "a", markers=("group_id",))] + with pytest.raises(DPADataError, match="only some of"): + _systems_are_grouped(("train_systems", systems)) + + +def test_grouped_train_ungrouped_valid_raises(tmp_path): + train = [_make_system(tmp_path, "train", markers=_MARKERS)] + valid = [_make_system(tmp_path, "valid", markers=None)] + with pytest.raises(DPADataError, match="Inconsistent grouped markers"): + _systems_are_grouped(("train_systems", train), ("valid_systems", valid)) + + +def test_grouped_train_and_valid_returns_true(tmp_path): + train = [_make_system(tmp_path, "train", markers=_MARKERS)] + valid = [_make_system(tmp_path, "valid", markers=_MARKERS)] + assert ( + _systems_are_grouped(("train_systems", train), ("valid_systems", valid)) is True + ) + + +def test_no_systems_returns_false(): + assert _systems_are_grouped(("train_systems", [])) is False diff --git a/source/tests/dpa_adapt/test_grouped_end_to_end.py b/source/tests/dpa_adapt/test_grouped_end_to_end.py new file mode 100644 index 0000000000..a4b8d3813b --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_end_to_end.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""End-to-end checks for heterogeneous grouped assemblies. + +These cover the full data path for groups whose components differ in size and +composition (e.g. OER O*/OH*/OOH*): the assembly writer pads every frame up to +the group's max atom count and emits the DeePMD ``mixed_type`` layout, and the +grouped model must aggregate the padded frames without letting the virtual +padding atoms corrupt the pooled embedding. + +The ``deepmd.pt`` integration tests need the compiled ``deepmd.lib`` extension +and therefore only run in a full build (they skip cleanly otherwise); the +standalone pooling test runs anywhere ``torch`` is available. +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + +from dpa_adapt import ( + Assembly, + ComponentSpec, + PoolMask, +) + + +def _require_real_torch(): + """Return real torch, skipping if a sibling test leaked a mock stub. + + Several ``dpa_adapt`` tests do ``sys.modules.setdefault("torch", MagicMock())`` + at import time, which shadows real torch for the rest of the session when + they are collected first. These tests need the genuine library. + """ + torch = pytest.importorskip("torch") + if type(torch).__module__.startswith("unittest.mock"): + pytest.skip("real torch shadowed by a leaked mock-torch stub in this session") + torch.set_default_device("cpu") + return torch + + +def _require_dataloader_backend(): + _require_real_torch() + # importorskip binds the names via module attributes (always defined) and + # skips cleanly when the compiled deepmd.lib extension is unavailable. + reason = "needs the compiled deepmd.lib extension" + dataloader = pytest.importorskip("deepmd.pt.utils.dataloader", reason=reason) + grouped = pytest.importorskip("deepmd.pt.utils.grouped", reason=reason) + data = pytest.importorskip("deepmd.utils.data", reason=reason) + return ( + dataloader.DpLoaderSet, + grouped.group_data_requirements, + data.DataRequirementItem, + ) + + +def _require_group_property_backend(): + torch = _require_real_torch() + reason = "needs the compiled deepmd.lib extension" + model = pytest.importorskip( + "deepmd.pt.model.model.group_property_model", reason=reason + ) + task = pytest.importorskip("deepmd.pt.model.task.group_property", reason=reason) + return torch, model.GroupPropertyModel, task.GroupPropertyFittingNet + + +def _heterogeneous_builder() -> Assembly: + builder = Assembly(target="target", type_map=["C", "O", "H"]) + group = builder.group(key="oer", label=1.5) + # component with 2 atoms -> padded to 3 with one virtual atom + group.add_component( + ComponentSpec.from_arrays( + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["C", "O"], box=np.eye(3) * 20 + ) + ) + # component with 3 atoms, capping H excluded from pooling + group.add_component( + ComponentSpec.from_arrays( + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0], [2.0, 0.0, 0.0]], + ["C", "O", "H"], + box=np.eye(3) * 20, + pool_mask=PoolMask.exclude_indices([2]), + ) + ) + return builder + + +def test_writer_to_dataloader_preserves_padding_and_group(tmp_path) -> None: + DpLoaderSet, group_data_requirements, DataRequirementItem = ( + _require_dataloader_backend() + ) + + builder = _heterogeneous_builder() + result = builder.write(tmp_path) + system = str(tmp_path / result["systems"][0]) + + data = DpLoaderSet([system], batch_size=8, type_map=["C", "O", "H"], shuffle=False) + req = [DataRequirementItem("target", ndof=1, atomic=False, must=True)] + req.extend(group_data_requirements()) + data.add_data_requirement(req) + + batch = next(iter(data.dataloaders[0])) + + atype = batch["atype"] + assert atype.shape == (2, 3) + # the smaller component keeps a trailing virtual atom (type -1) + assert atype[0].tolist() == [0, 1, -1] + assert atype[1].tolist() == [0, 1, 2] + # the whole group lands in a single batch (group completeness) + assert batch["group_id"].reshape(-1).tolist() == [0, 0] + # pool_mask stays aligned with coord: padding atom and excluded H are masked + assert batch["pool_mask"].reshape(2, 3).tolist() == [ + [1.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + ] + # padding atom coords are the large non-physical Task-D offset (not 0,0,0) + pad_coord = batch["coord"].reshape(2, 3, 3)[0, 2] + assert (pad_coord.abs() > 20.0).all() + + +def test_group_property_model_pool_is_nan_safe_for_virtual_atoms() -> None: + torch, GroupPropertyModel, GroupPropertyFittingNet = ( + _require_group_property_backend() + ) + + dim = 4 + + class _NaNOnVirtualDescriptor(torch.nn.Module): + """Descriptor stub that returns NaN for virtual (type < 0) atoms. + + This is the worst case for the masked mean pool: if padding rows were + multiplied by a zero ``pool_mask`` (``0 * NaN``) the frame embedding + would be NaN. The model must instead drop those rows entirely. + """ + + def get_rcut(self) -> float: + return 6.0 + + def get_sel(self) -> list[int]: + return [8] + + def mixed_types(self) -> bool: + return True + + def forward( + self, + extended_coord, + extended_atype, + nlist, + mapping=None, + charge_spin=None, + ): + nframes, nall = extended_atype.shape + desc = torch.ones(nframes, nall, dim, dtype=extended_coord.dtype) + desc = torch.where( + (extended_atype < 0)[:, :, None], + torch.full_like(desc, float("nan")), + desc, + ) + return desc, None, None, None, None + + fitting = GroupPropertyFittingNet( + ntypes=3, + dim_descrpt=dim, + property_name="target", + task_dim=1, + neuron=[4], + ) + model = GroupPropertyModel( + descriptor=_NaNOnVirtualDescriptor(), + fitting=fitting, + type_map=["C", "O", "H"], + ) + + coord = torch.tensor( + [ + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0], [0.0, 0.0, 0.0]], # 2 real + 1 virtual + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0], [2.0, 0.0, 0.0]], # 3 real + ] + ) + atype = torch.tensor([[0, 1, -1], [0, 1, 2]]) + group_id = torch.tensor([[0], [0]]) + weight = torch.tensor([[0.5], [0.5]]) + pool_mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) + + out = model( + coord, + atype, + box=None, + group_id=group_id, + weight=weight, + pool_mask=pool_mask, + ) + # the virtual padding atom must not poison the pooled embedding/prediction + assert torch.isfinite(out["frame_embedding"]).all() + assert torch.isfinite(out["target"]).all() + # the whole group aggregates to a single prediction row + assert out["target"].shape[0] == 1 + + +def test_group_property_model_honors_descriptor_mixed_types_contract( + monkeypatch, +) -> None: + """``forward`` must build the neighbor list with the descriptor's own + ``mixed_types()`` value, not a hard-coded ``True``, so descriptors that + require type-distinguished neighbor lists get one. + """ + torch, GroupPropertyModel, GroupPropertyFittingNet = ( + _require_group_property_backend() + ) + import deepmd.pt.model.model.group_property_model as gpm_module + + class _NonMixedDescriptor(torch.nn.Module): + def get_rcut(self) -> float: + return 6.0 + + def get_sel(self) -> list[int]: + return [8] + + def mixed_types(self) -> bool: + return False + + def forward( + self, + extended_coord, + extended_atype, + nlist, + mapping=None, + charge_spin=None, + ): + nframes, nall = extended_atype.shape + return ( + torch.ones(nframes, nall, 4, dtype=extended_coord.dtype), + None, + None, + None, + None, + ) + + fitting = GroupPropertyFittingNet( + ntypes=3, + dim_descrpt=4, + property_name="target", + task_dim=1, + neuron=[4], + ) + model = GroupPropertyModel( + descriptor=_NonMixedDescriptor(), + fitting=fitting, + type_map=["C", "O", "H"], + ) + assert model.mixed_types() is False + + seen_mixed_types = [] + real_fn = gpm_module.extend_input_and_build_neighbor_list + + def _spy(*args, **kwargs): + seen_mixed_types.append(kwargs.get("mixed_types")) + return real_fn(*args, **kwargs) + + monkeypatch.setattr(gpm_module, "extend_input_and_build_neighbor_list", _spy) + + coord = torch.tensor([[[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]]]) + atype = torch.tensor([[0, 1]]) + group_id = torch.tensor([[0]]) + weight = torch.tensor([[1.0]]) + pool_mask = torch.tensor([[1.0, 1.0]]) + + model( + coord, + atype, + box=None, + group_id=group_id, + weight=weight, + pool_mask=pool_mask, + ) + + assert seen_mixed_types == [False] + + +def test_masked_mean_pool_is_nan_safe() -> None: + """Standalone check of the model's NaN-safe masked mean pool. + + Mirrors the aggregation in ``GroupPropertyModel.forward``: excluded atoms + (``pool_mask`` == 0) -- padding/virtual atoms or capping H -- must be dropped + before the weighted sum, so a non-finite descriptor on those rows cannot + propagate through ``0 * NaN``. A naive ``descriptor * pool_mask`` would. + """ + torch = _require_real_torch() + + # atom 2 in frame 0 is excluded and carries a NaN descriptor + descriptor = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0], [float("nan"), float("nan")]], + [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], + ], + requires_grad=True, + ) + pool_mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) + + # naive form poisons the pooled embedding + naive = (descriptor * pool_mask[:, :, None]).sum(dim=1) + assert torch.isnan(naive[0]).any() + + # hardened form (as used by the model) stays finite and correct + denom = pool_mask.sum(dim=1).clamp_min(1.0) + keep = pool_mask[:, :, None] > 0 + masked = torch.where(keep, descriptor, torch.zeros_like(descriptor)) + frame_embedding = (masked * pool_mask[:, :, None]).sum(dim=1) / denom[:, None] + + assert torch.isfinite(frame_embedding).all() + # frame 0 mean over the two kept atoms: ([1,2] + [3,4]) / 2 = [2, 3] + assert torch.allclose(frame_embedding[0], torch.tensor([2.0, 3.0])) + + # gradient flows to kept atoms and is zero (not NaN) on the excluded row + frame_embedding.sum().backward() + assert torch.isfinite(descriptor.grad).all() + assert torch.equal(descriptor.grad[0, 2], torch.zeros(2)) + + +def test_group_property_model_consumes_per_group_fparam() -> None: + """numb_fparam widens the fitting input; the model concats the per-group + fparam AFTER aggregation, so it bypasses the weighted sum over frames. + """ + torch, GroupPropertyModel, GroupPropertyFittingNet = ( + _require_group_property_backend() + ) + + dim, fdim = 4, 2 + + fitting = GroupPropertyFittingNet( + ntypes=2, + dim_descrpt=dim, + property_name="y", + task_dim=1, + neuron=[4], + numb_fparam=fdim, + ) + assert fitting.get_dim_fparam() == fdim + first_linear = next(m for m in fitting.network if isinstance(m, torch.nn.Linear)) + assert first_linear.in_features == dim + fdim # descriptor dim + fparam dim + + class _ConstDescriptor(torch.nn.Module): + def get_rcut(self) -> float: + return 6.0 + + def get_sel(self) -> list[int]: + return [8] + + def mixed_types(self) -> bool: + return True + + def forward( + self, extended_coord, extended_atype, nlist, mapping=None, charge_spin=None + ): + nframes, nall = extended_atype.shape + desc = torch.ones(nframes, nall, dim, dtype=extended_coord.dtype) + return desc, None, None, None, None + + model = GroupPropertyModel( + descriptor=_ConstDescriptor(), fitting=fitting, type_map=["A", "B"] + ) + coord = torch.rand(2, 3, 3, dtype=torch.float64) + atype = torch.tensor([[0, 1, 0], [0, 1, 1]]) + group_id = torch.tensor([[0], [0]]) # both frames -> one group + weight = torch.tensor([[0.5], [0.5]]) + pool_mask = torch.ones(2, 3) + fparam = torch.tensor([[1.0, 2.0], [1.0, 2.0]], dtype=torch.float64) # per-group + + out = model( + coord, + atype, + box=None, + fparam=fparam, + group_id=group_id, + weight=weight, + pool_mask=pool_mask, + ) + assert out["y"].shape[0] == 1 # aggregates to one group prediction + assert torch.isfinite(out["y"]).all() + + +def test_group_property_model_rejects_inconsistent_group_fparam() -> None: + torch, GroupPropertyModel, GroupPropertyFittingNet = ( + _require_group_property_backend() + ) + + dim, fdim = 4, 2 + fitting = GroupPropertyFittingNet( + ntypes=2, + dim_descrpt=dim, + property_name="y", + task_dim=1, + neuron=[4], + numb_fparam=fdim, + ) + + class _ConstDescriptor(torch.nn.Module): + def get_rcut(self) -> float: + return 6.0 + + def get_sel(self) -> list[int]: + return [8] + + def mixed_types(self) -> bool: + return True + + def forward( + self, extended_coord, extended_atype, nlist, mapping=None, charge_spin=None + ): + nframes, nall = extended_atype.shape + desc = torch.ones(nframes, nall, dim, dtype=extended_coord.dtype) + return desc, None, None, None, None + + model = GroupPropertyModel( + descriptor=_ConstDescriptor(), fitting=fitting, type_map=["A", "B"] + ) + with pytest.raises(ValueError, match="fparam must be constant"): + model( + torch.rand(2, 3, 3, dtype=torch.float64), + torch.tensor([[0, 1, 0], [0, 1, 1]]), + box=None, + fparam=torch.tensor([[1.0, 2.0], [9.0, 2.0]], dtype=torch.float64), + group_id=torch.tensor([[0], [0]]), + weight=torch.tensor([[0.5], [0.5]]), + pool_mask=torch.ones(2, 3), + ) diff --git a/source/tests/dpa_adapt/test_grouped_finetuner.py b/source/tests/dpa_adapt/test_grouped_finetuner.py new file mode 100644 index 0000000000..bc411de3c3 --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_finetuner.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from __future__ import ( + annotations, +) + +import json +from pathlib import ( + Path, +) +from unittest.mock import ( + patch, +) + +import numpy as np + +from dpa_adapt.finetuner import ( + DPAFineTuner, +) +from source.tests.dpa_adapt.test_finetuner_strategies import ( + _FULL_TYPE_MAP, + _fake_ckpt_sd, + _make_system_dirs, + _mock_dp_train, +) + + +def test_grouped_training_strategy_uses_group_property_config(monkeypatch, tmp_path): + import torch + + monkeypatch.setattr(torch, "load", lambda *a, **kw: _fake_ckpt_sd()) + ckpt = tmp_path / "fake.pt" + ckpt.write_bytes(b"") + out_dir = tmp_path / "out" + systems = _make_system_dirs(tmp_path, formulas=("GroupedTrain",), n=2) + valid_systems = _make_system_dirs(tmp_path, formulas=("GroupedValid",), n=1) + for sid, sysdir in enumerate(systems + valid_systems): + set_dir = Path(sysdir) / "set.000" + np.save(set_dir / "group_id.npy", np.array([sid, sid], dtype=np.int64)) + np.save(set_dir / "weight.npy", np.array([0.5, 0.5], dtype=float)) + np.save(set_dir / "pool_mask.npy", np.ones((2, 2), dtype=float)) + + model = DPAFineTuner( + pretrained=str(ckpt), + strategy="finetune", + property_name="overpotential", + max_steps=20, + output_dir=str(out_dir), + ) + with patch("subprocess.run", side_effect=_mock_dp_train(str(out_dir))): + result = model.fit(train_data=systems, valid_data=valid_systems) + + assert result is not None + cfg = json.loads((out_dir / "input.json").read_text()) + assert cfg["model"]["type_map"] == _FULL_TYPE_MAP + assert cfg["model"]["fitting_net"]["type"] == "group_property" + assert cfg["model"]["fitting_net"]["property_name"] == "overpotential" + # Grouped head defaults to GELU (tanh saturates on the un-normalized + # embedding+fparam input and collapses predictions to a constant). + assert cfg["model"]["fitting_net"]["activation_function"] == "gelu" + assert cfg["loss"]["type"] == "group_property" + + +def test_grouped_target_alias_and_auto_fparam_dim(monkeypatch, tmp_path): + """target= aliases property_name; fit(train=/valid=) work; fparam_dim auto.""" + import torch + + monkeypatch.setattr(torch, "load", lambda *a, **kw: _fake_ckpt_sd()) + ckpt = tmp_path / "fake.pt" + ckpt.write_bytes(b"") + out_dir = tmp_path / "out" + systems = _make_system_dirs(tmp_path, formulas=("GroupedTrain",), n=2) + valid_systems = _make_system_dirs(tmp_path, formulas=("GroupedValid",), n=1) + for sid, sysdir in enumerate(systems + valid_systems): + set_dir = Path(sysdir) / "set.000" + np.save(set_dir / "group_id.npy", np.array([sid, sid], dtype=np.int64)) + np.save(set_dir / "weight.npy", np.array([0.5, 0.5], dtype=float)) + np.save(set_dir / "pool_mask.npy", np.ones((2, 2), dtype=float)) + # per-group side features -> fparam_dim should be auto-detected as 3 + np.save(set_dir / "fparam.npy", np.ones((2, 3), dtype=float)) + + model = DPAFineTuner( + pretrained=str(ckpt), + strategy="finetune", + target="overpotential", # alias for property_name + max_steps=20, + output_dir=str(out_dir), + ) + assert model.property_name == "overpotential" + + with patch("subprocess.run", side_effect=_mock_dp_train(str(out_dir))): + model.fit(train=systems, valid=valid_systems) # train=/valid= aliases + + cfg = json.loads((out_dir / "input.json").read_text()) + assert cfg["model"]["fitting_net"]["type"] == "group_property" + assert cfg["model"]["fitting_net"]["property_name"] == "overpotential" + # fparam_dim auto-inferred from set.*/fparam.npy -> numb_fparam wired into head + assert cfg["model"]["fitting_net"]["numb_fparam"] == 3 diff --git a/source/tests/dpa_adapt/test_grouped_hardening.py b/source/tests/dpa_adapt/test_grouped_hardening.py new file mode 100644 index 0000000000..9a705c510f --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_hardening.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Locally-runnable hardening checks (no compiled deepmd.lib). + +- Task C: all-masked frame rejected at Assembly construction time. +- Task C: masked-mean formula (รท mask sum, not natoms) โ€” torch mirror of the model. +- Task B: group_reduce mean-vs-sum โ€” torch mirror of the model reduction. +- Task F: shuffled-batch group-label aggregation invariance โ€” torch mirror of the loss. + +The model/loss/descriptor wiring versions live in source/tests/pt/ and need the +native extension (Bohrium/CI). +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + + +def _require_cpu_torch(): + torch = pytest.importorskip("torch") + torch.set_default_device("cpu") + return torch + + +def test_zero_mask_rejected_at_construction(): + from dpa_adapt.data.errors import ( + DPADataError, + ) + from dpa_adapt.grouped import ( + ComponentSpec, + ) + + comp = ComponentSpec.from_arrays( + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["H", "O"], pool_mask=[0.0, 0.0] + ) + with pytest.raises(DPADataError, match="all-zero"): + comp.validate() + + +def test_zero_mask_rejected_through_writer(tmp_path): + from dpa_adapt.data.errors import ( + DPADataError, + ) + from dpa_adapt.grouped import ( + Assembly, + ) + + a = Assembly(target="y", type_map=["H", "O"]) + g = a.group(label=1.0) + g.add([[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["H", "O"], pool_mask=[0.0, 0.0]) + with pytest.raises(DPADataError, match="all-zero"): + a.write(tmp_path / "out") + + +def test_masked_mean_divides_by_mask_sum(): + torch = _require_cpu_torch() + # frame 0: atoms 0,1 kept (atom 2 masked); frame 1: only atom 0 kept + descriptor = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0], [100.0, 100.0]], + [[10.0, 10.0], [999.0, 999.0], [999.0, 999.0]], + ] + ) + pool_mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 0.0, 0.0]]) + + # mirror GroupPropertyModel.forward masked mean + mask_sum = pool_mask.sum(dim=1) + assert not bool((mask_sum == 0).any()) + denom = mask_sum.clamp_min(1.0) + keep = pool_mask[:, :, None] > 0 + desc = torch.where(keep, descriptor, torch.zeros_like(descriptor)) + frame_emb = (desc * pool_mask[:, :, None]).sum(dim=1) / denom[:, None] + + # frame0 = ([1,2]+[3,4])/2 = [2,3]; frame1 = [10,10]/1 = [10,10] + assert torch.allclose(frame_emb, torch.tensor([[2.0, 3.0], [10.0, 10.0]])) + # NOT divided by natoms (=3): that would give [4/3, 6/3]=[1.333,2] for frame0 + assert not torch.allclose(frame_emb[0], torch.tensor([4.0 / 3, 6.0 / 3])) + + +def _reduce(frame_emb, weight, group_id, mode): + """Mirror of GroupPropertyModel.forward frame->group reduction.""" + torch = _require_cpu_torch() + + order, inverse = torch.unique(group_id, sorted=True, return_inverse=True) + n = order.shape[0] + ge = torch.zeros((n, frame_emb.shape[1])) + ge.index_add_(0, inverse, frame_emb * weight[:, None]) + if mode == "mean": + ws = torch.zeros((n, 1)) + ws.index_add_(0, inverse, weight[:, None]) + ge = ge / ws.clamp_min(1e-12) + return ge + + +def test_group_reduce_mean_vs_sum(): + torch = _require_cpu_torch() + frame_emb = torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]) # 2 identical frames + weight = torch.ones(2) + group_id = torch.tensor([0, 0]) # one group + + mean = _reduce(frame_emb, weight, group_id, "mean") + summ = _reduce(frame_emb, weight, group_id, "sum") + + assert torch.allclose(mean, frame_emb[:1]) # mean-reduce == the frame embedding + assert torch.allclose(summ, 2.0 * frame_emb[:1]) # sum-reduce == 2 x + + +def _group_labels(frame_label, group_inverse, n_groups): + """Mirror of GroupPropertyLoss._group_labels (uses shared inverse).""" + torch = _require_cpu_torch() + + gl = frame_label.new_zeros((n_groups, frame_label.shape[1])) + gl[group_inverse] = frame_label + assert torch.allclose(gl[group_inverse], frame_label) # consistency + return gl + + +def test_shuffled_batch_group_labels_invariance(): + torch = _require_cpu_torch() + # groups: 0->0.0, 1->10.0, 2->20.0; frames arrive interleaved + order_sorted = torch.tensor([0, 1, 2]) + + def group_labels_for(perm): + group_id = torch.tensor([2, 0, 2, 1, 0])[perm] + frame_label = torch.tensor([[20.0], [0.0], [20.0], [10.0], [0.0]])[perm] + gorder, inverse = torch.unique(group_id, sorted=True, return_inverse=True) + assert torch.equal(gorder, order_sorted) + return _group_labels(frame_label, inverse, gorder.shape[0]) + + ref = group_labels_for(torch.arange(5)) + shuffled = group_labels_for(torch.tensor([3, 0, 4, 1, 2])) + assert torch.allclose(ref, shuffled) + assert torch.allclose( + ref, torch.tensor([[0.0], [10.0], [20.0]]) + ) # sorted by group id + + +def test_padding_coords_are_non_physical(tmp_path): + """Task D data-writer: padding atoms get a large, spread-out offset (not 0,0,0).""" + from dpa_adapt.grouped import ( + Assembly, + ) + + a = Assembly(target="y", type_map=["C", "O", "H"]) + g = a.group(label=1.0) + g.add([[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["C", "O"], box=np.eye(3) * 20) # 2 atoms + g.add( + [[0.0, 0.0, 0.0], [1.1, 0.0, 0.0], [2.0, 0.0, 0.0]], + ["C", "O", "H"], + box=np.eye(3) * 20, + ) # 3 atoms -> frame 0 padded to 3 + res = a.write(tmp_path / "out") + set_dir = tmp_path / "out" / res["systems"][0] / "set.000" + coord = np.load(set_dir / "coord.npy").reshape(2, 3, 3) + real = np.load(set_dir / "real_atom_types.npy") + pad = real[0] < 0 + assert pad.sum() == 1 # frame 0 has one padding atom + assert ( + np.linalg.norm(coord[0][pad][0]) > 20.0 + ) # far outside the 20 A box, not origin + + +def test_write_disambiguates_groups_with_colliding_sanitized_names(tmp_path): + """Two group keys that sanitize to the same directory name must not + overwrite one another; each gets its own system dir and both labels + survive in the manifest. + """ + from dpa_adapt.grouped import ( + Assembly, + ) + + a = Assembly(target="y", type_map=["H", "O"]) + g1 = a.group(key="group!!", label=1.0) + g1.add([[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["H", "O"]) + g2 = a.group(key="group??", label=2.0) + g2.add([[0.0, 0.0, 0.0], [1.1, 0.0, 0.0]], ["H", "O"]) + + res = a.write(tmp_path / "out") + + systems = res["systems"] + assert len(systems) == len(set(systems)) == 2 + + labels = sorted( + float(np.load(tmp_path / "out" / s / "set.000" / "y.npy").reshape(-1)[0]) + for s in systems + ) + assert labels == [1.0, 2.0] diff --git a/source/tests/dpa_adapt/test_polymer_builder.py b/source/tests/dpa_adapt/test_polymer_builder.py new file mode 100644 index 0000000000..0d449d23ee --- /dev/null +++ b/source/tests/dpa_adapt/test_polymer_builder.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""End-to-end test for PolymerBuilder (CSV -> grouped DeePMD npy). + +Needs RDKit for SMILES->3D; skips cleanly otherwise. No compiled deepmd +backend required. +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + +pytest.importorskip("rdkit") + +from dpa_adapt.grouped._polymer import ( + PolymerBuilder, +) + +_CSV = """\ +reference;SMILES_start_group;SMILES_end_group;SMILES_repeating_unitA;molpercent_repeating_unitA;SMILES_repeating_unitB;molpercent_repeating_unitB;Mn;polymer_concentration_wpercent;additive1;additive1_concentration_molar;pH;cloud_point +r1;[C](C)(C)C#N;[C](C)(C)C#N;[CH2][CH](C(=O)NC(C)C);0.8;[CH2][CH](C(=O)O);0.2;11500;0.01;NaCl;0.01;7;32.1 +r2;;;[CH2][CH](C(=O)NC(C)C);1.0;;;20000;0.02;;;7;45.0 +""" + + +def _write_csv(tmp_path): + path = tmp_path / "cp.csv" + path.write_text(_CSV, encoding="utf-8") + return path + + +def test_from_csv_writes_grouped_polymer_systems(tmp_path): + csv = _write_csv(tmp_path) + out = tmp_path / "data" + result = PolymerBuilder.from_csv(csv, target="cloud_point", decimal=".").write(out) + + assert result["n_groups"] == 2 + F = result["fparam_dim"] + # schema = [mn_log, conc, pH, salt:NaCl] + assert F == 4 + + # ---- first polymer: 2 units (0.8/0.2) + 2 ends ---- + sys0 = out / result["systems"][0] + set0 = sys0 / "set.000" + + coord = np.load(set0 / "coord.npy") + real = np.load(set0 / "real_atom_types.npy") + pool = np.load(set0 / "pool_mask.npy") + weight = np.load(set0 / "weight.npy") + gid = np.load(set0 / "group_id.npy") + fparam = np.load(set0 / "fparam.npy") + label = np.load(set0 / "cloud_point.npy") + + nframes = 4 # unitA, unitB, start, end + natoms = real.shape[1] + assert coord.shape == (nframes, natoms * 3) + assert real.shape == (nframes, natoms) + assert pool.shape == (nframes, natoms) + + # one group per polymer -> constant group_id within the system + assert gid.shape == (nframes,) + assert len(set(gid.tolist())) == 1 + + # weights: repeating units keep their mole fraction; ends share the rest + assert weight[0] == pytest.approx(0.8) + assert weight[1] == pytest.approx(0.2) + assert weight[2] == weight[3] # both ends share the same weight + + # mixed_type padding: shorter frames carry -1 virtual atoms, pool_mask == (real>=0) + assert (real < 0).any() + np.testing.assert_array_equal(pool, (real >= 0).astype(pool.dtype)) + + # per-group fparam: constant across the group's frames, width F + assert fparam.shape == (nframes, F) + assert np.allclose(fparam, fparam[0]) + + # label repeated across frames + assert label.shape == (nframes, 1) + assert np.allclose(label, 32.1) + + # type_map is the element union across all monomers + tmap = (sys0 / "type_map.raw").read_text().split() + assert set(tmap) >= {"C", "H", "O", "N"} + + +def test_valid_split_reuses_training_scaler(tmp_path): + csv = _write_csv(tmp_path) + train = PolymerBuilder.from_csv(csv, target="cloud_point", decimal=".") + res_train = train.write(tmp_path / "train") + + # a second builder standardized with the training scaler + valid = PolymerBuilder.from_csv(csv, target="cloud_point", decimal=".") + res_valid = valid.write(tmp_path / "valid", scaler=res_train["scaler"]) + + assert res_valid["fparam_dim"] == res_train["fparam_dim"] + # same scaler -> identical standardized fparam for the same row + f_train = np.load( + tmp_path / "train" / res_train["systems"][0] / "set.000" / "fparam.npy" + ) + f_valid = np.load( + tmp_path / "valid" / res_valid["systems"][0] / "set.000" / "fparam.npy" + ) + np.testing.assert_allclose(f_train, f_valid) + + +def test_add_api_direct(tmp_path): + builder = PolymerBuilder(target="cloud_point") + builder.add( + units={"[CH2][CH](C(=O)NC(C)C)": 0.7, "[CH2][CH](C(=O)O)": 0.3}, + ends=["[C](C)(C)C#N"], + mol_weight=12000, + fparam={"pH": 7, "salts": {"KCl": 0.05}}, + target=28.0, + ) + result = builder.write(tmp_path / "d") + assert result["n_groups"] == 1 + assert result["fparam_dim"] == 3 # mn_log + pH + salt:KCl diff --git a/source/tests/dpa_adapt/test_pooling.py b/source/tests/dpa_adapt/test_pooling.py new file mode 100644 index 0000000000..dfc0c6f169 --- /dev/null +++ b/source/tests/dpa_adapt/test_pooling.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Task A: composable pooling primitives (dpa_adapt). + +parse_pooling + the strategy guard are pure-Python and run anywhere; the +byte-identical layout check needs torch (available without deepmd.lib). +""" + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest + +from dpa_adapt.finetuner import ( + POOLING_PRIMITIVES, + DPAFineTuner, + parse_pooling, +) + + +def test_parse_pooling_legacy_strings(): + # the 4 historical strings map to their exact primitive tuples + assert parse_pooling("mean") == ("mean",) + assert parse_pooling("sum") == ("sum",) + assert parse_pooling("mean+std") == ("mean", "std") + assert parse_pooling("mean+std+max+min") == ("mean", "std", "max", "min") + + +def test_parse_pooling_canonical_order_and_dedup(): + # output order follows POOLING_PRIMITIVES regardless of input order + assert parse_pooling("std+mean") == ("mean", "std") + assert parse_pooling("min+max+mean") == ("mean", "max", "min") + assert parse_pooling(["std", "mean", "std"]) == ("mean", "std") # dedup + assert parse_pooling("mean+sum") == ("mean", "sum") # legal, not redundant + + +def test_parse_pooling_list_input(): + assert parse_pooling(["mean", "std"]) == ("mean", "std") + assert parse_pooling(("max", "mean")) == ("mean", "max") + + +def test_parse_pooling_unknown_token_raises(): + with pytest.raises(ValueError, match="unknown pooling primitive"): + parse_pooling("mean+median") + with pytest.raises(ValueError, match="unknown pooling primitive"): + parse_pooling(["mean", "p95"]) + + +def test_parse_pooling_empty_raises(): + with pytest.raises(ValueError, match="empty pooling spec"): + parse_pooling("") + with pytest.raises(ValueError, match="empty pooling spec"): + parse_pooling([]) + + +def test_pooling_strategy_guard(): + # composite/non-mean pooling with a training strategy raises, pointing to `intensive` + for strat in ("frozen_head", "finetune", "mft"): + with pytest.raises(ValueError, match="intensive"): + DPAFineTuner(pretrained="x", strategy=strat, pooling="mean+std") + # plain mean is a harmless default for training strategies (no raise) + DPAFineTuner(pretrained="x", strategy="finetune", pooling="mean") + # frozen_sklearn accepts any composite + DPAFineTuner(pretrained="x", strategy="frozen_sklearn", pooling="mean+std+max+min") + + +def test_pool_descriptor_byte_identical_to_legacy(): + torch = pytest.importorskip("torch") + torch.set_default_device("cpu") + from dpa_adapt.finetuner import ( + _pool_descriptor, + ) + + rng = np.random.default_rng(0) + descrpt = torch.tensor(rng.standard_normal((3, 5, 4))) + + def legacy(spec): + if spec == "mean": + return descrpt.mean(dim=1) + if spec == "sum": + return descrpt.sum(dim=1) + if spec == "mean+std": + m = descrpt.mean(dim=1) + s = torch.nan_to_num(descrpt.std(dim=1), nan=0.0) + return torch.cat([m, s], dim=-1) + if spec == "mean+std+max+min": + m = descrpt.mean(dim=1) + s = torch.nan_to_num(descrpt.std(dim=1), nan=0.0) + return torch.cat( + [m, s, descrpt.max(dim=1).values, descrpt.min(dim=1).values], dim=-1 + ) + raise AssertionError(spec) + + for spec in ("mean", "sum", "mean+std", "mean+std+max+min"): + new = _pool_descriptor(descrpt, parse_pooling(spec)) + assert torch.equal(new, legacy(spec)), spec + + +def test_pool_descriptor_is_nan_safe_for_masked_atoms(): + """A non-finite descriptor on a masked (virtual/padding) atom must not + poison the pooled feature: ``descrpt * mask`` keeps ``0 * NaN == NaN``, + which would corrupt the whole frame instead of just dropping that atom. + """ + torch = pytest.importorskip("torch") + torch.set_default_device("cpu") + from dpa_adapt.finetuner import ( + _pool_descriptor, + ) + + # frame 0: atoms 0,1 kept ([1,2],[3,4]); atom 2 excluded and carries NaN. + # frame 1: all 3 atoms kept, no NaN. + descrpt = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0], [float("nan"), float("nan")]], + [[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]], + ] + ) + mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]) + + for spec in ("mean", "sum", "mean+std", "mean+std+max+min"): + out = _pool_descriptor(descrpt, parse_pooling(spec), mask=mask) + assert torch.isfinite(out).all(), spec + + mean = _pool_descriptor(descrpt, parse_pooling("mean"), mask=mask) + # frame 0 mean over the two kept atoms: ([1,2]+[3,4])/2 = [2,3] + assert torch.allclose(mean[0], torch.tensor([2.0, 3.0])) + # frame 1 unaffected: mean over all 3 kept atoms + assert torch.allclose(mean[1], torch.tensor([2.0, 2.0])) + + summ = _pool_descriptor(descrpt, parse_pooling("sum"), mask=mask) + assert torch.allclose(summ[0], torch.tensor([4.0, 6.0])) + + +def test_pooling_primitives_canonical_constant(): + # guards against accidental reordering that would shift feature columns + assert POOLING_PRIMITIVES == ("mean", "sum", "std", "max", "min") diff --git a/source/tests/pt/test_group_property.py b/source/tests/pt/test_group_property.py new file mode 100644 index 0000000000..191d99a394 --- /dev/null +++ b/source/tests/pt/test_group_property.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from __future__ import ( + annotations, +) + +import numpy as np +import pytest +import torch + +pytest.importorskip("deepmd.lib") + +from deepmd.pt.loss.group_property import ( + GroupPropertyLoss, +) +from deepmd.pt.utils.dataloader import ( + DpLoaderSet, + GroupCompleteBatchSampler, + GroupDistributedBatchSampler, +) +from deepmd.pt.utils.grouped import ( + GROUP_ID_KEY, + GROUP_WEIGHT_KEY, + POOL_MASK_KEY, + distributed_grouped_frame_batches, + group_data_requirements, + load_group_ids_for_system, + normalize_group_id_tensor, + normalize_pool_mask_tensor, + normalize_weight_tensor, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.path import ( + DPPath, +) + + +def _flatten_batches(batches: list[list[int]]) -> list[int]: + return [frame for batch in batches for frame in batch] + + +class ToyGroupedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.backbone = torch.nn.Linear(3, 4, bias=False) + self.head = torch.nn.Linear(4, 1) + self.var_name = "target" + + def forward( + self, + coord, + atype, + box=None, + do_atomic_virial=False, + fparam=None, + aparam=None, + charge_spin=None, + group_id=None, + weight=None, + pool_mask=None, + ): + del box, do_atomic_virial, fparam, aparam, charge_spin + nframes, natoms = atype.shape + coord = coord.reshape(nframes, natoms, 3) + descriptor = self.backbone(coord) + pool_mask = normalize_pool_mask_tensor(pool_mask, nframes, natoms).to( + descriptor.dtype + ) + frame_embedding = (descriptor * pool_mask[:, :, None]).sum( + dim=1 + ) / pool_mask.sum(dim=1).clamp_min(1.0)[:, None] + group_id = normalize_group_id_tensor(group_id, nframes) + weight = normalize_weight_tensor(weight, nframes).to(descriptor.dtype).detach() + group_order, inverse = torch.unique(group_id, sorted=True, return_inverse=True) + group_embedding = torch.zeros( + group_order.shape[0], frame_embedding.shape[1], dtype=frame_embedding.dtype + ) + group_embedding.index_add_(0, inverse, frame_embedding * weight[:, None]) + return { + "target": self.head(group_embedding), + "group_id": group_order, + "unique_group_ids": group_order, + "group_inverse": inverse, + } + + +def test_group_property_loss_trains_backbone_and_detaches_weight() -> None: + torch.manual_seed(7) + coord = torch.tensor( + [ + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + [[1.2, 0.1, 0.0], [0.0, 1.1, 0.0]], + [[0.0, 0.0, 1.0], [1.0, 1.0, 0.0]], + [[0.1, 0.0, 1.2], [1.1, 0.9, 0.0]], + ] + ) + atype = torch.zeros((4, 2), dtype=torch.long) + weight = torch.tensor([[0.5], [0.5], [0.5], [0.5]], requires_grad=True) + label = { + "target": torch.tensor([[1.0], [1.0], [2.0], [2.0]]), + GROUP_ID_KEY: torch.tensor([[0], [0], [1], [1]], dtype=torch.long), + GROUP_WEIGHT_KEY: weight, + POOL_MASK_KEY: torch.ones((4, 2, 1)), + f"find_{GROUP_ID_KEY}": torch.ones(1), + } + input_dict = { + "coord": coord, + "atype": atype, + "box": None, + "do_atomic_virial": False, + "fparam": None, + "aparam": None, + "charge_spin": None, + } + model = ToyGroupedModel() + loss_fn = GroupPropertyLoss(task_dim=1, var_name="target", loss_func="mse") + opt = torch.optim.SGD(model.parameters(), lr=0.05) + before = model.backbone.weight.detach().clone() + losses = [] + for _ in range(80): + opt.zero_grad() + _, loss, _ = loss_fn(input_dict, model, label, natoms=2) + losses.append(float(loss.detach())) + loss.backward() + opt.step() + assert losses[-1] < losses[0] + assert not torch.allclose(before, model.backbone.weight.detach()) + assert weight.grad is None + + +def test_group_property_loss_rejects_inconsistent_group_labels() -> None: + model = ToyGroupedModel() + loss_fn = GroupPropertyLoss(task_dim=1, var_name="target", loss_func="mse") + label = { + "target": torch.tensor([[1.0], [1.5]]), + GROUP_ID_KEY: torch.tensor([[0], [0]], dtype=torch.long), + GROUP_WEIGHT_KEY: torch.ones((2, 1)), + POOL_MASK_KEY: torch.ones((2, 1, 1)), + f"find_{GROUP_ID_KEY}": torch.ones(1), + } + input_dict = { + "coord": torch.ones((2, 1, 3)), + "atype": torch.zeros((2, 1), dtype=torch.long), + "box": None, + "do_atomic_virial": False, + "fparam": None, + "aparam": None, + "charge_spin": None, + } + try: + loss_fn(input_dict, model, label, natoms=1) + except ValueError as exc: + assert "Inconsistent target labels" in str(exc) + else: + raise AssertionError("expected inconsistent group labels to raise") + + +def test_group_complete_batch_sampler_keeps_groups_intact() -> None: + sampler = GroupCompleteBatchSampler( + np.array([0, 0, 1, 1, 2]), max_frames=2, shuffle=False + ) + assert list(sampler) == [[0, 1], [2, 3], [4]] + + +def test_distributed_group_batches_assign_whole_groups_to_one_rank() -> None: + group_ids = np.array([0, 0, 1, 1, 2, 3, 3, 4]) + rank0 = distributed_grouped_frame_batches( + group_ids, max_frames=3, num_replicas=2, rank=0, shuffle=False + ) + rank1 = distributed_grouped_frame_batches( + group_ids, max_frames=3, num_replicas=2, rank=1, shuffle=False + ) + + rank_frames = [set(_flatten_batches(rank0)), set(_flatten_batches(rank1))] + assert rank_frames[0].isdisjoint(rank_frames[1]) + assert rank_frames[0] | rank_frames[1] == set(range(len(group_ids))) + + owner_by_group = {} + for rank, frames in enumerate(rank_frames): + for frame in frames: + group = int(group_ids[frame]) + owner_by_group.setdefault(group, rank) + assert owner_by_group[group] == rank + + +def test_group_distributed_batch_sampler_keeps_groups_intact_per_rank() -> None: + group_ids = np.array([0, 0, 1, 1, 2, 3, 3, 4]) + samplers = [ + GroupDistributedBatchSampler( + group_ids, + max_frames=3, + num_replicas=2, + rank=rank, + shuffle=False, + ) + for rank in range(2) + ] + batches_by_rank = [list(sampler) for sampler in samplers] + frames_by_rank = [set(_flatten_batches(batches)) for batches in batches_by_rank] + + assert frames_by_rank[0].isdisjoint(frames_by_rank[1]) + assert frames_by_rank[0] | frames_by_rank[1] == set(range(len(group_ids))) + for rank, frames in enumerate(frames_by_rank): + for group in set(group_ids): + group_frames = {ii for ii, gid in enumerate(group_ids) if gid == group} + if frames & group_frames: + assert group_frames <= frames + + +def test_group_distributed_batch_sampler_uses_shared_seed_across_ranks() -> None: + group_ids = np.array([0, 0, 0, 1, 2, 2, 3, 4, 4, 4]) + samplers = [ + GroupDistributedBatchSampler( + group_ids, + max_frames=8, + num_replicas=2, + rank=rank, + shuffle=True, + seed=[11, 23, 37], + ) + for rank in range(2) + ] + frames_by_rank = [set(_flatten_batches(list(sampler))) for sampler in samplers] + + assert frames_by_rank[0].isdisjoint(frames_by_rank[1]) + assert frames_by_rank[0] | frames_by_rank[1] == set(range(len(group_ids))) + for frames in frames_by_rank: + for group in set(group_ids): + group_frames = {ii for ii, gid in enumerate(group_ids) if gid == group} + if frames & group_frames: + assert group_frames <= frames + + +def test_dploaderset_enables_group_batches_for_group_requirements(tmp_path) -> None: + system = tmp_path / "sys" + set_dir = system / "set.000" + set_dir.mkdir(parents=True) + (system / "type.raw").write_text("0\n0\n") + (system / "type_map.raw").write_text("H\n") + np.save(set_dir / "coord.npy", np.zeros((5, 6), dtype=np.float64)) + np.save(set_dir / "box.npy", np.tile(np.eye(3).reshape(1, 9), (5, 1))) + np.save(set_dir / "target.npy", np.array([[1.0], [1.0], [2.0], [2.0], [3.0]])) + np.save(set_dir / f"{GROUP_ID_KEY}.npy", np.array([0, 0, 1, 1, 2])) + np.save(set_dir / f"{GROUP_WEIGHT_KEY}.npy", np.ones(5)) + np.save(set_dir / f"{POOL_MASK_KEY}.npy", np.ones((5, 2))) + + data = DpLoaderSet([str(system)], batch_size=2, type_map=["H"], shuffle=False) + req = [DataRequirementItem("target", ndof=1, atomic=False, must=True)] + req.extend(group_data_requirements()) + data.add_data_requirement(req) + batches = [ + batch[GROUP_ID_KEY].reshape(-1).tolist() for batch in data.dataloaders[0] + ] + assert batches == [[0, 0], [1, 1], [2]] + + +def test_dploaderset_missing_group_id_falls_back_to_one_group_per_frame( + tmp_path, +) -> None: + """When group_id.npy is absent, the sampler must batch by ordinary + batch_size (one group per frame) -- the same "no explicit grouping" + semantics GroupPropertyLoss.forward falls back to (torch.arange(nframes)) + -- instead of treating the whole system as a single group, which would + ignore batch_size and (with fewer frames than ranks) break DDP. + """ + system = tmp_path / "sys" + set_dir = system / "set.000" + set_dir.mkdir(parents=True) + (system / "type.raw").write_text("0\n0\n") + (system / "type_map.raw").write_text("H\n") + np.save(set_dir / "coord.npy", np.zeros((5, 6), dtype=np.float64)) + np.save(set_dir / "box.npy", np.tile(np.eye(3).reshape(1, 9), (5, 1))) + np.save(set_dir / "target.npy", np.array([[1.0], [2.0], [3.0], [4.0], [5.0]])) + np.save(set_dir / f"{GROUP_WEIGHT_KEY}.npy", np.ones(5)) + np.save(set_dir / f"{POOL_MASK_KEY}.npy", np.ones((5, 2))) + # deliberately no group_id.npy + + data = DpLoaderSet([str(system)], batch_size=2, type_map=["H"], shuffle=False) + req = [DataRequirementItem("target", ndof=1, atomic=False, must=True)] + req.extend(group_data_requirements()) + data.add_data_requirement(req) + # Batch *composition* (frame count per batch) is what the sampler fallback + # controls; the raw GROUP_ID_KEY tensor content is separately defaulted to + # 0 by the data-requirement default and is not what this fallback governs + # (GroupPropertyLoss ignores that default via its own find_group_id check + # and recomputes one-group-per-frame itself). + batch_sizes = [batch["coord"].shape[0] for batch in data.dataloaders[0]] + # 5 frames, each its own group, batch_size=2 -> 2,2,1 (not one batch of 5) + assert batch_sizes == [2, 2, 1] + + +def test_load_group_ids_for_system_returns_none_when_missing(tmp_path) -> None: + set_dir = tmp_path / "set.000" + set_dir.mkdir(parents=True) + + class _FakeDataSystem: + dirs = (DPPath(str(set_dir)),) + + assert load_group_ids_for_system(_FakeDataSystem()) is None + + +def test_load_group_ids_for_system_reads_filesystem_and_hdf5(tmp_path) -> None: + """``load_group_ids_for_system`` must read group_id.npy through the + system's own DPPath ``dirs`` (backend-aware) so HDF5-backed systems are + supported the same way as on-disk ones, not silently reported as missing + because a plain filesystem glob can't see inside an HDF5 archive. + """ + # on-disk + set_dir = tmp_path / "os_sys" / "set.000" + set_dir.mkdir(parents=True) + np.save(set_dir / f"{GROUP_ID_KEY}.npy", np.array([0, 0, 1])) + + class _FakeOSDataSystem: + dirs = (DPPath(str(set_dir)),) + + ids = load_group_ids_for_system(_FakeOSDataSystem()) + assert ids is not None + assert ids.tolist() == [0, 0, 1] + + # HDF5-backed: writing and reading go through the same cached h5py.File + # handle (DPH5Path caches by (path, mode)), so use "a" consistently. + import h5py + + h5file = str(tmp_path / "sys.h5") + with h5py.File(h5file, "w") as f: + pass + h5_root = DPPath(h5file, "a") + h5_set_dir = h5_root / "set.000" + (h5_set_dir / f"{GROUP_ID_KEY}.npy").save_numpy(np.array([2, 2, 3])) + + class _FakeH5DataSystem: + dirs = (h5_set_dir,) + + ids = load_group_ids_for_system(_FakeH5DataSystem()) + assert ids is not None + assert ids.tolist() == [2, 2, 3] diff --git a/source/tests/pt/test_group_property_fitting_net.py b/source/tests/pt/test_group_property_fitting_net.py new file mode 100644 index 0000000000..2c259bd33c --- /dev/null +++ b/source/tests/pt/test_group_property_fitting_net.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""GroupPropertyFittingNet used to accept and silently ignore several +property-schema options (seed, numb_aparam, resnet_dt, ...) via a bare +``**kwargs``, and collapsed a per-layer ``trainable`` list to +``all(trainable)``. Config authors setting any of those got a model that +looked accepted but behaved differently than requested. Cover: an +unsupported option now raises instead of being swallowed, ``seed`` actually +makes initialization reproducible, and a per-layer ``trainable`` list +freezes exactly those layers. +""" + +from __future__ import ( + annotations, +) + +import pytest +import torch + +pytest.importorskip("deepmd.lib") + +from deepmd.pt.model.task.group_property import ( + GroupPropertyFittingNet, +) + + +def _make(**overrides): + kwargs = {"ntypes": 3, "dim_descrpt": 4, "property_name": "y", "neuron": [8]} + kwargs.update(overrides) + return GroupPropertyFittingNet(**kwargs) + + +def test_injected_type_and_mixed_types_are_accepted(): + # the generic model-building path always passes these two; they must not + # crash construction even though the fitting net itself doesn't use them. + _make(type="group_property", mixed_types=True) + + +def test_unsupported_property_options_are_rejected_not_ignored(): + for bad_kwarg in ( + "numb_aparam", + "default_fparam", + "dim_case_embd", + "resnet_dt", + "intensive", + "distinguish_types", + "some_future_typo", + ): + with pytest.raises(TypeError): + _make(**{bad_kwarg: True}) + + +def test_seed_makes_initialization_reproducible(): + fn1 = _make(seed=42) + fn2 = _make(seed=42) + w1 = fn1.network[0].weight.detach() + w2 = fn2.network[0].weight.detach() + assert torch.equal(w1, w2) + + +def test_different_seeds_give_different_initialization(): + fn1 = _make(seed=42) + fn2 = _make(seed=43) + assert not torch.equal( + fn1.network[0].weight.detach(), fn2.network[0].weight.detach() + ) + + +def test_unseeded_construction_advances_the_global_rng_stream(): + """Without an explicit seed, initialization must still draw from (and + advance) the caller's global RNG stream, same as an unseeded + torch.nn.Linear always has -- two back-to-back unseeded constructions + should not silently produce identical weights. + """ + torch.manual_seed(0) + fn_a = _make() + fn_b = _make() + assert not torch.equal( + fn_a.network[0].weight.detach(), fn_b.network[0].weight.detach() + ) + + +def test_trainable_list_freezes_only_the_named_layers(): + # neuron=[8] -> 2 Linear layers (hidden, output); freeze only the first. + fn = _make(trainable=[False, True]) + linear_layers = [m for m in fn.network if isinstance(m, torch.nn.Linear)] + assert len(linear_layers) == 2 + assert not any(p.requires_grad for p in linear_layers[0].parameters()) + assert all(p.requires_grad for p in linear_layers[1].parameters()) + + +def test_trainable_list_wrong_length_raises(): + with pytest.raises(ValueError, match="trainable"): + _make(trainable=[True]) + + +def test_trainable_bool_still_applies_uniformly(): + fn = _make(trainable=False) + assert not any(p.requires_grad for p in fn.parameters()) diff --git a/source/tests/pt/test_group_property_hardening.py b/source/tests/pt/test_group_property_hardening.py new file mode 100644 index 0000000000..6b5c9e8cc0 --- /dev/null +++ b/source/tests/pt/test_group_property_hardening.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt-backend hardening tests for the grouped property path (Tasks B-F). + +These require the compiled ``deepmd.lib`` extension and skip cleanly otherwise. +Run on Bohrium/CI with the DeePMD 3.1.3 env: + + pytest source/tests/pt/test_group_property_hardening.py -v +""" + +from __future__ import ( + annotations, +) + +import pytest + +pytest.importorskip("torch") +try: + import torch + + from deepmd.pt.loss.group_property import ( + GroupPropertyLoss, + ) + from deepmd.pt.model.model.group_property_model import ( + GroupPropertyModel, + ) + from deepmd.pt.model.task.group_property import ( + GroupPropertyFittingNet, + ) + from deepmd.pt.utils.grouped import ( + distributed_grouped_frame_batches, + ) +except ImportError as exc: # needs the compiled deepmd.lib extension + pytest.skip(f"deepmd backend unavailable: {exc}", allow_module_level=True) + + +class _StubDescriptor(torch.nn.Module): + """Per-atom feature = atom's local index (deterministic), so masked-mean / + reduce results are hand-checkable. Real atoms occupy the first ``natoms`` rows. + """ + + def __init__(self, dim: int = 2) -> None: + super().__init__() + self.dim = dim + + def get_rcut(self) -> float: + return 6.0 + + def get_sel(self) -> list[int]: + return [8] + + def mixed_types(self) -> bool: + return True + + def forward(self, ext_coord, ext_atype, nlist, mapping=None, charge_spin=None): + nframes, nall = ext_atype.shape + idx = torch.arange(nall, dtype=ext_coord.dtype).reshape(1, nall, 1) + desc = idx.expand(nframes, nall, self.dim).clone() + return desc, None, None, None, None + + +def _fitting(dim=2, task_dim=1, neuron=(4,), numb_fparam=0, group_reduce="mean"): + return GroupPropertyFittingNet( + ntypes=2, + dim_descrpt=dim, + property_name="y", + task_dim=task_dim, + neuron=list(neuron), + numb_fparam=numb_fparam, + group_reduce=group_reduce, + ) + + +def _model(fitting, dim=2): + return GroupPropertyModel( + descriptor=_StubDescriptor(dim), fitting=fitting, type_map=["A", "B"] + ) + + +def _coords(nframes, natoms): + return torch.rand(nframes, natoms, 3, dtype=torch.float64) + + +def test_group_property_model_passes_charge_spin_as_descriptor_fparam(): + class _ChargeSpinDescriptor(_StubDescriptor): + add_chg_spin_ebd = True + + def __init__(self, dim: int = 2) -> None: + super().__init__(dim) + self.default_chg_spin = [0.0, 1.0] + self.seen_fparam = None + + def forward(self, ext_coord, ext_atype, nlist, mapping=None, fparam=None): + assert fparam is not None + self.seen_fparam = fparam.detach().clone() + return super().forward(ext_coord, ext_atype, nlist, mapping=mapping) + + descriptor = _ChargeSpinDescriptor() + model = GroupPropertyModel( + descriptor=descriptor, + fitting=_fitting(), + type_map=["A", "B"], + ) + out = model( + _coords(2, 2), + torch.tensor([[0, 1], [0, 1]]), + box=None, + group_id=torch.tensor([[0], [1]]), + weight=torch.ones(2, 1), + pool_mask=torch.ones(2, 2), + ) + assert torch.isfinite(out["y"]).all() + assert descriptor.seen_fparam is not None + assert descriptor.seen_fparam.shape == (2, 2) + assert torch.allclose( + descriptor.seen_fparam, + torch.tensor([[0.0, 1.0], [0.0, 1.0]], dtype=torch.float64), + ) + + +# --------------------------------------------------------------------------- C +def test_masked_mean_model_divides_by_mask_sum(): + model = _model(_fitting()) + coord = _coords(1, 3) + atype = torch.tensor([[0, 1, 0]]) + pool_mask = torch.tensor([[1.0, 1.0, 0.0]]) # atom 2 excluded + out = model( + coord, + atype, + box=None, + group_id=torch.tensor([[0]]), + weight=torch.tensor([[1.0]]), + pool_mask=pool_mask, + ) + # per-atom features = [0,0],[1,1],[2,2]; masked mean over atoms 0,1 = [0.5,0.5] + assert torch.allclose( + out["frame_embedding"], torch.tensor([[0.5, 0.5]], dtype=torch.float64) + ) + + +def test_zero_mask_rejected_by_model(): + model = _model(_fitting()) + with pytest.raises(ValueError, match="all-zero pool_mask"): + model( + _coords(1, 3), + torch.tensor([[0, 1, 0]]), + box=None, + group_id=torch.tensor([[0]]), + weight=torch.tensor([[1.0]]), + pool_mask=torch.zeros(1, 3), + ) + + +# --------------------------------------------------------------------------- B +def test_group_reduce_mean_vs_sum_wiring(): + fitting = _fitting(group_reduce="mean") + model = _model(fitting) + atype = torch.tensor([[0, 1], [0, 1]]) + coord = _coords(2, 2) + kw = { + "box": None, + "pool_mask": torch.ones(2, 2), + "weight": torch.tensor([[1.0], [1.0]]), + } + + # two identical frames in one group: mean-reduce == a single-frame group + pred_two = model(coord, atype, group_id=torch.tensor([[0], [0]]), **kw)["y"] + pred_one = model( + coord[:1], + atype[:1], + box=None, + pool_mask=torch.ones(1, 2), + weight=torch.tensor([[1.0]]), + group_id=torch.tensor([[0]]), + )["y"] + assert torch.allclose(pred_two, pred_one) # mean is size-invariant + + fitting.group_reduce = "sum" # same weights, sum now doubles the embedding + pred_sum = model(coord, atype, group_id=torch.tensor([[0], [0]]), **kw)["y"] + assert not torch.allclose(pred_sum, pred_two) + + +def test_single_frame_group_equivalence(): + fitting = _fitting(group_reduce="mean") + model = _model(fitting) + coord = _coords(3, 2) + atype = torch.tensor([[0, 1], [0, 1], [0, 1]]) + out = model( + coord, + atype, + box=None, + group_id=torch.tensor([[0], [1], [2]]), # each frame its own group + weight=torch.ones(3, 1), + pool_mask=torch.ones(3, 2), + ) + # singleton groups + mean + weight 1 => grouping is a no-op: + # prediction == fitting applied to each frame embedding directly. + direct = fitting(out["frame_embedding"]) + assert torch.allclose(out["y"], direct) + + +# --------------------------------------------------------------------------- E +def test_group_output_bias_zero_initialized(): + fitting = _fitting() + last = fitting.network[-1] + assert isinstance(last, torch.nn.Linear) + assert torch.allclose(last.bias, torch.zeros_like(last.bias)) # route 3: zero-init + + +# --------------------------------------------------------------------------- F +def test_shuffled_batch_loss_invariance(): + torch.manual_seed(0) + fitting = _fitting() + model = _model(fitting) + loss_fn = GroupPropertyLoss(task_dim=1, var_name="y") + + natoms = 2 + nframes = 5 + coord = _coords(nframes, natoms) + atype = torch.zeros(nframes, natoms, dtype=torch.long) + group_id = torch.tensor([2, 0, 2, 1, 0]) + # consistent per-group labels + label_of = {0: 0.0, 1: 10.0, 2: 20.0} + y = torch.tensor([[label_of[int(g)]] for g in group_id], dtype=torch.float64) + + def run(perm): + inp = {"coord": coord[perm], "atype": atype[perm], "box": None} + lab = { + "y": y[perm], + "group_id": group_id[perm].reshape(-1, 1).double(), + "weight": torch.ones(nframes, 1).double(), + "pool_mask": torch.ones(nframes, natoms).double(), + } + _, loss, _ = loss_fn(inp, model, lab, natoms=natoms) + return loss + + ref = run(torch.arange(nframes)) + shuffled = run(torch.tensor([3, 0, 4, 1, 2])) + assert torch.allclose(ref, shuffled) + + +# --------------------------------------------------------------------------- serialize +def test_serialize_roundtrip_group_reduce_and_fparam(): + fitting = _fitting(numb_fparam=2, group_reduce="sum") + sd = fitting.serialize() + assert sd["group_reduce"] == "sum" + assert sd["numb_fparam"] == 2 + rebuilt = GroupPropertyFittingNet(**{k: v for k, v in sd.items() if k != "type"}) + assert rebuilt.group_reduce == "sum" + assert rebuilt.get_dim_fparam() == 2 + + +# --------------------------------------------------------------------------- D (nlist) +def test_padding_invariance(): + """Real descriptor: appending virtual atoms (atype=-1) must not change the + descriptors of real atoms, even when virtual atoms sit inside the cutoff. + """ + from deepmd.pt.model.descriptor import ( + DescrptSeA, + ) + from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, + ) + + desc = DescrptSeA(rcut=6.0, rcut_smth=5.0, sel=[20, 20], ntypes=2) + + def run(coord, atype, box): + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, atype, desc.get_rcut(), desc.get_sel(), mixed_types=True, box=box + ) + return desc(ext_coord, ext_atype, nlist, mapping=mapping)[0] + + torch.manual_seed(0) + base_coord = torch.rand(1, 8, 3, dtype=torch.float64) * 5.0 + base_atype = torch.tensor([[0, 1, 0, 1, 0, 1, 0, 1]]) + box = None # non-periodic + ref = run(base_coord, base_atype, box)[:, :8, :] + + # adversarial: virtual atoms placed INSIDE the cutoff of real atoms + pad = base_coord[:, :3, :].clone() # 3 virtual atoms among the real cloud + pad_coord = torch.cat([base_coord, pad], dim=1) + pad_atype = torch.cat([base_atype, torch.full((1, 3), -1)], dim=1) + got = run(pad_coord, pad_atype, box)[:, :8, :] + assert torch.allclose(ref, got, atol=1e-10) + + # periodic case: nlist masks atype<0 regardless of coord wrapping + box_p = (torch.eye(3, dtype=torch.float64) * 10.0).reshape(1, 9) + ref_p = run(base_coord, base_atype, box_p)[:, :8, :] + got_p = run(pad_coord, pad_atype, box_p)[:, :8, :] + assert torch.allclose(ref_p, got_p, atol=1e-10) + + +# --------------------------------------------------------------------------- DDP +def test_group_not_split_across_ranks(): + # 5 groups of varying size, world_size=2; assign groups (not frames) to ranks. + group_ids = torch.tensor([0, 0, 0, 1, 2, 2, 3, 4, 4, 4]).numpy() + seen: dict[int, int] = {} + for rank in range(2): + batches = distributed_grouped_frame_batches( + group_ids, num_replicas=2, rank=rank, max_frames=8 + ) + for batch in batches: + for frame_idx in batch: + gid = int(group_ids[frame_idx]) + assert seen.setdefault(gid, rank) == rank, ( + f"group {gid} split across ranks" + ) + # every group landed on exactly one rank, and all groups were covered + assert set(seen) == {0, 1, 2, 3, 4}