torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/generator/graph_ga/modeling_graph_ga.py

@@ -0,0 +1,310 @@
+import numpy as np
+import random
+import joblib
+from joblib import delayed
+from rdkit import Chem
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, Tuple, List, Callable
+from dataclasses import dataclass, field
+from warnings import warn
+
+from .crossover import crossover
+from .mutate import mutate
+from .oracle import Oracle
+
+from ...base import BaseMolecularGenerator
+
+@dataclass
+class GraphGAMolecularGenerator(BaseMolecularGenerator):
+    """This generator implements the Graph Genetic Algorithm for molecular generation.
+    
+    References
+    ----------
+    - A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of 
+    Median Molecules. Journal of Chemical Information and Computer Sciences. https://pubs.acs.org/doi/10.1021/ci034290p
+    - Implementation: https://github.com/wenhao-gao/mol_opt
+
+    Parameters
+    ----------
+    num_task : int, default=0
+        Number of properties to condition on. Set to 0 for unconditional generation.
+    population_size : int, default=100
+        Size of the population in each iteration.
+    offspring_size : int, default=50
+        Number of offspring molecules to generate in each iteration.
+    mutation_rate : float, default=0.0067
+        Probability of mutation occurring during reproduction.
+    n_jobs : int, default=1
+        Number of parallel jobs to run. -1 means using all processors.
+    iteration : int, default=5
+        Number of iterations for each target label (or random sample) to run the genetic algorithm.
+    verbose : bool, default=False
+        Whether to display progress bars and logs.
+    """
+
+    # GA parameters
+    num_task: int = 0
+    population_size: int = 100
+    offspring_size: int = 50
+    mutation_rate: float = 0.0067
+    n_jobs: int = 1
+    iteration: int = 5
+    
+    # Other parameters
+    verbose: bool = False
+    model_name: str = "GraphGAMolecularGenerator"
+    model_class = None
+    
+    def __post_init__(self):
+        super().__post_init__()
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return [
+            "num_task", "population_size", "offspring_size", "mutation_rate",
+            "n_jobs", "iteration", "verbose"
+        ]
+    
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        raise NotImplementedError("GraphGA does not support getting model parameters")
+
+    def save_to_local(self, path: str):
+        joblib.dump(self.oracle, path)
+        if self.verbose:
+            print(f"Saved oracle to {path}")
+
+    def load_from_local(self):
+        raise NotImplementedError(
+            "GraphGA does not support loading from local. "
+            "If you want to load the oracles saved through save_to_local, "
+            "you need to manually load the oracle from the path with joblib.load(path) "
+            "and pass it to the fit function."
+        )
+    
+    def save_to_hf(self, repo_id: str, task_id: str = "default"):
+        raise NotImplementedError("GraphGA does not support pushing to huggingface")
+    
+    def load_from_hf(self, repo_id: str, task_id: str = "default"):
+        raise NotImplementedError("GraphGA does not support loading from huggingface")
+    
+    def _setup_optimizers(self):
+        raise NotImplementedError("GraphGA does not support setting up optimizers")
+    
+    def _train_epoch(self, train_loader, optimizer):
+        raise NotImplementedError("GraphGA does not support training epochs")
+
+    def fit(
+        self,
+        X_train: List[str],
+        y_train: Optional[Union[List, np.ndarray]] = None,
+        oracle: Optional[List[Callable]] = None
+    ) -> "GraphGAMolecularGenerator":
+        """Fit the model to the training data.
+
+        Parameters
+        ----------
+        X_train : List[str]
+            Training data, which will be used as the initial population.
+        y_train : Optional[Union[List, np.ndarray]]
+            Training labels for conditional generation (num_task is not 0).
+        oracle : Optional[Callable]
+            Oracle used to score the generated molecules. If not provided, default oracles based on 
+            ``sklearn.ensemble.RandomForestRegressor`` are trained on the X_train and y_train.
+            
+            For a customized oracle, it should be a Callable object, i.e., ``oracle(X, y)``.
+            Please properly wrap your oracle to take two inputs:
+              - a list of ``rdkit.Chem.rdchem.Mol`` objects and 
+              - a (1, num_task) numpy array of target values that all the molecules in the list target to achieve. Take care of NaN values if any.
+            
+            Scores for different tasks should be aggregated, i.e., mean or sum. The return should be a list of scores (float).
+            Smaller scores mean closer to the target goal.
+            
+            Oracles are not needed for unconditional generation.
+
+        Returns
+        -------
+        self : GraphGAMolecularGenerator
+            Fitted model.
+        """
+        self.y_train = None
+        if oracle is not None:
+            self.oracle = oracle
+        else:
+            X_train, y_train = self._validate_inputs(X_train, y_train, num_task=self.num_task, return_rdkit_mol=False)
+            if y_train is not None:
+                warn("No oracles provided but y_train is provided, using default oracles (RandomForestRegressor)", UserWarning)
+                self.oracle = Oracle(num_task=self.num_task)
+                self.oracle.fit(X_train, y_train)
+                self.y_train = y_train
+            else:
+                assert self.num_task == 0, "No oracles or y_train provided but num_task is not 0"
+                self.oracle = None
+                    
+        self.X_train = X_train
+        self.is_fitted_ = True
+        return self
+
+    def _make_mating_pool(self, population_mol, population_scores, offspring_size: int):
+        """Create mating pool where smaller scores have higher selection probabilities."""
+        max_score = max(population_scores)
+        # Invert scores so that smaller scores become larger probabilities
+        inverted_scores = [max_score - s + 1e-6 for s in population_scores]  # Add small constant to avoid zeros
+        sum_scores = sum(inverted_scores)
+        population_probs = [p / sum_scores for p in inverted_scores]
+        mating_pool = np.random.choice(population_mol, p=population_probs, size=offspring_size, replace=True)
+        return mating_pool
+
+    def _reproduce(self, mating_pool, mutation_rate):
+        """Create new molecule through crossover and mutation."""
+        parent_a = random.choice(mating_pool)
+        parent_b = random.choice(mating_pool)
+        new_child = crossover(parent_a, parent_b)
+        if new_child is not None:
+            new_child = mutate(new_child, mutation_rate)
+        return new_child
+
+    def _sanitize_molecules(self, population_mol):
+        """Sanitize molecules by removing duplicates and invalid molecules."""
+        new_mol_list = []
+        smiles_set = set()
+        for mol in population_mol:
+            if mol is not None:
+                try:
+                    smiles = Chem.MolToSmiles(mol)
+                    if smiles is not None and smiles not in smiles_set:
+                        smiles_set.add(smiles)
+                        new_mol_list.append(mol)
+                except ValueError:
+                    pass
+        return new_mol_list
+    
+    def _get_score(self, mol_list, label):
+        if label is None:
+            return [1.0] * len(mol_list)  # For unconditional generation
+        return self.oracle(mol_list, label)
+
+    def generate(
+        self, 
+        labels: Optional[Union[List[List], np.ndarray]] = None,
+        num_samples: int = 32
+    ) -> List[str]:
+        """Generate molecules using genetic algorithm optimization."""
+        if not self.is_fitted_:
+            raise RuntimeError("Model must be fitted before generating")
+        
+        all_generated_mols = []
+        
+        if labels is not None:
+            try:
+                labels = np.array(labels).reshape(-1, self.num_task)
+            except:
+                raise ValueError(f"labels must be convertible to a numpy array with shape (-1, {self.num_task})")
+            
+            # Prepare all inputs for parallel processing
+            parallel_inputs = []
+            for i in range(labels.shape[0]):
+                label = labels[i:i+1]  # Keep as 2D array
+                
+                # Initialize population based on similarity to target label
+                if self.y_train is not None:
+                    population_mol = self._initialize_population_for_label(label)
+                else:
+                    population_idx = np.random.choice(len(self.X_train), min(self.population_size, len(self.X_train)))
+                    population_smiles = [self.X_train[idx] for idx in population_idx]
+                    population_mol = [Chem.MolFromSmiles(s) for s in population_smiles]
+                
+                parallel_inputs.append((population_mol, label))
+            
+            # Run GA for all labels in parallel with tqdm progress bar
+            if self.verbose:
+                results = joblib.Parallel(n_jobs=self.n_jobs)(
+                    delayed(self._run_generation)(pop_mol, lbl) 
+                    for pop_mol, lbl in tqdm(parallel_inputs, desc="Generating molecules")
+                )
+            else:
+                results = joblib.Parallel(n_jobs=self.n_jobs)(
+                    delayed(self._run_generation)(pop_mol, lbl) 
+                    for pop_mol, lbl in parallel_inputs
+                )
+            
+            # Convert results to SMILES in the original order
+            all_generated_mols = [Chem.MolToSmiles(mol) for mol in results]
+        else:
+            # Prepare all inputs for parallel processing
+            parallel_inputs = []
+            for _ in range(num_samples):
+                population_idx = np.random.choice(len(self.X_train), min(self.population_size, len(self.X_train)))
+                population_smiles = [self.X_train[idx] for idx in population_idx]
+                population_mol = [Chem.MolFromSmiles(s) for s in population_smiles]
+                parallel_inputs.append((population_mol, None))
+            
+            # Run GA for all samples in parallel with tqdm progress bar
+            if self.verbose:
+                results = joblib.Parallel(n_jobs=self.n_jobs)(
+                    delayed(self._run_generation)(pop_mol, lbl) 
+                    for pop_mol, lbl in tqdm(parallel_inputs, desc="Generating molecules", total=num_samples)
+                )
+            else:
+                results = joblib.Parallel(n_jobs=self.n_jobs)(
+                    delayed(self._run_generation)(pop_mol, lbl) 
+                    for pop_mol, lbl in parallel_inputs
+                )
+            
+            # Convert results to SMILES
+            all_generated_mols = [Chem.MolToSmiles(mol) for mol in results]
+        
+        return all_generated_mols
+    
+    def _initialize_population_for_label(self, label):
+        """Initialize population based on similarity to target label."""
+        similarities = []
+        
+        for i in range(len(self.X_train)):
+            sample_label = self.y_train[i]
+            similarity = -np.nansum((sample_label - label[0])**2)
+            similarities.append((i, similarity))
+        
+        if similarities:
+            similarities.sort(key=lambda x: x[1], reverse=True)
+            top_indices = [x[0] for x in similarities[:self.population_size]]
+        else:
+            top_indices = np.random.choice(len(self.X_train), min(self.population_size, len(self.X_train)))
+            
+        population_smiles = [self.X_train[i] for i in top_indices]
+        return [Chem.MolFromSmiles(s) for s in population_smiles]
+    
+    def _run_generation(self, population_mol, label):
+        """Run the genetic algorithm for a specific population and label."""
+        for generation_idx in range(self.iteration):
+            if label is not None:
+                population_scores = self._get_score(population_mol, label)
+            else:
+                population_scores = [1.0] * len(population_mol)  # For unconditional generation
+            
+            mating_pool = self._make_mating_pool(population_mol, population_scores, self.offspring_size)
+            
+            # Create offspring sequentially (parallelization is at the higher level now)
+            offspring_mol = []
+            for _ in range(self.offspring_size):
+                offspring = self._reproduce(mating_pool, self.mutation_rate)
+                offspring_mol.append(offspring)
+            
+            population_mol += offspring_mol
+            population_mol = self._sanitize_molecules(population_mol)
+
+            # Re-score the expanded population
+            if label is not None:
+                population_scores = self._get_score(population_mol, label)
+            else:
+                population_scores = [1.0] * len(population_mol)
+            
+            # Select top molecules for next generation
+            population_tuples = list(zip(population_scores, population_mol))
+            population_tuples = sorted(population_tuples, key=lambda x: x[0], reverse=False) # lower score is better
+            population_tuples = population_tuples[:self.population_size]
+            
+            population_mol = [t[1] for t in population_tuples]
+        
+        # Return the best molecule
+        return population_mol[0]
+    

torch_molecule/generator/graph_ga/mutate.py

@@ -0,0 +1,136 @@
+import random
+
+import numpy as np
+from rdkit import Chem, rdBase
+from rdkit.Chem import AllChem
+
+from .crossover import mol_ok, ring_OK
+
+rdBase.DisableLog('rdApp.error')
+
+def delete_atom():
+    choices = ['[*:1]~[D1:2]>>[*:1]', '[*:1]~[D2:2]~[*:3]>>[*:1]-[*:3]',
+               '[*:1]~[D3:2](~[*;!H0:3])~[*:4]>>[*:1]-[*:3]-[*:4]',
+               '[*:1]~[D4:2](~[*;!H0:3])(~[*;!H0:4])~[*:5]>>[*:1]-[*:3]-[*:4]-[*:5]',
+               '[*:1]~[D4:2](~[*;!H0;!H1:3])(~[*:4])~[*:5]>>[*:1]-[*:3](-[*:4])-[*:5]']
+    p = [0.25, 0.25, 0.25, 0.1875, 0.0625]
+
+    return np.random.choice(choices, p=p)
+
+
+def append_atom():
+    choices = [['single', ['C', 'N', 'O', 'F', 'S', 'Cl', 'Br'], 7 * [1.0 / 7.0]],
+               ['double', ['C', 'N', 'O'], 3 * [1.0 / 3.0]],
+               ['triple', ['C', 'N'], 2 * [1.0 / 2.0]]]
+    p_BO = [0.60, 0.35, 0.05]
+
+    index = np.random.choice(list(range(3)), p=p_BO)
+
+    BO, atom_list, p = choices[index]
+    new_atom = np.random.choice(atom_list, p=p)
+
+    if BO == 'single':
+        rxn_smarts = '[*;!H0:1]>>[*:1]X'.replace('X', '-' + new_atom)
+    if BO == 'double':
+        rxn_smarts = '[*;!H0;!H1:1]>>[*:1]X'.replace('X', '=' + new_atom)
+    if BO == 'triple':
+        rxn_smarts = '[*;H3:1]>>[*:1]X'.replace('X', '#' + new_atom)
+
+    return rxn_smarts
+
+
+def insert_atom():
+    choices = [['single', ['C', 'N', 'O', 'S'], 4 * [1.0 / 4.0]],
+               ['double', ['C', 'N'], 2 * [1.0 / 2.0]],
+               ['triple', ['C'], [1.0]]]
+    p_BO = [0.60, 0.35, 0.05]
+
+    index = np.random.choice(list(range(3)), p=p_BO)
+
+    BO, atom_list, p = choices[index]
+    new_atom = np.random.choice(atom_list, p=p)
+
+    if BO == 'single':
+        rxn_smarts = '[*:1]~[*:2]>>[*:1]X[*:2]'.replace('X', new_atom)
+    if BO == 'double':
+        rxn_smarts = '[*;!H0:1]~[*:2]>>[*:1]=X-[*:2]'.replace('X', new_atom)
+    if BO == 'triple':
+        rxn_smarts = '[*;!R;!H1;!H0:1]~[*:2]>>[*:1]#X-[*:2]'.replace('X', new_atom)
+
+    return rxn_smarts
+
+
+def change_bond_order():
+    choices = ['[*:1]!-[*:2]>>[*:1]-[*:2]', '[*;!H0:1]-[*;!H0:2]>>[*:1]=[*:2]',
+               '[*:1]#[*:2]>>[*:1]=[*:2]', '[*;!R;!H1;!H0:1]~[*:2]>>[*:1]#[*:2]']
+    p = [0.45, 0.45, 0.05, 0.05]
+
+    return np.random.choice(choices, p=p)
+
+
+def delete_cyclic_bond():
+    return '[*:1]@[*:2]>>([*:1].[*:2])'
+
+
+def add_ring():
+    choices = ['[*;!r;!H0:1]~[*;!r:2]~[*;!r;!H0:3]>>[*:1]1~[*:2]~[*:3]1',
+               '[*;!r;!H0:1]~[*!r:2]~[*!r:3]~[*;!r;!H0:4]>>[*:1]1~[*:2]~[*:3]~[*:4]1',
+               '[*;!r;!H0:1]~[*!r:2]~[*:3]~[*:4]~[*;!r;!H0:5]>>[*:1]1~[*:2]~[*:3]~[*:4]~[*:5]1',
+               '[*;!r;!H0:1]~[*!r:2]~[*:3]~[*:4]~[*!r:5]~[*;!r;!H0:6]>>[*:1]1~[*:2]~[*:3]~[*:4]~[*:5]~[*:6]1']
+    p = [0.05, 0.05, 0.45, 0.45]
+
+    return np.random.choice(choices, p=p)
+
+
+def change_atom(mol):
+    choices = ['#6', '#7', '#8', '#9', '#16', '#17', '#35']
+    p = [0.15, 0.15, 0.14, 0.14, 0.14, 0.14, 0.14]
+
+    X = np.random.choice(choices, p=p)
+    while not mol.HasSubstructMatch(Chem.MolFromSmarts('[' + X + ']')):
+        X = np.random.choice(choices, p=p)
+    Y = np.random.choice(choices, p=p)
+    while Y == X:
+        Y = np.random.choice(choices, p=p)
+
+    return '[X:1]>>[Y:1]'.replace('X', X).replace('Y', Y)
+
+
+def mutate(mol, mutation_rate):
+    if random.random() > mutation_rate:
+        return mol
+
+    try:
+        Chem.Kekulize(mol, clearAromaticFlags=True)
+    except ValueError:
+        return mol
+
+    p = [0.15, 0.14, 0.14, 0.14, 0.14, 0.14, 0.15]
+    for i in range(10):
+        rxn_smarts_list = 7 * ['']
+        rxn_smarts_list[0] = insert_atom()
+        rxn_smarts_list[1] = change_bond_order()
+        rxn_smarts_list[2] = delete_cyclic_bond()
+        rxn_smarts_list[3] = add_ring()
+        rxn_smarts_list[4] = delete_atom()
+        rxn_smarts_list[5] = change_atom(mol)
+        rxn_smarts_list[6] = append_atom()
+        rxn_smarts = np.random.choice(rxn_smarts_list, p=p)
+
+        # print 'mutation',rxn_smarts
+
+        rxn = AllChem.ReactionFromSmarts(rxn_smarts)
+
+        new_mol_trial = rxn.RunReactants((mol,))
+
+        new_mols = []
+        for m in new_mol_trial:
+            m = m[0]
+            # print Chem.MolToSmiles(mol),mol_ok(mol)
+            if mol_ok(m) and ring_OK(m):
+                new_mols.append(m)
+
+        if len(new_mols) > 0:
+            return random.choice(new_mols)
+
+    return None

torch_molecule/generator/graph_ga/oracle.py

@@ -0,0 +1,95 @@
+import numpy as np
+from sklearn.ensemble import RandomForestRegressor
+from rdkit import Chem
+from typing import List, Any
+from ...utils.graph.features import getmorganfingerprint
+
+
+class Oracle:
+    """The default Oracle class for scoring molecules in GraphGA.
+    
+    This class wraps predictive models (like RandomForestRegressor) to score molecules
+    based on their properties. It handles conversion of SMILES to fingerprints.
+    
+    Parameters
+    ----------
+    models : List[Any], optional
+        List of trained models that implement a predict method.
+        If None, RandomForestRegressors will be created when fit is called.
+    num_task : int, default=1
+        Number of properties to predict.
+    """
+    
+    def __init__(self, models=None, num_task=1):
+        self.models = models if models is not None else [RandomForestRegressor() for _ in range(num_task)]
+        self.num_task = num_task if models is None else len(models)
+        
+    def _convert_to_fingerprint(self, molecules):
+        """Convert SMILES or RDKit molecules to fingerprints."""
+        if isinstance(molecules[0], str):
+            return np.array([getmorganfingerprint(Chem.MolFromSmiles(mol)) for mol in molecules])
+        else:
+            return np.array([getmorganfingerprint(mol) for mol in molecules])
+    
+    def fit(self, X_train, y_train):
+        """Fit the underlying models with training data.
+        
+        Parameters
+        ----------
+        X_train : List[str] or List[RDKit.Mol]
+            Training molecules as SMILES strings or RDKit Mol objects.
+        y_train : np.ndarray
+            Training labels with shape (n_samples, num_task).
+        
+        Returns
+        -------
+        self : Oracle
+            Fitted oracle.
+        """
+        X_train_fp = self._convert_to_fingerprint(X_train)
+        
+        for i in range(self.num_task):
+            nan_mask = ~np.isnan(y_train[:, i])
+            y_train_ = y_train[:, i][nan_mask]
+            X_train_fp_ = X_train_fp[nan_mask]
+            self.models[i].fit(X_train_fp_, y_train_)
+            
+        return self
+    
+    def __call__(self, molecules, target_values):
+        """Score molecules based on their predicted properties.
+        
+        Parameters
+        ----------
+        molecules : List[str] or List[RDKit.Mol]
+            Molecules to score as SMILES strings or RDKit Mol objects.
+        target_values : np.ndarray,
+            Scores will be based on distance to these targets.
+            
+        Returns
+        -------
+        List[float]
+            Scores for each molecule.
+        """
+        fps = self._convert_to_fingerprint(molecules)
+        scores_list = []
+        
+        for i, fp in enumerate(fps):
+            if self.num_task == 1:
+                score = self.models[0].predict([fp])[0]
+                scores_list.append(float(score))
+            else:
+                mol_scores = []
+                for idx in range(self.num_task):
+                    pred = self.models[idx].predict([fp])[0]
+                    
+                    if target_values is not None and not np.isnan(target_values[0][idx]):
+                        # Lower score for values closer to target
+                        target = target_values[0][idx]
+                        dist = abs(float(pred) - target) / (abs(target) + 1e-8)
+                        mol_scores.append(dist)
+                        
+                score = np.nanmean(mol_scores)  # Lower is better when using distances
+                scores_list.append(float(score))
+
+        return scores_list

torch_molecule/generator/jtvae/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_jtvae import JTVAEMolecularGenerator
+
+__all__ = ['JTVAEMolecularGenerator']