fetch_compounds.py

#!/usr/bin/env python3
"""
Fetch ~1000 compounds from PubChem and generate Python code for database.py
"""
import requests
import json
import time

compounds = []
seen_smiles = set()

# Add the original 27 to avoid duplicates
original_smiles = {
    "C", "CC", "CCC", "CCCC", "CCCCC", "CCCCCC", "C1CCCCC1",
    "C=C", "CC=C", "C=CCC",
    "C#C", "CC#C",
    "c1ccccc1", "Cc1ccccc1", "Oc1ccccc1",
    "CO", "CCO", "CCCO", "CC(O)C",
    "C=O", "CC=O", "CC(=O)C",
    "O=CO", "CC(=O)O", "CCC(=O)O",
    "CN", "CCN"
}
seen_smiles.update(original_smiles)

def fetch_from_formula(formula):
    """Fetch CIDs for a chemical formula from PubChem"""
    try:
        url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/fastidentity/json"
        params = {
            "identity": formula,
            "type": "formula",
            "format": "json"
        }
        response = requests.get(url, params=params, timeout=5)
        if response.status_code == 200:
            data = response.json()
            if "IdentifierList" in data and "CID" in data["IdentifierList"]:
                cids = data["IdentifierList"]["CID"]
                return cids[:100]  # Limit to 100 per formula
    except Exception as e:
        print(f"  Error fetching {formula}: {e}")
    return []

def fetch_properties_batch(cids):
    """Fetch SMILES and formula for CIDs"""
    if not cids:
        return []
    
    try:
        cid_str = ",".join(map(str, cids[:100]))  # Batch up to 100
        url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid_str}/property/CanonicalSMILES,MolecularFormula,IUPACName/json"
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            if "properties" in data:
                results = []
                for prop in data["properties"]:
                    if "CanonicalSMILES" in prop and "MolecularFormula" in prop:
                        smiles = prop["CanonicalSMILES"].strip()
                        formula = prop["MolecularFormula"]
                        name = prop.get("IUPACName", "").strip()
                        if smiles and len(smiles) < 200:  # Skip huge molecules
                            results.append({
                                "name": name or formula,
                                "smiles": smiles,
                                "formula": formula
                            })
                return results
    except Exception as e:
        print(f"  Error fetching properties: {e}")
    
    return []

# Formulas to fetch - organic chemistry variety
formulas = [
    # Alkanes (C1-C20)
    "CH4", "C2H6", "C3H8", "C4H10", "C5H12", "C6H14", "C7H16", "C8H18", "C9H20", "C10H22",
    "C11H24", "C12H26", "C13H28", "C14H30", "C15H32",
    # Alkenes (C2-C15)
    "C2H4", "C3H6", "C4H8", "C5H10", "C6H12", "C7H14", "C8H16", "C9H18", "C10H20",
    "C11H22", "C12H24",
    # Alkynes (C2-C12)
    "C2H2", "C3H4", "C4H6", "C5H8", "C6H10", "C7H12", "C8H14", "C9H16", "C10H18",
    # Alcohols
    "CH3OH", "C2H5OH", "C3H7OH", "C4H9OH", "C5H11OH", "C6H13OH", "C7H15OH", "C8H17OH",
    "C9H19OH", "C10H21OH",
    # Ethers
    "C2H6O", "C3H8O", "C4H10O", "C5H12O", "C6H14O", "C7H16O", "C8H18O",
    # Aldehydes
    "CH2O", "C2H4O", "C3H6O", "C4H8O", "C5H10O", "C6H12O", "C7H14O",
    # Ketones
    "C3H6O", "C4H8O", "C5H10O", "C6H12O", "C7H14O", "C8H16O", "C9H18O", "C10H20O",
    # Carboxylic acids
    "CH2O2", "C2H4O2", "C3H6O2", "C4H8O2", "C5H10O2", "C6H12O2", "C7H14O2", "C8H16O2",
    # Amines
    "CH5N", "C2H7N", "C3H9N", "C4H11N", "C5H13N", "C6H15N", "C7H17N",
    # Amides
    "CH3NO", "C2H5NO", "C3H7NO", "C4H9NO", "C5H11NO",
    # Esters
    "C2H4O2", "C3H6O2", "C4H8O2", "C5H10O2", "C6H12O2", "C7H14O2", "C8H16O2",
    # Haloalkanes
    "CH3Cl", "C2H5Cl", "C3H7Cl", "CH3Br", "C2H5Br", "CH3I", "C2H5I",
    "CH2Cl2", "CHCl3", "CF4",
    # Aromatic compounds
    "C6H6", "C7H8", "C8H10", "C6H5OH", "C6H5NH2", "C6H5Cl", "C6H4Cl2",
    # Amino acids (common proteinogenic)
    "C2H5NO2", "C3H7NO2", "C5H11NO2", "C6H13NO2", "C3H7NO3", "C4H9NO3",
    "C3H7NO2S", "C5H11NO2S", "C4H7NO4", "C5H9NO4", "C4H8N2O3", "C5H10N2O3",
    "C6H14N2O2", "C6H14N4O2", "C6H9N3O2", "C9H11NO2", "C9H11NO3",
    "C11H12N2O2", "C5H9NO2",
    # Nucleic acids (nucleobases + monophosphates)
    "C5H5N5", "C5H5N5O", "C4H5N3O", "C5H6N2O2", "C4H4N2O2",
    "C10H14N5O7P", "C9H14N3O8P", "C10H14N5O8P", "C9H13N2O9P", "C10H15N2O8P",
]

print("Fetching compounds from PubChem...")
print("=" * 60)

for i, formula in enumerate(formulas):
    print(f"[{i+1}/{len(formulas)}] Fetching {formula:15s}...", end=" ", flush=True)
    
    cids = fetch_from_formula(formula)
    if cids:
        props = fetch_properties_batch(cids)
        added = 0
        for prop in props:
            smiles = prop["smiles"]
            if smiles not in seen_smiles:
                compounds.append({
                    "name": prop["name"][:100],  # Truncate long names
                    "smiles": smiles,
                    "category": formula.split("H")[0]  # Use formula base as category
                })
                seen_smiles.add(smiles)
                added += 1
        
        print(f"✓ {added} new compounds (total: {len(compounds)})")
    else:
        print("no results")
    
    time.sleep(0.3)  # Be nice to PubChem API

print("=" * 60)
print(f"Total unique compounds collected: {len(compounds)}")

# Generate Python code
code_lines = [
    "# --- 1000+ Compounds from PubChem (Preloaded into COMPOUNDS_DB) ---",
    "PUBCHEM_COMPOUNDS = [",
]

for compound in compounds:
    code_lines.append(f'    {{"name": {repr(compound["name"])}, "smiles": {repr(compound["smiles"])}, "category": ""}},')

code_lines.append("]")

# Save Python code
with open("/tmp/pubchem_compounds_code.py", "w") as f:
    f.write("\n".join(code_lines))

print(f"\nGenerated code saved to /tmp/pubchem_compounds_code.py")
print(f"Lines of code: {len(code_lines)}")

# Also save raw JSON for reference
with open("/tmp/pubchem_compounds.json", "w") as f:
    json.dump(compounds, f, indent=2)

print(f"JSON saved to /tmp/pubchem_compounds.json")