#!/usr/bin/env python
"""Helpers for ADMET-guided genetic molecule optimization demos.
The public visualization entry is render_optimization_history(...). The module
also keeps small validation and scoring helpers that are useful when building a
demo pipeline around RDKit, SA-Score, and ADMET-AI.
"""
from __future__ import annotations
import io
import json
import math
import os
import re
import shutil
from collections import Counter
from pathlib import Path
from typing import Any
matplotlib = None
pd = None
plt = None
__all__ = [
"standardize_smiles",
"canonicalize_smiles",
"classify_admet_columns",
"aggregate_admet_predictions",
"operation_detail_json",
"validate_generation_log",
"render_optimization_history",
]
_SVG_DECLARATION = re.compile(r"<\?xml[^>]*\?>", re.IGNORECASE)
_SVG_DOCTYPE = re.compile(r"]*(?:\[[\s\S]*?\]\s*)?>", re.IGNORECASE)
def _require_pandas() -> Any:
global pd
if pd is None:
try:
import pandas as pandas_module
except ImportError as exc:
raise RuntimeError(
"pandas is required; switch to the admet-sa-ga environment "
"or install the skill prerequisites"
) from exc
pd = pandas_module
return pd
def _require_plotting() -> tuple[Any, Any]:
global matplotlib, plt
if matplotlib is None or plt is None:
os.environ.setdefault("MPLCONFIGDIR", "/tmp/admet-sa-ga-mpl-cache")
os.environ.setdefault("XDG_CACHE_HOME", "/tmp/admet-sa-ga-cache")
Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
Path(os.environ["XDG_CACHE_HOME"]).mkdir(parents=True, exist_ok=True)
try:
import matplotlib as matplotlib_module
matplotlib_module.use("Agg")
import matplotlib.pyplot as pyplot_module
except ImportError as exc:
raise RuntimeError(
"matplotlib is required to render the optimization dashboard; "
"switch to the admet-sa-ga environment or install the skill prerequisites"
) from exc
matplotlib = matplotlib_module
plt = pyplot_module
return matplotlib, plt
def standardize_smiles(smiles: str) -> tuple[str | None, str]:
"""Return (canonical_smiles, failure_reason) after RDKit standardization."""
try:
from rdkit import Chem
except ImportError as exc:
raise RuntimeError(
"RDKit is required to standardize SMILES; switch to the "
"admet-sa-ga environment or install the skill prerequisites"
) from exc
mol = Chem.MolFromSmiles(str(smiles))
if mol is None:
return None, "parse_failed"
try:
Chem.SanitizeMol(mol)
except Exception:
return None, "sanitize_failed"
mol = Chem.RemoveHs(mol)
if mol.GetNumAtoms() == 0:
return None, "empty_molecule"
fragments = Chem.GetMolFrags(mol, asMols=True, sanitizeFrags=True)
if not fragments:
return None, "multi_fragment_failed"
organic = [
frag
for frag in fragments
if any(atom.GetAtomicNum() == 6 for atom in frag.GetAtoms())
]
mol = max(organic or fragments, key=lambda item: item.GetNumHeavyAtoms())
try:
Chem.SanitizeMol(mol)
return Chem.MolToSmiles(mol, canonical=True, isomericSmiles=True), ""
except Exception:
return None, "sanitize_failed"
def canonicalize_smiles(smiles: str) -> str:
"""Return a canonical SMILES when RDKit is installed, else a stripped string."""
value = (smiles or "").strip()
if not value:
raise ValueError("SMILES is empty")
try:
from rdkit import Chem # type: ignore
except ImportError:
return value
mol = Chem.MolFromSmiles(value)
if mol is None:
raise ValueError(f"invalid SMILES: {smiles!r}")
return Chem.MolToSmiles(mol, canonical=True)
def classify_admet_columns(
columns: list[str], admet_config: dict[str, Any]
) -> dict[str, str]:
"""Map ADMET-AI columns to positive, negative, or ignored score roles."""
positive = [key.lower() for key in admet_config.get("positive_keywords", [])]
negative = [key.lower() for key in admet_config.get("negative_keywords", [])]
mapping = {}
for column in columns:
name = column.lower()
if "drugbank_approved_percentile" in name:
mapping[column] = "ignored"
elif any(key in name for key in negative):
mapping[column] = "negative"
elif any(key in name for key in positive):
mapping[column] = "positive"
else:
mapping[column] = "ignored"
return mapping
def _numeric_01(value: Any) -> float | None:
if value is None:
return None
if isinstance(value, bool):
return 1.0 if value else 0.0
try:
number = float(value)
except (TypeError, ValueError):
text = str(value).strip().lower()
if text in {"true", "yes", "positive", "active", "high", "1"}:
return 1.0
if text in {"false", "no", "negative", "inactive", "low", "0"}:
return 0.0
return None
if math.isnan(number):
return None
if 0.0 <= number <= 1.0:
return number
return 1.0 / (1.0 + pow(2.718281828459045, -number))
def aggregate_admet_predictions(
predictions: dict[str, Any],
admet_config: dict[str, Any],
) -> tuple[float, list[str], dict[str, str]]:
"""Aggregate raw ADMET endpoint predictions into score, risk flags, mapping."""
mapping = classify_admet_columns(list(predictions.keys()), admet_config)
scores: list[float] = []
flags: list[str] = []
threshold = float(admet_config.get("risk_threshold", 0.5))
for column, role in mapping.items():
if role == "ignored":
continue
value = _numeric_01(predictions.get(column))
if value is None:
continue
if role == "positive":
scores.append(value)
if value < 1.0 - threshold:
flags.append(f"low_{column}")
else:
scores.append(1.0 - value)
if value >= threshold:
flags.append(f"high_{column}")
score = sum(scores) / len(scores) if scores else 0.5
return score, flags, mapping
def operation_detail_json(
operation: str,
operator_detail: Any,
child_canonical_smiles: str,
parent_ids: list[str],
parent_smiles: list[str],
) -> str:
"""Build the canonical operation_detail JSON string used in generation logs."""
payload = {
"operation": operation,
"operator_detail": operator_detail,
"parent_ids": parent_ids,
"parent_smiles": parent_smiles,
"child_canonical_smiles": child_canonical_smiles,
}
return json.dumps(payload, ensure_ascii=False, sort_keys=True)
def validate_generation_log(frame: pd.DataFrame) -> None:
"""Validate lineage invariants expected by the visualization workflow."""
_require_pandas()
required = {"molecule_id", "smiles", "generation", "operation", "parent", "parents"}
missing = sorted(required - set(frame.columns))
if missing:
raise ValueError(f"generation log missing required columns: {missing}")
duplicate_ids = [
key
for key, count in Counter(frame["molecule_id"].astype(str)).items()
if count > 1
]
duplicate_smiles = [
key for key, count in Counter(frame["smiles"].astype(str)).items() if count > 1
]
if duplicate_ids or duplicate_smiles:
raise ValueError(
"molecule_id and canonical smiles must be one-to-one; "
f"duplicate_ids={duplicate_ids[:5]}, duplicate_smiles={duplicate_smiles[:5]}"
)
mutation = frame[frame["operation"].astype(str).str.lower() == "mutation"].fillna(
""
)
if len(mutation) and (
(mutation["parent"].astype(str) == "").any()
or (mutation["parents"].astype(str) != "").any()
):
raise ValueError(
"mutation rows must set parent to one ID and leave parents empty"
)
crossover = frame[frame["operation"].astype(str).str.lower() == "crossover"].fillna(
""
)
if len(crossover):
bad_parent = (crossover["parent"].astype(str) != "").any()
bad_parents = (
crossover["parents"].astype(str).str.split(";").map(len) != 2
).any()
if bad_parent or bad_parents:
raise ValueError(
"crossover rows must leave parent empty and set parents to exactly two IDs"
)
def _clean_svg(svg: str) -> str:
svg = _SVG_DECLARATION.sub("", svg)
svg = _SVG_DOCTYPE.sub("", svg)
return svg.strip()
def _mol_svg(smiles: str, width: int = 150, height: int = 95) -> str:
try:
from rdkit import Chem
from rdkit.Chem import Draw
except ImportError:
return (
f''
)
mol = Chem.MolFromSmiles(str(smiles))
if mol is None:
return "