import pandas as pd
import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors, rdMolDescriptors, rdDepictor
from rdkit.Chem import Crippen, Descriptors3D
from rdkit.Chem import rdFingerprintGenerator
import warnings
warnings.filterwarnings("ignore")
from rdkit import RDLogger
RDLogger.DisableLog("rdApp.*")
import os
import json
import argparse
import multiprocessing as mp
from pathlib import Path
from typing import Dict, Optional, Tuple
# ----------------------------------------------------------------------
# Logging / RDKit hygiene
# ----------------------------------------------------------------------
# RDKit can be chatty; we silence logs above via RDLogger.DisableLog.
# We also suppress Python warnings (set above).
# ----------------------------------------------------------------------
# Wildcard ("*") handling utilities
# ----------------------------------------------------------------------
ATOMIC_NUM_AT = 85 # Astatine (At) used as a placeholder for wildcard atoms
def process_star_atoms(mol: Chem.Mol) -> Chem.Mol:
"""
Replace all wildcard atoms ("*" or atomicNum == 0) with Astatine (At, Z=85).
Rationale:
- Polymer SMILES often contain '*' to indicate attachment points.
- Many RDKit operations fail or sanitize differently with atomicNum == 0.
- Mapping '*' -> At allows sanitization and downstream featurization while
keeping a consistent placeholder identity.
"""
if mol is None:
return mol
for atom in mol.GetAtoms():
if atom.GetAtomicNum() == 0 or atom.GetSymbol() == "*":
atom.SetAtomicNum(ATOMIC_NUM_AT)
return mol
# ----------------------------------------------------------------------
# Per-polymer worker function
# ----------------------------------------------------------------------
def process_single_polymer(args) -> Tuple[Optional[Dict], Optional[Dict]]:
"""
Worker that processes one row (one polymer) and returns:
(polymer_data, failed_info)
polymer_data is a dict containing serialized multimodal outputs.
failed_info is a dict with index/smiles/error if anything fails.
"""
idx, row_dict, extractor = args
polymer_data = None
failed_info = None
try:
smiles = row_dict.get("psmiles", None)
source = row_dict.get("source", None)
if pd.isna(smiles) or not isinstance(smiles, str) or len(smiles.strip()) == 0:
failed_info = {"index": idx, "smiles": str(smiles), "error": "Empty or invalid SMILES"}
return polymer_data, failed_info
canonical_smiles = extractor.validate_and_standardize_smiles(smiles)
if canonical_smiles is None:
failed_info = {"index": idx, "smiles": smiles, "error": "Invalid SMILES or cannot be standardized"}
return polymer_data, failed_info
polymer_data = {
"original_index": idx,
"psmiles": canonical_smiles,
"source": source,
"smiles": canonical_smiles,
}
# Graph
try:
polymer_data["graph"] = extractor.generate_molecular_graph(canonical_smiles)
except Exception:
polymer_data["graph"] = {}
# Geometry
try:
polymer_data["geometry"] = extractor.optimize_3d_geometry(canonical_smiles)
except Exception:
polymer_data["geometry"] = {}
# Fingerprints
try:
polymer_data["fingerprints"] = extractor.calculate_morgan_fingerprints(canonical_smiles)
except Exception:
polymer_data["fingerprints"] = {}
return polymer_data, failed_info
except Exception as e:
failed_info = {"index": idx, "smiles": row_dict.get("psmiles", ""), "error": str(e)}
return polymer_data, failed_info
# ----------------------------------------------------------------------
# Main extractor class
# ----------------------------------------------------------------------
class AdvancedPolymerMultimodalExtractor:
"""
Multimodal extractor that reads a CSV of polymers and adds:
- graph: node/edge features + adjacency + summary graph features
- geometry: best 3D conformer (or fallback 2D coords) + 3D descriptors
- fingerprints: Morgan fingerprints (bitstrings + counts) for multiple radii
Output:
- _processed.csv (appended chunk-by-chunk)
- _failures.jsonl (one JSON per failure)
"""
def __init__(self, csv_file: str):
self.csv_file = str(csv_file)
# ------------------------------
# SMILES validation/standardization
# ------------------------------
def validate_and_standardize_smiles(self, smiles: str) -> Optional[str]:
"""
Parse, sanitize, replace '*' with At, and return canonical SMILES.
Returns None if parsing/sanitization fails.
"""
try:
if not smiles or pd.isna(smiles):
return None
mol = Chem.MolFromSmiles(smiles, sanitize=False)
if mol is None:
return None
mol = process_star_atoms(mol) # pass 1
Chem.SanitizeMol(mol)
mol = process_star_atoms(mol) # pass 2 (robust)
canonical_smiles = Chem.MolToSmiles(mol, canonical=True)
if not canonical_smiles:
return None
return canonical_smiles
except Exception:
return None
# ------------------------------
# Molecular graph (RDKit -> JSONable dict)
# ------------------------------
def generate_molecular_graph(self, smiles: str) -> Dict:
"""
Build a molecular graph representation with atom/bond features and
global graph descriptors.
"""
mol = Chem.MolFromSmiles(smiles)
mol = process_star_atoms(mol)
if mol is None:
return {}
# Explicit hydrogens for atom-level features
mol = Chem.AddHs(mol)
node_features = []
for atom in mol.GetAtoms():
node_features.append(
{
"atomic_num": atom.GetAtomicNum(),
"degree": atom.GetDegree(),
"formal_charge": atom.GetFormalCharge(),
"hybridization": int(atom.GetHybridization()),
"is_aromatic": atom.GetIsAromatic(),
"is_in_ring": atom.IsInRing(),
"chirality": int(atom.GetChiralTag()),
"mass": atom.GetMass(),
"valence": atom.GetTotalValence(),
"num_radical_electrons": atom.GetNumRadicalElectrons(),
}
)
edge_features = []
edge_indices = []
for bond in mol.GetBonds():
i = bond.GetBeginAtomIdx()
j = bond.GetEndAtomIdx()
edge_features.append(
{
"bond_type": int(bond.GetBondType()),
"is_aromatic": bond.GetIsAromatic(),
"is_in_ring": bond.IsInRing(),
"stereo": int(bond.GetStereo()),
"is_conjugated": bond.GetIsConjugated(),
}
)
# Undirected -> store both directions for GNN-style edge lists
edge_indices.extend([[i, j], [j, i]])
graph_features = {
"num_atoms": mol.GetNumAtoms(),
"num_bonds": mol.GetNumBonds(),
"num_rings": rdMolDescriptors.CalcNumRings(mol),
"molecular_weight": Descriptors.MolWt(mol),
"logp": Crippen.MolLogP(mol),
"tpsa": Descriptors.TPSA(mol),
"num_rotatable_bonds": Descriptors.NumRotatableBonds(mol),
"num_h_acceptors": rdMolDescriptors.CalcNumHBA(mol),
"num_h_donors": rdMolDescriptors.CalcNumHBD(mol),
}
adj = Chem.GetAdjacencyMatrix(mol).tolist()
return {
"node_features": node_features,
"edge_features": edge_features,
"edge_indices": edge_indices,
"graph_features": graph_features,
"adjacency_matrix": adj,
}
# ------------------------------
# 3D geometry (ETKDG + MMFF/UFF)
# ------------------------------
def optimize_3d_geometry(self, smiles: str, num_conformers: int = 10) -> Dict:
"""
Generate multiple conformers, optimize (MMFF if available else UFF),
and return the lowest-energy conformer coordinates + 3D descriptors.
If no conformer is generated/optimized, fall back to 2D coordinates.
"""
mol = Chem.MolFromSmiles(smiles)
if mol is None or mol.GetNumAtoms() > 200:
return {}
mol = process_star_atoms(mol)
mol_h = Chem.AddHs(mol)
# Atomic numbers aligned to coordinate ordering (mol_h atoms)
atomic_numbers = [atom.GetAtomicNum() for atom in mol_h.GetAtoms()]
try:
params = AllChem.ETKDGv3()
params.randomSeed = 42
conformer_ids = AllChem.EmbedMultipleConfs(mol_h, numConfs=num_conformers, params=params)
except Exception:
conformer_ids = []
best_conformer = None
best_energy = float("inf")
for conf_id in conformer_ids:
try:
mmff_ok = AllChem.MMFFHasAllMoleculeParams(mol_h)
if mmff_ok:
AllChem.MMFFOptimizeMolecule(mol_h, confId=conf_id)
props = AllChem.MMFFGetMoleculeProperties(mol_h)
ff = AllChem.MMFFGetMoleculeForceField(mol_h, props, confId=conf_id)
else:
AllChem.UFFOptimizeMolecule(mol_h, confId=conf_id)
ff = AllChem.UFFGetMoleculeForceField(mol_h, confId=conf_id)
energy = ff.CalcEnergy() if ff is not None else None
if energy is None or energy >= best_energy:
continue
conf = mol_h.GetConformer(conf_id)
coords = [
[conf.GetAtomPosition(i).x, conf.GetAtomPosition(i).y, conf.GetAtomPosition(i).z]
for i in range(mol_h.GetNumAtoms())
]
descriptors_3d = {}
try:
descriptors_3d = {
"asphericity": Descriptors3D.Asphericity(mol_h, confId=conf_id),
"eccentricity": Descriptors3D.Eccentricity(mol_h, confId=conf_id),
"inertial_shape_factor": Descriptors3D.InertialShapeFactor(mol_h, confId=conf_id),
"radius_of_gyration": Descriptors3D.RadiusOfGyration(mol_h, confId=conf_id),
"spherocity_index": Descriptors3D.SpherocityIndex(mol_h, confId=conf_id),
}
except Exception:
pass
best_conformer = {
"conformer_id": int(conf_id),
"coordinates": coords,
"atomic_numbers": atomic_numbers,
"energy": float(energy),
"descriptors_3d": descriptors_3d,
}
best_energy = energy
except Exception:
continue
if best_conformer is not None:
return {
"best_conformer": best_conformer,
"num_conformers_generated": int(len(conformer_ids)),
"converted_smiles": Chem.MolToSmiles(mol),
}
# Fallback: 2D coordinates
try:
rdDepictor.Compute2DCoords(mol)
coords_2d = mol.GetConformer().GetPositions().tolist()
atomic_numbers_2d = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
return {
"best_conformer": {
"conformer_id": -1,
"coordinates": coords_2d,
"atomic_numbers": atomic_numbers_2d,
"energy": None,
"descriptors_3d": {},
},
"num_conformers_generated": 0,
"converted_smiles": Chem.MolToSmiles(mol),
}
except Exception:
return {}
# ------------------------------
# Morgan fingerprints (multi-radius)
# ------------------------------
def calculate_morgan_fingerprints(self, smiles: str, radius: int = 3, n_bits: int = 2048) -> Dict:
"""
Compute Morgan fingerprints:
- bitstring (as list of '0'/'1' chars) at radius=radius
- counts (as dict) at radius=radius
Also includes all radii r in [1, radius-1].
"""
mol = Chem.MolFromSmiles(smiles)
mol = process_star_atoms(mol)
if mol is None:
return {}
fingerprints = {}
# Main radius
generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits)
fp_bitvect = generator.GetFingerprint(mol)
fingerprints[f"morgan_r{radius}_bits"] = list(fp_bitvect.ToBitString())
fingerprints[f"morgan_r{radius}_counts"] = dict(AllChem.GetMorganFingerprint(mol, radius).GetNonzeroElements())
# Additional radii
for r in range(1, radius):
gen = rdFingerprintGenerator.GetMorganGenerator(radius=r, fpSize=n_bits)
bitvect = gen.GetFingerprint(mol)
fingerprints[f"morgan_r{r}_bits"] = list(bitvect.ToBitString())
fingerprints[f"morgan_r{r}_counts"] = dict(AllChem.GetMorganFingerprint(mol, r).GetNonzeroElements())
return fingerprints
# ------------------------------
# Chunked parallel processing over CSV
# ------------------------------
def process_all_polymers_parallel(self, chunk_size: int = 100, num_workers: int = 40) -> str:
"""
Read the input CSV in chunks, fill missing multimodal columns, and process
only rows that are missing any of: graph/geometry/fingerprints.
Appends processed chunks to _processed.csv and failures to _failures.jsonl.
"""
chunk_iterator = pd.read_csv(self.csv_file, chunksize=chunk_size, engine="python")
for chunk in chunk_iterator:
# Ensure expected output columns exist and are object dtype (for JSON strings)
for col in ["graph", "geometry", "fingerprints"]:
if col not in chunk.columns:
chunk[col] = None
chunk[col] = chunk[col].astype(object)
# Only process rows missing any modality
chunk_to_process = chunk[chunk[["graph", "geometry", "fingerprints"]].isnull().any(axis=1)].copy()
# If all rows already done, just persist chunk and continue
if len(chunk_to_process) == 0:
self.save_chunk_to_csv(chunk)
continue
rows = list(chunk_to_process.iterrows())
argslist = [(i, row.to_dict(), self) for i, row in rows]
with mp.Pool(num_workers) as pool:
results = pool.map(process_single_polymer, argslist)
failed_list = []
for n, (output, fail) in enumerate(results):
idx = rows[n][0]
if output:
chunk.at[idx, "graph"] = json.dumps(output["graph"])
chunk.at[idx, "geometry"] = json.dumps(output["geometry"])
chunk.at[idx, "fingerprints"] = json.dumps(output["fingerprints"])
if fail:
failed_list.append(fail)
self.save_chunk_to_csv(chunk)
self.save_failed_to_json(failed_list)
return "Processing Done"
# ------------------------------
# Output helpers
# ------------------------------
def save_chunk_to_csv(self, chunk: pd.DataFrame) -> None:
"""
Append processed chunk to _processed.csv.
"""
out_csv = self.csv_file.replace(".csv", "_processed.csv")
if not os.path.exists(out_csv):
chunk.to_csv(out_csv, index=False, mode="w")
else:
chunk.to_csv(out_csv, index=False, mode="a", header=False)
def save_failed_to_json(self, failed_list) -> None:
"""
Append failures to _failures.jsonl (JSON lines).
"""
if not failed_list:
return
fail_json = self.csv_file.replace(".csv", "_failures.jsonl")
with open(fail_json, "a", encoding="utf-8") as f:
for fail in failed_list:
json.dump(fail, f)
f.write("\n")
def save_results(self, output_file: str = "polymer_multimodal_data.json"):
pass
def generate_summary_statistics(self) -> Dict:
return {}
# ----------------------------------------------------------------------
# CLI / entry-point helpers
# ----------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
"""
Command-line arguments:
--csv_file: path to input CSV (required)
--chunk_size: rows per chunk
--num_workers: multiprocessing workers
"""
parser = argparse.ArgumentParser(description="Polymer multimodal feature extraction (RDKit).")
parser.add_argument(
"--csv_file",
type=str,
default="/path/to/polymer_structures_unified.csv",
help="Path to the input CSV file containing at least a 'psmiles' column.",
)
parser.add_argument("--chunk_size", type=int, default=1000, help="Rows per chunk for streaming CSV processing.")
parser.add_argument("--num_workers", type=int, default=24, help="Number of parallel worker processes.")
return parser.parse_args()
def main() -> Tuple[AdvancedPolymerMultimodalExtractor, Optional[object]]:
"""
Script entry point.
Reads arguments, constructs the extractor, and runs chunked parallel processing.
"""
args = parse_args()
csv_file = args.csv_file
extractor = AdvancedPolymerMultimodalExtractor(csv_file)
try:
extractor.process_all_polymers_parallel(chunk_size=args.chunk_size, num_workers=args.num_workers)
except KeyboardInterrupt:
return extractor, None
except Exception as e:
print(f"CRASH! Error: {e}")
return extractor, None
print("\n=== Processing Complete ===")
return extractor, None
if __name__ == "__main__":
extractor, results = main()