evaluate.py

#!/usr/bin/env python
"""Evaluate a set of generated SMILES against a training corpus.

Reports the standard de novo generation benchmarks:

* **Validity**   -- fraction of generated strings that parse as valid molecules.
* **Uniqueness** -- fraction of valid molecules that are distinct (canonical SMILES).
* **Novelty**    -- fraction of unique valid molecules absent from the training set.
* **Internal diversity (IntDiv)** -- 1 - mean pairwise Tanimoto similarity
  (ECFP4 fingerprints) among the valid, unique generated set.

Falls back gracefully (validity assumed, diversity skipped) if RDKit is
not installed, matching the behaviour documented in the README.

Example
-------
    python evaluate.py --generated generated_molecules.smi --train_data data/sample_smiles.txt
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import List, Optional

try:
    from rdkit import Chem, DataStructs
    from rdkit.Chem import rdFingerprintGenerator

    RDKIT_AVAILABLE = True
    _MORGAN_GEN = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
except ImportError:
    RDKIT_AVAILABLE = False


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument(
        "--generated",
        type=str,
        required=True,
        help="Path to generated SMILES (one per line).",
    )
    p.add_argument(
        "--train_data",
        type=str,
        default=None,
        help="Training corpus, for novelty comparison.",
    )
    return p.parse_args()


def read_lines(path: str) -> List[str]:
    return [l.strip() for l in Path(path).read_text().splitlines() if l.strip()]


def canonicalize(smiles: str) -> Optional[str]:
    if not RDKIT_AVAILABLE:
        return smiles
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol)


def compute_validity(smiles_list: List[str]) -> tuple[float, List[str]]:
    """Returns (validity_fraction, list_of_valid_canonical_smiles)."""
    if not RDKIT_AVAILABLE:
        return 1.0, list(smiles_list)  # assumed valid, per README fallback behavior
    valid = []
    for smi in smiles_list:
        canon = canonicalize(smi)
        if canon is not None:
            valid.append(canon)
    validity = len(valid) / len(smiles_list) if smiles_list else 0.0
    return validity, valid


def compute_uniqueness(valid_smiles: List[str]) -> tuple[float, List[str]]:
    unique = list(dict.fromkeys(valid_smiles))  # preserves order, dedupes
    uniqueness = len(unique) / len(valid_smiles) if valid_smiles else 0.0
    return uniqueness, unique


def compute_novelty(
    unique_smiles: List[str], train_smiles: Optional[List[str]]
) -> Optional[float]:
    if train_smiles is None:
        return None
    train_canon = {canonicalize(s) for s in train_smiles}
    train_canon.discard(None)
    if not unique_smiles:
        return 0.0
    novel = [s for s in unique_smiles if s not in train_canon]
    return len(novel) / len(unique_smiles)


def compute_internal_diversity(unique_smiles: List[str]) -> Optional[float]:
    if not RDKIT_AVAILABLE or len(unique_smiles) < 2:
        return None
    fps = []
    for smi in unique_smiles:
        mol = Chem.MolFromSmiles(smi)
        if mol is not None:
            fps.append(_MORGAN_GEN.GetFingerprint(mol))

    if len(fps) < 2:
        return None

    sims = []
    for i in range(len(fps)):
        sims.extend(DataStructs.BulkTanimotoSimilarity(fps[i], fps[i + 1 :]))

    if not sims:
        return None
    mean_sim = sum(sims) / len(sims)
    return 1.0 - mean_sim


def main() -> None:
    args = parse_args()
    generated = read_lines(args.generated)
    train_smiles = read_lines(args.train_data) if args.train_data else None

    if not generated:
        print("No generated molecules to evaluate.")
        return

    validity, valid_smiles = compute_validity(generated)
    uniqueness, unique_smiles = compute_uniqueness(valid_smiles)
    novelty = compute_novelty(unique_smiles, train_smiles)
    int_div = compute_internal_diversity(unique_smiles)

    print(f"Generated:            {len(generated)}")
    print(
        f"Validity:             {validity:.1%}"
        + ("" if RDKIT_AVAILABLE else "  (assumed -- RDKit not installed)")
    )
    print(f"Uniqueness (of valid): {uniqueness:.1%}")
    if novelty is not None:
        print(f"Novelty (of unique):  {novelty:.1%}")
    else:
        print("Novelty:              n/a (no --train_data provided)")
    if int_div is not None:
        print(f"Internal diversity:   {int_div:.3f}")
    else:
        print(
            "Internal diversity:   n/a"
            + ("" if RDKIT_AVAILABLE else " (RDKit not installed)")
        )


if __name__ == "__main__":
    main()