from rdkit import Chem from rdkit.Chem import AllChem from typing import List, Optional, Set import random METAL_SWAP_OPTIONS = { "Pt": ["Pd", "Ni", "Cu", "Rh", "Ru"], "Pd": ["Pt", "Ni", "Cu", "Rh"], "Ni": ["Co", "Fe", "Cu", "Pd"], "Cu": ["Ag", "Au", "Ni", "Pd"], "Fe": ["Co", "Ni", "Ru"], "Ru": ["Rh", "Fe", "Co"], "Rh": ["Ru", "Pd", "Pt"], } LIGAND_VARIANTS = ["=O", "O", "C"] DOPING_ATOMS = ["N", "O", "S"] def generate_3d_molecule(smiles: str): """ Generates a 3D RDKit molecule from SMILES. Returns the molecule object with 3D coordinates or None if failed. Handles fallback for small molecules. """ mol = Chem.MolFromSmiles(smiles) if mol is None: return None mol = Chem.AddHs(mol) params = AllChem.ETKDGv3() params.randomSeed = 42 success = AllChem.EmbedMolecule(mol, params) if success != 0: # Fallback: Compute 2D coords if 3D fails AllChem.Compute2DCoords(mol) return mol try: AllChem.UFFOptimizeMolecule(mol, maxIters=200) except Exception: # Keep embedded geometry even if force field optimization fails. pass return mol def _is_valid_smiles(smiles: str) -> bool: """Return True if RDKit can parse SMILES.""" return Chem.MolFromSmiles(smiles) is not None def _sanitize_and_canonicalize(smiles: str) -> Optional[str]: """Return canonical sanitized SMILES or None for invalid molecules.""" try: mol = Chem.MolFromSmiles(smiles, sanitize=False) if mol is None: return None Chem.SanitizeMol(mol) return Chem.MolToSmiles(mol, canonical=True) except Exception: return None def _metal_symbols_in_molecule(mol: Chem.Mol) -> List[str]: """Return unique metal symbols from a molecule in insertion order.""" seen = set() metals = [] for atom in mol.GetAtoms(): symbol = atom.GetSymbol() if symbol in METAL_SWAP_OPTIONS and symbol not in seen: seen.add(symbol) metals.append(symbol) return metals def mutate_catalyst( base_smiles: str, num_variations: int = 5, mutation_mode: str = "all" ) -> List[str]: """ Generate mutated catalyst candidates from a base SMILES string. Mutation strategies: 1. Metal swapping (mutation_mode="metal_swap") 2. Ligand addition oxide/hydroxyl/methyl (mutation_mode="ligand_oxo" | "ligand_hydroxyl" | "ligand_methyl") 3. Peripheral-atom doping (mutation_mode="doping") Returns only unique, sanitized canonical SMILES strings. """ base_mol = Chem.MolFromSmiles(base_smiles) if base_mol is None: return [] base_canonical = _sanitize_and_canonicalize(base_smiles) valid_mutants: List[str] = [] seen: Set[str] = set() def add_candidate(candidate_smiles: str): canonical = _sanitize_and_canonicalize(candidate_smiles) if canonical is None: return if base_canonical is not None and canonical == base_canonical: return if canonical in seen: return seen.add(canonical) valid_mutants.append(canonical) do_all = mutation_mode == "all" # 1) Metal-swapping mutations using substructure replacement. if do_all or mutation_mode == "metal_swap": metals = _metal_symbols_in_molecule(base_mol) for metal in metals: query = Chem.MolFromSmarts(f"[{metal}]") if query is None: continue for replacement_symbol in METAL_SWAP_OPTIONS.get(metal, []): replacement = Chem.MolFromSmiles(f"[{replacement_symbol}]") if replacement is None: continue replaced = Chem.ReplaceSubstructs(base_mol, query, replacement, replaceAll=False) for mol in replaced: add_candidate(Chem.MolToSmiles(mol, canonical=True)) # 2) Ligand addition for single-metal centers. if base_mol.GetNumAtoms() == 1 and base_mol.GetAtomWithIdx(0).GetSymbol() in METAL_SWAP_OPTIONS: metal = base_mol.GetAtomWithIdx(0).GetSymbol() if do_all or mutation_mode == "ligand_oxo": add_candidate(f"[{metal}]=O") if do_all or mutation_mode == "ligand_hydroxyl": add_candidate(f"[{metal}]O") if do_all or mutation_mode == "ligand_methyl": add_candidate(f"[{metal}]C") # 3) Doping: replace first peripheral non-metal with common hetero dopants. if do_all or mutation_mode == "doping": atom_symbols = [atom.GetSymbol() for atom in base_mol.GetAtoms()] for idx, symbol in enumerate(atom_symbols): if symbol in METAL_SWAP_OPTIONS or symbol == "H": continue for dopant in DOPING_ATOMS: if dopant == symbol: continue editable = Chem.RWMol(base_mol) editable.GetAtomWithIdx(idx).SetAtomicNum(Chem.Atom(dopant).GetAtomicNum()) add_candidate(Chem.MolToSmiles(editable.GetMol(), canonical=True)) break if not valid_mutants: return [] return valid_mutants[:max(0, num_variations)] def mutate_catalyst_smiles(base_smiles: str, seed: Optional[int] = None) -> Optional[str]: """ Generate a simple catalyst mutation from a base catalyst SMILES. Mutations are chemistry-inspired but lightweight for fast exploration. """ if seed is not None: random.seed(seed) candidates = mutate_catalyst(base_smiles, num_variations=5) if not candidates: return None return random.choice(candidates) def generate_catalyst_candidates( reaction_type: str, num_candidates: int = 5, seed: Optional[int] = 42 ) -> List[str]: """ Generate candidate catalyst SMILES for a target reaction. Uses a seeded random exploration strategy over known catalyst families, plus lightweight mutation for novelty. """ if seed is not None: random.seed(seed) reaction_pools = { "H2_O2": ["[Pt]", "[Pd]", "[Rh]", "[Ru]", "[Ni]=O"], "N2_H2": ["[Fe]", "[Ru]", "[Co]", "[Fe]=O", "[Ni]"], "CO2_reduction": ["[Cu]", "[Ag]", "[Au]", "[Pd]", "[Ni]"], } base_pool = reaction_pools.get(reaction_type, ["[Pt]", "[Pd]", "[Fe]", "[Ni]", "[Cu]"]) generated = [] for _ in range(max(1, num_candidates * 3)): base = random.choice(base_pool) variants = mutate_catalyst(base, num_variations=2) if variants: for chosen in variants: if chosen not in generated and _is_valid_smiles(chosen): generated.append(chosen) if len(generated) >= num_candidates: break elif base not in generated and _is_valid_smiles(base): generated.append(base) if len(generated) >= num_candidates: break # Fallback to base pool if mutation did not produce enough unique candidates. for base in base_pool: if len(generated) >= num_candidates: break if base not in generated and _is_valid_smiles(base): generated.append(base) return generated[:num_candidates]