【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,12 +42,12 @@
 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.polymer_chemprop_dataset import PolymerChempropDataset
+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.utils import logger
@@ -64,9 +65,10 @@
     "HighLevelWaterDataset",
     "MSDnmrDataset",
     "MatbenchDataset",
-    "DensityDataset", 
+    "DensityDataset",
     "SmallDensityDataset",
     "OMol25Dataset",
+    "PolymerChempropDataset",
 ]
 
 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,42 @@ def pad_sequence(sequences, batch_first=False, padding_value=0):
             out_tensor[:length, i, ...] = tensor
 
     return out_tensor
+
+
+class PolymerChempropCollator:
+    """Collator for PolymerChempropDataset.
+
+    Batches MolGraph objects into BatchMolGraph and stacks targets.
+    """
+
+    def __call__(self, batch):
+        from ppmat.models.polymer_chemprop.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,
+            "labels": 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/polymer_chemprop_dataset.py

@@ -0,0 +1,110 @@
+# 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
+from typing import Dict
+from typing import List
+
+import numpy as np
+from paddle.io import Dataset
+
+
+class PolymerChempropDataset(Dataset):
+    """Dataset for polymer-chemprop model.
+
+    Loads CSV data with SMILES and target columns, converts to MolGraph.
+    """
+
+    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__()
+        self.path = path
+
+        # Build featurization config (lazy import to avoid hard dependency at module load)
+        from ppmat.models.polymer_chemprop.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.polymer_chemprop.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/models/__init__.py

@@ -35,13 +35,14 @@
 from ppmat.models.diffnmr.diffnmr import MolecularGraphFormer
 from ppmat.models.diffnmr.diffnmr import NMRNetCLIP
 from ppmat.models.dimenetpp.dimenetpp import DimeNetPlusPlus
+from ppmat.models.infgcn.infgcn import InfGCN
+from ppmat.models.mateno.mateno import MatENO
 from ppmat.models.mattergen.mattergen import MatterGen
 from ppmat.models.mattergen.mattergen import MatterGenWithCondition
 from ppmat.models.mattersim.m3gnet import M3GNet
 from ppmat.models.mattersim.m3gnet_graph_converter import M3GNetGraphConvertor
 from ppmat.models.megnet.megnet import MEGNetPlus
-from ppmat.models.infgcn.infgcn import InfGCN
-from ppmat.models.mateno.mateno import MatENO
+from ppmat.models.polymer_chemprop.model import PolymerChempropModel
 from ppmat.utils import download
 from ppmat.utils import logger
 from ppmat.utils import save_load
@@ -67,6 +68,7 @@
     "DiffNMR",
     "InfGCN",
     "MatENO",
+    "PolymerChempropModel",
 ]
 
 # Warning: The key of the dictionary must be consistent with the file name of the value

ppmat/models/polymer_chemprop/__init__.py

@@ -0,0 +1,3 @@
+from ppmat.models.polymer_chemprop.model import PolymerChempropModel
+
+__all__ = ["PolymerChempropModel"]