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

torch_molecule/generator/graph_dit/transformer.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import PlaceHolder
+from ...nn import MLP, AttentionWithNodeMask
+from ...nn.embedder import TimestepEmbedder, CategoricalEmbedder, ClusterContinuousEmbedder
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        max_node,
+        hidden_size=384,
+        num_layer=12,
+        num_head=16,
+        mlp_ratio=4.0,
+        dropout=0.,
+        drop_condition=0.1,
+        input_dim_X=118,
+        input_dim_E=5,
+        input_dim_y=None,
+        task_type=[], # a list of 'regression' or 'classification'
+    ):
+        super().__init__()
+        self.input_dim_y = input_dim_y
+        self.x_embedder = nn.Linear(input_dim_X + max_node * input_dim_E, hidden_size, bias=False)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+
+        if input_dim_y > 0 and len(task_type) > 0:
+            self.y_embedder_list = torch.nn.ModuleList()
+            for i in range(input_dim_y):
+                if task_type[i] == 'regression':
+                    self.y_embedder_list.append(ClusterContinuousEmbedder(1, hidden_size, drop_condition))
+                else:
+                    self.y_embedder_list.append(CategoricalEmbedder(2, hidden_size, drop_condition))
+        else:
+            self.y_embedder_list = None
+
+        self.blocks = nn.ModuleList(
+            [
+                AttentionBlock(hidden_size, num_head, mlp_ratio=mlp_ratio, dropout=dropout)
+                for _ in range(num_layer)
+            ]
+        )
+
+        self.final_layer = FinalLayer(
+            max_node=max_node,
+            hidden_size=hidden_size,
+            atom_type=input_dim_X,
+            bond_type=input_dim_E,
+            mlp_ratio=mlp_ratio,
+            num_head=num_head,
+        )
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        def _constant_init(module, i):
+            if isinstance(module, nn.Linear):
+                nn.init.constant_(module.weight, i)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, i)
+
+        self.apply(_basic_init)
+
+        for block in self.blocks :
+            _constant_init(block.adaLN_modulation[0], 0)
+        _constant_init(self.final_layer.adaLN_modulation[0], 0)
+
+    def forward(self, noisy_data, unconditioned):
+        X, E = noisy_data['X_t'].float(), noisy_data['E_t'].float()
+        X_in, E_in = X, E
+        node_mask, t =  noisy_data['node_mask'], noisy_data['t']
+        y = noisy_data['y_t'] if self.y_embedder_list is not None else None
+        assert (y is None) == (self.y_embedder_list is None), "y and y_embedder_list must both be None or both be not None"
+
+        if y is not None:
+            force_drop_id = torch.zeros_like(y.sum(-1))
+            force_drop_id[torch.isnan(y.sum(-1))] = 1
+            if unconditioned:
+                force_drop_id = torch.ones_like(y[:, 0])
+        
+        bs, n, _ = X.size()
+        h = torch.cat([X, E.reshape(bs, n, -1)], dim=-1)
+        h = self.x_embedder(h)
+        c = self.t_embedder(t)
+        if self.y_embedder_list is not None:
+            for i in range(self.input_dim_y):
+                c = c + self.y_embedder_list[i](y[:, i:i+1], self.training, force_drop_id)
+        
+        for i, block in enumerate(self.blocks):
+            h = block(h, c, node_mask)
+
+        X, E, y = self.final_layer(h, X_in, E_in, c, t, node_mask)
+        return PlaceHolder(X=X, E=E, y=y).mask(node_mask)
+    
+    def compute_loss(self, noisy_data, true_X, true_E, lw_X, lw_E, unconditioned=False):
+        pred = self.forward(noisy_data, unconditioned=unconditioned)
+        
+        # Reshape predictions and targets
+        true_X = torch.reshape(true_X, (-1, true_X.size(-1)))  # (bs * n, dx)
+        true_E = torch.reshape(true_E, (-1, true_E.size(-1)))  # (bs * n * n, de)
+        masked_pred_X = torch.reshape(pred.X, (-1, pred.X.size(-1)))  # (bs * n, dx)
+        masked_pred_E = torch.reshape(pred.E, (-1, pred.E.size(-1)))   # (bs * n * n, de)
+
+        # Remove masked rows
+        mask_X = (true_X != 0.).any(dim=-1)
+        mask_E = (true_E != 0.).any(dim=-1)
+
+        flat_true_X = true_X[mask_X, :]
+        flat_pred_X = masked_pred_X[mask_X, :]
+        flat_true_E = true_E[mask_E, :]
+        flat_pred_E = masked_pred_E[mask_E, :]
+        
+        # Calculate node and edge losses using cross entropy
+        loss_X = F.cross_entropy(flat_pred_X, torch.argmax(flat_true_X, dim=-1)) if true_X.numel() > 0 else 0.0
+        loss_E = F.cross_entropy(flat_pred_E, torch.argmax(flat_true_E, dim=-1)) if true_E.numel() > 0 else 0.0
+        loss = lw_X * loss_X + lw_E * loss_E
+        return loss, loss_X, loss_E
+
+class AttentionBlock(nn.Module):
+    def __init__(self, hidden_size, num_head, mlp_ratio=4.0, dropout=0.):
+        super().__init__()
+        self.dropout = dropout
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False)
+
+        self.attn = AttentionWithNodeMask(
+            hidden_size, num_head=num_head, qkv_bias=True, qk_norm=True, proj_drop=dropout
+        )
+        self.mlp = MLP(
+            in_features=hidden_size,
+            hidden_features=int(hidden_size * mlp_ratio),
+            drop=self.dropout,
+            use_bn=False,
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        
+    def forward(self, x, c, node_mask):
+        (
+            shift_msa,
+            scale_msa,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+        ) = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = x + gate_msa.unsqueeze(1) * modulate(self.norm1(self.attn(x, node_mask=node_mask)), shift_msa, scale_msa)
+        x = x + gate_mlp.unsqueeze(1) * modulate(self.norm2(self.mlp(x)), shift_mlp, scale_mlp)
+        return x
+
+class FinalLayer(nn.Module):
+    def __init__(self, max_node, hidden_size, atom_type, bond_type, mlp_ratio, num_head=None):
+        super().__init__()
+        self.atom_type = atom_type
+        self.bond_type = bond_type
+        final_size = atom_type + max_node * bond_type
+        self.XEdecoder = MLP(in_features=hidden_size, out_features=final_size, drop=0., use_bn=False)
+
+        self.norm_final = nn.LayerNorm(final_size, elementwise_affine=False)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * final_size, bias=True)
+        )
+
+    def forward(self, x, x_in, e_in, c, t, node_mask):
+        x_all = self.XEdecoder(x)
+        bs, n, _ = x_all.size()
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x_all = modulate(self.norm_final(x_all), shift, scale)
+        
+        atom_out = x_all[:, :, :self.atom_type]
+        atom_out = x_in + atom_out
+
+        bond_out = x_all[:, :, self.atom_type:].reshape(bs, n, n, self.bond_type)
+        bond_out = e_in + bond_out
+
+        ##### standardize adj_out
+        edge_mask = (~node_mask)[:, :, None] & (~node_mask)[:, None, :]
+        diag_mask = (
+            torch.eye(n, dtype=torch.bool)
+            .unsqueeze(0)
+            .expand(bs, -1, -1)
+            .type_as(edge_mask)
+        )
+        bond_out.masked_fill_(edge_mask[:, :, :, None], 0)
+        bond_out.masked_fill_(diag_mask[:, :, :, None], 0)
+        bond_out = 1 / 2 * (bond_out + torch.transpose(bond_out, 1, 2))
+        return atom_out, bond_out, None

torch_molecule/generator/graph_dit/utils.py

@@ -0,0 +1,220 @@
+import torch
+import json
+from rdkit import Chem
+import numpy as np
+from torch_geometric.utils import to_dense_adj, to_dense_batch, remove_self_loops
+
+def to_dense(x, edge_index, edge_attr, batch, max_num_nodes=None):
+    X, node_mask = to_dense_batch(x=x, batch=batch, max_num_nodes=max_num_nodes)
+    edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
+    if max_num_nodes is None:
+        max_num_nodes = X.size(1)
+    E = to_dense_adj(edge_index=edge_index, batch=batch, edge_attr=edge_attr, max_num_nodes=max_num_nodes)
+    E = encode_no_edge(E)
+    return PlaceHolder(X=X, E=E, y=None), node_mask
+
+def encode_no_edge(E):
+    assert len(E.shape) == 4
+    if E.shape[-1] == 0:
+        return E
+    no_edge = torch.sum(E, dim=3) == 0
+    first_elt = E[:, :, :, 0]
+    first_elt[no_edge] = 1
+    E[:, :, :, 0] = first_elt
+    diag = torch.eye(E.shape[1], dtype=torch.bool).unsqueeze(0).expand(E.shape[0], -1, -1)
+    E[diag] = 0
+    return E
+
+class PlaceHolder:
+    def __init__(self, X, E, y):
+        self.X = X
+        self.E = E
+        self.y = y
+
+    def type_as(self, x: torch.Tensor, categorical: bool = False):
+        """ Changes the device and dtype of X, E, y. """
+        self.X = self.X.type_as(x)
+        self.E = self.E.type_as(x)
+        if categorical:
+            self.y = self.y.type_as(x)
+        return self
+
+    def mask(self, node_mask, collapse=False):
+        x_mask = node_mask.unsqueeze(-1)          # bs, n, 1
+        e_mask1 = x_mask.unsqueeze(2)             # bs, n, 1, 1
+        e_mask2 = x_mask.unsqueeze(1)             # bs, 1, n, 1
+
+        if collapse:
+            self.X = torch.argmax(self.X, dim=-1)
+            self.E = torch.argmax(self.E, dim=-1)
+
+            self.X[node_mask == 0] = - 1
+            self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = - 1
+        else:
+            self.X = self.X * x_mask
+            self.E = self.E * e_mask1 * e_mask2
+            assert torch.allclose(self.E, torch.transpose(self.E, 1, 2))
+        return self
+
+def compute_dataset_info(smiles_or_mol_list, cache_path=None):
+    pt = Chem.GetPeriodicTable()
+    atom_name_list = []
+    atom_count_list = []
+    for i in range(2, 119):
+        atom_name_list.append(pt.GetElementSymbol(i))
+        atom_count_list.append(0)
+    atom_name_list.append('*')
+    atom_count_list.append(0)
+
+    bond_count_list = [0, 0, 0, 0, 0]
+    bond_type_to_index =  {Chem.BondType.SINGLE: 1, Chem.BondType.DOUBLE: 2, Chem.BondType.TRIPLE: 3, Chem.BondType.AROMATIC: 4}
+    tansition_E = np.zeros((118, 118, 5))
+
+    n_atom_list = []
+    n_bond_list = []
+
+    n_atoms_per_mol_count = {}
+    max_node = 0
+    for i, sms_or_mol in enumerate(smiles_or_mol_list):
+        if isinstance(sms_or_mol, str):
+            mol = Chem.MolFromSmiles(sms_or_mol)
+        else:
+            mol = sms_or_mol
+
+        n_atom = mol.GetNumHeavyAtoms()
+        n_bond = mol.GetNumBonds()
+        n_atom_list.append(n_atom)
+        n_bond_list.append(n_bond)
+        max_node = max(max_node, n_atom)
+
+        # Count atoms per molecule size
+        if n_atom in n_atoms_per_mol_count:
+            n_atoms_per_mol_count[n_atom] += 1
+        else:
+            n_atoms_per_mol_count[n_atom] = 1
+
+        cur_atom_count_arr = np.zeros(118)
+        for atom in mol.GetAtoms():
+            symbol = atom.GetSymbol()
+            if symbol == 'H':
+                continue
+            elif symbol == '*':
+                atom_count_list[-1] += 1
+                cur_atom_count_arr[-1] += 1
+            else:
+                atom_count_list[atom.GetAtomicNum()-2] += 1
+                cur_atom_count_arr[atom.GetAtomicNum()-2] += 1
+        
+        tansition_E_temp = np.zeros((118, 118, 5))
+        for bond in mol.GetBonds():
+            start_atom, end_atom = bond.GetBeginAtom(), bond.GetEndAtom()
+            if start_atom.GetSymbol() == 'H' or end_atom.GetSymbol() == 'H':
+                continue
+            
+            if start_atom.GetSymbol() == '*':
+                start_index = 117
+            else:
+                start_index = start_atom.GetAtomicNum() - 2
+            if end_atom.GetSymbol() == '*':
+                end_index = 117
+            else:
+                end_index = end_atom.GetAtomicNum() - 2
+
+            bond_type = bond.GetBondType()
+            bond_index = bond_type_to_index[bond_type]
+            bond_count_list[bond_index] += 2
+
+            tansition_E[start_index, end_index, bond_index] += 2
+            tansition_E[end_index, start_index, bond_index] += 2
+            tansition_E_temp[start_index, end_index, bond_index] += 2
+            tansition_E_temp[end_index, start_index, bond_index] += 2
+
+        bond_count_list[0] += n_atom * (n_atom - 1) - n_bond * 2
+        cur_tot_bond = cur_atom_count_arr.reshape(-1,1) * cur_atom_count_arr.reshape(1,-1) * 2 # 118 * 118
+        cur_tot_bond = cur_tot_bond - np.diag(cur_atom_count_arr) * 2 # 118 * 118
+        tansition_E[:, :, 0] += cur_tot_bond - tansition_E_temp.sum(axis=-1)
+        assert (cur_tot_bond > tansition_E_temp.sum(axis=-1)).sum() >= 0, f'i:{i}, sms:{sms}'
+    
+    # Create n_atoms_per_mol array with proper size
+    n_atoms_per_mol = [0] * (max_node + 1)
+    for n_atom, count in n_atoms_per_mol_count.items():
+        n_atoms_per_mol[n_atom] = count
+
+    n_atoms_per_mol = np.array(n_atoms_per_mol) / np.sum(n_atoms_per_mol)
+    n_atoms_per_mol = n_atoms_per_mol.tolist()
+
+    atom_count_list = np.array(atom_count_list) / np.sum(atom_count_list)
+    active_atoms = np.array(atom_name_list)[atom_count_list > 0]
+    active_atoms = active_atoms.tolist()
+    atom_count_list = atom_count_list.tolist()
+
+    bond_count_list = np.array(bond_count_list) / np.sum(bond_count_list)
+    bond_count_list = bond_count_list.tolist()
+
+    no_edge = np.sum(tansition_E, axis=-1) == 0
+    first_elt = tansition_E[:, :, 0]
+    first_elt[no_edge] = 1
+    tansition_E[:, :, 0] = first_elt
+    tansition_E = tansition_E / np.sum(tansition_E, axis=-1, keepdims=True)
+    
+    node_types = torch.Tensor(atom_count_list)
+    edge_types = torch.Tensor(bond_count_list)
+    active_index = (node_types > 0).nonzero().squeeze()
+
+    x_margins = node_types.float() / torch.sum(node_types)
+    x_margins = x_margins[active_index]
+    e_margins = edge_types.float() / torch.sum(edge_types)
+
+    xe_conditions = torch.Tensor(tansition_E)
+    xe_conditions = xe_conditions[active_index][:, active_index] 
+    
+    xe_conditions = xe_conditions.sum(dim=1)
+    ex_conditions = xe_conditions.t()
+    xe_conditions = xe_conditions / (xe_conditions.sum(dim=-1, keepdim=True) + 1e-10)
+    ex_conditions = ex_conditions / (ex_conditions.sum(dim=-1, keepdim=True) + 1e-10)
+    num_nodes_dist = DistributionNodes(torch.Tensor(n_atoms_per_mol))
+
+    meta_dict = {
+        'active_index': active_index,
+        'x_margins': x_margins,
+        'e_margins': e_margins,
+        'xe_conditions': xe_conditions,
+        'ex_conditions': ex_conditions,
+        'atom_decoder': active_atoms,
+        'num_nodes_dist': num_nodes_dist,
+        'max_node': max_node, 
+        }
+    if cache_path is not None:
+        with open(cache_path, "w") as f:
+            json.dump(meta_dict, f)
+
+    return meta_dict
+
+class DistributionNodes:
+    def __init__(self, histogram):
+        """ Compute the distribution of the number of nodes in the dataset, and sample from this distribution.
+            historgram: dict. The keys are num_nodes, the values are counts
+        """
+
+        if type(histogram) == dict:
+            max_n_nodes = max(histogram.keys())
+            prob = torch.zeros(max_n_nodes + 1)
+            for num_nodes, count in histogram.items():
+                prob[num_nodes] = count
+        else:
+            prob = histogram
+
+        self.prob = prob / prob.sum()
+        self.m = torch.distributions.Categorical(prob)
+
+    def sample_n(self, n_samples, device):
+        idx = self.m.sample((n_samples,))
+        return idx.to(device)
+
+    def log_prob(self, batch_n_nodes):
+        assert len(batch_n_nodes.size()) == 1
+        p = self.prob.to(batch_n_nodes.device)
+
+        probas = p[batch_n_nodes]
+        log_p = torch.log(probas + 1e-30)
+        return log_p

torch_molecule/generator/graph_ga/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_graph_ga import GraphGAMolecularGenerator
+
+__all__ = ['GraphGAMolecularGenerator']

torch_molecule/generator/graph_ga/crossover.py

@@ -0,0 +1,182 @@
+import random
+
+import numpy as np
+from rdkit import Chem, rdBase
+from rdkit.Chem import AllChem
+rdBase.DisableLog('rdApp.error')
+
+
+def cut(mol):
+    if not mol.HasSubstructMatch(Chem.MolFromSmarts('[*]-;!@[*]')):
+        return None
+
+    bis = random.choice(mol.GetSubstructMatches(Chem.MolFromSmarts('[*]-;!@[*]')))  # single bond not in ring
+
+    bs = [mol.GetBondBetweenAtoms(bis[0], bis[1]).GetIdx()]
+
+    fragments_mol = Chem.FragmentOnBonds(mol, bs, addDummies=True, dummyLabels=[(1, 1)])
+
+    try:
+        return Chem.GetMolFrags(fragments_mol, asMols=True, sanitizeFrags=True)
+    except ValueError:
+        return None
+
+    return None
+
+
+def cut_ring(mol):
+
+    for i in range(10):
+        if random.random() < 0.5:
+            if not mol.HasSubstructMatch(Chem.MolFromSmarts('[R]@[R]@[R]@[R]')):
+                return None
+            bis = random.choice(mol.GetSubstructMatches(Chem.MolFromSmarts('[R]@[R]@[R]@[R]')))
+            bis = ((bis[0], bis[1]), (bis[2], bis[3]),)
+        else:
+            if not mol.HasSubstructMatch(Chem.MolFromSmarts('[R]@[R;!D2]@[R]')):
+                return None
+            bis = random.choice(mol.GetSubstructMatches(Chem.MolFromSmarts('[R]@[R;!D2]@[R]')))
+            bis = ((bis[0], bis[1]), (bis[1], bis[2]),)
+
+        bs = [mol.GetBondBetweenAtoms(x, y).GetIdx() for x, y in bis]
+
+        fragments_mol = Chem.FragmentOnBonds(mol, bs, addDummies=True, dummyLabels=[(1, 1), (1, 1)])
+
+        try:
+            fragments = Chem.GetMolFrags(fragments_mol, asMols=True, sanitizeFrags=True)
+            if len(fragments) == 2:
+                return fragments
+        except ValueError:
+            return None
+
+    return None
+
+
+def ring_OK(mol):
+    if not mol.HasSubstructMatch(Chem.MolFromSmarts('[R]')):
+        return True
+
+    ring_allene = mol.HasSubstructMatch(Chem.MolFromSmarts('[R]=[R]=[R]'))
+
+    cycle_list = mol.GetRingInfo().AtomRings()
+    max_cycle_length = max([len(j) for j in cycle_list])
+    macro_cycle = max_cycle_length > 6
+
+    double_bond_in_small_ring = mol.HasSubstructMatch(Chem.MolFromSmarts('[r3,r4]=[r3,r4]'))
+
+    return not ring_allene and not macro_cycle and not double_bond_in_small_ring
+
+
+# TODO: set from main? calculate for dataset?
+average_size = 39.15
+size_stdev = 3.50
+
+
+def mol_ok(mol):
+    try:
+        Chem.SanitizeMol(mol)
+        target_size = size_stdev * np.random.randn() + average_size  # parameters set in GA_mol
+        if mol.GetNumAtoms() > 5 and mol.GetNumAtoms() < target_size:
+            return True
+        else:
+            return False
+    except ValueError:
+        return False
+
+
+def crossover_ring(parent_A, parent_B):
+    ring_smarts = Chem.MolFromSmarts('[R]')
+    if not parent_A.HasSubstructMatch(ring_smarts) and not parent_B.HasSubstructMatch(ring_smarts):
+        return None
+
+    rxn_smarts1 = ['[*:1]~[1*].[1*]~[*:2]>>[*:1]-[*:2]', '[*:1]~[1*].[1*]~[*:2]>>[*:1]=[*:2]']
+    rxn_smarts2 = ['([*:1]~[1*].[1*]~[*:2])>>[*:1]-[*:2]', '([*:1]~[1*].[1*]~[*:2])>>[*:1]=[*:2]']
+
+    for i in range(10):
+        fragments_A = cut_ring(parent_A)
+        fragments_B = cut_ring(parent_B)
+
+        if fragments_A is None or fragments_B is None:
+            return None
+
+        new_mol_trial = []
+        for rs in rxn_smarts1:
+            rxn1 = AllChem.ReactionFromSmarts(rs)
+            new_mol_trial = []
+            for fa in fragments_A:
+                for fb in fragments_B:
+                    if len(rxn1.RunReactants((fa, fb))) > 0:
+                        new_mol_trial.append(rxn1.RunReactants((fa, fb))[0])
+
+        new_mols = []
+        for rs in rxn_smarts2:
+            rxn2 = AllChem.ReactionFromSmarts(rs)
+            for m in new_mol_trial:
+                m = m[0]
+                if mol_ok(m):
+                    new_mols += list(rxn2.RunReactants((m,)))
+
+        new_mols2 = []
+        for m in new_mols:
+            m = m[0]
+            if mol_ok(m) and ring_OK(m):
+                new_mols2.append(m)
+
+        if len(new_mols2) > 0:
+            return random.choice(new_mols2)
+
+    return None
+
+
+def crossover_non_ring(parent_A, parent_B):
+
+    for i in range(10):
+        fragments_A = cut(parent_A)
+        fragments_B = cut(parent_B)
+        if fragments_A is None or fragments_B is None:
+            return None
+        rxn = AllChem.ReactionFromSmarts('[*:1]-[1*].[1*]-[*:2]>>[*:1]-[*:2]')
+        new_mol_trial = []
+        for fa in fragments_A:
+            for fb in fragments_B:
+                if len(rxn.RunReactants((fa, fb))) > 0:
+                    new_mol_trial.append(rxn.RunReactants((fa, fb))[0])
+
+        new_mols = []
+        for mol in new_mol_trial:
+            mol = mol[0]
+            if mol_ok(mol):
+                new_mols.append(mol)
+
+        if len(new_mols) > 0:
+            return random.choice(new_mols)
+
+    return None
+
+
+def crossover(parent_A, parent_B):
+    parent_smiles = [Chem.MolToSmiles(parent_A), Chem.MolToSmiles(parent_B)]
+    try:
+        Chem.Kekulize(parent_A, clearAromaticFlags=True)
+        Chem.Kekulize(parent_B, clearAromaticFlags=True)
+
+    except ValueError:
+        pass
+
+    for i in range(10):
+        if random.random() <= 0.5:
+            # print 'non-ring crossover'
+            new_mol = crossover_non_ring(parent_A, parent_B)
+            if new_mol is not None:
+                new_smiles = Chem.MolToSmiles(new_mol)
+                if new_smiles is not None and new_smiles not in parent_smiles:
+                    return new_mol
+        else:
+            # print 'ring crossover'
+            new_mol = crossover_ring(parent_A, parent_B)
+            if new_mol is not None:
+                new_smiles = Chem.MolToSmiles(new_mol)
+                if new_smiles is not None and new_smiles not in parent_smiles:
+                    return new_mol
+
+    return None