alphagrammar 0.1.0__tar.gz

alphagrammar-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Minghao Guo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

alphagrammar-0.1.0/PKG-INFO

@@ -0,0 +1,131 @@
+Metadata-Version: 2.1
+Name: alphagrammar
+Version: 0.1.0
+Summary: AlphaGrammar: Grammar-based molecular representation learning
+Author-email: Michael Sun <msun415@mit.edu>
+License: MIT License
+        
+        Copyright (c) 2022 Minghao Guo
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Keywords: molecule,grammar,SMILES,representation-learning,chemistry
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: torch
+Requires-Dist: rdkit
+Requires-Dist: networkx
+Requires-Dist: matplotlib
+Requires-Dist: tqdm
+Requires-Dist: torch_geometric
+
+# AlphaGrammar
+
+**AlphaGrammar** is a grammar-based molecular representation learning framework.
+It learns a hyperedge-replacement grammar (HRG) over molecular graphs via
+Monte Carlo Tree Search (MCMC) with a learned agent.
+
+## Installation
+
+```bash
+pip install -e /path/to/AlphaGrammar_pkg
+```
+
+Or from PyPI (once published):
+
+```bash
+pip install alphagrammar
+```
+
+### Pretrained models
+
+After installation, copy the pretrained model files to the package data directory:
+
+```bash
+DATA_DIR=$(python -c "import alphagrammar, os; print(os.path.join(os.path.dirname(alphagrammar.__file__), 'data'))")
+cp /path/to/AlphaGrammar/ckpts/vocab_epoch5.pkl $DATA_DIR/
+cp /path/to/AlphaGrammar/ckpts/best_agent_epoch0_R0.0000.pkl $DATA_DIR/
+cp /path/to/AlphaGrammar/GCN/supervised_contextpred.pth $DATA_DIR/
+```
+
+## Usage
+
+### Command-line interface
+
+```bash
+# Parse a single SMILES string (rollout mode)
+alphagrammar parse "CCO"
+
+# Parse from a file (one SMILES per line)
+alphagrammar parse molecules.smi
+
+# Parse with Bolinas parser (10-second timeout per molecule)
+alphagrammar parse "CCO" --timeout 10
+
+# Use only top-100 rules from vocab
+alphagrammar parse "CCO" --vocab_size 100
+```
+
+### Python API
+
+```python
+from alphagrammar import MoleculeDataset, _collate_mol_batch, RuleStats
+from alphagrammar.grammar_generation import MCMC_sampling
+from alphagrammar.agent import Agent
+import torch, pickle
+
+# Load pretrained models
+with open("data/vocab_epoch5.pkl", "rb") as f:
+    rule_stats = pickle.load(f)
+
+agent = Agent(feat_dim=300, hidden_size=256)
+agent.load_state_dict(torch.load("data/best_agent_epoch0_R0.0000.pkl"))
+agent.eval()
+
+# Build input dataset
+dataset = MoleculeDataset(["CCO", "c1ccccc1"], GNN_model_path="data/supervised_contextpred.pth")
+batch = [dataset[i] for i in range(len(dataset))]
+input_graphs_dict = _collate_mol_batch(batch)
+
+# Run MCMC sampling
+results, rules_per_mol, sequential_steps = MCMC_sampling(
+    "output_dir", agent, input_graphs_dict, MCMC_size=1, debug=True
+)
+```
+
+## Package structure
+
+```
+src/alphagrammar/
+├── __init__.py          # Public API
+├── cli.py               # argparse CLI entry point
+├── core.py              # Core functions (bolinas_evaluate, MoleculeDataset, RuleStats, ...)
+├── grammar_generation.py# MCMC sampling and grammar generation
+├── agent.py             # Neural agent (policy network)
+├── hrg_td_parser_undirected.py  # Bolinas HRG parser
+├── private/             # Internal hypergraph / grammar data structures
+├── fuseprop/            # Molecular fragmentation utilities
+├── GCN/                 # Graph neural network feature extraction
+├── data/                # Pretrained model files (copy here after install)
+└── commands/
+    └── parse.py         # 'alphagrammar parse' subcommand
+```

alphagrammar-0.1.0/README.md

@@ -0,0 +1,91 @@
+# AlphaGrammar
+
+**AlphaGrammar** is a grammar-based molecular representation learning framework.
+It learns a hyperedge-replacement grammar (HRG) over molecular graphs via
+Monte Carlo Tree Search (MCMC) with a learned agent.
+
+## Installation
+
+```bash
+pip install -e /path/to/AlphaGrammar_pkg
+```
+
+Or from PyPI (once published):
+
+```bash
+pip install alphagrammar
+```
+
+### Pretrained models
+
+After installation, copy the pretrained model files to the package data directory:
+
+```bash
+DATA_DIR=$(python -c "import alphagrammar, os; print(os.path.join(os.path.dirname(alphagrammar.__file__), 'data'))")
+cp /path/to/AlphaGrammar/ckpts/vocab_epoch5.pkl $DATA_DIR/
+cp /path/to/AlphaGrammar/ckpts/best_agent_epoch0_R0.0000.pkl $DATA_DIR/
+cp /path/to/AlphaGrammar/GCN/supervised_contextpred.pth $DATA_DIR/
+```
+
+## Usage
+
+### Command-line interface
+
+```bash
+# Parse a single SMILES string (rollout mode)
+alphagrammar parse "CCO"
+
+# Parse from a file (one SMILES per line)
+alphagrammar parse molecules.smi
+
+# Parse with Bolinas parser (10-second timeout per molecule)
+alphagrammar parse "CCO" --timeout 10
+
+# Use only top-100 rules from vocab
+alphagrammar parse "CCO" --vocab_size 100
+```
+
+### Python API
+
+```python
+from alphagrammar import MoleculeDataset, _collate_mol_batch, RuleStats
+from alphagrammar.grammar_generation import MCMC_sampling
+from alphagrammar.agent import Agent
+import torch, pickle
+
+# Load pretrained models
+with open("data/vocab_epoch5.pkl", "rb") as f:
+    rule_stats = pickle.load(f)
+
+agent = Agent(feat_dim=300, hidden_size=256)
+agent.load_state_dict(torch.load("data/best_agent_epoch0_R0.0000.pkl"))
+agent.eval()
+
+# Build input dataset
+dataset = MoleculeDataset(["CCO", "c1ccccc1"], GNN_model_path="data/supervised_contextpred.pth")
+batch = [dataset[i] for i in range(len(dataset))]
+input_graphs_dict = _collate_mol_batch(batch)
+
+# Run MCMC sampling
+results, rules_per_mol, sequential_steps = MCMC_sampling(
+    "output_dir", agent, input_graphs_dict, MCMC_size=1, debug=True
+)
+```
+
+## Package structure
+
+```
+src/alphagrammar/
+├── __init__.py          # Public API
+├── cli.py               # argparse CLI entry point
+├── core.py              # Core functions (bolinas_evaluate, MoleculeDataset, RuleStats, ...)
+├── grammar_generation.py# MCMC sampling and grammar generation
+├── agent.py             # Neural agent (policy network)
+├── hrg_td_parser_undirected.py  # Bolinas HRG parser
+├── private/             # Internal hypergraph / grammar data structures
+├── fuseprop/            # Molecular fragmentation utilities
+├── GCN/                 # Graph neural network feature extraction
+├── data/                # Pretrained model files (copy here after install)
+└── commands/
+    └── parse.py         # 'alphagrammar parse' subcommand
+```

alphagrammar-0.1.0/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "alphagrammar"
+version = "0.1.0"
+description = "AlphaGrammar: Grammar-based molecular representation learning"
+readme = "README.md"
+requires-python = ">=3.8"
+license = {file = "LICENSE"}
+authors = [
+  { name = "Michael Sun", email = "msun415@mit.edu" }
+]
+keywords = ["molecule", "grammar", "SMILES", "representation-learning", "chemistry"]
+dependencies = [
+  "numpy",
+  "pandas",
+  "torch",
+  "rdkit",
+  "networkx",
+  "matplotlib",
+  "tqdm",
+  "torch_geometric",
+]
+
+[project.scripts]
+alphagrammar = "alphagrammar.cli:main"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["alphagrammar*"]
+
+[tool.setuptools.package-data]
+alphagrammar = ["data/*", "data/.gitkeep"]

alphagrammar-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

alphagrammar-0.1.0/src/alphagrammar/GCN/__init__.py

@@ -0,0 +1 @@
+# GCN subpackage for AlphaGrammar

alphagrammar-0.1.0/src/alphagrammar/GCN/batch.py

@@ -0,0 +1,228 @@
+import torch
+from torch_geometric.data import Data, Batch
+
+class BatchMasking(Data):
+    r"""A plain old python object modeling a batch of graphs as one big
+    (dicconnected) graph. With :class:`torch_geometric.data.Data` being the
+    base class, all its methods can also be used here.
+    In addition, single graphs can be reconstructed via the assignment vector
+    :obj:`batch`, which maps each node to its respective graph identifier.
+    """
+
+    def __init__(self, batch=None, **kwargs):
+        super(BatchMasking, self).__init__(**kwargs)
+        self.batch = batch
+
+    @staticmethod
+    def from_data_list(data_list):
+        r"""Constructs a batch object from a python list holding
+        :class:`torch_geometric.data.Data` objects.
+        The assignment vector :obj:`batch` is created on the fly."""
+        keys = [set(data.keys) for data in data_list]
+        keys = list(set.union(*keys))
+        assert 'batch' not in keys
+
+        batch = BatchMasking()
+
+        for key in keys:
+            batch[key] = []
+        batch.batch = []
+
+        cumsum_node = 0
+        cumsum_edge = 0
+
+        for i, data in enumerate(data_list):
+            num_nodes = data.num_nodes
+            batch.batch.append(torch.full((num_nodes, ), i, dtype=torch.long))
+            for key in data.keys:
+                item = data[key]
+                if key in ['edge_index', 'masked_atom_indices']:
+                    item = item + cumsum_node
+                elif key  == 'connected_edge_indices':
+                    item = item + cumsum_edge
+                batch[key].append(item)
+
+            cumsum_node += num_nodes
+            cumsum_edge += data.edge_index.shape[1]
+
+        for key in keys:
+            batch[key] = torch.cat(
+                batch[key], dim=data_list[0].cat_dim(key, batch[key][0]))
+        batch.batch = torch.cat(batch.batch, dim=-1)
+        return batch.contiguous()
+
+    def cumsum(self, key, item):
+        r"""If :obj:`True`, the attribute :obj:`key` with content :obj:`item`
+        should be added up cumulatively before concatenated together.
+        .. note::
+            This method is for internal use only, and should only be overridden
+            if the batch concatenation process is corrupted for a specific data
+            attribute.
+        """
+        return key in ['edge_index', 'face', 'masked_atom_indices', 'connected_edge_indices']
+
+    @property
+    def num_graphs(self):
+        """Returns the number of graphs in the batch."""
+        return self.batch[-1].item() + 1
+
+class BatchAE(Data):
+    r"""A plain old python object modeling a batch of graphs as one big
+    (dicconnected) graph. With :class:`torch_geometric.data.Data` being the
+    base class, all its methods can also be used here.
+    In addition, single graphs can be reconstructed via the assignment vector
+    :obj:`batch`, which maps each node to its respective graph identifier.
+    """
+
+    def __init__(self, batch=None, **kwargs):
+        super(BatchAE, self).__init__(**kwargs)
+        self.batch = batch
+
+    @staticmethod
+    def from_data_list(data_list):
+        r"""Constructs a batch object from a python list holding
+        :class:`torch_geometric.data.Data` objects.
+        The assignment vector :obj:`batch` is created on the fly."""
+        keys = [set(data.keys) for data in data_list]
+        keys = list(set.union(*keys))
+        assert 'batch' not in keys
+
+        batch = BatchAE()
+
+        for key in keys:
+            batch[key] = []
+        batch.batch = []
+
+        cumsum_node = 0
+
+        for i, data in enumerate(data_list):
+            num_nodes = data.num_nodes
+            batch.batch.append(torch.full((num_nodes, ), i, dtype=torch.long))
+            for key in data.keys:
+                item = data[key]
+                if key in ['edge_index', 'negative_edge_index']:
+                    item = item + cumsum_node
+                batch[key].append(item)
+
+            cumsum_node += num_nodes
+
+        for key in keys:
+            batch[key] = torch.cat(
+                batch[key], dim=batch.cat_dim(key))
+        batch.batch = torch.cat(batch.batch, dim=-1)
+        return batch.contiguous()
+
+    @property
+    def num_graphs(self):
+        """Returns the number of graphs in the batch."""
+        return self.batch[-1].item() + 1
+
+    def cat_dim(self, key):
+        return -1 if key in ["edge_index", "negative_edge_index"] else 0
+
+
+class BatchSubstructContext(Data):
+    r"""A plain old python object modeling a batch of graphs as one big
+    (dicconnected) graph. With :class:`torch_geometric.data.Data` being the
+    base class, all its methods can also be used here.
+    In addition, single graphs can be reconstructed via the assignment vector
+    :obj:`batch`, which maps each node to its respective graph identifier.
+    """
+
+    """
+    Specialized batching for substructure context pair!
+    """
+
+    def __init__(self, batch=None, **kwargs):
+        super(BatchSubstructContext, self).__init__(**kwargs)
+        self.batch = batch
+
+    @staticmethod
+    def from_data_list(data_list):
+        r"""Constructs a batch object from a python list holding
+        :class:`torch_geometric.data.Data` objects.
+        The assignment vector :obj:`batch` is created on the fly."""
+        #keys = [set(data.keys) for data in data_list]
+        #keys = list(set.union(*keys))
+        #assert 'batch' not in keys
+
+        batch = BatchSubstructContext()
+        keys = ["center_substruct_idx", "edge_attr_substruct", "edge_index_substruct", "x_substruct", "overlap_context_substruct_idx", "edge_attr_context", "edge_index_context", "x_context"]
+
+        for key in keys:
+            #print(key)
+            batch[key] = []
+
+        #batch.batch = []
+        #used for pooling the context
+        batch.batch_overlapped_context = []
+        batch.overlapped_context_size = []
+
+        cumsum_main = 0
+        cumsum_substruct = 0
+        cumsum_context = 0
+
+        i = 0
+        
+        for data in data_list:
+            #If there is no context, just skip!!
+            if hasattr(data, "x_context"):
+                num_nodes = data.num_nodes
+                num_nodes_substruct = len(data.x_substruct)
+                num_nodes_context = len(data.x_context)
+
+                #batch.batch.append(torch.full((num_nodes, ), i, dtype=torch.long))
+                batch.batch_overlapped_context.append(torch.full((len(data.overlap_context_substruct_idx), ), i, dtype=torch.long))
+                batch.overlapped_context_size.append(len(data.overlap_context_substruct_idx))
+
+                ###batching for the main graph
+                #for key in data.keys:
+                #    if not "context" in key and not "substruct" in key:
+                #        item = data[key]
+                #        item = item + cumsum_main if batch.cumsum(key, item) else item
+                #        batch[key].append(item)
+                
+                ###batching for the substructure graph
+                for key in ["center_substruct_idx", "edge_attr_substruct", "edge_index_substruct", "x_substruct"]:
+                    item = data[key]
+                    item = item + cumsum_substruct if batch.cumsum(key, item) else item
+                    batch[key].append(item)
+                
+
+                ###batching for the context graph
+                for key in ["overlap_context_substruct_idx", "edge_attr_context", "edge_index_context", "x_context"]:
+                    item = data[key]
+                    item = item + cumsum_context if batch.cumsum(key, item) else item
+                    batch[key].append(item)
+
+                cumsum_main += num_nodes
+                cumsum_substruct += num_nodes_substruct   
+                cumsum_context += num_nodes_context
+                i += 1
+
+        for key in keys:
+            batch[key] = torch.cat(
+                batch[key], dim=batch.cat_dim(key))
+        #batch.batch = torch.cat(batch.batch, dim=-1)
+        batch.batch_overlapped_context = torch.cat(batch.batch_overlapped_context, dim=-1)
+        batch.overlapped_context_size = torch.LongTensor(batch.overlapped_context_size)
+
+        return batch.contiguous()
+
+    def cat_dim(self, key):
+        return -1 if key in ["edge_index", "edge_index_substruct", "edge_index_context"] else 0
+
+    def cumsum(self, key, item):
+        r"""If :obj:`True`, the attribute :obj:`key` with content :obj:`item`
+        should be added up cumulatively before concatenated together.
+        .. note::
+            This method is for internal use only, and should only be overridden
+            if the batch concatenation process is corrupted for a specific data
+            attribute.
+        """
+        return key in ["edge_index", "edge_index_substruct", "edge_index_context", "overlap_context_substruct_idx", "center_substruct_idx"]
+
+    @property
+    def num_graphs(self):
+        """Returns the number of graphs in the batch."""
+        return self.batch[-1].item() + 1

alphagrammar-0.1.0/src/alphagrammar/GCN/feature_extract.py

@@ -0,0 +1,38 @@
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from rdkit import Chem
+from tqdm import tqdm
+import numpy as np
+# from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+from .model import GNN, GNN_feature
+from .loader import mol_to_graph_data_obj_simple
+
+
+class feature_extractor():
+    def __init__(self, pretrained_model_path):
+        self.pretrained_model_path = pretrained_model_path
+
+    def preprocessing(self, graph_mol):
+        # rdkit_mol = AllChem.MolFromSmiles(graph_sml)
+        # kekulize the molecule to distinguish the aromatic bonds
+        data = mol_to_graph_data_obj_simple(Chem.MolFromSmiles(Chem.MolToSmiles(graph_mol)))
+        # print("data.x, data.edge_index, data.edge_attr:", data.x, data.edge_index, data.edge_attr)
+        # print(Chem.MolToSmiles(graph_mol))
+        # import pdb; pdb.set_trace()
+        return data
+
+    def extract(self, graph_mol):
+        model = GNN_feature(num_layer=5, emb_dim=300, num_tasks=1, JK='last', drop_ratio=0, graph_pooling='mean', gnn_type='gin')
+        model.from_pretrained(self.pretrained_model_path)
+        # model.cuda(device=0)
+        model.eval()
+        graph_data = self.preprocessing(graph_mol)
+        # graph_data = graph_data.cuda(device=0)
+        with torch.no_grad():
+            node_features = model(graph_data.x, graph_data.edge_index, graph_data.edge_attr)
+        del model
+        return node_features
+