quantui 0.5.1__py3-none-any.whl

quantui/structure_providers.py

@@ -0,0 +1,277 @@
+"""Unified structure-resolver chain with offline fallback.
+
+Resolution order (first hit wins):
+
+1. **Local RDKit** for SMILES / InChI input — offline, no network.
+2. **Bundled library** exact hit (formula key) — offline, instant.
+3. **PubChem** (hardened client).
+4. **NCI CACTUS** resolver.
+5. **Bundled-library fuzzy fallback** (name/description substring) — the
+   last-resort offline answer so the search box is never a dead end, even with
+   no network.
+
+Every resolver returns a normalized :class:`ResolvedStructure` so callers
+(and the disambiguation UI) treat all sources uniformly. The
+``source`` field records which resolver answered, so the UI can be honest about
+provenance.
+
+The bundled-library steps currently search ``config.MOLECULE_LIBRARY`` (the 20
+presets). The library lives in an indexed package-data store, reached only
+through :func:`_library_exact` / :func:`_library_fuzzy`, so the chain is
+unaffected by how it is stored.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+from . import cactus, config
+from .pubchem import (
+    MoleculeNotFoundError,
+    PubChemAPIError,
+    classify_query,
+    fetch_structure,
+    search_pubchem_candidates,
+)
+
+logger = logging.getLogger(__name__)
+
+# Elements that count as "heavy" exclusions when tallying heavy atoms.
+_HYDROGEN = "H"
+
+
+@dataclass
+class ResolvedStructure:
+    """A structure resolved by the provider chain, normalized across sources."""
+
+    xyz: str
+    source: str  # "rdkit-smiles" | "rdkit-inchi" | "library" | "pubchem" | "cactus" | "library-offline-fallback"
+    formula: str = "?"
+    num_atoms: int = 0
+    num_heavy_atoms: int = 0
+    charge: int = 0
+    multiplicity: int = 1
+    molecular_weight: Optional[float] = None
+    conformer_origin: str = ""
+    identifiers: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def is_offline(self) -> bool:
+        return self.source in (
+            "library",
+            "library-offline-fallback",
+        ) or self.source.startswith("rdkit-")
+
+
+def _from_metadata(xyz: str, meta: Dict[str, Any], *, source: str) -> ResolvedStructure:
+    """Build a ResolvedStructure from a ``(xyz, metadata)`` resolver result."""
+    identifiers = {
+        k: meta[k]
+        for k in ("pubchem_cid", "smiles", "canonical_smiles", "inchi", "query_type")
+        if k in meta
+    }
+    return ResolvedStructure(
+        xyz=xyz,
+        source=source,
+        formula=meta.get("formula", "?"),
+        num_atoms=int(meta.get("num_atoms", 0) or 0),
+        num_heavy_atoms=int(meta.get("num_heavy_atoms", 0) or 0),
+        charge=int(meta.get("charge", 0) or 0),
+        molecular_weight=meta.get("molecular_weight"),
+        conformer_origin=meta.get("conformer_origin", meta.get("source", "")),
+        identifiers=identifiers,
+    )
+
+
+def _entry_to_xyz(name: str, entry: Dict[str, Any]) -> str:
+    """Format a ``MOLECULE_LIBRARY`` entry as an XYZ string."""
+    atoms = entry["atoms"]
+    coords = entry["coordinates"]
+    desc = entry.get("description", "")
+    lines = [str(len(atoms)), f"{name}: {desc}".strip()]
+    for sym, (x, y, z) in zip(atoms, coords):
+        lines.append(f"{sym:3s} {float(x):12.6f} {float(y):12.6f} {float(z):12.6f}")
+    return "\n".join(lines)
+
+
+def _entry_to_resolved(
+    name: str, entry: Dict[str, Any], *, source: str
+) -> ResolvedStructure:
+    atoms = entry["atoms"]
+    return ResolvedStructure(
+        xyz=_entry_to_xyz(name, entry),
+        source=source,
+        formula=name,
+        num_atoms=len(atoms),
+        num_heavy_atoms=sum(1 for a in atoms if a != _HYDROGEN),
+        charge=int(entry.get("charge", 0)),
+        multiplicity=int(entry.get("multiplicity", 1)),
+        conformer_origin="library",
+        identifiers={"library_key": name},
+    )
+
+
+def _library_exact(query: str) -> Optional[ResolvedStructure]:
+    """Exact (case-insensitive) match on a bundled-library formula key."""
+    q = query.strip().lower()
+    for key, entry in config.MOLECULE_LIBRARY.items():
+        if q == key.lower():
+            logger.info(f"Resolved '{query}' from bundled library (exact key '{key}')")
+            return _entry_to_resolved(key, entry, source="library")
+    return None
+
+
+def _library_fuzzy(query: str) -> Optional[ResolvedStructure]:
+    """Loose offline fallback: substring match over keys + descriptions.
+
+    Used only when the network resolvers are unreachable or all miss, so a
+    looser match is acceptable (and clearly labelled as a fallback to the user).
+    """
+    q = query.strip().lower()
+    if not q:
+        return None
+    for key, entry in config.MOLECULE_LIBRARY.items():
+        haystack = f"{key} {entry.get('description', '')}".lower()
+        if q == key.lower() or q in haystack:
+            logger.info(f"Offline fallback matched '{query}' to library entry '{key}'")
+            return _entry_to_resolved(key, entry, source="library-offline-fallback")
+    return None
+
+
+def resolve_structure(
+    query: str,
+    *,
+    conformer_3d: bool = True,
+    allow_network: bool = True,
+) -> ResolvedStructure:
+    """Resolve ``query`` through the provider chain. Raises on total failure.
+
+    Raises:
+        MoleculeNotFoundError: nothing in the chain could resolve the query.
+        ValueError: empty query, or RDKit failed to parse a SMILES/InChI.
+    """
+    qtype = classify_query(query)
+    logger.info(f"Resolving '{query}' (type={qtype}, network={allow_network})")
+
+    # 1. Local RDKit for SMILES / InChI — no network, no library needed.
+    if qtype in ("smiles", "inchi"):
+        xyz, meta = fetch_structure(query, conformer_3d=conformer_3d)
+        return _from_metadata(xyz, meta, source=meta["source"])
+
+    # 2. Exact bundled-library hit (offline, instant) — e.g. "H2O", "C6H6".
+    hit = _library_exact(query)
+    if hit is not None:
+        return hit
+
+    # 3/4. Network resolvers: PubChem, then CACTUS. A genuine miss
+    # (MoleculeNotFoundError) or a transport error (PubChemAPIError) both fall
+    # through to the next resolver, then to the offline fallback.
+    errors: List[Tuple[str, Exception]] = []
+    if allow_network:
+        try:
+            xyz, meta = fetch_structure(query, conformer_3d=conformer_3d)
+            return _from_metadata(xyz, meta, source="pubchem")
+        except (MoleculeNotFoundError, PubChemAPIError) as exc:
+            logger.info(f"PubChem did not resolve '{query}': {exc}")
+            errors.append(("pubchem", exc))
+
+        try:
+            xyz, meta = cactus.fetch_from_cactus(query, conformer_3d=conformer_3d)
+            return _from_metadata(xyz, meta, source="cactus")
+        except (MoleculeNotFoundError, PubChemAPIError) as exc:
+            logger.info(f"CACTUS did not resolve '{query}': {exc}")
+            errors.append(("cactus", exc))
+
+    # 5. Offline fuzzy fallback against the bundled library.
+    hit = _library_fuzzy(query)
+    if hit is not None:
+        return hit
+
+    tried = ", ".join(name for name, _ in errors) or "offline only"
+    raise MoleculeNotFoundError(
+        f"Could not resolve '{query}' (tried: {tried}, bundled library)"
+    )
+
+
+def search_candidates(query: str) -> List[Dict[str, Any]]:
+    """Return disambiguation candidates for a multi-match name/formula query.
+
+    Only name/formula queries can be ambiguous (SMILES/InChI/CID resolve to a
+    single structure, and the local library is exact). Returns ``[]`` for those
+    types, and ``[]`` on any network failure so the caller falls back to the
+    full single-result chain (PubChem → CACTUS → offline library).
+    """
+    if classify_query(query) not in ("name", "formula"):
+        return []
+    try:
+        return search_pubchem_candidates(query)
+    except (MoleculeNotFoundError, PubChemAPIError):
+        return []
+
+
+def resolve_structure_with_message(
+    query: str,
+) -> Tuple[Optional[str], str, Optional[str], bool]:
+    """Chain-backed resolve that also reports provenance.
+
+    Returns ``(xyz_or_None, message, source_or_None, is_offline)``. ``source``
+    is the raw provider key (``"pubchem"`` / ``"cactus"`` / ``"library"`` /
+    ``"library-offline-fallback"`` / ``"rdkit-smiles"`` / ``"rdkit-inchi"``) so
+    the UI can label the loaded molecule by where it *actually* came from
+    rather than always saying "PubChem". ``is_offline`` is True when the
+    structure was produced without a network call (local RDKit or the bundled
+    library) — the UI surfaces a no-network note for the offline *fallback*
+    (network was attempted and failed).
+    """
+    try:
+        result = resolve_structure(query, conformer_3d=True)
+    except (MoleculeNotFoundError, ValueError) as exc:
+        return (
+            None,
+            (
+                f"❌ Could not resolve '{query}'.\n"
+                f"   {exc}\n"
+                f"   Try a different name, a SMILES (e.g. CC(=O)O), a CAS "
+                f"number, or check spelling.\n"
+                f"   Search manually at: https://pubchem.ncbi.nlm.nih.gov/"
+            ),
+            None,
+            False,
+        )
+    except Exception as exc:  # pragma: no cover - unexpected
+        logger.error(f"Unexpected error resolving '{query}': {exc}", exc_info=True)
+        return None, f"❌ Error resolving '{query}': {exc}", None, False
+
+    source_label = {
+        "rdkit-smiles": "generated locally from SMILES",
+        "rdkit-inchi": "generated locally from InChI",
+        "library": "the bundled library (offline)",
+        "library-offline-fallback": "the bundled library (offline fallback — "
+        "network resolvers were unavailable)",
+        "pubchem": "PubChem",
+        "cactus": "NCI CACTUS",
+    }.get(result.source, result.source)
+
+    embedded = (
+        " (2D structure embedded by RDKit)"
+        if result.conformer_origin == "rdkit-embedded"
+        else ""
+    )
+    mw = f"{result.molecular_weight:.2f} g/mol" if result.molecular_weight else "—"
+    message = (
+        f"✓ Resolved '{query}' via {source_label}.\n"
+        f"  Formula: {result.formula}\n"
+        f"  Atoms: {result.num_atoms} ({result.num_heavy_atoms} heavy)\n"
+        f"  Molecular weight: {mw}{embedded}"
+    )
+    return result.xyz, message, result.source, result.is_offline
+
+
+def student_friendly_resolve(query: str) -> Tuple[Optional[str], str]:
+    """Chain-backed, student-friendly resolve. Drop-in for the UI handler.
+
+    Returns ``(xyz_string_or_None, message)``. Thin wrapper over
+    :func:`resolve_structure_with_message` (which also reports provenance).
+    """
+    xyz, message, _source, _is_offline = resolve_structure_with_message(query)
+    return xyz, message

quantui/tddft_calc.py

@@ -0,0 +1,307 @@
+"""
+TD-DFT excited-state calculation using PySCF.
+
+Computes vertical excitation energies and oscillator strengths using
+time-dependent density functional theory (TD-DFT).  For Hartree-Fock
+methods (RHF/UHF), falls back to TDHF (equivalent to CIS) and notes
+this in the output.
+
+Platform notes
+--------------
+Requires PySCF — Linux / macOS / WSL only.
+
+Educational value
+-----------------
+* Students see which wavelengths a molecule absorbs (UV-Vis spectrum).
+* Oscillator strengths indicate which transitions are optically allowed
+  (bright, f > ~0.01) versus dark (f ≈ 0).
+* Teaches the connection between electronic structure and spectroscopy.
+* Comparing TD-DFT results for different functionals shows how the
+  choice of functional affects excitation energies.
+
+Typical usage
+-------------
+>>> from quantui.tddft_calc import run_tddft_calc
+>>> result = run_tddft_calc(molecule, method="B3LYP", basis="6-31G")
+>>> for e, f in zip(result.excitation_energies_ev, result.oscillator_strengths):
+...     print(f"  E = {e:.3f} eV,  f = {f:.4f}")
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from dataclasses import dataclass, field
+from typing import IO, Any, List, Optional
+
+from .molecule import Molecule
+from .session_calc import HARTREE_TO_EV
+
+logger = logging.getLogger(__name__)
+
+# Planck × speed of light in eV·nm  (h·c = 1239.84 eV·nm)
+_EV_TO_NM: float = 1239.84193
+
+
+# ============================================================================
+# Result dataclass
+# ============================================================================
+
+
+@dataclass
+class TDDFTResult:
+    """Structured output from a TD-DFT excited-state calculation.
+
+    Attributes:
+        energy_hartree: Ground-state SCF energy in Hartrees.
+        homo_lumo_gap_ev: HOMO-LUMO gap in eV from the ground-state SCF,
+            or ``None``.
+        converged: ``True`` if the ground-state SCF converged.
+        n_iterations: Number of ground-state SCF macro-iterations.
+        method: DFT functional or HF method used.
+        basis: Basis set.
+        formula: Hill-notation molecular formula.
+        excitation_energies_ev: Vertical excitation energies in eV.
+        oscillator_strengths: Oscillator strengths (dimensionless).
+            Bright (optically allowed) transitions have f > ~0.01.
+        nstates: Number of excited states requested.
+    """
+
+    energy_hartree: float
+    homo_lumo_gap_ev: Optional[float]
+    converged: bool
+    n_iterations: int
+    method: str
+    basis: str
+    formula: str
+    excitation_energies_ev: List[float] = field(default_factory=list)
+    oscillator_strengths: List[float] = field(default_factory=list)
+    nstates: int = 10
+
+    @property
+    def energy_ev(self) -> float:
+        """Ground-state SCF energy in electronvolts."""
+        return self.energy_hartree * HARTREE_TO_EV
+
+    def wavelengths_nm(self) -> List[float]:
+        """Return excitation wavelengths in nm (λ = 1239.84 / E_eV)."""
+        return [
+            _EV_TO_NM / e if e > 0 else float("inf")
+            for e in self.excitation_energies_ev
+        ]
+
+
+# ============================================================================
+# Main function
+# ============================================================================
+
+
+def run_tddft_calc(
+    molecule: Molecule,
+    method: str = "B3LYP",
+    basis: str = "STO-3G",
+    nstates: int = 10,
+    progress_stream: Optional[IO[str]] = None,
+) -> TDDFTResult:
+    """Run a TD-DFT excited-state calculation to obtain UV-Vis absorption data.
+
+    Converges the ground-state SCF, then runs the time-dependent response
+    equations to compute the requested number of vertical excitation energies
+    and their oscillator strengths.
+
+    When *method* is ``'RHF'`` or ``'UHF'``, the function uses TDHF (CIS)
+    rather than TD-DFT and writes a note to *progress_stream*.  For a proper
+    UV-Vis simulation, a DFT functional such as ``'B3LYP'`` or ``'PBE0'`` is
+    strongly recommended.
+
+    Args:
+        molecule: Validated :class:`~quantui.molecule.Molecule`.
+        method: DFT functional (e.g. ``'B3LYP'``, ``'PBE0'``,
+            ``'CAM-B3LYP'``) or ``'RHF'``/``'UHF'`` for TDHF.
+            Default: ``'B3LYP'``.
+        basis: Basis set name.  Default: ``'STO-3G'``.
+        nstates: Number of excited states to compute.  Default: 10.
+        progress_stream: Optional writable text stream for live PySCF output.
+
+    Returns:
+        :class:`TDDFTResult` with excitation energies and oscillator strengths.
+
+    Raises:
+        ImportError: If PySCF is not installed.
+        RuntimeError: If the ground-state SCF calculation fails.  If the
+            TD calculation fails, excitation lists are empty and a warning
+            is written to progress_stream — no exception is raised.
+    """
+    # Post-HF methods (MP2/CCSD/CCSD(T)) have no special-casing below —
+    # without this guard, method='CCSD' silently falls into the DFT
+    # branch (sets mf.xc = "CCSD") and fails deep inside PySCF with a
+    # cryptic "LibXCFunctional: name 'CCSD' not found" instead of a clear
+    # message. TD-DFT/TDHF is not defined for these methods here.
+    from . import config as _config
+
+    if method.strip().upper() in _config.POST_HF_METHODS:
+        raise ValueError(
+            f"'{method}' is a post-HF method and cannot be used for "
+            "TD-DFT/UV-Vis — use RHF/UHF (TDHF) or a DFT functional "
+            "instead."
+        )
+
+    try:
+        from pyscf import dft, gto, scf
+    except ImportError as exc:
+        raise ImportError(
+            "PySCF is not installed — cannot run TD-DFT.\n"
+            "PySCF requires Linux, macOS, or WSL."
+        ) from exc
+
+    stream: IO[str] = progress_stream if progress_stream is not None else sys.stdout
+
+    # See quantui/c_stderr.py — captures fd-2 stderr
+    # from libcint / BLAS / LAPACK / TDA solver C code and relays to
+    # ``stream`` on exit. POSIX-only; no-op on Windows.
+    from quantui.c_stderr import capture_c_stderr
+
+    with capture_c_stderr(stream):
+        return _run_tddft_calc_body(
+            molecule=molecule,
+            method=method,
+            basis=basis,
+            nstates=nstates,
+            progress_stream=progress_stream,
+            _dft=dft,
+            _gto=gto,
+            _scf=scf,
+            stream=stream,
+        )
+
+
+def _run_tddft_calc_body(
+    *,
+    molecule: Molecule,
+    method: str,
+    basis: str,
+    nstates: int,
+    progress_stream: Optional[IO[str]],
+    _dft: Any,
+    _gto: Any,
+    _scf: Any,
+    stream: IO[str],
+) -> TDDFTResult:
+    """Inner body of :func:`run_tddft_calc` (split out for stderr-capture wrap)."""
+    dft, gto, scf = _dft, _gto, _scf
+
+    # ── Build Mole object ────────────────────────────────────────────────────
+    mol = gto.Mole()
+    mol.atom = molecule.to_pyscf_format()
+    mol.basis = basis
+    mol.charge = molecule.charge
+    mol.spin = molecule.multiplicity - 1
+    mol.verbose = 4
+    mol.stdout = stream
+    mol.build()
+
+    # ── SCF ──────────────────────────────────────────────────────────────────
+    method_upper = method.upper()
+    using_hf = method_upper in ("RHF", "UHF")
+
+    if method_upper == "RHF":
+        mf = scf.RHF(mol)
+    elif method_upper == "UHF":
+        mf = scf.UHF(mol)
+    else:
+        # Route through resolve_xc + maybe_apply_d3 so
+        # methods like wB97X-D (PySCF rejects "wb97x-d") map cleanly.
+        from .session_calc import maybe_apply_d3, resolve_xc
+
+        mf = dft.RKS(mol) if mol.spin == 0 else dft.UKS(mol)
+        mf.xc = resolve_xc(method)
+        mf = maybe_apply_d3(mf, method, progress_stream=progress_stream)
+
+    if using_hf and progress_stream is not None:
+        try:
+            progress_stream.write(
+                "\nNote: Using TDHF (CIS) for excited states — RHF/UHF was selected.\n"
+                "For a proper TD-DFT UV-Vis spectrum, use a DFT functional\n"
+                "such as B3LYP or PBE0 in the Method dropdown.\n\n"
+            )
+        except Exception:  # noqa: BLE001 — cleanup (stream may be closed)
+            pass
+
+    # Cooperative cancel between SCF cycles.
+    from .cancellation import attach_scf_cancel_callback, cancel_check_from_stream
+    from .log_utils import emit_status
+
+    attach_scf_cancel_callback(mf, cancel_check_from_stream(stream))
+
+    emit_status(stream, "Running SCF (ground state)…")
+    try:
+        energy_hartree = float(mf.kernel())
+    except Exception as exc:
+        raise RuntimeError(
+            f"SCF failed for {molecule.get_formula()} ({method}/{basis}): {exc}"
+        ) from exc
+
+    converged = bool(getattr(mf, "converged", False))
+    n_iterations = int(getattr(mf, "cycles", -1))
+
+    # ── HOMO-LUMO gap (non-fatal) ────────────────────────────────────────────
+    homo_lumo_gap_ev: Optional[float] = None
+    try:
+        import numpy as _np
+
+        mo_occ = mf.mo_occ
+        mo_energy = mf.mo_energy
+        if isinstance(mo_energy, (list, _np.ndarray)) and hasattr(
+            mo_energy[0], "__len__"
+        ):
+            mo_e_ref, mo_occ_ref = mo_energy[0], mo_occ[0]
+        else:
+            mo_e_ref, mo_occ_ref = mo_energy, mo_occ
+        n_occ = int((_np.array(mo_occ_ref) > 0).sum())
+        if 0 < n_occ < len(mo_e_ref):
+            homo_lumo_gap_ev = float(
+                (mo_e_ref[n_occ] - mo_e_ref[n_occ - 1]) * HARTREE_TO_EV
+            )
+    except Exception as exc:
+        logger.debug("HOMO-LUMO gap extraction failed in TD-DFT calc: %s", exc)
+
+    # ── TD-DFT / TDHF ────────────────────────────────────────────────────────
+    excitation_energies_ev: List[float] = []
+    oscillator_strengths: List[float] = []
+
+    try:
+        emit_status(
+            stream,
+            f"Solving {'TDHF (CIS)' if using_hf else 'TD-DFT'} "
+            f"excited states ({nstates})…",
+        )
+        td = mf.TDHF() if using_hf else mf.TDDFT()
+        td.nstates = nstates
+        td.verbose = 3
+        td.stdout = stream
+        td.kernel()
+
+        excitation_energies_ev = [float(e) * HARTREE_TO_EV for e in td.e]
+        osc = td.oscillator_strength()
+        oscillator_strengths = [float(f) for f in osc]
+
+    except Exception as exc:
+        logger.warning("TD-DFT/TDHF calculation failed: %s", exc)
+        if progress_stream is not None:
+            try:
+                progress_stream.write(f"\n⚠ TD-DFT failed: {exc}\n")
+            except Exception:  # noqa: BLE001 — cleanup (stream may be closed)
+                pass
+
+    return TDDFTResult(
+        energy_hartree=energy_hartree,
+        homo_lumo_gap_ev=homo_lumo_gap_ev,
+        converged=converged,
+        n_iterations=n_iterations,
+        method=method,
+        basis=basis,
+        formula=molecule.get_formula(),
+        excitation_energies_ev=excitation_energies_ev,
+        oscillator_strengths=oscillator_strengths,
+        nstates=nstates,
+    )