revise grin compatibility

Author: liugangcode

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torch_molecule"
-version = "0.1.6.post1"
+version = "0.1.7"
 description = "Deep learning packages for molecular discovery with a simple sklearn-style interface"
 authors = [{name = "Gang Liu", email = "gliu7@nd.edu"}]
 readme = "README.md"

tests/predictor/grin.py

@@ -42,7 +42,7 @@ def test_grin_predictor():
     model = GRINMolecularPredictor(
         num_task=1,
         task_type="regression",
-        polymer_train_augmentation=3,
+        repetition_augmentation=3,
         num_layer=3,
         hidden_size=128,
         batch_size=4,
@@ -58,10 +58,10 @@ def test_grin_predictor():
 
     # 2.2. Prediction test on different repeat times
     print("\n=== Testing model prediction with polymer train augmentation ===")
-    predictions = model.predict(smiles_list, test_augmentation=3)
+    predictions = model.predict(smiles_list, test_time_augmentation=3)
     print(f"Prediction shape: {predictions['prediction'].shape}")
     print(f"Prediction for new polymer repeat times 3: {predictions['prediction']}")
-    predictions = model.predict(smiles_list, test_augmentation=5)
+    predictions = model.predict(smiles_list, test_time_augmentation=5)
     print(f"Prediction for new polymer repeat times 5: {predictions['prediction']}")
 
     # 3. Auto-fitting test with custom parameters

torch_molecule/predictor/grin/modeling_grin.py

@@ -22,13 +22,13 @@ class GRINMolecularPredictor(GNNMolecularPredictor):
 
     References
     ----------
-    - Learning Repetition-Invariant Representations for Polymer Informatics.
+    - Learning Repetition-Invariant Representations for Polymer Informatics. NeurIPS 2025
       https://arxiv.org/pdf/2505.10726
 
     Parameters
     ----------
-    polymer_train_augmentation : int, default=None
-        Number of times to repeat the polymer for training.
+    repetition_augmentation : bool, default=False
+        Whether to enable polymer augmentation for training.
     l1_penalty : float, default=1e-3
         Weight for the L1 penalty.
     epochs_to_penalize : int, default=100
@@ -86,7 +86,7 @@ class GRINMolecularPredictor(GNNMolecularPredictor):
     def __init__(
         self,
         # GRIN-specific parameters
-        polymer_train_augmentation: Optional[int] = None,
+        repetition_augmentation: bool = False,
         l1_penalty: float = 1e-3,
         epochs_to_penalize: int = 100,
         # Core model parameters
@@ -148,11 +148,21 @@ def __init__(
         )
 
         # GRIN-specific parameters
-        self.polymer_train_augmentation = polymer_train_augmentation
+        self.repetition_augmentation = repetition_augmentation
         self.l1_penalty = l1_penalty
         self.epochs_to_penalize = epochs_to_penalize
         self.model_class = GRIN
 
+        # Check CombineMols dependency if polymer augmentation is enabled
+        if self.repetition_augmentation:
+            try:
+                from CombineMols.CombineMols import CombineMols
+            except ImportError:
+                raise ImportError(
+                    "CombineMols is required for repetition augmentation for polymer. "
+                    "Please install it using: pip install CombineMols"
+                )
+
 
     @staticmethod
     def _get_param_names() -> List[str]:
@@ -210,8 +220,8 @@ def fit(
         optimizer, scheduler = self._setup_optimizers()
 
         # Prepare datasets and loaders
-        if self.polymer_train_augmentation is not None:
-            X_train_aug, y_train_aug = SmilesRepeat(self.polymer_train_augmentation).repeat(X_train, y_train)
+        if self.repetition_augmentation:
+            X_train_aug, y_train_aug = SmilesRepeat(2).repeat(X_train, y_train)
             X_train = X_train + X_train_aug
             if y_train_aug is not None:
                 if isinstance(y_train, np.ndarray):
@@ -236,8 +246,8 @@ def fit(
                 UserWarning
             )
         else:
-            if self.polymer_train_augmentation is not None:
-                X_val_aug, y_val_aug = SmilesRepeat(self.polymer_train_augmentation).repeat(X_val, y_val)
+            if self.repetition_augmentation:
+                X_val_aug, y_val_aug = SmilesRepeat(2).repeat(X_val, y_val)
                 X_val = X_val + X_val_aug
                 if y_val_aug is not None:
                     if isinstance(y_val, np.ndarray):
@@ -376,15 +386,15 @@ def _train_epoch(self, train_loader, optimizer, epoch, global_pbar=None):
 
         return losses
 
-    def predict(self, X: List[str], test_augmentation: Optional[int] = None) -> Dict[str, np.ndarray]:
+    def predict(self, X: List[str], test_time_augmentation: bool = False) -> Dict[str, np.ndarray]:
         """Make predictions using the fitted model.
 
         Parameters
         ----------
         X : List[str]
             List of SMILES strings to make predictions for
-        test_augmentation : int, optional
-            Number of times to repeat the polymer for making predictions.
+        test_time_augmentation : bool, default=False
+            Whether to enable polymer augmentation for making predictions.
 
         Returns
         -------
@@ -394,8 +404,8 @@ def predict(self, X: List[str], test_augmentation: Optional[int] = None) -> Dict
 
         """
         self._check_is_fitted()
-        if test_augmentation is not None:
-            X_aug, _ = SmilesRepeat(test_augmentation).repeat(X)
+        if test_time_augmentation:
+            X_aug, _ = SmilesRepeat(2).repeat(X)
             X = X_aug
 
         # Convert to PyTorch Geometric format and create loader

torch_molecule/predictor/grin/utils.py

@@ -2,9 +2,13 @@
 import numpy as np
 
 from rdkit import Chem
-from CombineMols.CombineMols import CombineMols
 from rdkit.Chem import Draw
 
+try:
+    from CombineMols.CombineMols import CombineMols
+except ImportError:
+    pass
+
 class SmilesRepeat():
     def __init__(self, repeat_times) -> None:
         self.repeat_times = repeat_times
@@ -121,7 +125,7 @@ def edit_mol(self, ori_psmiles, des_psmiles) -> str:
         # Obtain connection info for future bonds/atoms remove/add
         info = self.get_connection_info(combo)
         if not info:
-            print("************************** No Star Mark! **************************")
+            print(f"************************** No Star Mark for polymer {ori_psmiles} **************************")
             return des_psmiles
 
         # add a new bond between two star symbols and discard these two stars
@@ -144,7 +148,7 @@ def star_edge(self, ori_psmiles) -> str:
         ori_mol = self.get_mol(ori_psmiles)
         info = self.get_connection_info(ori_mol)
         if not info or not info["neighbor"]['path']:
-            print(f"************************** No Star Mark: {ori_psmiles} **************************")
+            print(f"************************** No Star Mark for polymer {ori_psmiles} **************************")
             return ori_psmiles
 
         edsmiles = Chem.EditableMol(ori_mol)