caddforge-descriptors 0.11.0__py3-none-any.whl

cadd_descriptors/__init__.py

@@ -0,0 +1,48 @@
+"""CADDForge descriptors tool package."""
+
+from .ani_conformers import AniConformersConfig, AniConformersWorkflow
+from .ani_psa.pipeline import (
+    boltzmann_averaged_psa,
+    boltzmann_weights_from_energies,
+    compute_3d_psa_for_conformer,
+    compute_3d_psa_for_conformers,
+    pipeline_boltzmann_3d_psa,
+)
+from .ani_utils import (
+    ConformerSettings,
+    ani_energies_for_conformers,
+    cluster_conformers,
+    generate_conformers,
+    seed_everything,
+    select_top_n_by_ani,
+)
+from .ani_psa.workflow import AniPsaConfig, AniPsaWorkflow
+from .cluster_molecules import ClusterMoleculesConfig, ClusterMoleculesWorkflow
+from .phesa import PhesaConfig, PhesaWorkflow
+from .tool import AniConformersTool, AniPsaTool, ClusterMoleculesTool, PhesaTool
+
+__all__ = [
+    "AniPsaConfig",
+    "AniPsaWorkflow",
+    "AniConformersConfig",
+    "AniConformersWorkflow",
+    "AniPsaTool",
+    "AniConformersTool",
+    "ClusterMoleculesTool",
+    "PhesaTool",
+    "ClusterMoleculesConfig",
+    "ClusterMoleculesWorkflow",
+    "PhesaConfig",
+    "PhesaWorkflow",
+    "ConformerSettings",
+    "ani_energies_for_conformers",
+    "boltzmann_averaged_psa",
+    "boltzmann_weights_from_energies",
+    "cluster_conformers",
+    "compute_3d_psa_for_conformer",
+    "compute_3d_psa_for_conformers",
+    "generate_conformers",
+    "pipeline_boltzmann_3d_psa",
+    "seed_everything",
+    "select_top_n_by_ani",
+]

cadd_descriptors/ani_conformers/__init__.py

@@ -0,0 +1,11 @@
+"""ANI-driven conformer generation and ranking."""
+
+from .pipeline import extract_conformer_mol, pipeline_ranked_conformers
+from .workflow import AniConformersConfig, AniConformersWorkflow
+
+__all__ = [
+    "AniConformersConfig",
+    "AniConformersWorkflow",
+    "extract_conformer_mol",
+    "pipeline_ranked_conformers",
+]

cadd_descriptors/ani_conformers/pipeline.py

@@ -0,0 +1,72 @@
+"""ANI re-ranked conformer generation pipeline."""
+
+from __future__ import annotations
+
+from typing import Any
+import logging
+from rdkit import Chem
+import torch
+
+from cadd_descriptors.ani_utils import (
+    ConformerSettings,
+    cluster_conformers,
+    generate_conformers,
+    mmff_energies_for_conformers,
+    seed_everything,
+    select_top_n_by_ani,
+)
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "extract_conformer_mol",
+    "pipeline_ranked_conformers",
+]
+
+
+def extract_conformer_mol(mol: Chem.Mol, conf_id: int) -> Chem.Mol:
+    """Create a molecule copy containing only the specified conformer."""
+    mol_copy = Chem.Mol(mol)
+    mol_copy.RemoveAllConformers()
+    conf = Chem.Conformer(mol.GetConformer(conf_id))
+    mol_copy.AddConformer(conf, assignId=True)
+    return mol_copy
+
+
+def pipeline_ranked_conformers(
+    mol: Chem.Mol,
+    settings: ConformerSettings | None = None,
+    *,
+    device: torch.device | None = None,
+    model_name: str = "ANI2x",
+) -> dict[str, Any]:
+    """Run the ANI re-ranked conformer pipeline for a single molecule."""
+    seed_everything()
+    if settings is None:
+        settings = ConformerSettings()
+
+    mol = generate_conformers(mol, settings)
+    logger.info("Generated %d conformers", mol.GetNumConformers())
+
+    logger.info("Clustering conformers...")
+    clusters, rep_conf_ids = cluster_conformers(mol, settings.rmsd_cluster_thresh)
+    logger.info("  %d clusters; using %d representatives", len(clusters), len(rep_conf_ids))
+
+    logger.info("Re-ranking with TorchANI...")
+    top_conf_ids, top_energies_hartree = select_top_n_by_ani(
+        mol,
+        rep_conf_ids,
+        settings,
+        device=device,
+        model_name=model_name,
+    )
+
+    rdkit_energies = mmff_energies_for_conformers(mol, top_conf_ids)
+    return {
+        "mol": mol,
+        "clusters": clusters,
+        "rep_conf_ids": rep_conf_ids,
+        "top_conf_ids": top_conf_ids,
+        "top_ani_energies_hartree": top_energies_hartree,
+        "top_rdkit_energies_kcal_mol": rdkit_energies,
+    }

cadd_descriptors/ani_conformers/workflow.py

@@ -0,0 +1,390 @@
+"""Workflow and config for ANI re-ranked conformer generation."""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from pydantic import Field, model_validator
+from rdkit import Chem
+import torch
+
+from caddforge.config import BaseConfig
+from caddforge.core.job import Job, JobStatus
+from caddforge.core.job_registry import record_job, update_status
+from caddforge.core.result import Result
+from caddforge.core.workflow import Workflow
+
+from cadd_descriptors.ani_utils import ConformerSettings, mol_from_smiles
+from cadd_descriptors.inputs import load_smiles_file, load_structure_file, records_from_smiles_list
+
+from .pipeline import extract_conformer_mol, pipeline_ranked_conformers
+
+logger = logging.getLogger(__name__)
+
+
+class AniConformersConfig(BaseConfig):
+    """Validated configuration for ANI re-ranked conformer generation."""
+
+    @classmethod
+    def from_smiles(cls, smiles: str, **kwargs: Any) -> "AniConformersConfig":
+        """Create a config from a single SMILES while keeping other defaults."""
+        return cls(smiles=smiles, **kwargs)
+
+    @classmethod
+    def from_smiles_list(cls, smiles_list: list[str], **kwargs: Any) -> "AniConformersConfig":
+        """Create a config from an inline list of SMILES."""
+        return cls(smiles_list=smiles_list, **kwargs)
+
+    @classmethod
+    def from_smiles_file(cls, smiles_file: Path, **kwargs: Any) -> "AniConformersConfig":
+        """Create a config from a SMILES file (headerless or with a column)."""
+        return cls(smiles_file=smiles_file, **kwargs)
+
+    @classmethod
+    def from_structure_file(cls, structure_file: Path, **kwargs: Any) -> "AniConformersConfig":
+        """Create a config from an SDF file."""
+        return cls(structure_file=structure_file, **kwargs)
+
+    smiles: str | None = Field(None, description="Single SMILES string to evaluate.")
+    smiles_list: list[str] | None = Field(
+        default=None,
+        description="Inline list of SMILES strings to evaluate as a batch.",
+    )
+    smiles_file: Path | None = Field(
+        default=None,
+        description="Path to a file containing SMILES (one per line, or delimited with a column).",
+    )
+    structure_file: Path | None = Field(
+        default=None,
+        description="Path to an SDF file containing molecules.",
+    )
+    smiles_column: str = Field(
+        default="smiles",
+        description="Column name to read from a delimited SMILES file (CSV/TSV).",
+    )
+    num_confs: int = Field(500, ge=1, description="Number of conformers to generate.")
+    rmsd_cluster_thresh: float = Field(
+        0.5, gt=0, description="RMSD threshold (Angstrom) for Butina clustering."
+    )
+    max_mmff_iters: int = Field(
+        200, ge=1, description="MMFF optimization steps per conformer."
+    )
+    top_n: int = Field(50, ge=1, description="Top-N ANI conformers to keep.")
+    model_name: str = Field(
+        "ANI2x",
+        description="TorchANI model name to use (e.g., ANI2x).",
+    )
+    device: str | None = Field(
+        default="cpu",
+        description="Torch device string (cpu, cuda, cuda:0). Auto-selects if None.",
+    )
+    output_sdf_name: str = Field(
+        default="ranked_conformers.sdf",
+        description="Filename for the output SDF containing ranked conformers.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_inputs(self) -> "AniConformersConfig":
+        """Ensure exactly one input source is provided."""
+        provided = [
+            bool(self.smiles),
+            bool(self.smiles_list),
+            bool(self.smiles_file),
+            bool(self.structure_file),
+        ]
+        if sum(provided) != 1:
+            raise ValueError(
+                "Provide exactly one of: smiles, smiles_list, smiles_file, structure_file."
+            )
+        return self
+
+    def to_conformer_settings(self) -> ConformerSettings:
+        """Convert validated config into ConformerSettings."""
+        return ConformerSettings(
+            num_confs=self.num_confs,
+            rmsd_cluster_thresh=self.rmsd_cluster_thresh,
+            max_mmff_iters=self.max_mmff_iters,
+            top_n=self.top_n,
+        )
+
+    def resolved_inputs(self) -> list[dict[str, Any]]:
+        """Materialize the configured inputs into a list of records."""
+        if self.smiles:
+            return [{"smiles": self.smiles, "name": None, "mol": None}]
+        if self.smiles_list:
+            records = records_from_smiles_list(self.smiles_list)
+            return [{**rec, "mol": None} for rec in records]
+        if self.smiles_file:
+            records = load_smiles_file(self.smiles_file, self.smiles_column)
+            return [{**rec, "mol": None} for rec in records]
+        if self.structure_file:
+            return load_structure_file(self.structure_file)
+        raise ValueError("No inputs resolved.")
+
+    def input_source(self) -> str:
+        """Human-readable description of where inputs were loaded from."""
+        if self.smiles:
+            return "inline_smiles"
+        if self.smiles_list:
+            return "inline_smiles_list"
+        if self.smiles_file:
+            return f"file:{self.smiles_file}"
+        if self.structure_file:
+            return f"file:{self.structure_file}"
+        return "unknown"
+
+    @classmethod
+    def example_dict(cls) -> dict[str, Any]:
+        """Provide richer examples for batch and single-molecule usage."""
+        example = super().example_dict()
+        example.update(
+            {
+                "smiles": "CCO",
+                "smiles_list": ["CCO", "CCN"],
+                "smiles_file": "batch.smi",
+                "structure_file": "batch.sdf",
+                "smiles_column": "smiles",
+            }
+        )
+        return example
+
+
+class AniConformersWorkflow(Workflow):
+    """Workflow that generates ANI re-ranked conformers."""
+
+    tool_name = "ani-conformers"
+    section = "descriptors"
+
+    def __init__(self, config: AniConformersConfig, **_compat):
+        self.config = config
+        self._resolved_inputs: list[dict[str, Any]] | None = None
+
+    def _default_run_directory(self, job_id: str) -> Path:
+        if self.config.run_directory:
+            return Path(self.config.run_directory)
+        date_folder = datetime.now().astimezone().date().isoformat()
+        return Path("runs") / date_folder / job_id
+
+    def _input_records(self) -> list[dict[str, Any]]:
+        if self._resolved_inputs is None:
+            self._resolved_inputs = self.config.resolved_inputs()
+        return self._resolved_inputs
+
+    def prepare_job(self) -> Job:
+        job_id = uuid4().hex
+        run_dir = self._default_run_directory(job_id)
+        run_dir.mkdir(parents=True, exist_ok=True)
+
+        job = Job(
+            id=job_id,
+            tool_name=self.tool_name,
+            run_directory=run_dir,
+            metadata={
+                "input_source": self.config.input_source(),
+                "model_name": self.config.model_name,
+                "section": self.section,
+            },
+        )
+
+        self._write_config(run_dir)
+        job.save_metadata(run_dir / "metadata.json")
+        record_job(job, registry_path=self.config.registry_path)
+        return job
+
+    def _write_config(self, run_dir: Path) -> None:
+        config_path = run_dir / "config.json"
+        config_path.write_text(
+            json.dumps(self.config.model_dump(), indent=2, default=str),
+            encoding="utf-8",
+        )
+
+    def _configure_logging(self, run_dir: Path) -> None:
+        log_path = run_dir / "stdout.log"
+        handler = logging.FileHandler(log_path, mode="w", encoding="utf-8")
+        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+
+        tool_logger = logging.getLogger("cadd_descriptors.ani_conformers")
+        tool_logger.setLevel(logging.INFO)
+        tool_logger.handlers.clear()
+        tool_logger.addHandler(handler)
+        tool_logger.propagate = False
+
+    def _write_ranked_conformers(
+        self,
+        writer: Chem.SDWriter,
+        mol: Chem.Mol,
+        smiles: str,
+        name: str | None,
+        top_conf_ids: list[int],
+        ani_energies: list[float],
+        ani_energies_kcal: list[float],
+        rdkit_energies: list[float],
+    ) -> None:
+        for rank, (cid, ani_energy, rdkit_energy) in enumerate(
+            zip(top_conf_ids, ani_energies, rdkit_energies), start=1
+        ):
+            conf_mol = extract_conformer_mol(mol, cid)
+            conf_mol.SetProp("source_smiles", smiles)
+            if name:
+                conf_mol.SetProp("source_name", name)
+            conf_mol.SetProp("conformer_rank", str(rank))
+            conf_mol.SetProp("conf_id", str(cid))
+            conf_mol.SetProp("ani_energy_hartree", f"{ani_energy:.10f}")
+            conf_mol.SetProp("ani_energy_kcal_mol", f"{ani_energies_kcal[rank - 1]:.6f}")
+            conf_mol.SetProp("rdkit_energy_kcal_mol", f"{rdkit_energy:.6f}")
+            writer.write(conf_mol)
+
+    def run(self, job: Job) -> Result:
+        run_dir = job.run_directory or self._default_run_directory(job.id)
+        run_dir.mkdir(parents=True, exist_ok=True)
+        self._configure_logging(run_dir)
+        started_at = datetime.now().astimezone()
+        job.started_at = started_at
+        update_status(
+            job.id,
+            JobStatus.RUNNING,
+            started_at=started_at,
+            registry_path=self.config.registry_path,
+        )
+
+        results_dir = run_dir / "results"
+        results_dir.mkdir(parents=True, exist_ok=True)
+        results_json_path = results_dir / "results.json"
+        output_sdf_path = results_dir / self.config.output_sdf_name
+
+        inputs = self._input_records()
+        job.metadata.update(
+            {
+                "num_inputs": len(inputs),
+                "smiles_preview": [record["smiles"] for record in inputs[:3]],
+            }
+        )
+        job.save_metadata(run_dir / "metadata.json")
+        molecule_results: list[dict[str, Any]] = []
+        failures: list[dict[str, str]] = []
+        result_obj: Result | None = None
+        error: Exception | None = None
+
+        try:
+            settings = self.config.to_conformer_settings()
+            device = torch.device(self.config.device) if self.config.device else None
+
+            writer = Chem.SDWriter(str(output_sdf_path))
+            try:
+                for idx, record in enumerate(inputs, start=1):
+                    smiles = record.get("smiles") or ""
+                    name = record.get("name")
+                    mol_input = record.get("mol")
+                    try:
+                        mol = mol_input if isinstance(mol_input, Chem.Mol) else mol_from_smiles(smiles)
+                        if not smiles:
+                            smiles = Chem.MolToSmiles(mol)
+                        pipeline_result = pipeline_ranked_conformers(
+                            mol,
+                            settings=settings,
+                            device=device,
+                            model_name=self.config.model_name,
+                        )
+                        top_conf_ids = pipeline_result["top_conf_ids"]
+                        top_ani = pipeline_result["top_ani_energies_hartree"]
+                        top_rdkit = pipeline_result["top_rdkit_energies_kcal_mol"]
+                        hartree_to_kcal = 627.509474
+                        top_ani_kcal = [float(val) * hartree_to_kcal for val in top_ani]
+                        self._write_ranked_conformers(
+                            writer,
+                            pipeline_result["mol"],
+                            smiles,
+                            name,
+                            top_conf_ids,
+                            [float(val) for val in top_ani],
+                            top_ani_kcal,
+                            [float(val) for val in top_rdkit],
+                        )
+                        molecule_results.append(
+                            {
+                                "index": idx,
+                                "smiles": smiles,
+                                "name": name,
+                                "num_conformers_written": len(top_conf_ids),
+                                "top_conf_ids": [int(cid) for cid in top_conf_ids],
+                                "top_ani_energies_hartree": [float(val) for val in top_ani],
+                                "top_ani_energies_kcal_mol": top_ani_kcal,
+                                "top_rdkit_energies_kcal_mol": [float(val) for val in top_rdkit],
+                                "num_clusters": len(pipeline_result["clusters"]),
+                                "num_representatives": len(pipeline_result["rep_conf_ids"]),
+                            }
+                        )
+                    except Exception as exc:  # pragma: no cover - per-molecule guard
+                        logger.exception("Conformer workflow failed for %s: %s", smiles, exc)
+                        failures.append({"smiles": smiles, "error": str(exc)})
+            finally:
+                writer.close()
+
+            status = JobStatus.COMPLETED if not failures else JobStatus.FAILED
+            metrics: dict[str, float | int | str] = {
+                "molecules_succeeded": len(molecule_results),
+                "molecules_failed": len(failures),
+            }
+            if molecule_results:
+                avg_written = sum(entry["num_conformers_written"] for entry in molecule_results) / len(
+                    molecule_results
+                )
+                metrics["mean_conformers_written"] = float(avg_written)
+
+            results_payload = {"molecules": molecule_results, "failures": failures}
+            results_json_path.write_text(json.dumps(results_payload, indent=2), encoding="utf-8")
+
+            paths = {"results_json": results_json_path, "ranked_conformers_sdf": output_sdf_path}
+            result_obj = Result(
+                job_id=job.id,
+                tool_name=self.tool_name,
+                status=status,
+                metrics=metrics,
+                paths=paths,
+                extra=results_payload,
+            )
+        except Exception as exc:  # pragma: no cover - defensive guard
+            error = exc
+            logger.exception("Conformer workflow failed: %s", exc)
+            status = JobStatus.FAILED
+            result_obj = Result(
+                job_id=job.id,
+                tool_name=self.tool_name,
+                status=status,
+                metrics={},
+                paths={},
+                extra={"error": str(exc)},
+            )
+
+        job.metadata.update(
+            {
+                "num_inputs": len(inputs),
+                "molecules_succeeded": len(molecule_results),
+                "molecules_failed": len(failures),
+            }
+        )
+
+        finished_at = datetime.now().astimezone()
+        job.finished_at = finished_at
+        job.status = status
+        job.save_metadata(run_dir / "metadata.json")
+        update_status(
+            job.id,
+            status,
+            finished_at=finished_at,
+            registry_path=self.config.registry_path,
+        )
+
+        if result_obj is not None:
+            result_obj.save(results_dir / "result_record.json")
+
+        if error:
+            raise error
+
+        return result_obj

cadd_descriptors/ani_psa/__init__.py

@@ -0,0 +1,32 @@
+"""ANI-weighted PSA descriptor implementation."""
+
+from .pipeline import (
+    ConformerSettings,
+    ani_energies_for_conformers,
+    boltzmann_averaged_psa,
+    boltzmann_weights_from_energies,
+    cluster_conformers,
+    compute_3d_psa_for_conformer,
+    compute_3d_psa_for_conformers,
+    generate_conformers,
+    pipeline_boltzmann_3d_psa,
+    seed_everything,
+    select_top_n_by_ani,
+)
+from .workflow import AniPsaConfig, AniPsaWorkflow
+
+__all__ = [
+    "AniPsaConfig",
+    "AniPsaWorkflow",
+    "ConformerSettings",
+    "ani_energies_for_conformers",
+    "boltzmann_averaged_psa",
+    "boltzmann_weights_from_energies",
+    "cluster_conformers",
+    "compute_3d_psa_for_conformer",
+    "compute_3d_psa_for_conformers",
+    "generate_conformers",
+    "pipeline_boltzmann_3d_psa",
+    "seed_everything",
+    "select_top_n_by_ani",
+]