polymon 1.0.2__tar.gz → 1.0.3__tar.gz

{polymon-1.0.2 → polymon-1.0.3}/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.4
 Name: polymon
-Version: 1.0.2
+Version: 1.0.3
 Summary: PolyMon
 Author: PolyMon Group
+Requires-Dist: numpy<2
 Requires-Dist: mordredcommunity
 Requires-Dist: mordred==1.2.0
 Requires-Dist: rdkit==2023.09.6

polymon-1.0.3/README.md

@@ -0,0 +1,240 @@
+<img src="assets/polymon.png" alt="Polymon Icon">
+
+`PolyMon` is a unified framework for polymer property prediction. It combines traditional machine learning methods (Random Forest, XGBoost, LightGBM, CatBoost, TabPFN) with state-of-the-art deep learning models (Graph Neural Networks including GATv2, GIN, PNA, DimeNet++, and KAN-based architectures).
+
+<p align="center">
+  <img src="assets/framework.png" alt="framework" width="400">
+</p>
+
+## Features
+
+- **Multiple Model Types**: Support for both tabular ML models and graph-based deep learning models
+- **Flexible Training Strategies**: K-fold cross-validation, hyperparameter optimization, ensemble learning, multi-fidelity learning, and active learning
+- **Comprehensive Descriptors**: RDKit 2D/3D descriptors, ECFP fingerprints, Mordred descriptors, and graph-based representations
+- **Multiple Properties**: Predict glass transition temperature (Tg), fractional free volume (FFV), radius of gyration (Rg), density, and thermal conductivity (Tc)
+- **Easy-to-Use CLI**: Simple command-line interface for training, prediction, and active learning recommendations
+
+## Installation
+
+### Prerequisites
+
+This package requires `torch>=2.2.2` and `torch_geometric>=2.5.3`. We recommend installing PyTorch with CUDA support first.
+
+> **Note:** RDKit (a required dependency) is not yet compatible with NumPy 2.x. If you have NumPy 2.x installed, downgrade it with `pip install 'numpy<2'` before installing PolyMon.
+
+### Step 1: Install PyTorch and PyTorch Geometric
+
+```bash
+# For CUDA 11.8
+conda install -y pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 \
+                 pytorch-cuda=11.8 -c pytorch -c nvidia
+pip install torch_geometric
+pip install torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.3.0+cu118.html
+```
+
+### Step 2: Install PolyMon
+
+```bash
+pip install polymon
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/fate1997/polymon.git
+cd polymon
+pip install -e .
+```
+
+## Quick Start
+
+### Training a Model
+
+Train a tabular model (e.g., Random Forest) with RDKit 2D descriptors:
+
+```bash
+polymon train \
+    --raw-csv ./database/database.csv \
+    --sources Kaggle PI1070 PolyMetriX \
+    --labels Tg \
+    --feature-names rdkit2d \
+    --model rf \
+    --n-fold 5 \
+    --out-dir ./results
+```
+
+Train a graph neural network (GNN) model:
+
+```bash
+polymon train \
+    --raw-csv ./database/database.csv \
+    --sources Kaggle PI1070 PolyMetriX \
+    --labels Tg \
+    --model gatv2 \
+    --n-fold 5 \
+    --n-trials 15 \
+    --num-epochs 2500 \
+    --out-dir ./results
+```
+
+### Making Predictions
+
+```bash
+polymon predict \
+    --model-path ./results/gatv2/Tg/train/gatv2_Tg.pt \
+    --csv-path ./data/new_polymers.csv \
+    --smiles-column SMILES
+```
+
+### Active Learning Recommendations
+
+```bash
+polymon rec \
+    --pool-csv ./database/pool.csv \
+    --trained-model ./results/gatv2/Rg/train/gatv2_Rg-KFold.pt \
+    --acquisition uncertainty \
+    --sample-size 20 \
+    --save-path recommended.csv
+```
+
+## Available Models
+
+### Tabular Models (for use with `--feature-names`)
+
+| Model | CLI Name | Description |
+|-------|----------|-------------|
+| Random Forest | `rf` | Ensemble of decision trees |
+| XGBoost | `xgb` | Gradient boosting framework |
+| LightGBM | `lgbm` | Light gradient boosting machine |
+| CatBoost | `catboost` | Gradient boosting on decision trees |
+| TabPFN | `tabpfn` | Prior-data trained network |
+
+### Deep Learning Models
+
+| Model | CLI Name | Description |
+|-------|----------|-------------|
+| GATv2 | `gatv2` | Graph Attention Network v2 |
+| GIN | `gin` | Graph Isomorphism Network |
+| PNA | `pna` | Principal Neighbourhood Aggregation |
+| AttentiveFP | `attentivefp` | Attention-based molecular fingerprinting |
+| DimeNet++ | `dimenetpp` | Directional message passing |
+| GPS | `gps` | Graph Positional Encoding network |
+| KAN-GATv2 | `fastkan_gatv2` | Kolmogorov-Arnold Network + GATv2 |
+| KAN-GPS | `kan_gps` | Kolmogorov-Arnold Network + GPS |
+
+## Available Descriptors
+
+### For Tabular Models (`--feature-names`)
+
+- `rdkit2d`: RDKit 2D molecular descriptors
+- `ecfp4`: Extended Connectivity Fingerprints (ECFP4)
+- `mordred`: 1800+ Mordred descriptors
+- `maccs`: MACCS keys
+- `xenonpy_desc`: XenonPy elemental composition descriptors
+
+### For Graph Models
+
+Graph models automatically use molecular graph features. Additional descriptors can be added via `--descriptors`:
+- `rdkit2d`, `ecfp4`, `mordred`, `maccs`, `xenonpy_desc`
+- `oligomer_rdkit2d`, `oligomer_mordred`, `oligomer_ecfp4` (for oligomer representations)
+
+## Target Properties
+
+| Property | Symbol | Unit | Description |
+|----------|--------|------|-------------|
+| Glass Transition Temperature | `Tg` | K | Temperature at which polymer transitions from glassy to rubbery |
+| Fractional Free Volume | `FFV` | - | Fraction of volume not occupied by polymer chains |
+| Radius of Gyration | `Rg` | Å | Measure of polymer chain size |
+| Density | `Density` | g/cm³ | Mass per unit volume |
+| Thermal Conductivity | `Tc` | W/m·K | Heat transfer capability |
+
+## Advanced Usage
+
+### Hyperparameter Optimization
+
+```bash
+polymon train \
+    --labels Tg \
+    --model gatv2 \
+    --n-trials 15 \
+    --n-fold 5 \
+    --raw-csv ./database/database.csv
+```
+
+### Multi-Fidelity Learning (Fine-tuning)
+
+```bash
+# Train on low-fidelity data first
+polymon train \
+    --labels Density \
+    --model gatv2 \
+    --sources MD-simulation \
+    --run-production
+
+# Fine-tune on high-fidelity data
+polymon train \
+    --labels Density \
+    --model gatv2 \
+    --sources Experimental \
+    --finetune \
+    --pretrained-model ./results/gatv2/Density/production/gatv2_Density.pt \
+    --finetune-csv-path ./database/experimental.csv
+```
+
+### Ensemble Learning
+
+```bash
+polymon train \
+    --labels Rg \
+    --model gatv2 \
+    --n-estimator 10 \
+    --ensemble-type voting \
+    --raw-csv ./database/database.csv
+```
+
+### Delta-Learning with Empirical Estimators
+
+```bash
+polymon train \
+    --labels Density \
+    --model gatv2 \
+    --train-residual \
+    --estimator-name Density-IBM \
+    --raw-csv ./database/database.csv
+```
+
+## Python API
+
+For more advanced usage, you can use the Python API directly:
+
+```python
+from polymon.model.base import ModelWrapper
+
+# Load a trained model
+model = ModelWrapper.from_file('results/gatv2/Tg/train/gatv2_Tg.pt')
+
+# Make predictions
+predictions = model.predict(['*C*', '*CC*', '*CCC*'])
+print(predictions)
+```
+
+## Citation
+
+If you use PolyMon in your research, please cite:
+
+```bibtex
+@article{polymon2024,
+  title={PolyMon: A Unified Framework for Polymer Property Prediction},
+  author={Ren, Gaopeng and Yang, Yijie and Zhou, Jiajun and Jelfs, Kim E.},
+  journal={Journal of Chemical Information and Modeling},
+  year={2024}
+}
+```
+
+## License
+
+This project is licensed under the MIT License.
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.

polymon-1.0.3/polymon/__init__.py

@@ -0,0 +1 @@
+__version__ = '1.0.3'

polymon-1.0.3/polymon/cli/main.py

@@ -0,0 +1,363 @@
+"""PolyMon Command-Line Interface.
+
+This module provides the main CLI entry point for the PolyMon framework.
+It supports three main modes:
+    - train: Train ML/DL models for polymer property prediction
+    - rec: Active learning - recommend molecules for labeling
+    - predict: Make predictions on new polymer data
+
+Example:
+    Train a model:
+        $ polymon train --labels Tg --model rf --n-fold 5
+
+    Make predictions:
+        $ polymon predict --model-path model.pt --csv-path data.csv --smiles-column SMILES
+
+    Recommend molecules:
+        $ polymon rec --pool-csv pool.csv --trained-model model.pt --acquisition uncertainty
+"""
+
+import argparse
+
+from polymon.cli.train_dl import main as main_dl
+from polymon.cli.train_ml import MODELS
+from polymon.cli.train_ml import main as main_ml
+from polymon.cli.predict import main as main_predict
+from polymon.cli.recommend import main as main_recommend
+
+
+def parse_args():
+    """Parse command-line arguments for the PolyMon CLI.
+
+    Returns:
+        argparse.Namespace: Parsed arguments with the following attributes:
+            - mode (str): Subcommand to run ('train', 'rec', or 'predict')
+            - ... (additional arguments depend on the subcommand)
+    """
+    parser = argparse.ArgumentParser(
+        description='PolyMon - Polymer Property Prediction Framework',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    subparsers = parser.add_subparsers(dest='mode', required=True)
+
+    # ============================================================================
+    # Train subcommand
+    # ============================================================================
+    train_parser = subparsers.add_parser(
+        'train',
+        help='Train a ML/DL model for polymer property prediction',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    train_parser.add_argument(
+        '--raw-csv',
+        type=str,
+        default='database/database.csv',
+        help='Path to the raw CSV file containing polymer data'
+    )
+    train_parser.add_argument(
+        '--sources',
+        type=str,
+        nargs='+',
+        default=['Kaggle'],
+        help='Data sources to filter from the dataset (e.g., Kaggle, PI1070, PolyMetriX)'
+    )
+    train_parser.add_argument(
+        '--tag',
+        type=str,
+        default='debug',
+        help='Tag for organizing this training run'
+    )
+    train_parser.add_argument(
+        '--labels',
+        nargs='+',
+        required=True,
+        choices=['Tg', 'FFV', 'Density', 'Rg', 'Tc'],
+        help='Target property/properties to predict'
+    )
+    train_parser.add_argument(
+        '--feature-names',
+        type=str,
+        nargs='+',
+        default=['rdkit2d'],
+        help='Feature names for tabular models (rdkit2d, ecfp4, mordred, maccs, xenonpy_desc)'
+    )
+    train_parser.add_argument(
+        '--n-trials',
+        type=int,
+        default=None,
+        help='Number of trials for hyperparameter optimization (enables optimization if set)'
+    )
+    train_parser.add_argument(
+        '--out-dir',
+        type=str,
+        default='./results',
+        help='Directory to save training results'
+    )
+    train_parser.add_argument(
+        '--hparams-from',
+        type=str,
+        default=None,
+        help='Path to hyperparameters file (.json, .pt, or .pkl) to reuse from previous run'
+    )
+    train_parser.add_argument(
+        '--n-fold',
+        type=int,
+        default=1,
+        help='Number of folds for cross-validation (use 1 for single split)'
+    )
+    train_parser.add_argument(
+        '--split-mode',
+        type=str,
+        default='random',
+        choices=['random', 'source', 'scaffold'],
+        help='Data splitting strategy: random, by source, or by molecular scaffold'
+    )
+    train_parser.add_argument(
+        '--seed',
+        type=int,
+        default=42,
+        help='Random seed for reproducibility'
+    )
+    train_parser.add_argument(
+        '--remove-hydrogens',
+        action='store_true',
+        help='Remove hydrogens from molecular graphs (reduces computation)'
+    )
+    train_parser.add_argument(
+        '--descriptors',
+        type=str,
+        nargs='+',
+        default=None,
+        help='Additional descriptors to concatenate with graph features for DL models'
+    )
+    train_parser.add_argument(
+        '--model',
+        type=str,
+        default='rf',
+        help='Model type: tabular (rf, xgb, lgbm, catboost, tabpfn) or GNN (gatv2, gin, pna, etc.)'
+    )
+    train_parser.add_argument(
+        '--hidden-dim',
+        type=int,
+        default=32,
+        help='Hidden dimension for neural network models'
+    )
+    train_parser.add_argument(
+        '--num-layers',
+        type=int,
+        default=3,
+        help='Number of layers for neural network models'
+    )
+    train_parser.add_argument(
+        '--batch-size',
+        type=int,
+        default=128,
+        help='Batch size for training neural networks'
+    )
+    train_parser.add_argument(
+        '--lr',
+        type=float,
+        default=1e-3,
+        help='Learning rate for neural network training'
+    )
+    train_parser.add_argument(
+        '--num-epochs',
+        type=int,
+        default=2500,
+        help='Maximum number of training epochs'
+    )
+    train_parser.add_argument(
+        '--early-stopping-patience',
+        type=int,
+        default=250,
+        help='Patience for early stopping (epochs without improvement)'
+    )
+    train_parser.add_argument(
+        '--device',
+        type=str,
+        default='cuda',
+        help='Device for training (cuda or cpu)'
+    )
+    train_parser.add_argument(
+        '--run-production',
+        action='store_true',
+        help='Run in production mode (95:5 train:val split, no test set)'
+    )
+    train_parser.add_argument(
+        '--finetune',
+        action='store_true',
+        help='Fine-tune a pretrained model on new data'
+    )
+    train_parser.add_argument(
+        '--finetune-csv-path',
+        type=str,
+        default=None,
+        help='Path to CSV file for fine-tuning data'
+    )
+    train_parser.add_argument(
+        '--pretrained-model',
+        type=str,
+        default=None,
+        help='Path to pretrained model for fine-tuning'
+    )
+    train_parser.add_argument(
+        '--n-estimator',
+        type=int,
+        default=1,
+        help='Number of estimators for ensemble learning (>1 enables ensemble mode)'
+    )
+    train_parser.add_argument(
+        '--additional-features',
+        type=str,
+        nargs='+',
+        default=None,
+        help='Additional graph features (monomer, periodic_bond, etc.)'
+    )
+    train_parser.add_argument(
+        '--skip-train',
+        action='store_true',
+        help='Skip training step (use with --n-estimator to only build ensemble)'
+    )
+    train_parser.add_argument(
+        '--low-fidelity-model',
+        type=str,
+        default=None,
+        help='Path to low-fidelity model for residual learning'
+    )
+    train_parser.add_argument(
+        '--estimator-name',
+        type=str,
+        default=None,
+        help='Name of empirical estimator for delta-learning (e.g., Density-IBM, Rg-monomer)'
+    )
+    train_parser.add_argument(
+        '--emb-model',
+        type=str,
+        default=None,
+        help='Path to embedding model for knowledge transfer'
+    )
+    train_parser.add_argument(
+        '--ensemble-type',
+        type=str,
+        default='voting',
+        choices=['voting', 'bagging', 'gradient_boosting', 'snapshot', 'soft_gradient_boosting'],
+        help='Type of ensemble to use'
+    )
+    train_parser.add_argument(
+        '--train-residual',
+        action='store_true',
+        help='Train on residuals from base estimator or low-fidelity model'
+    )
+    train_parser.add_argument(
+        '--normalizer-type',
+        type=str,
+        default='normalizer',
+        choices=['normalizer', 'log_normalizer', 'none'],
+        help='Type of label normalization'
+    )
+    train_parser.add_argument(
+        '--augmentation',
+        action='store_true',
+        help='Use data augmentation (oligomer building)'
+    )
+
+    # ============================================================================
+    # Recommend subcommand
+    # ============================================================================
+    recommend_parser = subparsers.add_parser(
+        'rec',
+        help='Recommend molecules for active learning',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    recommend_parser.add_argument(
+        '--pool-csv',
+        type=str,
+        required=True,
+        help='Path to CSV file with candidate molecules (must have SMILES column)'
+    )
+    recommend_parser.add_argument(
+        '--trained-model',
+        type=str,
+        required=True,
+        help='Path to trained model (.pt or .pkl file)'
+    )
+    recommend_parser.add_argument(
+        '--model-type',
+        type=str,
+        default='kfold',
+        choices=['kfold', 'ensemble'],
+        help='Type of model: kfold or ensemble'
+    )
+    recommend_parser.add_argument(
+        '--acquisition',
+        type=str,
+        default='uncertainty',
+        choices=['epig', 'uncertainty', 'random'],
+        help='Acquisition function: epig (expected improvement), uncertainty, or random'
+    )
+    recommend_parser.add_argument(
+        '--sample-size',
+        type=int,
+        default=100,
+        help='Number of molecules to recommend'
+    )
+    recommend_parser.add_argument(
+        '--save-path',
+        type=str,
+        default=None,
+        help='Path to save recommended molecules (CSV format)'
+    )
+
+    # ============================================================================
+    # Predict subcommand
+    # ============================================================================
+    predict_parser = subparsers.add_parser(
+        'predict',
+        help='Make predictions on new polymer data',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    predict_parser.add_argument(
+        '--model-path',
+        type=str,
+        required=True,
+        help='Path to trained model (.pt or .pkl file)'
+    )
+    predict_parser.add_argument(
+        '--csv-path',
+        type=str,
+        required=True,
+        help='Path to CSV file with molecules to predict'
+    )
+    predict_parser.add_argument(
+        '--smiles-column',
+        type=str,
+        required=True,
+        help='Name of column containing SMILES strings'
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    """Main entry point for the PolyMon CLI.
+
+    Parses command-line arguments and routes to the appropriate subcommand handler:
+    - train: Routes to ML or DL training based on model type
+    - rec: Active learning recommendations
+    - predict: Inference on new data
+    """
+    args = parse_args()
+    if args.mode == 'train':
+        # Route to appropriate training function
+        if args.model in MODELS.keys():
+            main_ml(args)
+        else:
+            main_dl(args)
+    elif args.mode == 'rec':
+        main_recommend(args)
+    elif args.mode == 'predict':
+        main_predict(args)
+
+
+if __name__ == '__main__':
+    main()

{polymon-1.0.2 → polymon-1.0.3}/polymon/cli/merge.py

@@ -365,4 +365,66 @@ def main(args: argparse.Namespace):
 
 if __name__ == '__main__':
     args = arg_parser()
-    main(args)
+    main(args)
+
+
+
+def score(
+    train_file: str, 
+    model_file: str, 
+    smiles_cluster_file: str = None, 
+    model_type: str = 'ensemble', 
+    device: str = 'cuda', 
+    ordered_tasks: List[str] = ['Rg', 'Density', 'Bulk_modulus', 'FFV', 'PLD', 'CLD'],
+    acquisition_function: str = 'uncertainty',
+    all_clusters_file: str = 'all_clusters.npy',
+    n_sample: int = 50,
+):
+    if smiles_cluster_file is not None:
+        smiles_cluster_df = pd.read_csv(smiles_cluster_file)
+        cluster_smiles = smiles_cluster_df.groupby('cluster_id')['SMILES'].apply(list).to_dict()
+    else:
+        all_smiles = np.load(all_clusters_file, allow_pickle=True).item()
+        
+    train_df = pd.read_csv(train_file)
+    train_df = train_df[train_df['Source'].isin(INITIAL_SOURCES)]
+    train_smiles = train_df['SMILES'].tolist()
+    scorer = Acquisition(
+        acquisition_function= acquisition_function,
+        model_file = model_file,
+        model_type=model_type,
+        device=device,
+        ordered_tasks=ordered_tasks,
+    )
+    sampled_scores = defaultdict(list)
+    for cluster_id, smiles in tqdm(cluster_smiles.items(), desc='Scoring'):
+        score = scorer.score(
+            pool_smiles = smiles,
+            train_smiles = train_smiles,
+        )
+        sampled_scores[cluster_id] = score
+    
+    for cluster_id, scores in sampled_scores.items():
+        if isinstance(scores, list):
+            scores_np = []
+            for s in scores:
+                if hasattr(s, 'detach'):  # torch.Tensor
+                    s = s.detach().cpu().numpy()
+                scores_np.append(s)
+            sampled_scores[cluster_id] = np.array(scores_np)
+        elif hasattr(scores, 'detach'):
+            sampled_scores[cluster_id] = scores.detach().cpu().numpy()
+
+    all_scores = [score for scores in sampled_scores.values() for score in scores]
+    threshold = np.quantile(all_scores, 0.95)
+    best_cluster = max(
+        sampled_scores,
+        key=lambda c: sum(s > threshold for s in sampled_scores[c]) / len(sampled_scores[c])
+    )
+    all_clusters = np.load(all_clusters_file, allow_pickle=True)
+    cluster_smiles = all_clusters.item()[best_cluster]
+    cluster_scores = scorer.score(cluster_smiles, train_smiles)
+    cluster_scores = cluster_scores.detach().cpu()
+    top_n_sample = torch.topk(cluster_scores, n_sample).indices.tolist()
+    top_n_smiles = [cluster_smiles[i] for i in top_n_sample]
+    return dict(zip(top_n_smiles, cluster_scores[top_n_sample]))