src/pino/thermo/geometry.py

import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import Descriptors3D


def extract_3d_shape_features(smiles: str) -> list[float]:
    """
    Generates a 3D conformation for a SMILES string and extracts geometric shape
    and chirality descriptors to capture stereochemical variations.

    Returns a 7-D vector:
        [Asphericity, Eccentricity, InertialShapeFactor, SpherocityIndex,
         RadiusOfGyration, PBF, ChiralScore]
    """
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return [0.0] * 7

    # Step 1: Add hydrogens for accurate physical volume and surface estimations.
    mol = Chem.AddHs(mol)

    # Step 2: Generate 3D conformation using ETKDG v3.
    params = AllChem.ETKDGv3()
    params.randomSeed = 42
    conf_id = AllChem.EmbedMolecule(mol, params)

    if conf_id < 0:
        # Fallback to empty features if 3D coordinate embedding fails.
        return [0.0] * 7

    # Step 3: Optimize geometry using MMFF94 force field.
    try:
        AllChem.MMFFOptimizeMolecule(mol, confId=conf_id)
    except Exception:
        pass  # Keep raw embedded coordinates if minimization fails.

    # Step 4: Extract 3D shape descriptors.
    desc_3d = Descriptors3D.CalcMolDescriptors3D(mol, confId=conf_id)

    feature_vector = [
        float(desc_3d.get("Asphericity", 0.0)),
        float(desc_3d.get("Eccentricity", 0.0)),
        float(desc_3d.get("InertialShapeFactor", 0.0)),
        float(desc_3d.get("SpherocityIndex", 0.0)),
        float(desc_3d.get("RadiusOfGyration", 0.0)),
        float(desc_3d.get("PBF", 0.0)),  # Plane of best fit.
    ]

    # Step 5: Score chiral centers to differentiate mirror reflections.
    chiral_centers = Chem.FindMolChiralCenters(mol, includeUnassigned=True)
    chiral_score = sum(
        1.0 if center[1] == "R" else -1.0 for center in chiral_centers
    )
    feature_vector.append(float(chiral_score))

    return feature_vector