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 8671a1e94e..aea23ae42b 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, ) @@ -355,6 +358,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..c1b9f79df3 --- /dev/null +++ b/deepmd/pt/model/model/group_property_model.py @@ -0,0 +1,430 @@ +# 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 + + def set_case_embd(self, case_idx: int) -> None: + """Set the case (branch) embedding of this model by the given + case_idx, used to distinguish branches that share a descriptor in + multi-task training. GroupPropertyModel does not go through + make_model(), so it does not inherit the generic model -> + atomic_model -> fitting_net proxy chain and forwards directly. + """ + self.fitting_net.set_case_embd(case_idx) + + @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..b59a301c72 --- /dev/null +++ b/deepmd/pt/model/task/group_property.py @@ -0,0 +1,281 @@ +# 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, + fparam_neuron: list[int] | None = None, + dim_case_embd: 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) + self.fparam_neuron = list(fparam_neuron or []) + if self.fparam_neuron and self.numb_fparam <= 0: + raise ValueError("fparam_neuron requires numb_fparam > 0.") + self.dim_case_embd = int(dim_case_embd) + if self.dim_case_embd > 0: + # One-hot branch identifier for multi-task training, where several + # fitting nets (e.g. an MFT aux ener head and this group_property + # head) share one descriptor and need to stay distinguishable; see + # set_case_embd. Concatenated onto the group embedding in + # forward(), so it counts toward the first Linear layer's input + # width below. + self.register_buffer( + "case_embd", + torch.zeros(self.dim_case_embd, dtype=self.prec, device=env.DEVICE), + ) + else: + self.case_embd = None + 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) can either be concatenated directly + # to the aggregated group embedding (default/backward-compatible path), + # or first encoded through a small fparam-only branch. The branch helps + # when low-dimensional experimental conditions (salt/concentration/pH) + # would otherwise be numerically drowned by a high-dimensional structural + # embedding in the first fused Linear layer. + fparam_out_dim = ( + self.fparam_neuron[-1] if self.fparam_neuron else self.numb_fparam + ) + dims = [ + dim_descrpt + fparam_out_dim + self.dim_case_embd, + *self.neuron, + task_dim, + ] + + def _build_mlp( + dims: list[int], *, activate_last: bool = False + ) -> 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 activate_last or ii < len(dims) - 2: + layers.append(_activation(self.activation_function)) + return layers + + def _build_layers() -> tuple[list[torch.nn.Module], list[torch.nn.Module]]: + fparam_layers = ( + _build_mlp([self.numb_fparam, *self.fparam_neuron], activate_last=True) + if self.fparam_neuron + else [] + ) + return fparam_layers, _build_mlp(dims) + + 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. + fparam_layers, 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) + fparam_layers, layers = _build_layers() + self.fparam_network = torch.nn.Sequential(*fparam_layers).to(env.DEVICE) + 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: + group_embedding = group_embedding.to(self.prec) + if self.fparam_neuron: + descrpt = group_embedding[:, : self.dim_descrpt] + fparam = group_embedding[:, self.dim_descrpt :] + if fparam.shape[-1] != self.numb_fparam: + raise ValueError( + f"expected {self.numb_fparam} fparam columns, got {fparam.shape[-1]}" + ) + group_embedding = torch.cat([descrpt, self.fparam_network(fparam)], dim=-1) + if self.dim_case_embd > 0: + case_embd = self.case_embd.to( + dtype=group_embedding.dtype, device=group_embedding.device + ) + case_embd = case_embd.unsqueeze(0).expand(group_embedding.shape[0], -1) + group_embedding = torch.cat([group_embedding, case_embd], dim=-1) + return self.network(group_embedding) + + 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 set_case_embd(self, case_idx: int) -> None: + """ + Set the case (branch) embedding of this fitting net by the given + case_idx, concatenated with the aggregated group embedding and fed + into the fitting net. Used to keep multiple fitting nets + distinguishable when they share one descriptor in multi-task + training. + """ + self.case_embd = torch.eye( + self.dim_case_embd, dtype=self.prec, device=env.DEVICE + )[case_idx] + + 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, + "fparam_neuron": self.fparam_neuron, + "dim_case_embd": self.dim_case_embd, + "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 8dd68c8cdb..5d526decee 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -47,6 +47,7 @@ EnergyHessianStdLoss, EnergySpinLoss, EnergyStdLoss, + GroupPropertyLoss, PopulationLoss, PropertyLoss, TaskLoss, @@ -899,8 +900,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 @@ -2546,6 +2562,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 91dfab4824..920cee290e 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2891,6 +2891,75 @@ 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``, + ``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. + + ``dim_case_embd`` IS supported (unlike the fields above): it lets a + group_property head share a descriptor with another fitting net in + multi-task training (e.g. MFT jointly training a downstream + group_property head and an auxiliary ener head) by giving each branch a + distinct one-hot case embedding. + """ + 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)." + ) + doc_fparam_neuron = ( + "Hidden sizes for an optional fparam-only branch before fusing fparam " + "with the aggregated group embedding. Empty list keeps the historical " + "direct concatenation path." + ) + _unsupported = { + "numb_aparam", + "default_fparam", + "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( + "fparam_neuron", + list[int], + optional=True, + default=[], + doc=doc_fparam_neuron, + ) + ) + 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." @@ -4941,6 +5010,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/input_formats.md b/doc/dpa_adapt/input_formats.md index fe68c878f4..1414d775c9 100644 --- a/doc/dpa_adapt/input_formats.md +++ b/doc/dpa_adapt/input_formats.md @@ -144,3 +144,97 @@ dpaad data convert --input "calcs/**/OUTCAR" --output ./npy_root --fmt vasp/outc For example, `calcs/run1/OUTCAR` is written as `npy_root/run1/OUTCAR/`. When `--strict` is set, the first conversion error fails immediately. Without it, errors are skipped and logged in the manifest. + +## Grouped `deepmd/npy` Systems + +Grouped data is still ordinary `deepmd/npy` data, with marker arrays under each +`set.*` directory. `frozen_head` and `finetune` auto-detect grouped training when +all training and validation `set.*` directories carry the complete marker set. + +| File | Shape | Dtype | Meaning | +| --------------- | ------------------- | ------------------ | ----------------------------------------------------------------------------------------------- | +| `group_id.npy` | `(nframes,)` | integer | Group index for each frame. Frames with the same id are pooled into one group-level prediction. | +| `weight.npy` | `(nframes,)` | float | Frame/component weight used during group pooling. Use `1.0` for unweighted groups. | +| `pool_mask.npy` | `(nframes, natoms)` | float or bool-like | Atom mask for descriptor pooling. Values > 0 are included; values 0 are excluded. | + +Optional context features use the existing DeePMD frame-parameter convention: + +| File | Shape | Meaning | +| ------------ | -------------------- | --------------------------------------------------------------------------------------------------- | +| `fparam.npy` | `(nframes, nfparam)` | Per-group side features copied to each frame in the group. Rows within one group must be identical. | + +Labels remain standard per-frame label arrays, for example +`set.*/overpotential.npy` or `set.*/cloud_point.npy`. Within one group, label +rows must be identical because loss is computed once per group. + +### Mark existing systems + +Use `dpa-adapt data mark-groups` when the structures and labels already exist +and only grouped markers are missing: + +```bash +# One DeePMD system directory is one group. +dpa-adapt data mark-groups --input oer/dpdata --target overpotential \ + --group-by system + +# Several groups merged into one system; identical labels define groups. +dpa-adapt data mark-groups --input merged_system --target property \ + --group-by label + +# Every consecutive 3 frames form one group. +dpa-adapt data mark-groups --input triplets --group-by 3 +``` + +Equivalent Python API: + +```python +from dpa_adapt import mark_groups + +mark_groups("oer/dpdata", target="overpotential", group_by="system") +``` + +`mark_groups` searches recursively for directories that directly contain +`set.*`. It writes the complete grouped marker contract in every set: +`group_id.npy`, `weight.npy`, and `pool_mask.npy`. When `real_atom_types.npy` is +present, `pool_mask.npy` is derived as `real_atom_types >= 0`; otherwise an +all-one mask is written. `weight.npy` defaults to explicit `1.0` values unless +`--weight` / `weight=` is provided. Existing marker files are preserved unless +`overwrite=True` or `--overwrite` is used. + +For `frozen_head` and `finetune`, avoid partial marker sets: every grouped +`set.*` in both training and validation data must contain all three marker files. +The training preflight rejects partial marker sets to prevent mixed grouped and +ungrouped systems. + +### Low-level grouped writer + +For converters that already have component arrays in memory, use +`dpa_adapt.Assembly`. Each group becomes one DeePMD system; components inside +the group become frames. Components may have different atom counts and are +padded to the largest component in that group. Padding atoms use +`real_atom_types = -1` and `pool_mask = 0`. + +```python +from dpa_adapt import Assembly, PoolMask + +assembly = Assembly(target="cloud_point") +group = assembly.group(key="PNIPAM-co-AA", label=32.1, fparam={"mn_log": 4.6}) +group.add( + coords=repeat_unit_xyz, + symbols=["C", "C", "N", "O"], + weight=0.9, + pool_mask=PoolMask.all(), + role="repeat_unit", +) +group.add( + coords=end_group_xyz, + symbols=["C", "H", "H", "H"], + weight=0.1, + role="end_group", +) +assembly.write("cloud_point_grouped", overwrite=True) +``` + +The writer also creates `manifest.json` with roles, component metadata, type map, +fparam schema, and provenance. Training consumes the `.npy` arrays; the manifest +is for traceability. diff --git a/doc/dpa_adapt/overview.md b/doc/dpa_adapt/overview.md index b742ebefb7..a269bc11d8 100644 --- a/doc/dpa_adapt/overview.md +++ b/doc/dpa_adapt/overview.md @@ -25,6 +25,94 @@ The strategy is the core choice. All four share the same pre-trained DPA backbon | `finetune` | End-to-end full parameter fine-tuning | Large (>10k) | Maximum accuracy on large datasets | | `mft` | Multi-task co-training (property + force field) | Small / low-data | Mitigating representation collapse | +## Grouped property targets + +Most property datasets have one label per structure. Some scientific tasks have +one label for a **group** of structures instead: for example, an OER catalyst +label shared by the clean, `O*`, `OH*`, and `OOH*` slabs; or a polymer property +defined by several repeating-unit and end-group components. DPA-ADAPT supports +this layout with grouped markers stored in ordinary `deepmd/npy` systems. + +Grouped training is enabled automatically when every training and validation +`set.*` directory contains `group_id.npy`, `weight.npy`, and `pool_mask.npy`. +The same high-level `DPAFineTuner` API is used; grouped data switches the +underlying DeePMD fitting net and loss to `group_property`. + +```python +from dpa_adapt import DPAFineTuner, mark_groups + +# Existing deepmd/npy data: add grouped markers in place. +# A common OER layout is "one system directory == one labeled group". +mark_groups( + "oer/dpdata", + target="overpotential", + group_by="system", +) + +model = DPAFineTuner( + pretrained="DPA-3.1-3M", + strategy="finetune", + property_name="overpotential", + task_dim=1, + init_branch="SPICE2", + max_steps=50_000, +) +model.fit(train_data="oer/train/*", valid_data="oer/valid/*") +``` + +The grouped marker API is intentionally small: + +```python +mark_groups( + data, + target="property", + group_by="system", # "system" | "label" | positive integer group size + weight=None, # None writes explicit weight.npy = 1.0 + overwrite=False, + dry_run=False, +) +``` + +- `group_by="system"` assigns all frames in one system to one group. +- `group_by="label"` groups frames with identical rows in `set.*/.npy`. +- `group_by=N` groups every consecutive `N` frames. +- `pool_mask.npy` is always written explicitly. It is derived from + `real_atom_types.npy >= 0` when available; otherwise it is an all-one mask. +- `weight.npy` is always written explicitly. `weight=None` writes `1.0` for + every frame; pass `weight=` to set another constant value. +- For `frozen_head` and `finetune`, every grouped `set.*` directory must contain + the complete marker set (`group_id.npy`, `weight.npy`, and `pool_mask.npy`). + Partial marker sets are rejected so grouped and ungrouped data cannot be mixed + silently. + +For converter implementations or tests that already have arrays in memory, +`Assembly` is the low-level writer: + +```python +from dpa_adapt import Assembly, PoolMask + +assembly = Assembly(target="cloud_point") +group = assembly.group(key="polymer_0", label=32.1, fparam={"mn_log": 4.5}) +group.add( + coords=repeat_unit_xyz, + symbols=["C", "C", "O"], + weight=0.8, + pool_mask=PoolMask.all(), + role="repeat_unit", +) +group.add( + coords=end_group_xyz, + symbols=["C", "H", "H", "H"], + weight=0.2, + role="end_group", +) +assembly.write("polymer_grouped", overwrite=True) +``` + +`Assembly` writes each group as one DeePMD system and stores only the tensors +needed by training in `set.*/`. Scientific provenance such as roles, source +paths, and fparam schema is stored in `manifest.json`. + ### frozen_sklearn — CPU-only, scikit-learn predictor Freezes the DPA backbone as a feature extractor and fits a scikit-learn @@ -100,6 +188,11 @@ pred = model.predict(data="/data/test") metrics = model.evaluate(data="/data/test") # .mae, .rmse, .r2 ``` +When `train_data` and `valid_data` contain grouped markers, `frozen_head` and +`finetune` auto-detect grouped mode. `fitting_net_params` still works as an +override; otherwise grouped fine-tuning uses a GELU activation by default for +the grouped property head. + | Parameter | Type | Default | Description | | -------------------- | ---------------- | ------------------ | ------------------------------------------------------------------------------------- | | `pretrained` | `str` | `"DPA-3.1-3M"` | Checkpoint path or built-in name | @@ -183,6 +276,90 @@ metrics = model.evaluate(data="/data/test") # .mae, .rmse, .r2 | `downstream_batch_size` | `int` or `None` | `None` | Batch size for downstream head; auto-selected if `None` | | `fitting_net_params` | `dict` or `None` | `None` | Overrides for the **aux** fitting net; downstream uses property defaults. `None` = auto-read from checkpoint. | +## Training-time regularizer + +`Regularizer` defines extra training-time losses on the current downstream +training path. It is shared by ordinary property training and grouped property +training. It does not duplicate MFT: if the regularization signal comes from an +external auxiliary dataset and an auxiliary head, use `strategy="mft"`. + +```python +from dpa_adapt import DPAFineTuner, Regularizer + +model = DPAFineTuner( + strategy="finetune", + target="cloud_point", + regularizer=Regularizer(descriptor_anchor=0.0), +) +``` + +The first API term is `descriptor_anchor`, which will penalize drift between +the current descriptor and the frozen base descriptor on the same downstream +structures. Non-zero `descriptor_anchor` is reserved for the descriptor-anchor +training backend and currently fails loudly instead of silently doing nothing. +Use MFT today for auxiliary-task regularization. + +```text +MFT: + external auxiliary task + auxiliary head + regularizes the shared descriptor through aux_data + +Regularizer: + extra loss on the current downstream training path + no auxiliary dataset or auxiliary head +``` + +## Post-training calibration + +`Calibrator` defines post-training prediction correction. It is fitted after the +model has already been trained, and it does not participate in backward. + +The main user-facing API is `model.calibrate(...)`: + +```python +model.calibrate( + data="polymer_grouped/calib", + method="ridge", + alpha=3.0, + features=["prediction", "fparam", "group_stats"], +) + +pred = model.predict("polymer_grouped/test") +``` + +Conceptually: + +```text +trained model prediction +-> calibrator +-> corrected prediction +``` + +Implemented feature sources are: + +- `"prediction"`: raw model prediction columns. +- `"fparam"`: columns from `set.*/fparam.npy`, aligned per frame for ordinary + data and per group for grouped data. Grouped calibration requires fparam to be + constant within each group. +- `"group_stats"`: `n_frames`, weight sum/mean/std/min/max, effective sample + count, and L2 weight norm from grouped `group_id.npy` and `weight.npy`. + Ordinary data uses the one-frame, unit-weight equivalent. + +`Calibrator` is useful when the finetuned model has learned a reasonable trend, +but its absolute scale or OOD boundary behavior is biased. Unlike `Regularizer`, +it does not change the descriptor or fitting net; it only changes the final +prediction rule. + +Advanced users may construct a reusable calibrator object explicitly: + +```python +from dpa_adapt import Calibrator + +calibrator = Calibrator(method="ridge", alpha=3.0, features=["prediction", "fparam", "group_stats"]) +calibrator.fit(model, data="polymer_grouped/calib") +pred = calibrator.predict(model, data="polymer_grouped/test") +``` + ## Data preparation DPA-ADAPT trains on `deepmd/npy` data. Use `dpa-adapt data convert` (or the Python @@ -244,7 +421,7 @@ dpa-adapt data convert --input "calcs/**/OUTCAR" --output ./npy_root --fmt vasp/ Lower-level helpers: ```python -from dpa_adapt import convert, attach_labels, check_data +from dpa_adapt import convert, attach_labels, check_data, mark_groups convert("OUTCAR", "./npy", fmt="vasp/outcar") convert("calcs/**/OUTCAR", "./npy_root", fmt="vasp/outcar") @@ -257,6 +434,9 @@ labels = np.load("labels.npy") # shape (n_systems,) attach_labels("./npy/", head="bandgap", values=labels) check_data("/data/system") # → list[Issue] + +# Add grouped markers to existing deepmd/npy systems. +mark_groups("./npy_root", target="overpotential", group_by="system") ``` For the full option list and supported dpdata formats, see @@ -345,6 +525,8 @@ from dpa_adapt import ( check_data, # data sanity checks attach_labels, # inject label arrays load_dataset, # label-filtered data loading + Regularizer, # training-time downstream-batch regularizer config + Calibrator, # post-training prediction correction ) ``` diff --git a/dpa_adapt/__init__.py b/dpa_adapt/__init__.py index 5af6edff26..df0be326d0 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"), @@ -25,6 +33,8 @@ "MFTFineTuner": (".mft", "MFTFineTuner"), "DPAPredictor": (".predictor", "DPAPredictor"), "DPATrainer": (".trainer", "DPATrainer"), + "Regularizer": (".regularizer", "Regularizer"), + "Calibrator": (".calibrator", "Calibrator"), } __all__ = list(_LAZY) diff --git a/dpa_adapt/calibrator.py b/dpa_adapt/calibrator.py new file mode 100644 index 0000000000..5a353932b1 --- /dev/null +++ b/dpa_adapt/calibrator.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Post-training prediction calibration.""" + +from __future__ import annotations + +from dataclasses import ( + dataclass, + field, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, +) + +if TYPE_CHECKING: + from collections.abc import ( + Sequence, + ) + +import numpy as np + +_GROUP_STAT_NAMES = ( + "n_frames", + "weight_sum", + "weight_mean", + "weight_std", + "weight_min", + "weight_max", + "weight_eff_n", + "weight_l2", +) + + +@dataclass +class Calibrator: + """Post-training prediction correction. + + Supported feature sources: + + - ``"prediction"``: raw model prediction columns. + - ``"fparam"``: per-frame or per-group ``set.*/fparam.npy`` columns. + - ``"group_stats"``: grouped marker statistics derived from + ``group_id.npy`` and ``weight.npy``; ordinary non-grouped data gets the + one-frame/unweighted equivalent. + """ + + method: str = "ridge" + alpha: float = 1.0 + features: Sequence[str] = ("prediction",) + model: object | None = field(default=None, init=False, repr=False) + task_dim: int | None = field(default=None, init=False) + feature_names_: tuple[str, ...] = field(default=(), init=False) + + def __post_init__(self) -> None: + self.method = str(self.method).lower() + if self.method not in {"ridge", "linear"}: + raise ValueError( + f"Calibrator method must be 'ridge' or 'linear'; got {self.method!r}." + ) + self.features = tuple(self.features) + if not self.features: + raise ValueError("Calibrator requires at least one feature.") + supported = {"prediction", "fparam", "group_stats"} + unsupported = sorted(set(self.features) - supported) + if unsupported: + raise ValueError( + f"Unsupported Calibrator feature sources: {unsupported}. " + f"Supported sources are {sorted(supported)}." + ) + if self.alpha < 0.0: + raise ValueError(f"alpha must be non-negative; got {self.alpha}.") + + def fit_from_arrays( + self, + predictions: np.ndarray, + labels: np.ndarray, + *, + data: str | list[str] | None = None, + fmt: str | None = None, + ) -> Calibrator: + x, names = self._build_features(predictions, data=data, fmt=fmt) + y = np.asarray(labels, dtype=float) + if y.ndim == 1: + y = y.reshape(-1, 1) + if x.shape[0] != y.shape[0]: + raise ValueError( + f"Calibration row mismatch: features have {x.shape[0]} rows, " + f"labels have {y.shape[0]} rows." + ) + + from sklearn.linear_model import ( + LinearRegression, + Ridge, + ) + + if self.method == "ridge": + self.model = Ridge(alpha=float(self.alpha)) + else: + self.model = LinearRegression() + self.model.fit(x, y) + self.task_dim = int(y.shape[1]) + self.feature_names_ = tuple(names) + return self + + def fit( + self, model: object, data: str | list[str], fmt: str | None = None + ) -> Calibrator: + result = model.evaluate(data, fmt=fmt, calibrated=False) + return self.fit_from_arrays( + result.predictions, + result.labels, + data=data, + fmt=fmt, + ) + + def predict_from_arrays( + self, + predictions: np.ndarray, + *, + data: str | list[str] | None = None, + fmt: str | None = None, + ) -> np.ndarray: + if self.model is None: + raise RuntimeError("Calibrator.predict() called before fit().") + x, names = self._build_features(predictions, data=data, fmt=fmt) + if self.feature_names_ and tuple(names) != self.feature_names_: + raise ValueError( + "Calibration feature columns differ from fit-time columns.\n" + f" fit : {list(self.feature_names_)}\n" + f" predict : {names}" + ) + out = np.asarray(self.model.predict(x), dtype=float) + task_dim = self.task_dim or (1 if out.ndim == 1 else out.shape[-1]) + return out.reshape(-1, task_dim) + + def predict( + self, model: object, data: str | list[str], fmt: str | None = None + ) -> np.ndarray: + result = model.predict(data, fmt=fmt, calibrated=False) + return self.predict_from_arrays(result.predictions, data=data, fmt=fmt) + + def _build_features( + self, + predictions: np.ndarray, + *, + data: str | list[str] | None, + fmt: str | None, + ) -> tuple[np.ndarray, list[str]]: + pred = np.asarray(predictions, dtype=float) + if pred.ndim == 1: + pred = pred.reshape(-1, 1) + pred = pred.reshape(pred.shape[0], -1) + + parts: list[np.ndarray] = [] + names: list[str] = [] + if "prediction" in self.features: + parts.append(pred) + names.extend([f"prediction_{i}" for i in range(pred.shape[1])]) + + need_data = any(source in self.features for source in ("fparam", "group_stats")) + rows = None + if need_data: + if data is None: + raise ValueError( + "Calibrator features 'fparam' and 'group_stats' require data." + ) + rows = _feature_rows_from_data(data, fmt=fmt, expected_rows=pred.shape[0]) + + if "fparam" in self.features: + assert rows is not None + if rows.fparam is None: + raise ValueError( + "Calibrator feature 'fparam' requested, but no fparam.npy was found." + ) + parts.append(rows.fparam) + names.extend([f"fparam_{i}" for i in range(rows.fparam.shape[1])]) + + if "group_stats" in self.features: + assert rows is not None + parts.append(rows.group_stats) + names.extend(_GROUP_STAT_NAMES) + + if not parts: + raise ValueError("No calibration features were selected.") + return np.concatenate(parts, axis=1), names + + +@dataclass(frozen=True) +class _FeatureRows: + fparam: np.ndarray | None + group_stats: np.ndarray + + +def _feature_rows_from_data( + data: str | list[str], *, fmt: str | None, expected_rows: int +) -> _FeatureRows: + from dpa_adapt.data.loader import ( + _get_source, + load_data, + ) + + systems = load_data(data, fmt=fmt) + grouped_records = [] + frame_fparams = [] + frame_stats = [] + any_grouped = False + any_fparam = False + next_group_id = 0 + + for system in systems: + source = _get_source(system) + if source is None: + raise ValueError( + "Calibration fparam/group_stats require filesystem-backed deepmd/npy data." + ) + source_path = Path(source) + local_to_global: dict[int, int] = {} + for set_dir in sorted(source_path.glob("set.*")): + nframes, natoms = _set_shape(set_dir) + gid_path = set_dir / "group_id.npy" + if gid_path.is_file(): + group_ids = np.load(gid_path).reshape(-1).astype(np.int64) + if group_ids.shape != (nframes,): + raise ValueError( + f"{gid_path} shape {group_ids.shape}; expected ({nframes},)." + ) + any_grouped = True + else: + group_ids = np.arange(nframes, dtype=np.int64) + + weight_path = set_dir / "weight.npy" + weights = ( + np.asarray(np.load(weight_path), dtype=float).reshape(-1) + if weight_path.is_file() + else np.ones(nframes, dtype=float) + ) + if weights.shape != (nframes,): + raise ValueError( + f"{weight_path} shape {weights.shape}; expected ({nframes},)." + ) + + fp_path = set_dir / "fparam.npy" + fparam = None + if fp_path.is_file(): + fparam = np.asarray(np.load(fp_path), dtype=float).reshape(nframes, -1) + any_fparam = True + + for frame in range(nframes): + frame_stats.append(_stats_for_weights(weights[frame : frame + 1])) + if fparam is not None: + frame_fparams.append(fparam[frame]) + elif any_fparam: + raise ValueError( + "fparam.npy must be present for every set when Calibrator uses fparam." + ) + + for local_gid in group_ids: + if int(local_gid) not in local_to_global: + local_to_global[int(local_gid)] = next_group_id + next_group_id += 1 + for frame, local_gid in enumerate(group_ids): + grouped_records.append( + ( + local_to_global[int(local_gid)], + float(weights[frame]), + None if fparam is None else fparam[frame], + ) + ) + + if any_grouped: + group_ids = sorted({record[0] for record in grouped_records}) + group_stats = [] + group_fparams = [] + for gid in group_ids: + records = [record for record in grouped_records if record[0] == gid] + weights = np.asarray([record[1] for record in records], dtype=float) + group_stats.append(_stats_for_weights(weights)) + fps = [record[2] for record in records if record[2] is not None] + if fps: + arr = np.asarray(fps, dtype=float) + first = arr[0] + if not np.allclose(arr, first[None, :], atol=1e-8): + raise ValueError( + f"fparam must be constant within grouped calibration data; group {gid} differs." + ) + group_fparams.append(first) + elif any_fparam: + raise ValueError( + "fparam.npy must be present for every grouped frame when Calibrator uses fparam." + ) + rows = _FeatureRows( + fparam=np.asarray(group_fparams, dtype=float) if any_fparam else None, + group_stats=np.asarray(group_stats, dtype=float), + ) + if rows.group_stats.shape[0] == expected_rows: + return rows + + frame_rows = _FeatureRows( + fparam=np.asarray(frame_fparams, dtype=float) if any_fparam else None, + group_stats=np.asarray(frame_stats, dtype=float), + ) + if frame_rows.group_stats.shape[0] == expected_rows: + return frame_rows + + raise ValueError( + f"Could not align calibration features to predictions: got " + f"{frame_rows.group_stats.shape[0]} frame rows" + + (f" and {rows.group_stats.shape[0]} group rows" if any_grouped else "") + + f", expected {expected_rows}." + ) + + +def _set_shape(set_dir: Path) -> tuple[int, int]: + real_path = set_dir / "real_atom_types.npy" + if real_path.is_file(): + arr = np.load(real_path, mmap_mode="r") + return int(arr.shape[0]), int(arr.shape[1]) + coord = np.load(set_dir / "coord.npy", mmap_mode="r") + return int(coord.shape[0]), int(coord.shape[1] // 3) + + +def _stats_for_weights(weights: np.ndarray) -> np.ndarray: + w = np.asarray(weights, dtype=float).reshape(-1) + l2_sq = float(np.sum(w**2)) + weight_sum = float(np.sum(w)) + eff_n = float((weight_sum**2) / l2_sq) if l2_sq > 1e-12 else 0.0 + return np.asarray( + [ + float(len(w)), + weight_sum, + float(np.mean(w)) if len(w) else 0.0, + float(np.std(w)) if len(w) else 0.0, + float(np.min(w)) if len(w) else 0.0, + float(np.max(w)) if len(w) else 0.0, + eff_n, + float(np.sqrt(l2_sq)), + ], + dtype=float, + ) diff --git a/dpa_adapt/cli.py b/dpa_adapt/cli.py index 10d098715f..ebeb9c1441 100644 --- a/dpa_adapt/cli.py +++ b/dpa_adapt/cli.py @@ -352,6 +352,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 # --------------------------------------------------------------------------- @@ -368,6 +414,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, } @@ -674,6 +721,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/config/manager.py b/dpa_adapt/config/manager.py index 2d267a57a1..d6ec9f0023 100644 --- a/dpa_adapt/config/manager.py +++ b/dpa_adapt/config/manager.py @@ -9,19 +9,40 @@ resolve_dp_command, ) +_PROPERTY_LIKE_DOWNSTREAM_TYPES = ("property", "group_property") + + +def _aux_dim_case_embd(t: Any) -> int: + """dim_case_embd required on the DOWNSTREAM head to share a descriptor + with the aux branch. + + Read from the aux branch's own (ckpt-derived) ``fitting_net_params`` + rather than hardcoded: it must equal the branch count of whatever + multi-task checkpoint the aux branch was itself pretrained as part of + (deepmd-kit's multi-task trainer requires every model_dict branch to + declare the same dim_case_embd -- see + deepmd.pt.train.training.get_case_embd_config). That count is 31 for + DPA-3.1-3M but differs for other checkpoints (e.g. 23 for the + OMol25/Organic_Reactions/ODAC23 checkpoint); hardcoding 31 silently + mismatches every checkpoint that isn't DPA-3.1-3M. Falls back to 0 (no + case embedding, matching the aux branch) when the aux fitting_net has + none, e.g. a single-task-pretrained checkpoint. + """ + return (getattr(t, "fitting_net_params", None) or {}).get("dim_case_embd", 0) + + # Default property-head architecture for MFT DOWNSTREAM when # downstream_task_type="property". Mirrors DPATrainer.DEFAULT_FITTING_NET -# (trainer.py L64-70) plus dim_case_embd=31, which the DPA-3.1-3M ckpt -# requires for the case-embedding layer in multi-task mode. (DPATrainer is -# single-task and doesn't need this field; in MFT the descriptor is shared -# across branches so the property head must declare it.) +# (trainer.py L64-70). dim_case_embd is added dynamically by +# _build_property_fitting_net via _aux_dim_case_embd, not baked in here: +# DPATrainer is single-task and doesn't need this field at all, while MFT's +# correct value depends on which checkpoint is being finetuned. _PROPERTY_FITTING_NET_BASE = { "type": "property", "neuron": [240, 240, 240], "activation_function": "tanh", "resnet_dt": True, "precision": "float32", - "dim_case_embd": 31, } @@ -40,6 +61,9 @@ def _build_property_fitting_net(t: Any) -> dict: "seed": t.seed, } ) + dim_case_embd = _aux_dim_case_embd(t) + if dim_case_embd: + fn["dim_case_embd"] = dim_case_embd if getattr(t, "fparam_dim", 0) > 0: fn["numb_fparam"] = t.fparam_dim return fn @@ -59,6 +83,57 @@ def _build_property_loss() -> dict: } +# Default group_property-head architecture for MFT DOWNSTREAM when +# downstream_task_type="group_property". Independent of +# _PROPERTY_FITTING_NET_BASE (not a trimmed copy of it): GroupPropertyFittingNet +# is a small standalone MLP, not built on GeneralFitting, so several +# property-schema fields (resnet_dt, intensive, ...) don't exist on it and +# dargs strict-mode rejects them outright rather than ignoring them -- see +# deepmd.utils.argcheck.fitting_group_property. dim_case_embd is added +# dynamically by _build_group_property_fitting_net via _aux_dim_case_embd, +# not baked in here, for the same reason as the property head. +_GROUP_PROPERTY_FITTING_NET_BASE = { + "type": "group_property", + "neuron": [240, 240, 240], + "activation_function": "gelu", + "precision": "float32", +} + + +def _build_group_property_fitting_net(t: Any) -> dict: + """Construct a group_property fitting_net dict from a tuner's params.""" + fn = dict(_GROUP_PROPERTY_FITTING_NET_BASE) + fn.update( + { + "property_name": t.property_name, + "task_dim": t.task_dim, + "group_reduce": getattr(t, "group_reduce", "mean"), + "seed": t.seed, + } + ) + dim_case_embd = _aux_dim_case_embd(t) + if dim_case_embd: + fn["dim_case_embd"] = dim_case_embd + if getattr(t, "fparam_dim", 0) > 0: + fn["numb_fparam"] = t.fparam_dim + return fn + + +def _build_group_property_loss() -> dict: + """group_property-task loss for DOWNSTREAM. + + deepmd.utils.argcheck.loss_group_property() reuses loss_property()'s + schema verbatim, so this only differs from _build_property_loss() by + ``type``. + """ + return { + "type": "group_property", + "loss_func": "mse", + "metric": ["mae", "rmse"], + "beta": 1.0, + } + + _ENER_LOSS = { "type": "ener", "start_pref_e": 0.2, @@ -81,26 +156,36 @@ def build(self) -> dict: if getattr(t, "fitting_net_params", None) else {"type": "ener"} ) - # DOWNSTREAM branch: ener (legacy, sensitivity-analysis callers) or - # property (paper-faithful BOOM eval). Default 'ener' for back-compat - # with FakeTuners and existing callers that don't set the attr. + # DOWNSTREAM branch: ener (legacy, sensitivity-analysis callers), + # property (paper-faithful BOOM eval), or group_property (grouped/ + # assembly targets, e.g. OER overpotential). Default 'ener' for + # back-compat with FakeTuners and existing callers that don't set + # the attr. downstream_task_type = getattr(t, "downstream_task_type", "ener") - is_property = downstream_task_type == "property" - # Branch key for the downstream head. Paper qm9_gap/mft uses "property"; - # legacy ener mode keeps "DOWNSTREAM" so mp_data sensitivity-analysis - # configs stay byte-for-byte unchanged (renaming would break the branch - # name in their already-trained ckpts). - downstream_key = "property" if is_property else "DOWNSTREAM" - if is_property: + # Both property and group_property get a fresh, RANDOM-initialized + # downstream head sized by property_name/task_dim and follow the + # qm9_gap paper-alignment defaults below; only legacy ener mode + # reuses the aux branch's own fitting_net/finetune_head. + is_random_downstream = downstream_task_type in _PROPERTY_LIKE_DOWNSTREAM_TYPES + # Branch key for the downstream head. Paper qm9_gap/mft uses the task + # type itself ("property" / "group_property"); legacy ener mode keeps + # "DOWNSTREAM" so mp_data sensitivity-analysis configs stay + # byte-for-byte unchanged (renaming would break the branch name in + # their already-trained ckpts). + downstream_key = downstream_task_type if is_random_downstream else "DOWNSTREAM" + if downstream_task_type == "property": downstream_fitting_net = _build_property_fitting_net(t) downstream_loss = _build_property_loss() + elif downstream_task_type == "group_property": + downstream_fitting_net = _build_group_property_fitting_net(t) + downstream_loss = _build_group_property_loss() else: downstream_fitting_net = aux_fitting_net downstream_loss = dict(_ENER_LOSS) - # Paper qm9_gap/mft alignment is applied ONLY in property mode. The - # legacy ener path (mp_data sensitivity analysis) stays byte-for-byte - # unchanged. + # Paper qm9_gap/mft alignment is applied to both property and + # group_property downstream modes. The legacy ener path (mp_data + # sensitivity analysis) stays byte-for-byte unchanged. descriptor = { "type": "dpa3", "repflow": { @@ -130,7 +215,9 @@ def build(self) -> dict: "optim_update": True, "use_exp_switch": True, }, - "activation_function": "silut:3.0" if is_property else "custom_silu:3.0", + "activation_function": "silut:3.0" + if is_random_downstream + else "custom_silu:3.0", "precision": "float32", "use_tebd_bias": False, "concat_output_tebd": False, @@ -139,15 +226,16 @@ def build(self) -> dict: "trainable": True, "use_econf_tebd": False, } - if is_property: + if is_random_downstream: descriptor["repflow"]["fix_stat_std"] = 0.3 - # MFT branch heads. In property mode the paper pins finetune_head: - # the aux head loads from its named branch, the downstream property - # head is RANDOM-initialized (paper Eq 12). Legacy ener mode keeps the - # original layout (no finetune_head on aux; downstream = aux branch), - # including key order, so the emitted JSON is byte-for-byte unchanged. - if is_property: + # MFT branch heads. In property/group_property mode the paper pins + # finetune_head: the aux head loads from its named branch, the + # downstream head is RANDOM-initialized (paper Eq 12). Legacy ener + # mode keeps the original layout (no finetune_head on aux; downstream + # = aux branch), including key order, so the emitted JSON is + # byte-for-byte unchanged. + if is_random_downstream: aux_head = { "type_map": "type_map", "descriptor": "dpa3_descriptor", @@ -176,22 +264,23 @@ def build(self) -> dict: decay_steps = ( t.decay_steps if getattr(t, "decay_steps", None) is not None - else (1000 if is_property else 5000) + else (1000 if is_random_downstream else 5000) ) # Per-branch batch sizes: explicit override wins, then paper defaults - # for property mode, then the single batch_size for legacy ener mode. + # for property/group_property mode, then the single batch_size for + # legacy ener mode. aux_batch = getattr(t, "aux_batch_size", None) or ( - "auto:128" if is_property else t.batch_size + "auto:128" if is_random_downstream else t.batch_size ) downstream_batch = getattr(t, "downstream_batch_size", None) or ( - "auto:512" if is_property else t.batch_size + "auto:512" if is_random_downstream else t.batch_size ) # Paper default 0.5/0.5; aux_prob (default 0.5) controls the split, the # downstream share is the complement. Legacy keeps downstream at 1.0. aux_prob = float(t.aux_prob) if not 0.0 <= aux_prob <= 1.0: raise ValueError(f"aux_prob must be in [0, 1]; got {t.aux_prob!r}.") - downstream_prob = (1.0 - aux_prob) if is_property else 1.0 + downstream_prob = (1.0 - aux_prob) if is_random_downstream else 1.0 aux_systems = t.aux_data if isinstance(t.aux_data, list) else [t.aux_data] train_systems = ( @@ -233,7 +322,7 @@ def build(self) -> dict: "batch_size": downstream_batch, } - if is_property: + if is_random_downstream: # Paper qm9_gap: gradient clipping at 5.0. training["gradient_max_norm"] = 5.0 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..5c253264a8 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()) @@ -826,6 +999,9 @@ class DPAFineTuner: strategies. For ``frozen_sklearn``, fparam columns are standardized and concatenated to the descriptor via ``ConditionManager``. Default 0 (disabled). + regularizer : Regularizer, dict, or None + Training-time downstream-batch regularizer. Does not accept + MFT-style aux_data; use ``strategy='mft'`` for external aux tasks. output_dir : str Directory for ``input.json``, checkpoints, and logs. save_freq, disp_freq : int @@ -865,6 +1041,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", @@ -878,6 +1055,7 @@ def __init__( loss_function: str = "mse", fitting_net_params: dict | None = None, fparam_dim: int = 0, + regularizer: object | None = None, output_dir: str = "./dpa_output", save_freq: int = 10_000, disp_freq: int = 1_000, @@ -889,16 +1067,28 @@ 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) + from dpa_adapt.regularizer import ( + Regularizer, + ) + + self.regularizer = Regularizer.from_config(regularizer) self.strategy = strategy @@ -960,6 +1150,7 @@ def __init__( self._device = None # set when model is first loaded self._checkpoint_type_map = [] # set by _load_descriptor_model self._condition_manager = None + self.calibrator = None # ------------------------------------------------------------------ # Frozen-sklearn pipeline helpers (thin delegators) @@ -1331,13 +1522,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. @@ -1356,7 +1550,9 @@ def fit( Element symbols. Auto-inferred from the checkpoint and data ``type_map.raw`` when not provided. target_key : str, optional - (frozen_sklearn) Label key, e.g. ``"energy"``. + Label key, e.g. ``"energy"``. For ``frozen_sklearn`` this selects + labels to load. For training strategies it is a fit-time alias for + ``property_name``. labels : np.ndarray, optional (frozen_sklearn) Pre-computed labels. fmt : str, optional @@ -1365,9 +1561,35 @@ 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).") + + self.regularizer.require_supported_backend() + if self.strategy == "frozen_sklearn": return self._fit_sklearn(train_data, type_map, target_key, labels, fmt) + if labels is not None: + raise ValueError( + "labels is only valid when strategy='frozen_sklearn'; " + f"got strategy={self.strategy!r}." + ) + if target_key is not None: + if not isinstance(target_key, str): + raise ValueError( + "training strategies require a single string target_key; " + f"got {target_key!r}." + ) + self.property_name = target_key + if self._mft is not None: + self._mft.property_name = target_key + if self.strategy == "mft": if aux_data is None: raise ValueError( @@ -1455,6 +1677,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,7 +1753,75 @@ def _fit_sklearn( p._condition_manager = self._condition_manager p._fitted = True - def predict(self, data: str | list[str], fmt: str | None = None) -> DotDict: + 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, + *, + calibrated: bool = True, + ) -> DotDict: """ Predict with the adapted model. @@ -1542,20 +1842,48 @@ def predict(self, data: str | list[str], fmt: str | None = None) -> DotDict: ``predictions`` : np.ndarray, shape (n_frames, task_dim) """ if self.strategy in {"frozen_head", "finetune"}: - return self._run_training_predict(data, fmt=fmt) + result = self._run_training_predict(data, fmt=fmt) + return self._apply_calibrator( + result, calibrated=calibrated, data=data, fmt=fmt + ) if self.strategy == "mft": if fmt is not None: raise ValueError( "fmt is not supported for mft predict(); " "provide deepmd/npy system directories." ) - return self._ensure_mft().predict(data) + result = self._ensure_mft().predict(data) + return self._apply_calibrator( + result, calibrated=calibrated, data=data, fmt=fmt + ) if not self._fitted: raise RuntimeError( "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 self._apply_calibrator( + DotDict({"predictions": predictions}), + calibrated=calibrated, + data=data, + fmt=fmt, + ) + systems = load_data(data, fmt=fmt) features = self._extract_features(systems) @@ -1575,9 +1903,57 @@ def predict(self, data: str | list[str], fmt: str | None = None) -> DotDict: raw = self.predictor.predict(features) predictions = np.asarray(raw).reshape(-1, self._task_dim) - return DotDict({"predictions": predictions}) + return self._apply_calibrator( + DotDict({"predictions": predictions}), + calibrated=calibrated, + data=data, + fmt=fmt, + ) - def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: + def _apply_calibrator( + self, + result: DotDict, + *, + calibrated: bool, + data: str | list[str] | None = None, + fmt: str | None = None, + ) -> DotDict: + if calibrated and self.calibrator is not None: + raw_predictions = np.asarray(result.predictions) + result["raw_predictions"] = raw_predictions + result["predictions"] = self.calibrator.predict_from_arrays( + raw_predictions, + data=data, + fmt=fmt, + ) + return result + + def calibrate( + self, + data: str | list[str], + *, + method: str = "ridge", + alpha: float = 1.0, + features: list[str] | tuple[str, ...] = ("prediction",), + fmt: str | None = None, + ) -> object: + """Fit a post-training prediction calibrator and attach it to the model.""" + from dpa_adapt.calibrator import ( + Calibrator, + ) + + calibrator = Calibrator(method=method, alpha=alpha, features=features) + calibrator.fit(self, data=data, fmt=fmt) + self.calibrator = calibrator + return calibrator + + def evaluate( + self, + data: str | list[str], + fmt: str | None = None, + *, + calibrated: bool = True, + ) -> DotDict: """ Predict on ``data`` and compute evaluation metrics against stored labels. @@ -1597,9 +1973,14 @@ def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: """ if self.strategy in {"frozen_head", "finetune"}: result = self._run_training_predict(data, fmt=fmt) + result = self._apply_calibrator( + result, calibrated=calibrated, data=data, fmt=fmt + ) labels = result.labels predictions = result.predictions err = predictions - labels + result["mae"] = float(np.mean(np.abs(err))) + result["rmse"] = float(np.sqrt(np.mean(err**2))) ss_res = np.sum(err**2) ss_tot = np.sum((labels - labels.mean()) ** 2) result["r2"] = float(1.0 - ss_res / ss_tot) if ss_tot > 0 else float("nan") @@ -1614,20 +1995,42 @@ def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: if getattr(mft, "downstream_task_type", "property") == "ener": return DotDict(mft.evaluate(data)) result = mft.predict(data) + result = self._apply_calibrator( + result, calibrated=calibrated, data=data, fmt=fmt + ) labels = result.labels predictions = result.predictions err = predictions - labels + result["mae"] = float(np.mean(np.abs(err))) + result["rmse"] = float(np.sqrt(np.mean(err**2))) ss_res = np.sum(err**2) ss_tot = np.sum((labels - labels.mean()) ** 2) result["r2"] = float(1.0 - ss_res / ss_tot) if ss_tot > 0 else float("nan") return result - result = self.predict(data, fmt=fmt) + result = self.predict(data, fmt=fmt, calibrated=calibrated) 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( @@ -1718,6 +2121,7 @@ def freeze(self, output_path: str = "frozen_model.pth") -> str: "pooling": self.pooling, "condition_manager": self._condition_manager, "fparam_dim": self.fparam_dim, + "calibrator": self.calibrator, } output_path = str(output_path) 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..6822c45f72 --- /dev/null +++ b/dpa_adapt/grouped/_convert.py @@ -0,0 +1,394 @@ +# 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`` / ``weight`` / ``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 explicit ``weight.npy`` and ``pool_mask.npy`` the dataset is only + partially marked as grouped, making validation and downstream tooling + ambiguous even though the backend has numerical defaults. + +This module fills that gap **in place**, deriving ``pool_mask`` from the +mixed-type ``real_atom_types`` when present (otherwise all ones), assigning +``group_id`` by a configurable policy, and writing an explicit frame weight, +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, ``weight`` shape ``(nframes,)`` float64, ``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 grouped 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 + Constant ``weight.npy`` value written for every frame. When ``None``, + ``1.0`` is written explicitly so grouped data always carries the full + ``group_id`` / ``weight`` / ``pool_mask`` contract. + 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" + if _should_write(pool_mask_path, overwrite): + pool_mask = ( + (real_types >= 0).astype(np.float64) + if real_types is not None + else np.ones((nframes, natoms), dtype=np.float64) + ) + if not dry_run: + np.save(pool_mask_path, pool_mask) + result.wrote_pool_mask = True + + if _should_write(set_dir / f"{WEIGHT_KEY}.npy", overwrite): + if not dry_run: + np.save( + set_dir / f"{WEIGHT_KEY}.npy", + np.full( + (nframes,), + 1.0 if weight is None else 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..11cf34b608 --- /dev/null +++ b/dpa_adapt/grouped/_polymer.py @@ -0,0 +1,458 @@ +# 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:" +CATEGORY_PREFIX = "cat:" + + +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] + categories: dict[str, str] + 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 {} + categories_raw = conds.pop("categories", None) or {} + scalars = { + str(k): float(v) + for k, v in conds.items() + if v is not None and not _isnan(v) + } + categories = { + f"{name}={value}": str(value) + for name, value in categories_raw.items() + if name is not None and value is not None and not _isnan(value) + } + 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, + categories=categories, + 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), + } + mw = cell(row, "Mw") + if not _isnan(mw): + fparam["mw_log"] = math.log10(float(mw)) + pdi = cell(row, "PDI") + if not _isnan(pdi): + fparam["pdi"] = float(pdi) + if not _isnan(conc): + fparam["conc"] = float(conc) + + categories: dict[str, str] = {} + for col in ( + "polymer_type", + "polymer_type_style", + "polymer_architecture", + "polymerisation_type", + "mass_characterisation_method", + "mass_characterisation_standart", + "def_type", + "tacticity", + ): + value = cell(row, col) + if value is not None and not _isnan(value): + categories[col] = str(value) + if categories: + fparam["categories"] = categories + + 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) + builder.set_fparam_schema({"name": name, "default": 0.0} for name in schema) + 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: mass scalars, scalar fields, categories, salts.""" + columns: list[str] = [MN_KEY] + scalar_keys: set[str] = set() + category_names: set[str] = set() + salt_names: set[str] = set() + for row in self._rows: + scalar_keys.update(row.scalars) + category_names.update(row.categories) + salt_names.update(row.salts) + columns.extend(sorted(key for key in scalar_keys if key != MN_KEY)) + columns.extend(f"{CATEGORY_PREFIX}{name}" for name in sorted(category_names)) + 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(CATEGORY_PREFIX): + vec[i] = 1.0 if col[len(CATEGORY_PREFIX) :] in row.categories 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/mft.py b/dpa_adapt/mft.py index 731c9623ff..32b7d0eb30 100644 --- a/dpa_adapt/mft.py +++ b/dpa_adapt/mft.py @@ -22,6 +22,12 @@ _LOG = logging.getLogger("dpa_adapt.mft") +# DOWNSTREAM branch types that get a fresh, randomly-initialized head sized +# by property_name/task_dim (as opposed to "ener", which reuses the aux +# branch's own fitting_net dict). Shared by __init__ validation and +# MFTConfigManager's paper-alignment branching. +_PROPERTY_LIKE_DOWNSTREAM_TYPES = ("property", "group_property") + class MFTFineTuner: """ @@ -64,27 +70,41 @@ class MFTFineTuner: Pass an explicit dict only if you need to override the checkpoint's config (e.g. for experiments). downstream_task_type : str - Either ``"property"`` (intensive scalar head, e.g. HOMO/LUMO, the - default) or ``"ener"`` (force-field head, legacy mode). Selects how - the DOWNSTREAM branch's fitting_net and loss are built: + One of ``"property"`` (intensive scalar head, e.g. HOMO/LUMO, the + default), ``"group_property"`` (grouped/assembly head, e.g. an + OER O*/OH*/OOH* overpotential pooled over several structures), or + ``"ener"`` (force-field head, legacy mode). Selects how the + DOWNSTREAM branch's fitting_net and loss are built: * ``"property"`` — DOWNSTREAM gets a fresh ``type: property`` fitting_net (using ``property_name``, ``task_dim``, ``intensive``) and a property-style MSE loss with no force/virial prefs. This is what arXiv:2601.08486 Table 3 / Fig 2 reports for HOMO/LUMO. + * ``"group_property"`` — DOWNSTREAM gets a fresh + ``type: group_property`` fitting_net (using ``property_name``, + ``task_dim``, ``group_reduce``) and a group_property MSE loss. + Training data must carry the grouped markers (``group_id.npy``, + ``weight.npy``, ``pool_mask.npy``; see ``dpa_adapt.grouped``); + ``intensive`` does not apply (the group-level output bias is + always zero-initialized, not set from label statistics). * ``"ener"`` — DOWNSTREAM reuses the aux fitting_net dict and an ener-style loss with force/virial prefs. This is the legacy mode used by earlier mp_data sensitivity-analysis MFT experiments. property_name : str, optional - Required when ``downstream_task_type="property"``. Name of the - per-system property file (e.g. ``"homo"`` reads ``set.*/homo.npy``). - Must be a valid Python identifier. + Required when ``downstream_task_type`` is ``"property"`` or + ``"group_property"``. Name of the per-system property file (e.g. + ``"homo"`` reads ``set.*/homo.npy``). Must be a valid Python + identifier. task_dim : int Output dimensionality of the property head. Default ``1``. intensive : bool Whether the property is intensive (mean-pool) or extensive (sum). Default ``True`` (correct for HOMO/LUMO and most molecular - properties). + properties). Ignored when ``downstream_task_type="group_property"``. + group_reduce : str + Frame-embedding -> group-embedding reduction for + ``downstream_task_type="group_property"``: ``"mean"`` (weighted + average, default) or ``"sum"``. Ignored otherwise. learning_rate : float Initial learning rate. stop_lr : float @@ -119,6 +139,7 @@ def __init__( property_name: str | None = None, task_dim: int = 1, intensive: bool = True, + group_reduce: str = "mean", learning_rate: float = 1e-3, stop_lr: float = 1e-5, decay_steps: int | None = None, # None → auto: 1000 for property, 5000 for ener @@ -133,20 +154,27 @@ def __init__( save_freq: int = 10000, disp_freq: int = 1000, ) -> None: - if downstream_task_type not in ("ener", "property"): + if downstream_task_type not in ("ener", *_PROPERTY_LIKE_DOWNSTREAM_TYPES): raise ValueError( - f"downstream_task_type must be 'ener' or 'property'; " - f"got {downstream_task_type!r}." + "downstream_task_type must be 'ener', 'property', or " + f"'group_property'; got {downstream_task_type!r}." ) - if downstream_task_type == "property": + if downstream_task_type in _PROPERTY_LIKE_DOWNSTREAM_TYPES: if not isinstance(property_name, str) or not property_name.isidentifier(): raise ValueError( - "property_name is required when " - "downstream_task_type='property' and must be a valid " + "property_name is required when downstream_task_type is " + "'property' or 'group_property' and must be a valid " f"Python identifier; got {property_name!r}." ) if not isinstance(task_dim, int) or task_dim < 1: raise ValueError(f"task_dim must be an int >= 1; got {task_dim!r}.") + if downstream_task_type == "group_property" and group_reduce not in ( + "mean", + "sum", + ): + raise ValueError( + f"group_reduce must be 'mean' or 'sum'; got {group_reduce!r}." + ) validate_fparam_dim(fparam_dim) try: aux_prob = float(aux_prob) @@ -168,6 +196,7 @@ def __init__( self.property_name = property_name self.task_dim = task_dim self.intensive = intensive + self.group_reduce = group_reduce self.learning_rate = learning_rate self.stop_lr = stop_lr self.decay_steps = decay_steps @@ -452,14 +481,13 @@ def fit( @property def _downstream_head(self) -> str: - """Branch/head name of the downstream task. Paper property mode uses - "property" (matching MFTConfigManager); legacy ener mode keeps - "DOWNSTREAM". + """Branch/head name of the downstream task. Paper property/ + group_property modes use the task type itself as the branch key + (matching MFTConfigManager); legacy ener mode keeps "DOWNSTREAM". """ + task_type = getattr(self, "downstream_task_type", "ener") return ( - "property" - if getattr(self, "downstream_task_type", "ener") == "property" - else "DOWNSTREAM" + task_type if task_type in _PROPERTY_LIKE_DOWNSTREAM_TYPES else "DOWNSTREAM" ) def _freeze_ckpt(self) -> str: @@ -549,6 +577,43 @@ def _resolve_test_data(test_data: str | list[str]) -> list[str]: raise RuntimeError(f"test_data {test_data!r} resolved to 0 systems.") return unique + @staticmethod + def _check_no_multi_frame_groups(systems: list[str]) -> None: + """Refuse group_property evaluate()/predict() on genuine multi-frame + groups. + + ``dp --pt test`` (the generic ``DeepEval``/``DeepProperty`` path) + never threads ``group_id``/``weight``/``pool_mask`` through to the + model, so ``GroupPropertyModel.forward`` falls back to treating + every frame as its own one-frame group. That fallback is a correct + no-op when every real group is exactly one frame, but silently wrong + for assemblies spanning multiple frames (e.g. an OER O*/OH*/OOH* + overpotential group) -- the reported MAE/RMSE or per-row predictions + would not match what the model was trained to produce. This is a + limitation of ``dp --pt test`` itself, not specific to MFT. + """ + for sys_path in systems: + for set_dir in sorted(_glob.glob(os.path.join(sys_path, "set.*"))): + group_id_path = os.path.join(set_dir, "group_id.npy") + if not os.path.isfile(group_id_path): + continue + group_ids = np.load(group_id_path).reshape(-1) + _, counts = np.unique(group_ids, return_counts=True) + if np.any(counts > 1): + raise RuntimeError( + "MFT evaluate()/predict() cannot score a " + "group_property head against " + f"{set_dir}: it contains a group spanning more than " + "one frame, but `dp --pt test` evaluates frame-by-" + "frame and silently treats every frame as its own " + "group, so the reported numbers would not match what " + "the model was trained on. This is a known gap in " + "`dp --pt test`'s group_property support (not " + "specific to MFT); scoring multi-frame assemblies " + "requires a group-aware evaluator, which does not " + "exist yet." + ) + def evaluate(self, test_data: str | list[str]) -> dict: """ Evaluate the downstream head of the MFT checkpoint via ``dp --pt test``. @@ -587,10 +652,19 @@ def evaluate(self, test_data: str | list[str]) -> dict: The DeePMD-kit output labels the unit as ``eV`` regardless of the actual training units; callers using Hartree-trained checkpoints should treat the returned numbers as Hartree. - """ - frozen_path = self._freeze_ckpt() + For ``downstream_task_type="group_property"``, ``dp --pt test`` never + threads ``group_id``/``weight``/``pool_mask`` through its evaluation + path, so it silently scores every test frame as its own one-frame + group. This is a no-op for single-frame groups but wrong for genuine + multi-frame assemblies; ``evaluate()`` refuses to run against test + data containing one (see ``_check_no_multi_frame_groups``). + """ systems = self._resolve_test_data(test_data) + if self._downstream_head == "group_property": + self._check_no_multi_frame_groups(systems) + + frozen_path = self._freeze_ckpt() os.makedirs(self.output_dir, exist_ok=True) datafile = os.path.join(self.output_dir, "test_systems.txt") @@ -628,14 +702,18 @@ def predict(self, test_data: str | list[str]) -> DotDict: ``-d`` to ``dp --pt test`` and parses the generated property detail files so callers get frame-level labels and predictions. """ - if self._downstream_head != "property": + if self._downstream_head not in _PROPERTY_LIKE_DOWNSTREAM_TYPES: raise RuntimeError( - "MFT predict() is only supported for downstream_task_type='property'. " - "Energy-mode MFT can still use evaluate() for aggregate metrics." + "MFT predict() is only supported for downstream_task_type " + "'property' or 'group_property'. Energy-mode MFT can still " + "use evaluate() for aggregate metrics." ) - frozen_path = self._freeze_ckpt() systems = self._resolve_test_data(test_data) + if self._downstream_head == "group_property": + self._check_no_multi_frame_groups(systems) + + frozen_path = self._freeze_ckpt() os.makedirs(self.output_dir, exist_ok=True) datafile = os.path.join(self.output_dir, "predict_systems.txt") diff --git a/dpa_adapt/predictor.py b/dpa_adapt/predictor.py index 970bef7088..975204eb31 100644 --- a/dpa_adapt/predictor.py +++ b/dpa_adapt/predictor.py @@ -131,6 +131,7 @@ def __init__(self, model_path: str, n_committee: int = 1) -> None: self._pooling = bundle["pooling"] self._condition_manager = bundle.get("condition_manager") self._fparam_dim = bundle.get("fparam_dim", 0) + self._calibrator = bundle.get("calibrator") self.n_committee = n_committee # Detect estimator type from the final pipeline step. @@ -265,6 +266,8 @@ def predict( data: str | list[str], fmt: str | None = None, return_uncertainty: bool = False, + *, + calibrated: bool = True, ) -> DotDict: """ Run inference on ``data``. @@ -292,16 +295,66 @@ def predict( features = self._extract_and_condition(data, fmt) if return_uncertainty: - return self._predict_with_uncertainty(features) + result = self._predict_with_uncertainty(features) + return self._apply_calibrator( + result, calibrated=calibrated, data=data, fmt=fmt + ) if self.n_committee > 1: preds = np.array([e.predict(features) for e in self.estimators_]) preds = preds.reshape(self.n_committee, -1, self._task_dim) - return DotDict({"predictions": np.mean(preds, axis=0)}) + return self._apply_calibrator( + DotDict({"predictions": np.mean(preds, axis=0)}), + calibrated=calibrated, + data=data, + fmt=fmt, + ) raw = self._predictor.predict(features) predictions = np.asarray(raw).reshape(-1, self._task_dim) - return DotDict({"predictions": predictions}) + return self._apply_calibrator( + DotDict({"predictions": predictions}), + calibrated=calibrated, + data=data, + fmt=fmt, + ) + + def _apply_calibrator( + self, + result: DotDict, + *, + calibrated: bool, + data: str | list[str] | None = None, + fmt: str | None = None, + ) -> DotDict: + if calibrated and self._calibrator is not None: + raw_predictions = np.asarray(result.predictions) + result["raw_predictions"] = raw_predictions + result["predictions"] = self._calibrator.predict_from_arrays( + raw_predictions, + data=data, + fmt=fmt, + ) + return result + + def calibrate( + self, + data: str | list[str], + *, + method: str = "ridge", + alpha: float = 1.0, + features: list[str] | tuple[str, ...] = ("prediction",), + fmt: str | None = None, + ) -> object: + """Fit and attach a post-training prediction calibrator.""" + from dpa_adapt.calibrator import ( + Calibrator, + ) + + calibrator = Calibrator(method=method, alpha=alpha, features=features) + calibrator.fit(self, data=data, fmt=fmt) + self._calibrator = calibrator + return calibrator def _predict_with_uncertainty(self, features: np.ndarray) -> DotDict: """Per-estimator uncertainty dispatch.""" @@ -343,7 +396,13 @@ def _predict_with_uncertainty(self, features: np.ndarray) -> DotDict: f"with n_committee={self.n_committee}." ) - def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: + def evaluate( + self, + data: str | list[str], + fmt: str | None = None, + *, + calibrated: bool = True, + ) -> DotDict: """ Predict on ``data`` and compute evaluation metrics against stored labels. @@ -368,7 +427,7 @@ def evaluate(self, data: str | list[str], fmt: str | None = None) -> DotDict: _load_labels, ) - result = self.predict(data, fmt=fmt) + result = self.predict(data, fmt=fmt, calibrated=calibrated) predictions = result.predictions systems = load_data(data, fmt=fmt) diff --git a/dpa_adapt/regularizer.py b/dpa_adapt/regularizer.py new file mode 100644 index 0000000000..57633459e0 --- /dev/null +++ b/dpa_adapt/regularizer.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Training-time regularizer configuration for DPA-ADAPT.""" + +from __future__ import annotations + +from collections.abc import ( + Mapping, +) +from dataclasses import ( + dataclass, +) +from typing import ( + Any, +) + + +@dataclass(frozen=True) +class Regularizer: + """Extra downstream-batch losses added during training. + + ``Regularizer`` deliberately does not model MFT auxiliary data. External + auxiliary tasks belong to ``strategy="mft"``; this object is reserved for + losses on the current downstream training path. + """ + + descriptor_anchor: float = 0.0 + + @classmethod + def from_config(cls, value: Regularizer | Mapping[str, Any] | None) -> Regularizer: + if value is None: + return cls() + if isinstance(value, cls): + return value + if not isinstance(value, Mapping): + raise TypeError( + "regularizer must be a Regularizer, a dict, or None; " + f"got {type(value).__name__}." + ) + unsupported = sorted(set(value) - {"descriptor_anchor"}) + if unsupported: + raise ValueError( + "Regularizer does not accept MFT-style auxiliary fields or " + f"unknown options: {unsupported}. Use strategy='mft' for aux_data." + ) + return cls(descriptor_anchor=float(value.get("descriptor_anchor", 0.0))) + + def __post_init__(self) -> None: + if self.descriptor_anchor < 0.0: + raise ValueError( + f"descriptor_anchor must be non-negative; got {self.descriptor_anchor}." + ) + + @property + def enabled(self) -> bool: + return self.descriptor_anchor > 0.0 + + def require_supported_backend(self) -> None: + """Fail loudly until a backend implements these training-time losses.""" + if self.enabled: + raise NotImplementedError( + "Regularizer(descriptor_anchor=...) is part of the public API " + "contract, but the dp --pt training backend is not wired yet. " + "Use strategy='mft' for auxiliary-task regularization, or leave " + "regularizer unset until descriptor-anchor loss is implemented." + ) 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/examples/group_property/README.md b/examples/group_property/README.md new file mode 100644 index 0000000000..21b4fd843c --- /dev/null +++ b/examples/group_property/README.md @@ -0,0 +1,53 @@ +# Group Property Example + +This example demonstrates the `group_property` architecture. One supervised +sample is a group of component structures rather than a single frame. DeePMD +embeds each component, masks padded atoms with `pool_mask.npy`, weights +components with `weight.npy`, and predicts one group-level `target_property`. + +The bundled dataset is intentionally tiny: two training groups and one +validation group generated from three real rows of a public component-based +molecular dataset. + +## Directory layout + +```text +examples/group_property/ +|-- data/ +| |-- grouped_assembly_subset.csv +| |-- train/ # 2 grouped training systems +| |-- valid/ # 1 grouped validation system +| |-- train_systems.txt +| `-- valid_systems.txt +|-- scripts/ +| |-- prepare_data.py # regenerate data/train and data/valid +| `-- run_dpa_adapt_api.py # dpa-adapt API route +|-- train/ +| `-- input_torch.json # DeePMD group_property training input +`-- README.md +``` + +## Run + +To run the DeePMD training example: + +```bash +cd train +dp --pt train input_torch.json +``` + +To run the dpa-adapt API route: + +```bash +python scripts/run_dpa_adapt_api.py +``` + +To regenerate the grouped DeePMD data from the CSV subset: + +```bash +python scripts/prepare_data.py +``` + +Set `DPA_GROUP_PROPERTY_CSV=/path/to/dataset.csv` and optionally +`DPA_GROUP_PROPERTY_ROWS` / `DPA_GROUP_PROPERTY_VALID` to use a larger compatible +CSV split. diff --git a/examples/group_property/data/grouped_assembly_subset.csv b/examples/group_property/data/grouped_assembly_subset.csv new file mode 100644 index 0000000000..313b7c4444 --- /dev/null +++ b/examples/group_property/data/grouped_assembly_subset.csv @@ -0,0 +1,4 @@ +reference;polymer_type;polymer_type_style;polymer_architecture;polymerisation_type;SMILES_start_group;SMILES_end_group;SMILES_repeating_unitA;molpercent_repeating_unitA;SMILES_repeating_unitB;molpercent_repeating_unitB;SMILES_repeating_unitC;molpercent_repeating_unitC;SMILES_repeating_unitD;molpercent_repeating_unitD;SMILES_repeating_unitE;molpercent_repeating_unitE;Mn;Mw;PDI;mass_characterisation_method;mass_characterisation_standart;polymer_concentration_mass_conc;polymer_concentration_wpercent;additive1;additive1_concentration_molar;additive1_concentration_weight_percent;additive2;additive2_concentration_molar;additive2_concentration_weight_percent;target_property;N/A;def_type;pH;identifier;comment;tacticity;rating +10.1016/j.fuel.2016.10.075;random;;linear;FRP;[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;20700;1,79;SEC(DMAc);PMMA;;0,01;;;;;;;32,1;;A;;PNIPAM-co-AA;;; +10.1016/j.fuel.2016.10.075;random;;linear;FRP;[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,125;[CH2][CH](C(=O)NC1CC1);0,075;;;;;11700;22700;1,94;SEC(DMAc);PMMA;;0,01;;;;;;;43;;A;;PNIPAM-cycloprop-7.5;;; +10.1016/j.fuel.2016.10.075;random;;linear;FRP;[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,125;[CH2][CH](C(=O)NC1CCC1);0,075;;;;;12000;23100;1,93;SEC(DMAc);PMMA;;0,01;;;;;;;44;;A;;PNIPAM-cyclobut-7.5;;; diff --git a/examples/group_property/data/train/fparam_scaler.json b/examples/group_property/data/train/fparam_scaler.json new file mode 100644 index 0000000000..5433794cdb --- /dev/null +++ b/examples/group_property/data/train/fparam_scaler.json @@ -0,0 +1,43 @@ +{ + "columns": [ + "mn_log", + "conc", + "mw_log", + "pH", + "pdi", + "cat:def_type=A", + "cat:mass_characterisation_method=SEC(DMAc)", + "cat:mass_characterisation_standart=PMMA", + "cat:polymer_architecture=linear", + "cat:polymer_type=random", + "cat:polymerisation_type=FRP" + ], + "stats": { + "mean": [ + 4.064441851049887, + 0.01, + 4.3359981013250195, + 7.0, + 1.865, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0 + ], + "std": [ + 0.0037440106962747244, + 1.0, + 0.02002775586810257, + 1.0, + 0.07499999999999996, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0 + ] + } +} diff --git a/examples/group_property/data/train/manifest.json b/examples/group_property/data/train/manifest.json new file mode 100644 index 0000000000..c670d7a00a --- /dev/null +++ b/examples/group_property/data/train/manifest.json @@ -0,0 +1,168 @@ +{ + "fparam_schema": [ + { + "default": 0.0, + "name": "mn_log" + }, + { + "default": 0.0, + "name": "conc" + }, + { + "default": 0.0, + "name": "mw_log" + }, + { + "default": 0.0, + "name": "pH" + }, + { + "default": 0.0, + "name": "pdi" + }, + { + "default": 0.0, + "name": "cat:def_type=A" + }, + { + "default": 0.0, + "name": "cat:mass_characterisation_method=SEC(DMAc)" + }, + { + "default": 0.0, + "name": "cat:mass_characterisation_standart=PMMA" + }, + { + "default": 0.0, + "name": "cat:polymer_architecture=linear" + }, + { + "default": 0.0, + "name": "cat:polymer_type=random" + }, + { + "default": 0.0, + "name": "cat:polymerisation_type=FRP" + } + ], + "groups": [ + { + "components": [ + { + "block": "rep", + "frame": 0, + "pool_mask_excluded": [], + "weight": 0.8 + }, + { + "block": "rep", + "frame": 1, + "pool_mask_excluded": [], + "weight": 0.2 + }, + { + "block": "end", + "frame": 2, + "pool_mask_excluded": [], + "weight": 0.008938430581571325 + }, + { + "block": "end", + "frame": 3, + "pool_mask_excluded": [], + "weight": 0.008938430581571325 + } + ], + "fparam": { + "cat:def_type=A": 0.0, + "cat:mass_characterisation_method=SEC(DMAc)": 0.0, + "cat:mass_characterisation_standart=PMMA": 0.0, + "cat:polymer_architecture=linear": 0.0, + "cat:polymer_type=random": 0.0, + "cat:polymerisation_type=FRP": 0.0, + "conc": 0.0, + "mn_log": -1.0000000000001186, + "mw_log": -0.9999999999999778, + "pH": 0.0, + "pdi": -1.0 + }, + "group_id": 0, + "key": "row_0", + "label": [ + 32.1 + ], + "metadata": {}, + "system": "systems/row_0" + }, + { + "components": [ + { + "block": "rep", + "frame": 0, + "pool_mask_excluded": [], + "weight": 0.8 + }, + { + "block": "rep", + "frame": 1, + "pool_mask_excluded": [], + "weight": 0.125 + }, + { + "block": "rep", + "frame": 2, + "pool_mask_excluded": [], + "weight": 0.075 + }, + { + "block": "end", + "frame": 3, + "pool_mask_excluded": [], + "weight": 0.009122953842692424 + }, + { + "block": "end", + "frame": 4, + "pool_mask_excluded": [], + "weight": 0.009122953842692424 + } + ], + "fparam": { + "cat:def_type=A": 0.0, + "cat:mass_characterisation_method=SEC(DMAc)": 0.0, + "cat:mass_characterisation_standart=PMMA": 0.0, + "cat:polymer_architecture=linear": 0.0, + "cat:polymer_type=random": 0.0, + "cat:polymerisation_type=FRP": 0.0, + "conc": 0.0, + "mn_log": 0.9999999999998814, + "mw_log": 1.0000000000000222, + "pH": 0.0, + "pdi": 1.0 + }, + "group_id": 1, + "key": "row_1", + "label": [ + 43.0 + ], + "metadata": {}, + "system": "systems/row_1" + } + ], + "property_name": "target_property", + "schema": "dpa_adapt.assembly.v1", + "system_dir": "systems", + "tensor_fields": { + "fparam": "fparam", + "group_id": "group_id", + "label": "target_property", + "pool_mask": "pool_mask", + "weight": "weight" + }, + "type_map": [ + "C", + "O", + "N", + "H" + ] +} diff --git a/examples/group_property/data/train/systems/row_0/set.000/box.npy b/examples/group_property/data/train/systems/row_0/set.000/box.npy new file mode 100644 index 0000000000..7fed1617df Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/box.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/coord.npy b/examples/group_property/data/train/systems/row_0/set.000/coord.npy new file mode 100644 index 0000000000..17cfab5509 Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/coord.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/fparam.npy b/examples/group_property/data/train/systems/row_0/set.000/fparam.npy new file mode 100644 index 0000000000..de6177429a Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/fparam.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/group_id.npy b/examples/group_property/data/train/systems/row_0/set.000/group_id.npy new file mode 100644 index 0000000000..87743c0468 Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/group_id.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/pool_mask.npy b/examples/group_property/data/train/systems/row_0/set.000/pool_mask.npy new file mode 100644 index 0000000000..5441333e9f Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/pool_mask.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/real_atom_types.npy b/examples/group_property/data/train/systems/row_0/set.000/real_atom_types.npy new file mode 100644 index 0000000000..94c513f1de Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/real_atom_types.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/target_property.npy b/examples/group_property/data/train/systems/row_0/set.000/target_property.npy new file mode 100644 index 0000000000..8a3ee682c5 Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/target_property.npy differ diff --git a/examples/group_property/data/train/systems/row_0/set.000/weight.npy b/examples/group_property/data/train/systems/row_0/set.000/weight.npy new file mode 100644 index 0000000000..f628c25e69 Binary files /dev/null and b/examples/group_property/data/train/systems/row_0/set.000/weight.npy differ diff --git a/examples/group_property/data/train/systems/row_0/type.raw b/examples/group_property/data/train/systems/row_0/type.raw new file mode 100644 index 0000000000..6d2f65f0ff --- /dev/null +++ b/examples/group_property/data/train/systems/row_0/type.raw @@ -0,0 +1,19 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 diff --git a/examples/group_property/data/train/systems/row_0/type_map.raw b/examples/group_property/data/train/systems/row_0/type_map.raw new file mode 100644 index 0000000000..c6ee287c17 --- /dev/null +++ b/examples/group_property/data/train/systems/row_0/type_map.raw @@ -0,0 +1,4 @@ +C +O +N +H diff --git a/examples/group_property/data/train/systems/row_1/set.000/box.npy b/examples/group_property/data/train/systems/row_1/set.000/box.npy new file mode 100644 index 0000000000..b9e595190b Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/box.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/coord.npy b/examples/group_property/data/train/systems/row_1/set.000/coord.npy new file mode 100644 index 0000000000..89ad7bb3eb Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/coord.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/fparam.npy b/examples/group_property/data/train/systems/row_1/set.000/fparam.npy new file mode 100644 index 0000000000..5f97077153 Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/fparam.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/group_id.npy b/examples/group_property/data/train/systems/row_1/set.000/group_id.npy new file mode 100644 index 0000000000..3b3a388202 Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/group_id.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/pool_mask.npy b/examples/group_property/data/train/systems/row_1/set.000/pool_mask.npy new file mode 100644 index 0000000000..18106c6d7b Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/pool_mask.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/real_atom_types.npy b/examples/group_property/data/train/systems/row_1/set.000/real_atom_types.npy new file mode 100644 index 0000000000..50a21c8e50 Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/real_atom_types.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/target_property.npy b/examples/group_property/data/train/systems/row_1/set.000/target_property.npy new file mode 100644 index 0000000000..c97a61aff4 Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/target_property.npy differ diff --git a/examples/group_property/data/train/systems/row_1/set.000/weight.npy b/examples/group_property/data/train/systems/row_1/set.000/weight.npy new file mode 100644 index 0000000000..26d9906793 Binary files /dev/null and b/examples/group_property/data/train/systems/row_1/set.000/weight.npy differ diff --git a/examples/group_property/data/train/systems/row_1/type.raw b/examples/group_property/data/train/systems/row_1/type.raw new file mode 100644 index 0000000000..6d2f65f0ff --- /dev/null +++ b/examples/group_property/data/train/systems/row_1/type.raw @@ -0,0 +1,19 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 diff --git a/examples/group_property/data/train/systems/row_1/type_map.raw b/examples/group_property/data/train/systems/row_1/type_map.raw new file mode 100644 index 0000000000..c6ee287c17 --- /dev/null +++ b/examples/group_property/data/train/systems/row_1/type_map.raw @@ -0,0 +1,4 @@ +C +O +N +H diff --git a/examples/group_property/data/train_systems.txt b/examples/group_property/data/train_systems.txt new file mode 100644 index 0000000000..e8b2f2bc7c --- /dev/null +++ b/examples/group_property/data/train_systems.txt @@ -0,0 +1,2 @@ +train/systems/row_0 +train/systems/row_1 diff --git a/examples/group_property/data/valid/fparam_scaler.json b/examples/group_property/data/valid/fparam_scaler.json new file mode 100644 index 0000000000..5433794cdb --- /dev/null +++ b/examples/group_property/data/valid/fparam_scaler.json @@ -0,0 +1,43 @@ +{ + "columns": [ + "mn_log", + "conc", + "mw_log", + "pH", + "pdi", + "cat:def_type=A", + "cat:mass_characterisation_method=SEC(DMAc)", + "cat:mass_characterisation_standart=PMMA", + "cat:polymer_architecture=linear", + "cat:polymer_type=random", + "cat:polymerisation_type=FRP" + ], + "stats": { + "mean": [ + 4.064441851049887, + 0.01, + 4.3359981013250195, + 7.0, + 1.865, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0 + ], + "std": [ + 0.0037440106962747244, + 1.0, + 0.02002775586810257, + 1.0, + 0.07499999999999996, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0 + ] + } +} diff --git a/examples/group_property/data/valid/manifest.json b/examples/group_property/data/valid/manifest.json new file mode 100644 index 0000000000..755ddbec4e --- /dev/null +++ b/examples/group_property/data/valid/manifest.json @@ -0,0 +1,120 @@ +{ + "fparam_schema": [ + { + "default": 0.0, + "name": "mn_log" + }, + { + "default": 0.0, + "name": "conc" + }, + { + "default": 0.0, + "name": "mw_log" + }, + { + "default": 0.0, + "name": "pH" + }, + { + "default": 0.0, + "name": "pdi" + }, + { + "default": 0.0, + "name": "cat:def_type=A" + }, + { + "default": 0.0, + "name": "cat:mass_characterisation_method=SEC(DMAc)" + }, + { + "default": 0.0, + "name": "cat:mass_characterisation_standart=PMMA" + }, + { + "default": 0.0, + "name": "cat:polymer_architecture=linear" + }, + { + "default": 0.0, + "name": "cat:polymer_type=random" + }, + { + "default": 0.0, + "name": "cat:polymerisation_type=FRP" + } + ], + "groups": [ + { + "components": [ + { + "block": "rep", + "frame": 0, + "pool_mask_excluded": [], + "weight": 0.8 + }, + { + "block": "rep", + "frame": 1, + "pool_mask_excluded": [], + "weight": 0.125 + }, + { + "block": "rep", + "frame": 2, + "pool_mask_excluded": [], + "weight": 0.075 + }, + { + "block": "end", + "frame": 3, + "pool_mask_excluded": [], + "weight": 0.008963771089468452 + }, + { + "block": "end", + "frame": 4, + "pool_mask_excluded": [], + "weight": 0.008963771089468452 + } + ], + "fparam": { + "cat:def_type=A": 0.0, + "cat:mass_characterisation_method=SEC(DMAc)": 0.0, + "cat:mass_characterisation_standart=PMMA": 0.0, + "cat:polymer_architecture=linear": 0.0, + "cat:polymer_type=random": 0.0, + "cat:polymerisation_type=FRP": 0.0, + "conc": 0.0, + "mn_log": 3.9367929724142536, + "mw_log": 1.378780465918503, + "pH": 0.0, + "pdi": 0.8666666666666665 + }, + "group_id": 0, + "key": "row_2", + "label": [ + 44.0 + ], + "metadata": {}, + "system": "systems/row_2" + } + ], + "property_name": "target_property", + "schema": "dpa_adapt.assembly.v1", + "system_dir": "systems", + "tensor_fields": { + "fparam": "fparam", + "group_id": "group_id", + "label": "target_property", + "pool_mask": "pool_mask", + "weight": "weight" + }, + "type_map": [ + "C", + "O", + "N", + "H" + ] +} diff --git a/examples/group_property/data/valid/systems/row_2/set.000/box.npy b/examples/group_property/data/valid/systems/row_2/set.000/box.npy new file mode 100644 index 0000000000..b9e595190b Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/box.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/coord.npy b/examples/group_property/data/valid/systems/row_2/set.000/coord.npy new file mode 100644 index 0000000000..0a6d221459 Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/coord.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/fparam.npy b/examples/group_property/data/valid/systems/row_2/set.000/fparam.npy new file mode 100644 index 0000000000..4f4ed889f7 Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/fparam.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/group_id.npy b/examples/group_property/data/valid/systems/row_2/set.000/group_id.npy new file mode 100644 index 0000000000..9258843872 Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/group_id.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/pool_mask.npy b/examples/group_property/data/valid/systems/row_2/set.000/pool_mask.npy new file mode 100644 index 0000000000..6c9473da73 Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/pool_mask.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/real_atom_types.npy b/examples/group_property/data/valid/systems/row_2/set.000/real_atom_types.npy new file mode 100644 index 0000000000..f1137dbd89 Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/real_atom_types.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/target_property.npy b/examples/group_property/data/valid/systems/row_2/set.000/target_property.npy new file mode 100644 index 0000000000..6dc56a31ac Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/target_property.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/set.000/weight.npy b/examples/group_property/data/valid/systems/row_2/set.000/weight.npy new file mode 100644 index 0000000000..259fec950c Binary files /dev/null and b/examples/group_property/data/valid/systems/row_2/set.000/weight.npy differ diff --git a/examples/group_property/data/valid/systems/row_2/type.raw b/examples/group_property/data/valid/systems/row_2/type.raw new file mode 100644 index 0000000000..d677b495ec --- /dev/null +++ b/examples/group_property/data/valid/systems/row_2/type.raw @@ -0,0 +1,20 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 diff --git a/examples/group_property/data/valid/systems/row_2/type_map.raw b/examples/group_property/data/valid/systems/row_2/type_map.raw new file mode 100644 index 0000000000..c6ee287c17 --- /dev/null +++ b/examples/group_property/data/valid/systems/row_2/type_map.raw @@ -0,0 +1,4 @@ +C +O +N +H diff --git a/examples/group_property/data/valid_systems.txt b/examples/group_property/data/valid_systems.txt new file mode 100644 index 0000000000..21dd57b30a --- /dev/null +++ b/examples/group_property/data/valid_systems.txt @@ -0,0 +1 @@ +valid/systems/row_2 diff --git a/examples/group_property/scripts/prepare_data.py b/examples/group_property/scripts/prepare_data.py new file mode 100644 index 0000000000..e26f820dd6 --- /dev/null +++ b/examples/group_property/scripts/prepare_data.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Prepare a tiny grouped assembly-property dataset. + +The bundled CSV contains three real rows from a public component-based +molecular dataset. Each row is one supervised sample. Its component SMILES are +converted to structures, stored as DeePMD frames, and pooled through +``group_id.npy``, ``weight.npy``, and ``pool_mask.npy``. +""" + +from __future__ import ( + annotations, +) + +import logging +import os +import shutil +from pathlib import ( + Path, +) + +from dpa_adapt.grouped._polymer import ( + PolymerBuilder, +) + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEMO_DIR = Path(__file__).resolve().parent.parent +DATA_DIR = DEMO_DIR / "data" +DEFAULT_CSV = DATA_DIR / "grouped_assembly_subset.csv" +TARGET = "target_property" + +N_ROWS = int(os.environ.get("DPA_GROUP_PROPERTY_ROWS", "3")) +N_VALID = int(os.environ.get("DPA_GROUP_PROPERTY_VALID", "1")) + + +def _source_csv() -> Path: + env_path = os.environ.get("DPA_GROUP_PROPERTY_CSV") + candidates = [Path(env_path)] if env_path else [] + candidates.append(DEFAULT_CSV) + for path in candidates: + if path and path.is_file(): + return path + raise FileNotFoundError( + "Could not find grouped_assembly_subset.csv. Set DPA_GROUP_PROPERTY_CSV " + f"or place a compatible dataset at {DEFAULT_CSV}." + ) + + +def _clone_with_rows(source: PolymerBuilder, rows: list) -> PolymerBuilder: + builder = PolymerBuilder( + target=source.target, type_map=source.type_map, seed=source.seed + ) + builder._rows = rows + return builder + + +def _write_list(path: Path, systems: list[str], split: str) -> None: + path.write_text( + "".join(f"{split}/{system}\n" for system in systems), encoding="utf-8" + ) + + +def _ensure_text_newline(path: Path) -> None: + if path.is_file(): + path.write_text( + path.read_text(encoding="utf-8").rstrip() + "\n", encoding="utf-8" + ) + + +def _normalize_generated_json(split_dir: Path) -> None: + old = split_dir / "polymer_scaler.json" + scaler = split_dir / "fparam_scaler.json" + if old.is_file(): + old.rename(scaler) + _ensure_text_newline(scaler) + _ensure_text_newline(split_dir / "manifest.json") + + +def main() -> None: + csv_path = _source_csv() + logger.info("reading grouped assembly CSV: %s", csv_path) + builder = PolymerBuilder.from_csv(csv_path, target=TARGET) + rows = builder._rows[:N_ROWS] + if len(rows) <= N_VALID: + raise ValueError(f"Need more than {N_VALID} usable rows; got {len(rows)}.") + + for split in ("train", "valid"): + split_dir = DATA_DIR / split + if split_dir.exists(): + shutil.rmtree(split_dir) + + train_rows = rows[:-N_VALID] + valid_rows = rows[-N_VALID:] + + train = _clone_with_rows(builder, train_rows) + train_res = train.write(DATA_DIR / "train", overwrite=True) + _normalize_generated_json(DATA_DIR / "train") + + valid = _clone_with_rows(builder, valid_rows) + valid.type_map = train_res["type_map"] + valid_res = valid.write( + DATA_DIR / "valid", + overwrite=True, + scaler=train_res["scaler"], + ) + _normalize_generated_json(DATA_DIR / "valid") + + _write_list(DATA_DIR / "train_systems.txt", train_res["systems"], "train") + _write_list(DATA_DIR / "valid_systems.txt", valid_res["systems"], "valid") + + logger.info( + "wrote train groups: %d -> %s", train_res["n_groups"], DATA_DIR / "train" + ) + logger.info( + "wrote valid groups: %d -> %s", valid_res["n_groups"], DATA_DIR / "valid" + ) + logger.info("fparam_dim: %d", train_res["fparam_dim"]) + + +if __name__ == "__main__": + main() diff --git a/examples/group_property/scripts/run_dpa_adapt_api.py b/examples/group_property/scripts/run_dpa_adapt_api.py new file mode 100644 index 0000000000..ca4cc7ff79 --- /dev/null +++ b/examples/group_property/scripts/run_dpa_adapt_api.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Train a lightweight grouped embedding regressor through the dpa-adapt API.""" + +from __future__ import ( + annotations, +) + +import logging +from pathlib import ( + Path, +) + +import prepare_data + +from dpa_adapt import ( + DPAFineTuner, +) + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEMO_DIR = Path(__file__).resolve().parent.parent +DATA = DEMO_DIR / "data" +TARGET = "target_property" + +if not (DATA / "train").is_dir() or not (DATA / "valid").is_dir(): + prepare_data.main() + +model = DPAFineTuner( + pretrained="DPA-3.1-3M", + model_branch="Domains_Drug", + strategy="frozen_sklearn", + predictor="linear", + seed=42, +) +model.fit(train_data=str(DATA / "train" / "systems" / "*"), target_key=TARGET) + +metrics = model.evaluate(data=str(DATA / "valid" / "systems" / "*")) +logger.info("MAE = %.4f", metrics.mae) +logger.info("RMSE = %.4f", metrics.rmse) +logger.info("R2 = %.4f", metrics.r2) diff --git a/examples/group_property/train/input_torch.json b/examples/group_property/train/input_torch.json new file mode 100644 index 0000000000..70cf7ba0c6 --- /dev/null +++ b/examples/group_property/train/input_torch.json @@ -0,0 +1,85 @@ +{ + "_comment": "Grouped assembly-property training example. Component frames in each system are pooled by group_id/weight/pool_mask and predict target_property.", + "model": { + "type_map": [ + "C", + "O", + "N", + "H" + ], + "descriptor": { + "type": "dpa1", + "sel": 120, + "rcut_smth": 0.5, + "rcut": 6.0, + "neuron": [ + 25, + 50, + 100 + ], + "tebd_dim": 8, + "axis_neuron": 16, + "type_one_side": true, + "attn": 128, + "attn_layer": 0, + "attn_dotr": true, + "attn_mask": false, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": true, + "temperature": 1.0 + }, + "fitting_net": { + "type": "group_property", + "property_name": "target_property", + "task_dim": 1, + "numb_fparam": 11, + "group_reduce": "mean", + "neuron": [ + 240, + 240, + 240 + ], + "activation_function": "gelu", + "precision": "float32", + "seed": 1 + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 1000, + "start_lr": 0.0002, + "stop_lr": 3.51e-08 + }, + "loss": { + "type": "group_property", + "metric": [ + "mae", + "rmse" + ], + "loss_func": "mse", + "beta": 1.0 + }, + "training": { + "training_data": { + "systems": [ + "../data/train/systems/row_0", + "../data/train/systems/row_1" + ], + "batch_size": 1 + }, + "validation_data": { + "systems": [ + "../data/valid/systems/row_2" + ], + "batch_size": 1, + "numb_batch": 1 + }, + "numb_steps": 1000, + "gradient_max_norm": 5.0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 500 + } +} 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..c19dccc52f --- /dev/null +++ b/source/tests/common/test_argcheck_group_property.py @@ -0,0 +1,151 @@ +# 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``, ``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. + +``dim_case_embd`` was originally removed for the same reason, but +GroupPropertyFittingNet later gained real support for it (a one-hot branch +identifier, concatenated onto the group embedding) so that a group_property +head can share a descriptor with another fitting net in multi-task training +(e.g. MFT). It is back in the schema and is exercised below instead of +being asserted-absent. +""" + +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", + "resnet_dt", + "intensive", + "distinguish_types", + ): + assert unsupported not in names + + +def test_dim_case_embd_argument_is_declared(): + names = [arg.name for arg in fitting_group_property()] + assert "dim_case_embd" in names + + +def test_dim_case_embd_passes_strict_validation(): + out = _normalize_and_check( + {"type": "group_property", "property_name": "y", "dim_case_embd": 31} + ) + assert out["dim_case_embd"] == 31 + + +def test_dim_case_embd_defaults_to_zero_when_omitted(): + out = _normalize_and_check({"type": "group_property", "property_name": "y"}) + assert out["dim_case_embd"] == 0 + + +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..3461fcb2b9 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. + _make_system(tmp_path, "s1", natoms=3, nframes=2) + _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 d332153060..0bc5f74c6c 100644 --- a/source/tests/dpa_adapt/test_finetuner_strategies.py +++ b/source/tests/dpa_adapt/test_finetuner_strategies.py @@ -500,3 +500,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..884f47e86a --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_convert.py @@ -0,0 +1,166 @@ +# 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_full_group_marker_contract(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 and r.wrote_weight + + 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] + + weight = np.load(set_dir / "weight.npy") + assert weight.dtype == np.float64 + assert weight.shape == (3,) + assert weight.tolist() == [1.0, 1.0, 1.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 the full grouped marker contract + for gid in tmp_path.rglob("group_id.npy"): + set_dir = gid.parent + assert np.load(gid).tolist() == [0, 0, 0] + assert (set_dir / "weight.npy").is_file() + assert (set_dir / "pool_mask.npy").is_file() + + +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_writes_all_one_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 r.wrote_pool_mask + np.testing.assert_array_equal( + np.load(set_dir / "pool_mask.npy"), + np.ones((3, 3), dtype=np.float64), + ) + + +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 + and results[0].wrote_weight + ) + assert not (set_dir / "group_id.npy").exists() + assert not (set_dir / "pool_mask.npy").exists() + assert not (set_dir / "weight.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..829d3035fb --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_end_to_end.py @@ -0,0 +1,489 @@ +# 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), + ) + + +def test_group_property_model_set_case_embd_proxies_to_fitting_net() -> None: + """GroupPropertyModel does not go through make_model() (unlike + EnergyModel/PropertyModel), so it does not inherit the generic + model -> atomic_model -> fitting_net set_case_embd chain for free. The + multi-task trainer calls ``model.set_case_embd(idx)`` on every branch + when any branch declares ``dim_case_embd`` (e.g. an MFT run sharing a + descriptor between a group_property head and an ener head); this must + reach the real fitting net, not silently no-op. + """ + torch, GroupPropertyModel, GroupPropertyFittingNet = ( + _require_group_property_backend() + ) + + dim = 4 + fitting = GroupPropertyFittingNet( + ntypes=2, + dim_descrpt=dim, + property_name="y", + task_dim=1, + neuron=[4], + dim_case_embd=3, + ) + + 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"] + ) + + model.set_case_embd(2) + assert torch.equal(fitting.case_embd, torch.eye(3)[2]) + + coord = torch.rand(2, 3, 3, dtype=torch.float64) + atype = torch.tensor([[0, 1, 0], [0, 1, 1]]) + out = model( + coord, + atype, + box=None, + group_id=torch.tensor([[0], [1]]), + weight=torch.tensor([[1.0], [1.0]]), + pool_mask=torch.ones(2, 3), + ) + assert torch.isfinite(out["y"]).all() + assert out["y"].shape[0] == 2 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..5dd752ecc5 --- /dev/null +++ b/source/tests/dpa_adapt/test_grouped_finetuner.py @@ -0,0 +1,131 @@ +# 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 _write_group_markers(set_dir: Path, group_id: int, *, natoms: int = 2) -> None: + np.save(set_dir / "group_id.npy", np.array([group_id, group_id], dtype=np.int64)) + np.save(set_dir / "weight.npy", np.ones(2, dtype=float)) + np.save(set_dir / "pool_mask.npy", np.ones((2, natoms), dtype=float)) + + +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" + _write_group_markers(set_dir, sid) + + 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_training_strategy_target_key_aliases_property_name(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" + _write_group_markers(set_dir, sid) + + model = DPAFineTuner( + pretrained=str(ckpt), + strategy="finetune", + max_steps=20, + output_dir=str(out_dir), + ) + with patch("subprocess.run", side_effect=_mock_dp_train(str(out_dir))): + model.fit( + train_data=systems, + valid_data=valid_systems, + target_key="overpotential", + ) + + cfg = json.loads((out_dir / "input.json").read_text()) + assert model.property_name == "overpotential" + assert cfg["model"]["fitting_net"]["property_name"] == "overpotential" + + +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" + _write_group_markers(set_dir, sid) + # 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_mft_grouped_config.py b/source/tests/dpa_adapt/test_mft_grouped_config.py new file mode 100644 index 0000000000..75772c90d0 --- /dev/null +++ b/source/tests/dpa_adapt/test_mft_grouped_config.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for MFT downstream_task_type='group_property'. + +Mirrors test_mft_property_task.py's coverage of the paper-faithful +downstream_task_type='property' branch, but for the grouped/assembly +variant: a fresh group_property fitting_net + group_property loss for the +DOWNSTREAM head (sharing a descriptor with the aux ener branch, exactly like +property mode), while the aux branch keeps its ener fitting_net pulled from +the ckpt. + +GroupPropertyFittingNet is a small standalone MLP, not built on +GeneralFitting like the property head is, so several property-schema +fields (resnet_dt, intensive, distinguish_types, numb_aparam) don't exist +on it -- dargs strict-mode rejects them outright, so the config builder +must not carry them over from the property path (see +_build_group_property_fitting_net, independent of +_build_property_fitting_net). +""" + +from __future__ import ( + annotations, +) + +import os +from pathlib import ( + Path, +) +from typing import ( + ClassVar, +) + +import numpy as np +import pytest + +from dpa_adapt.config.manager import ( + MFTConfigManager, +) +from dpa_adapt.mft import ( + MFTFineTuner, +) + +DUMMY_TYPE_MAP = ["H", "C", "N", "O"] + + +class _FakeGroupPropertyTuner: + """Tuner-shaped object configured for downstream_task_type='group_property'. + + Bypasses MFTFineTuner.__init__ so tests don't need a real ckpt. + """ + + pretrained = "/share/DPA-3.1-3M.pt" + aux_branch = "SPICE2" + aux_prob = 0.5 + type_map: ClassVar[list[str]] = DUMMY_TYPE_MAP + # aux fitting_net pulled from ckpt — an ener config (the actual SPICE2 + # head). dim_case_embd=31 mirrors what a real DPA-3.1-3M checkpoint's aux + # branch fitting_net carries (it was itself trained as one of 31 branches). + fitting_net_params: ClassVar[dict[str, object]] = { + "type": "ener", + "neuron": [240, 240, 240], + "dim_case_embd": 31, + } + downstream_task_type = "group_property" + property_name = "overpotential" + task_dim = 1 + group_reduce = "sum" + learning_rate = 1e-3 + stop_lr = 1e-5 + max_steps = 1000 + batch_size = "auto:32" + seed = 42 + output_dir = "/tmp/mft_group_property_test" + save_freq = 500 + disp_freq = 100 + train_data = "/data/oer_downstream" + aux_data = "/data/spice2" + valid_data = None + + +def _build() -> dict: + return MFTConfigManager(_FakeGroupPropertyTuner()).build() + + +# --------------------------------------------------------------------------- +# Config building +# --------------------------------------------------------------------------- + + +def test_group_property_task_uses_task_type_as_branch_key(): + config = _build() + assert "group_property" in config["model"]["model_dict"] + assert "SPICE2" in config["model"]["model_dict"] + + +def test_group_property_task_config_has_group_property_fitting_net(): + config = _build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["type"] == "group_property" + assert fn["property_name"] == "overpotential" + assert fn["task_dim"] == 1 + + +def test_group_property_task_group_reduce_propagated(): + config = _build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["group_reduce"] == "sum" + + +def test_group_property_task_group_reduce_defaults_to_mean(): + # A tuner-shaped object that genuinely has no group_reduce attribute + # (rather than mutating the shared _FakeGroupPropertyTuner class), so + # _build_group_property_fitting_net's getattr(t, "group_reduce", "mean") + # exercises its real default. + attrs = { + k: v + for k, v in vars(_FakeGroupPropertyTuner).items() + if not k.startswith("__") and k != "group_reduce" + } + tuner = type("_TunerWithoutGroupReduce", (), attrs)() + config = MFTConfigManager(tuner).build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["group_reduce"] == "mean" + + +def test_group_property_task_config_has_group_property_loss(): + config = _build() + loss = config["loss_dict"]["group_property"] + assert loss["type"] == "group_property" + assert loss["loss_func"] == "mse" + + +def test_group_property_task_dim_case_embd_matches_aux_branch_requirement(): + """The descriptor is shared with the aux branch; deepmd-kit's multi-task + trainer requires every model_dict branch to declare the same + dim_case_embd (deepmd.pt.train.training.get_case_embd_config), so the + group_property head must declare it too, exactly like the property head + does -- read from the aux branch's own fitting_net_params, not hardcoded + (a checkpoint other than DPA-3.1-3M has a different branch count; see + test_dim_case_embd_is_read_from_aux_fitting_net_not_hardcoded below). + """ + config = _build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["dim_case_embd"] == 31 + + +def test_dim_case_embd_is_read_from_aux_fitting_net_not_hardcoded(): + """Regression pin: a checkpoint with a different branch count (e.g. the + 23-branch OMol25/Organic_Reactions/ODAC23 checkpoint used for the + cloud-point OER grouped run) must produce a matching, not hardcoded-31, + dim_case_embd on the group_property downstream head. Hardcoding 31 here + silently mismatches every checkpoint that isn't DPA-3.1-3M. + """ + t = _FakeGroupPropertyTuner() + t.fitting_net_params = { + "type": "ener", + "neuron": [240, 240, 240], + "dim_case_embd": 23, + } + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["dim_case_embd"] == 23 + + +def test_dim_case_embd_omitted_when_aux_fitting_net_has_none(): + """A single-task-pretrained (non-multitask) checkpoint's aux fitting_net + has no dim_case_embd at all; the downstream head must not invent one. + """ + t = _FakeGroupPropertyTuner() + t.fitting_net_params = {"type": "ener", "neuron": [240, 240, 240]} + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert "dim_case_embd" not in fn + + +def test_group_property_task_no_unsupported_fitting_fields(): + """GroupPropertyFittingNet has no wiring for these property-schema + fields; dargs strict mode rejects them outright rather than ignoring + them (deepmd.utils.argcheck.fitting_group_property), so the config + builder must never emit them. + """ + config = _build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + unsupported = {"resnet_dt", "intensive", "distinguish_types", "numb_aparam"} + assert not (unsupported & fn.keys()) + + +def test_group_property_task_finetune_head_random_for_downstream(): + config = _build() + downstream = config["model"]["model_dict"]["group_property"] + assert downstream["finetune_head"] == "RANDOM" + + +def test_group_property_task_aux_branch_keeps_ener_fitting_net(): + config = _build() + aux = config["model"]["model_dict"]["SPICE2"] + assert aux["fitting_net"]["type"] == "ener" + assert aux["finetune_head"] == "SPICE2" + + +def test_group_property_task_aux_branch_keeps_ener_loss(): + config = _build() + assert config["loss_dict"]["SPICE2"]["type"] == "ener" + + +def test_group_property_task_model_prob_splits_by_aux_prob(): + config = _build() + prob = config["training"]["model_prob"] + assert prob["SPICE2"] == pytest.approx(0.5) + assert prob["group_property"] == pytest.approx(0.5) + + +def test_group_property_task_gradient_clipping_matches_paper_alignment(): + config = _build() + assert config["training"]["gradient_max_norm"] == 5.0 + + +def test_group_property_task_descriptor_matches_property_alignment(): + """group_property is a "random downstream head" mode just like property, + so it must follow the same paper-alignment descriptor tweaks (silut + activation, fix_stat_std) rather than the legacy ener-mode descriptor. + """ + config = _build() + desc = config["model"]["shared_dict"]["dpa3_descriptor"] + assert desc["activation_function"] == "silut:3.0" + assert desc["repflow"]["fix_stat_std"] == 0.3 + + +def test_group_property_task_multidim_task_dim(): + t = _FakeGroupPropertyTuner() + t.task_dim = 3 + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["task_dim"] == 3 + + +def test_group_property_task_fparam_dim_propagated(): + t = _FakeGroupPropertyTuner() + t.fparam_dim = 2 + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["group_property"]["fitting_net"] + assert fn["numb_fparam"] == 2 + + +# --------------------------------------------------------------------------- +# MFTFineTuner.__init__ validation +# --------------------------------------------------------------------------- + + +def test_group_property_task_requires_property_name(monkeypatch): + import torch + + monkeypatch.setattr( + torch, + "load", + lambda *a, **kw: { + "model": { + "_extra_state": { + "model_params": { + "model_dict": {"SPICE2": {"fitting_net": {"type": "ener"}}} + } + } + } + }, + ) + with pytest.raises(ValueError, match="property_name"): + MFTFineTuner( + pretrained="/does/not/exist.pt", + aux_branch="SPICE2", + downstream_task_type="group_property", + ) + + +def test_group_property_task_invalid_group_reduce_raises(): + with pytest.raises(ValueError, match="group_reduce"): + MFTFineTuner( + pretrained="/does/not/exist.pt", + aux_branch="SPICE2", + downstream_task_type="group_property", + property_name="overpotential", + group_reduce="max", + ) + + +def test_group_property_task_group_reduce_stored(monkeypatch): + import torch + + monkeypatch.setattr( + torch, + "load", + lambda *a, **kw: { + "model": { + "_extra_state": { + "model_params": { + "model_dict": {"SPICE2": {"fitting_net": {"type": "ener"}}} + } + } + } + }, + ) + t = MFTFineTuner( + pretrained="/does/not/exist.pt", + aux_branch="SPICE2", + downstream_task_type="group_property", + property_name="overpotential", + group_reduce="sum", + ) + assert t.group_reduce == "sum" + + +def test_downstream_head_is_group_property(): + t = MFTFineTuner.__new__(MFTFineTuner) + t.downstream_task_type = "group_property" + assert t._downstream_head == "group_property" + + +# --------------------------------------------------------------------------- +# evaluate()/predict(): multi-frame-group safety guard +# +# `dp --pt test` never threads group_id/weight/pool_mask through, so a frozen +# group_property head silently scores every test frame as its own one-frame +# group. Harmless when every real group is one frame; silently wrong for +# genuine multi-frame assemblies. evaluate()/predict() must refuse the +# latter rather than return numbers that look plausible but aren't. +# --------------------------------------------------------------------------- + + +def _make_group_property_finetuner(tmp_path): + ft = MFTFineTuner.__new__(MFTFineTuner) + ft.pretrained = str(tmp_path / "dummy.pt") + ft.aux_branch = "SPICE2" + ft.aux_prob = 0.5 + ft.type_map = DUMMY_TYPE_MAP + ft.fitting_net_params = {} + ft.downstream_task_type = "group_property" + ft.property_name = "overpotential" + ft.task_dim = 1 + ft.group_reduce = "mean" + ft.learning_rate = 1e-3 + ft.stop_lr = 1e-5 + ft.max_steps = 100 + ft.batch_size = "auto:32" + ft.seed = 42 + ft.output_dir = str(tmp_path / "out") + ft.save_freq = 10 + ft.disp_freq = 10 + ft.train_data = None + ft.aux_data = None + ft.valid_data = None + os.makedirs(ft.output_dir, exist_ok=True) + return ft + + +def _write_group_id(set_dir: Path, group_ids: list[int]) -> None: + set_dir.mkdir(parents=True, exist_ok=True) + np.save(set_dir / "group_id.npy", np.asarray(group_ids, dtype=np.int64)) + + +def test_check_no_multi_frame_groups_passes_when_every_group_is_one_frame(tmp_path): + sys_dir = tmp_path / "sys_ok" + _write_group_id(sys_dir / "set.000", [0, 1, 2]) + ft = _make_group_property_finetuner(tmp_path) + ft._check_no_multi_frame_groups([str(sys_dir)]) # must not raise + + +def test_check_no_multi_frame_groups_passes_when_group_id_absent(tmp_path): + sys_dir = tmp_path / "sys_no_marker" + (sys_dir / "set.000").mkdir(parents=True) + ft = _make_group_property_finetuner(tmp_path) + ft._check_no_multi_frame_groups([str(sys_dir)]) # must not raise + + +def test_check_no_multi_frame_groups_raises_for_genuine_multi_frame_group(tmp_path): + sys_dir = tmp_path / "sys_oer" + # group 0 spans two frames (e.g. O* + OH* of one OER assembly) + _write_group_id(sys_dir / "set.000", [0, 0, 1]) + ft = _make_group_property_finetuner(tmp_path) + with pytest.raises(RuntimeError, match="spanning more than one frame"): + ft._check_no_multi_frame_groups([str(sys_dir)]) + + +def test_evaluate_rejects_multi_frame_groups_before_touching_subprocess(tmp_path): + """The guard must fire before any dp --pt freeze/test subprocess call, + so evaluate() fails fast without a spurious freeze/test invocation. + """ + from unittest.mock import ( + patch, + ) + + ft = _make_group_property_finetuner(tmp_path) + (Path(ft.output_dir) / "model.ckpt-100.pt").write_bytes(b"") + sys_dir = tmp_path / "sys_oer" + _write_group_id(sys_dir / "set.000", [0, 0, 1]) + + with ( + patch("subprocess.run", side_effect=AssertionError("must not be called")), + pytest.raises(RuntimeError, match="spanning more than one frame"), + ): + ft.evaluate(str(sys_dir)) + + +def test_predict_allows_group_property_head(): + t = MFTFineTuner.__new__(MFTFineTuner) + t.downstream_task_type = "ener" + with pytest.raises(RuntimeError, match="'property' or 'group_property'"): + t.predict("some/data") + t.downstream_task_type = "group_property" + # Advances past the head-type guard; next it calls _freeze_ckpt(), which + # fails for an unrelated reason (no output_dir/checkpoint set up here) -- + # proving the group_property RuntimeError above is gone. + with pytest.raises(Exception) as exc_info: + t.predict("some/data") + assert "only supported for downstream_task_type" not in str(exc_info.value) diff --git a/source/tests/dpa_adapt/test_mft_property_task.py b/source/tests/dpa_adapt/test_mft_property_task.py index e5fa5f3045..0f461e68b0 100644 --- a/source/tests/dpa_adapt/test_mft_property_task.py +++ b/source/tests/dpa_adapt/test_mft_property_task.py @@ -38,10 +38,13 @@ class _FakePropertyTuner: aux_branch = "SPICE2" aux_prob = 0.5 type_map: ClassVar[list[str]] = ["H", "C", "N", "O"] - # aux fitting_net pulled from ckpt — an ener config (the actual SPICE2 head) + # aux fitting_net pulled from ckpt — an ener config (the actual SPICE2 + # head). dim_case_embd=31 mirrors what a real DPA-3.1-3M checkpoint's aux + # branch fitting_net carries (it was itself trained as one of 31 branches). fitting_net_params: ClassVar[dict[str, object]] = { "type": "ener", "neuron": [240, 240, 240], + "dim_case_embd": 31, } downstream_task_type = "property" property_name = "homo" @@ -106,10 +109,46 @@ def test_property_task_config_has_property_fitting_net(): assert fn["neuron"] == [240, 240, 240] assert fn["activation_function"] == "tanh" assert fn["seed"] == 42 - # Required for DPA-3.1-3M multi-task case-embedding layer. + # dim_case_embd is read from the aux branch's own (ckpt-derived) + # fitting_net_params, not hardcoded — see + # test_dim_case_embd_is_read_from_aux_fitting_net_not_hardcoded below. assert fn["dim_case_embd"] == 31 +def test_dim_case_embd_is_read_from_aux_fitting_net_not_hardcoded(): + """dim_case_embd must match the aux branch's own pretrained checkpoint + (i.e. the number of branches THAT checkpoint was originally trained + with), not a hardcoded constant. A checkpoint other than DPA-3.1-3M + (e.g. a 23-branch checkpoint) has a different aux dim_case_embd, and + hardcoding 31 would build a downstream head with the wrong case-embedding + width -- get_case_embd_config's all-branches-must-match check would then + reject the config, or (with resuming=True) the mismatched shapes would + fail to load. + """ + t = _FakePropertyTuner() + t.fitting_net_params = { + "type": "ener", + "neuron": [240, 240, 240], + "dim_case_embd": 23, + } + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["property"]["fitting_net"] + assert fn["dim_case_embd"] == 23 + + +def test_dim_case_embd_omitted_when_aux_fitting_net_has_none(): + """A single-task-pretrained (non-multitask) checkpoint's aux fitting_net + has no dim_case_embd at all; the downstream head must not invent one + (which would fail get_case_embd_config's all-branches-match check + against the aux branch's real dim_case_embd of 0). + """ + t = _FakePropertyTuner() + t.fitting_net_params = {"type": "ener", "neuron": [240, 240, 240]} + config = MFTConfigManager(t).build() + fn = config["model"]["model_dict"]["property"]["fitting_net"] + assert "dim_case_embd" not in fn + + def test_property_task_config_has_property_loss(): """DOWNSTREAM loss must be type='property' with mse + mae/rmse metrics.""" config = MFTConfigManager(_FakePropertyTuner()).build() @@ -164,7 +203,11 @@ def test_property_task_aux_branch_keeps_ener_fitting_net(): config = MFTConfigManager(_FakePropertyTuner()).build() aux_fn = config["model"]["model_dict"]["SPICE2"]["fitting_net"] assert aux_fn["type"] == "ener" - assert aux_fn == {"type": "ener", "neuron": [240, 240, 240]} + assert aux_fn == { + "type": "ener", + "neuron": [240, 240, 240], + "dim_case_embd": 31, + } def test_property_task_aux_branch_keeps_ener_loss(): 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..4416386ca1 --- /dev/null +++ b/source/tests/dpa_adapt/test_polymer_builder.py @@ -0,0 +1,158 @@ +# 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_write_with_scaler_aligns_different_fparam_schema(tmp_path): + train = PolymerBuilder(target="cloud_point") + train.add( + units={"[CH2][CH](C(=O)NC(C)C)": 1.0}, + mol_weight=12000, + fparam={"conc": 0.01, "pH": 7.0, "salts": {"NaCl": 0.1}}, + target=32.0, + ) + train.add( + units={"[CH2][CH](C(=O)NC(C)C)": 1.0}, + mol_weight=24000, + fparam={"conc": 0.03, "pH": 8.0, "salts": {"NaCl": 0.2}}, + target=40.0, + ) + res_train = train.write(tmp_path / "train") + + aux = PolymerBuilder(target="cloud_point") + aux.add( + units={"[CH2][CH](C(=O)NC(C)C)": 1.0}, + mol_weight=18000, + # ``mw_log`` is extra and ``salt:NaCl`` is missing; the training scaler + # must still define the written fparam width and column order. + fparam={"mw_log": 4.4, "pH": 7.5}, + target=50.0, + ) + res_aux = aux.write(tmp_path / "aux", scaler=res_train["scaler"]) + + assert res_aux["fparam_dim"] == res_train["fparam_dim"] + assert res_aux["scaler"]["columns"] == res_train["scaler"]["columns"] + assert res_aux["scaler"]["columns"] == ["mn_log", "conc", "pH", "salt:NaCl"] + + set_dir = tmp_path / "aux" / res_aux["systems"][0] / "set.000" + fparam = np.load(set_dir / "fparam.npy") + assert fparam.shape[1] == 4 + + manifest = (tmp_path / "aux" / "manifest.json").read_text(encoding="utf-8") + assert '"name": "salt:NaCl"' in manifest + + +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/dpa_adapt/test_regularizer_calibrator.py b/source/tests/dpa_adapt/test_regularizer_calibrator.py new file mode 100644 index 0000000000..69fd13288a --- /dev/null +++ b/source/tests/dpa_adapt/test_regularizer_calibrator.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) +from unittest.mock import ( + patch, +) + +if TYPE_CHECKING: + from pathlib import ( + Path, + ) + +import numpy as np +import pytest + +from dpa_adapt import ( + Calibrator, + DPAFineTuner, + DPAPredictor, + Regularizer, +) + +from source.tests.dpa_adapt.test_predictor import ( + _make_npy_system, + _mock_extract_features, + _mock_load_descriptor_model, +) + + +def test_regularizer_rejects_mft_aux_fields() -> None: + with pytest.raises(ValueError, match="strategy='mft'"): + Regularizer.from_config({"aux_data": "mp_traj/dpdata"}) + + +def test_regularizer_nonzero_fails_loudly_until_backend_exists(tmp_path: Path) -> None: + system = tmp_path / "sys" + system.mkdir() + _make_npy_system(system) + model = DPAFineTuner( + pretrained="fake.pt", + predictor="linear", + regularizer={"descriptor_anchor": 0.01}, + ) + + with pytest.raises(NotImplementedError, match="descriptor-anchor"): + model.fit(str(system), target_key="energy") + + +def test_calibrator_prediction_only_ridge() -> None: + calibrator = Calibrator(method="ridge", alpha=0.0) + raw = np.array([[1.0], [2.0], [3.0], [4.0]]) + labels = 2.0 * raw + 1.0 + + calibrator.fit_from_arrays(raw, labels) + pred = calibrator.predict_from_arrays(np.array([[5.0], [6.0]])) + + np.testing.assert_allclose(pred, np.array([[11.0], [13.0]]), atol=1e-10) + + +def test_calibrator_uses_fparam_and_group_stats(tmp_path: Path) -> None: + system = tmp_path / "sys" + system.mkdir() + _make_npy_system(system, n_frames=4) + fparam = np.array([[0.0, 1.0], [1.0, 1.0], [2.0, 1.0], [3.0, 1.0]]) + np.save(system / "set.000" / "fparam.npy", fparam) + + raw = np.array([[1.0], [1.5], [2.0], [2.5]]) + labels = raw + fparam[:, :1] * 0.5 + 2.0 + + calibrator = Calibrator( + method="ridge", + alpha=0.0, + features=["prediction", "fparam", "group_stats"], + ) + calibrator.fit_from_arrays(raw, labels, data=str(system)) + pred = calibrator.predict_from_arrays(raw, data=str(system)) + + assert "fparam_0" in calibrator.feature_names_ + assert "weight_eff_n" in calibrator.feature_names_ + np.testing.assert_allclose(pred, labels, atol=1e-10) + + +def test_calibrator_aligns_group_stats_to_group_predictions(tmp_path: Path) -> None: + system = tmp_path / "grouped" + system.mkdir() + _make_npy_system(system, n_frames=4) + set_dir = system / "set.000" + np.save(set_dir / "group_id.npy", np.array([0, 0, 1, 1])) + np.save(set_dir / "weight.npy", np.array([0.25, 0.75, 1.0, 3.0])) + np.save(set_dir / "pool_mask.npy", np.ones((4, 2))) + np.save( + set_dir / "fparam.npy", + np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]), + ) + + raw = np.array([[2.0], [4.0]]) + labels = np.array([[10.0], [20.0]]) + calibrator = Calibrator( + method="ridge", + alpha=0.0, + features=["prediction", "fparam", "group_stats"], + ) + + calibrator.fit_from_arrays(raw, labels, data=str(system)) + pred = calibrator.predict_from_arrays(raw, data=str(system)) + + assert pred.shape == (2, 1) + assert "n_frames" in calibrator.feature_names_ + assert "weight_l2" in calibrator.feature_names_ + + +def test_finetuner_calibrate_applies_and_freezes(tmp_path: Path) -> None: + system = tmp_path / "sys" + system.mkdir() + _make_npy_system(system, n_frames=5) + + with ( + patch.object( + DPAFineTuner, "_load_descriptor_model", _mock_load_descriptor_model + ), + patch.object(DPAFineTuner, "_extract_features", _mock_extract_features), + ): + model = DPAFineTuner(pretrained="fake.pt", predictor="linear") + model.fit(str(system), target_key="energy") + raw = model.predict(str(system), calibrated=False).predictions + model.calibrate(str(system), method="ridge", alpha=0.0) + calibrated = model.predict(str(system)).predictions + + labels = np.arange(5).reshape(-1, 1) + raw_rmse = float(np.sqrt(np.mean((raw - labels) ** 2))) + calibrated_rmse = float(np.sqrt(np.mean((calibrated - labels) ** 2))) + + assert model.calibrator is not None + assert calibrated_rmse <= raw_rmse + np.testing.assert_allclose( + model.predict(str(system)).raw_predictions, + raw, + ) + + frozen = model.freeze(str(tmp_path / "model.pth")) + pred = DPAPredictor(frozen) + loaded = pred.predict(str(system)) + + assert hasattr(loaded, "raw_predictions") + assert loaded.predictions.shape == (5, 1) 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..0e0ce67714 --- /dev/null +++ b/source/tests/pt/test_group_property_fitting_net.py @@ -0,0 +1,157 @@ +# 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", + "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()) + + +def test_fparam_branch_encodes_side_features_before_fusion(): + fn = _make(numb_fparam=3, fparam_neuron=[5], neuron=[7]) + assert fn.fparam_neuron == [5] + assert isinstance(fn.fparam_network[0], torch.nn.Linear) + assert fn.fparam_network[0].in_features == 3 + assert fn.fparam_network[0].out_features == 5 + assert fn.network[0].in_features == 4 + 5 + out = fn(torch.zeros(2, 4 + 3)) + assert out.shape == (2, 1) + + +def test_fparam_branch_requires_fparam_columns(): + with pytest.raises(ValueError, match="fparam_neuron"): + _make(numb_fparam=0, fparam_neuron=[5]) + + +def test_dim_case_embd_defaults_to_disabled(): + fn = _make() + assert fn.dim_case_embd == 0 + assert fn.case_embd is None + assert fn.network[0].in_features == 4 # dim_descrpt only + + +def test_dim_case_embd_widens_first_layer_and_inits_zero(): + fn = _make(dim_case_embd=5) + assert fn.network[0].in_features == 4 + 5 # dim_descrpt + dim_case_embd + assert fn.case_embd.shape == (5,) + assert torch.equal(fn.case_embd, torch.zeros(5)) + + +def test_set_case_embd_produces_one_hot_row(): + fn = _make(dim_case_embd=4) + fn.set_case_embd(2) + assert torch.equal(fn.case_embd, torch.eye(4)[2]) + + +def test_set_case_embd_changes_forward_output(): + fn = _make(dim_case_embd=3, seed=42) + group_embedding = torch.rand(2, 4) + out_before = fn(group_embedding).detach().clone() + fn.set_case_embd(1) + out_after = fn(group_embedding).detach().clone() + assert not torch.allclose(out_before, out_after) + + +def test_dim_case_embd_combines_with_fparam_neuron(): + fn = _make(numb_fparam=3, fparam_neuron=[5], dim_case_embd=4, neuron=[7]) + # dim_descrpt(4) + fparam_neuron out(5) + dim_case_embd(4) + assert fn.network[0].in_features == 4 + 5 + 4 + fn.set_case_embd(0) + out = fn(torch.zeros(2, 4 + 3)) + assert out.shape == (2, 1) + + +def test_serialize_round_trips_dim_case_embd(): + fn = _make(dim_case_embd=6) + assert fn.serialize()["dim_case_embd"] == 6 + assert _make().serialize()["dim_case_embd"] == 0 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}