【Hackathon 10th Spring No.7】wD-MPNN模型复现

.pre-commit-config.yaml

@@ -16,6 +16,7 @@ repos:
     rev: "v0.0.272"
     hooks:
       - id: ruff
+        args: ["--line-length=120"]
         exclude: '(jointContribution|legacy)/.*'
 
   - repo: https://github.com/pre-commit/pre-commit-hooks

ppmat/datasets/__init__.py

@@ -30,6 +30,7 @@
 from paddle.io import DistributedBatchSampler  # noqa
 
 from ppmat.datasets import collate_fn
+from ppmat.datasets.density_dataset import DensityDataset
 from ppmat.datasets.high_level_water_dataset import HighLevelWaterDataset
 from ppmat.datasets.jarvis_dataset import JarvisDataset
 from ppmat.datasets.matbench_dataset import MatbenchDataset
@@ -41,14 +42,14 @@
 from ppmat.datasets.mptrj_dataset import MPTrjDataset
 from ppmat.datasets.msd_nmr_dataset import MSDnmrDataset
 from ppmat.datasets.msd_nmr_dataset import MSDnmrinfos
-from ppmat.datasets.density_dataset import DensityDataset
-from ppmat.datasets.small_density_dataset import SmallDensityDataset 
 from ppmat.datasets.num_atom_crystal_dataset import NumAtomsCrystalDataset
 from ppmat.datasets.oc20_s2ef_dataset import OC20S2EFDataset  # noqa
-from ppmat.datasets.qm9_dataset import QM9Dataset # noqa
 from ppmat.datasets.omol25_dataset import OMol25Dataset
+from ppmat.datasets.qm9_dataset import QM9Dataset  # noqa
+from ppmat.datasets.small_density_dataset import SmallDensityDataset
 from ppmat.datasets.split_mptrj_data import none_to_zero
 from ppmat.datasets.transform import build_transforms
+from ppmat.datasets.wd_mpnn_dataset import WDMPNNDataset
 from ppmat.utils import logger
 
 __all__ = [
@@ -64,9 +65,10 @@
     "HighLevelWaterDataset",
     "MSDnmrDataset",
     "MatbenchDataset",
-    "DensityDataset", 
+    "DensityDataset",
     "SmallDensityDataset",
     "OMol25Dataset",
+    "WDMPNNDataset",
 ]
 
 INFO_CLASS_REGISTRY: Dict[str, type] = {
@@ -98,6 +100,7 @@ def term_mp(sig_num, frame):
     print("main proc {} exit, kill process group " "{}".format(pid, pgid))
     os.killpg(pgid, signal.SIGKILL)
 
+
 def set_signal_handlers():
     """
     Set up signal handlers for safe process group termination.
@@ -107,7 +110,7 @@ def set_signal_handlers():
     2. The current process is the process group leader
 
     This allows safe termination of the entire process group via:
-    - Ctrl+C (SIGINT) 
+    - Ctrl+C (SIGINT)
     - Termination signals (SIGTERM)
 
     Safety Notes:

ppmat/datasets/collate_fn.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import numbers
+import warnings
 from collections.abc import Mapping
 from collections.abc import Sequence
 from typing import Any
@@ -24,7 +25,6 @@
 import numpy as np
 import paddle
 import pgl
-import warnings
 
 from ppmat.datasets.custom_data_type import ConcatData
 from ppmat.datasets.custom_data_type import ConcatNumpyWarper
@@ -108,7 +108,9 @@ def __init__(
         self.sampling_mode = sampling_mode.lower()
         self.uniform_random_offset = bool(uniform_random_offset)
         self.sampling_seed = sampling_seed
-        self._rng = np.random.default_rng(sampling_seed) if sampling_seed is not None else None
+        self._rng = (
+            np.random.default_rng(sampling_seed) if sampling_seed is not None else None
+        )
         self.clip_max = clip_max
         self.importance_sampling = bool(importance_sampling)
         self.importance_threshold = importance_threshold
@@ -154,11 +156,18 @@ def __call__(self, batch):
                         extreme_mask = dense_vals >= self.extreme_threshold
                         extreme_idx = total_idx[extreme_mask]
                         # ensure extreme is subset of high
-                        extreme_idx = np.intersect1d(extreme_idx, high_idx, assume_unique=True)
+                        extreme_idx = np.intersect1d(
+                            extreme_idx, high_idx, assume_unique=True
+                        )
                     mid_idx = np.setdiff1d(high_idx, extreme_idx, assume_unique=True)
 
-                    high_quota = min(target_samples, max(0, int(target_samples * self.importance_ratio)))
-                    extreme_quota = min(target_samples, max(0, int(target_samples * self.extreme_ratio)))
+                    high_quota = min(
+                        target_samples,
+                        max(0, int(target_samples * self.importance_ratio)),
+                    )
+                    extreme_quota = min(
+                        target_samples, max(0, int(target_samples * self.extreme_ratio))
+                    )
 
                     extreme_take = min(len(extreme_idx), extreme_quota)
                     indices_extreme = (
@@ -178,11 +187,15 @@ def __call__(self, batch):
                     selected = np.concatenate([indices_extreme, indices_mid])
                     remaining = target_samples - len(selected)
                     if remaining > 0:
-                        low_candidates = np.setdiff1d(total_idx, selected, assume_unique=False)
+                        low_candidates = np.setdiff1d(
+                            total_idx, selected, assume_unique=False
+                        )
                         if len(low_candidates) == 0:
                             low_candidates = total_idx
                         replace_low = remaining > len(low_candidates)
-                        indices_low = np.random.choice(low_candidates, remaining, replace=replace_low)
+                        indices_low = np.random.choice(
+                            low_candidates, remaining, replace=replace_low
+                        )
                         indices = np.concatenate([selected, indices_low])
                     else:
                         indices = selected
@@ -192,14 +205,22 @@ def __call__(self, batch):
                             if self._rng is None:
                                 self._rng = np.random.default_rng()
                             step = (total - 1) / max(target_samples - 1, 1)
-                            offset = float(self._rng.uniform(0, max(step, 1.0))) if step > 0 else 0.0
+                            offset = (
+                                float(self._rng.uniform(0, max(step, 1.0)))
+                                if step > 0
+                                else 0.0
+                            )
                             idx = offset + step * np.arange(target_samples)
                             indices = np.clip(np.round(idx).astype(int), 0, total - 1)
                         else:
-                            indices = np.linspace(0, total - 1, num=target_samples, dtype=int)
+                            indices = np.linspace(
+                                0, total - 1, num=target_samples, dtype=int
+                            )
                     elif self.sampling_mode == "random":
                         replace = target_samples > total
-                        indices = np.random.choice(total, target_samples, replace=replace)
+                        indices = np.random.choice(
+                            total, target_samples, replace=replace
+                        )
                     else:
                         raise ValueError(
                             f"Unsupported sampling_mode '{self.sampling_mode}'. "
@@ -208,16 +229,16 @@ def __call__(self, batch):
                 indices.sort()
                 sampled_density.append(d[indices])
                 sampled_grid.append(coord[indices])
-                mask.append(
-                    paddle.ones_like(x=sampled_density[-1], dtype="float32")
-                )
+                mask.append(paddle.ones_like(x=sampled_density[-1], dtype="float32"))
             densities = paddle.stack(x=sampled_density, axis=0)
             grid_coord = paddle.stack(x=sampled_grid, axis=0)
             mask = paddle.stack(x=mask, axis=0)
 
         densities = densities * mask
         if self.clip_max is not None:
-            densities = paddle.clip(densities, min=self.padding_value, max=self.clip_max)
+            densities = paddle.clip(
+                densities, min=self.padding_value, max=self.clip_max
+            )
         return {
             "density": densities,
             "density_mask": mask,
@@ -282,6 +303,7 @@ def __call__(self, batch):
             "infos": list(infos),
         }
 
+
 # utils DensityCollator
 def pad_sequence(sequences, batch_first=False, padding_value=0):
     max_len = max([int(s.shape[0]) for s in sequences])  # 确保转换为Python整数
@@ -302,3 +324,51 @@ def pad_sequence(sequences, batch_first=False, padding_value=0):
             out_tensor[:length, i, ...] = tensor
 
     return out_tensor
+
+
+class WDMPNNCollator:
+    """Collator for WDMPNNDataset.
+
+    Batches MolGraph objects into BatchMolGraph and stacks targets.
+    """
+
+    def __init__(self, label_names: List[str] = None):
+        """
+        Initialize WDMPNNCollator.
+
+        :param label_names: List of label names to use as keys in the output dict.
+                           If None, defaults to ["labels"].
+        """
+        self.label_names = label_names if label_names is not None else ["labels"]
+
+    def __call__(self, batch):
+        from ppmat.models.wd_mpnn.featurization import BatchMolGraph
+
+        num_molecules = len(batch[0]["mol_graphs"])
+
+        # Transpose: group mol_graphs by molecule slot
+        batch_graphs = []
+        for i in range(num_molecules):
+            graphs = [sample["mol_graphs"][i] for sample in batch]
+            batch_graphs.append(BatchMolGraph(graphs))
+
+        targets = np.stack([s["targets"] for s in batch])
+        target_mask = np.stack([s["target_mask"] for s in batch])
+
+        result = {
+            "batch_graphs": batch_graphs,
+            self.label_names[0]: paddle.to_tensor(targets, dtype="float32"),
+            "label_mask": paddle.to_tensor(target_mask, dtype="float32"),
+            "smiles": [s["smiles"] for s in batch],
+        }
+
+        # Handle optional features
+        if batch[0]["features"] is not None:
+            features = np.stack([s["features"] for s in batch])
+            result["features"] = paddle.to_tensor(features, dtype="float32")
+        else:
+            result["features"] = None
+
+        result["atom_descriptors_batch"] = None
+
+        return result

ppmat/datasets/wd_mpnn_dataset.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import csv
+import os.path as osp
+from typing import Dict
+from typing import List
+
+import numpy as np
+from paddle.io import Dataset
+
+from ppmat.utils import download
+from ppmat.utils import logger
+
+
+class WDMPNNDataset(Dataset):
+    """Dataset for polymer-chemprop model.
+
+    Loads CSV data with SMILES and target columns, converts to MolGraph.
+    """
+
+    url_bace = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/Custom_Poly/bace.csv"  # noqa
+    md5_bace = "b8962e1adc9a83d2d1d706a4224a8f0b"
+
+    url_delaney = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/Custom_Poly/delaney.csv"  # noqa
+    md5_delaney = "b300f12938e848e51d2a9f200486ff9e"
+
+    # Mapping from filename stem to (url, md5) attribute names
+    _dataset_urls = {
+        "bace": ("url_bace", "md5_bace"),
+        "delaney": ("url_delaney", "md5_delaney"),
+    }
+
+    def __init__(
+        self,
+        path: str,
+        smiles_columns: List[str] = None,
+        target_columns: List[str] = None,
+        featurization_config: Dict = None,
+        max_data_size: int = None,
+    ):
+        super().__init__()
+
+        if not osp.exists(path):
+            logger.message("The dataset is not found. Will download it now.")
+            basename = osp.basename(path)
+            stem = osp.splitext(basename)[0]
+            if stem in self._dataset_urls:
+                url_attr, md5_attr = self._dataset_urls[stem]
+                root_path = download.get_datasets_path_from_url(
+                    getattr(self, url_attr), getattr(self, md5_attr)
+                )
+                path = osp.join(root_path, basename)
+            else:
+                logger.warning(
+                    f"Dataset file '{basename}' is not available for "
+                    f"auto-download. Available datasets: "
+                    f"{list(self._dataset_urls.keys())}"
+                )
+
+        self.path = path
+
+        # Build featurization config (lazy import to avoid hard dependency at module load)
+        from ppmat.models.wd_mpnn.featurization import Featurization_parameters
+
+        if featurization_config is not None:
+            self.feat_config = Featurization_parameters(**featurization_config)
+        else:
+            self.feat_config = Featurization_parameters()
+
+        # Read CSV
+        with open(path) as f:
+            reader = csv.DictReader(f)
+            header = reader.fieldnames
+
+            # Default: first column is SMILES
+            if smiles_columns is None:
+                smiles_columns = [header[0]]
+            self.smiles_columns = smiles_columns
+
+            # Default: all non-SMILES columns are targets
+            if target_columns is None:
+                target_columns = [c for c in header if c not in smiles_columns]
+            self.target_columns = target_columns
+
+            rows = list(reader)
+
+        if max_data_size is not None:
+            rows = rows[:max_data_size]
+
+        self.smiles_list = []
+        self.targets_list = []
+        self.target_masks = []
+
+        for row in rows:
+            smiles = [row[col] for col in self.smiles_columns]
+            targets = []
+            mask = []
+            for col in self.target_columns:
+                val = row.get(col, "")
+                if val == "" or val is None:
+                    targets.append(0.0)
+                    mask.append(0.0)
+                else:
+                    targets.append(float(val))
+                    mask.append(1.0)
+            self.smiles_list.append(smiles)
+            self.targets_list.append(np.array(targets, dtype=np.float32))
+            self.target_masks.append(np.array(mask, dtype=np.float32))
+
+    def __len__(self):
+        return len(self.smiles_list)
+
+    def __getitem__(self, idx):
+        smiles = self.smiles_list[idx]
+
+        from ppmat.models.wd_mpnn.featurization import MolGraph
+
+        # Build MolGraph for each molecule
+        mol_graphs = []
+        for smi in smiles:
+            mol_graph = MolGraph(smi, config=self.feat_config)
+            mol_graphs.append(mol_graph)
+
+        return {
+            "mol_graphs": mol_graphs,
+            "targets": self.targets_list[idx],
+            "target_mask": self.target_masks[idx],
+            "features": None,
+            "smiles": smiles,
+        }

ppmat/metrics/__init__.py

@@ -16,13 +16,15 @@
 
 import paddle  # noqa
 
+from ppmat.metrics.auc_metric import AucMetric
 from ppmat.metrics.csp_metric import CSPMetric
 from ppmat.metrics.diffnmr_streaming_adapter import DiffNMRStreamingAdapter
 
 __all__ = [
     "build_metric",
     "CSPMetric",
     "DiffNMRStreamingAdapter",
+    "AucMetric",
     # "DiffNMRMetric",
     # "NLL", "CrossEntropyMetric", "SumExceptBatchMetric", "SumExceptBatchKL",
 ]