Source code for vibeqc.io.molden

"""Molden-format writer for vibe-qc SCF results.

Molden files carry the geometry, the basis set, and the molecular orbitals
in a single plain-text document. Any orbital viewer that understands the
format (Molden itself, Jmol, Avogadro, MolView, IQmol, ...) can visualise
the occupied and virtual MOs without further processing.

Supports RHF, UHF, RKS, and UKS results. Restricted results produce one
[MO] block; unrestricted results produce two (alpha followed by beta).

Basis-function ordering
-----------------------

Molden and libint disagree on the ordering of real spherical harmonics
within a shell of angular momentum L >= 1. We reorder MO coefficients
at write time so the emitted file matches Molden's convention:

- p (L=1): molden (px, py, pz) <- libint (py, pz, px)
- d (L=2): molden (d0, d+1, d-1, d+2, d-2) <- libint (d-2, d-1, d0, d+1, d+2)
- f, g, h: molden (m=0, +1, -1, +2, -2, ..., +L, -L) <- libint (m = -L..+L)

vibe-qc forces ``set_pure(true)`` on the BasisSet at construction, so every
shell is spherical-harmonic in this writer — Cartesian d/f are not emitted.
"""

from __future__ import annotations

import os
from typing import List, Sequence, Tuple

import numpy as np

from .._vibeqc_core import BasisSet, Molecule


# Element-symbol table; indexed by atomic number (0 is unused).
_ELEMENTS: Tuple[str, ...] = (
    "X",
    "H",  "He",
    "Li", "Be", "B",  "C",  "N",  "O",  "F",  "Ne",
    "Na", "Mg", "Al", "Si", "P",  "S",  "Cl", "Ar",
    "K",  "Ca", "Sc", "Ti", "V",  "Cr", "Mn", "Fe", "Co", "Ni",
    "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr",
    "Rb", "Sr", "Y",  "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd",
    "Ag", "Cd", "In", "Sn", "Sb", "Te", "I",  "Xe",
)


def _element(z: int) -> str:
    if 0 < z < len(_ELEMENTS):
        return _ELEMENTS[z]
    return f"Z{z}"


# Molden-spec shell-type letters (lowercase).
_SHELL_LETTER = ("s", "p", "d", "f", "g", "h", "i", "k", "l")


def _primitive_normalisation(alpha: float, L: int) -> float:
    """Normalisation constant N_i such that (N_i · x^lx y^ly z^lz e^{-α r²})
    has unit overlap for (lx+ly+lz = L).

    The formula

        N = (2α/π)^{3/4} · (4α)^{L/2} / √((2L−1)!!)

    matches the convention used by libint2 / Gaussian / Molpro / Turbomole.
    Used here to strip libint's built-in primitive normalisation out of the
    contraction coefficients so that what lands in the molden file is the
    raw basis-set-file value (which is what molden readers assume).
    """
    import math

    double_factorial = 1.0
    for k in range(1, 2 * L, 2):
        double_factorial *= k
    return (
        (2.0 * alpha / math.pi) ** 0.75
        * (4.0 * alpha) ** (L / 2.0)
        / math.sqrt(double_factorial)
    )


def _pure_molden_to_libint(L: int) -> List[int]:
    """Index permutation: molden_order[j] = libint slot of molden basis j.

    For L=0 a single s function, identity.

    For L=1 libint's ``pure`` p-shell indexes m = -1, 0, +1 which in real
    solid harmonics is (py, pz, px). Molden wants (px, py, pz).

    For L >= 2 libint uses m = -L, -L+1, ..., 0, ..., +L, while molden uses
    (m=0, +1, -1, +2, -2, ..., +L, -L).
    """
    if L == 0:
        return [0]
    if L == 1:
        return [2, 0, 1]  # px(m=+1=libint 2), py(m=-1=libint 0), pz(m=0=libint 1)
    idx = [L]  # m = 0
    for k in range(1, L + 1):
        idx.append(L + k)  # m = +k
        idx.append(L - k)  # m = -k
    return idx


def _build_ao_permutation(basis: BasisSet) -> np.ndarray:
    """Global AO permutation: ``C_molden = C_libint[ao_perm, :]``.

    Each shell contributes 2L+1 consecutive entries; we walk shells in
    libint's native order and stitch per-shell spherical-harmonic
    reorderings together.
    """
    perm: List[int] = []
    offset = 0
    for shell in basis.shells():
        if not shell.pure:
            raise ValueError(
                "write_molden: Cartesian shells are not supported yet "
                f"(shell L={shell.l} pure={shell.pure}). vibeqc forces "
                "set_pure(true) at construction — this error means the "
                "BasisSet was created outside the normal path."
            )
        local = _pure_molden_to_libint(shell.l)
        perm.extend(offset + i for i in local)
        offset += len(local)
    return np.asarray(perm, dtype=np.intp)


def _write_header(f, title: str) -> None:
    f.write("[Molden Format]\n")
    f.write("[Title]\n")
    # Molden tolerates a blank title line; keep it non-empty to be safe.
    f.write(f" {title or 'vibe-qc output'}\n")


def _write_atoms(f, mol: Molecule) -> None:
    f.write("[Atoms] (AU)\n")
    for idx, atom in enumerate(mol.atoms, start=1):
        sym = _element(atom.Z)
        x, y, z = atom.xyz[0], atom.xyz[1], atom.xyz[2]
        f.write(
            f" {sym:<4s} {idx:5d} {atom.Z:5d} "
            f"{x:20.12f} {y:20.12f} {z:20.12f}\n"
        )


def _write_gto(f, basis: BasisSet, n_atoms: int) -> None:
    """[GTO] block: per-atom shell listings, in atom index order."""
    # Group shells by atom (libint emits shells atom-by-atom but we don't
    # rely on that — sort explicitly).
    shells_by_atom: List[List] = [[] for _ in range(n_atoms)]
    for shell in basis.shells():
        shells_by_atom[shell.atom_index].append(shell)

    f.write("[GTO]\n")
    for atom_idx in range(n_atoms):
        # Trailing zero is the historical "scaling factor" field — always 0.
        f.write(f"  {atom_idx + 1} 0\n")
        for shell in shells_by_atom[atom_idx]:
            letter = _SHELL_LETTER[shell.l]
            n_prim = len(shell.exponents)
            f.write(f" {letter}   {n_prim}  1.00\n")
            for alpha, c_eff in zip(shell.exponents, shell.coefficients):
                # libint stores primitive-normalised coefficients; molden
                # readers expect the raw basis-set-file values and apply
                # primitive normalisation themselves.
                c_raw = c_eff / _primitive_normalisation(alpha, shell.l)
                f.write(f"    {alpha:18.10E}    {c_raw:18.10E}\n")
        f.write("\n")


def _write_pure_flags(f, basis: BasisSet) -> None:
    """Molden spec: ``[5D]`` / ``[7F]`` / ``[9G]`` announce spherical."""
    lmax = 0
    for shell in basis.shells():
        if shell.l > lmax:
            lmax = shell.l
    # Emit the flags relevant to the basis. [5D] implies [7F] [9G] ...
    # in some readers, but writing them explicitly is the safe bet.
    if lmax >= 2:
        f.write("[5D]\n")
    if lmax >= 3:
        f.write("[7F]\n")
    if lmax >= 4:
        f.write("[9G]\n")


def _write_mo_block(
    f,
    mo_energies: np.ndarray,
    mo_coeffs: np.ndarray,   # shape (nbasis, nmo), columns are MOs
    occupations: Sequence[float],
    spin: str,               # "Alpha" or "Beta"
    ao_perm: np.ndarray,
) -> None:
    """Write one [MO] block — a list of orbitals with energy / spin / occ."""
    n_ao, n_mo = mo_coeffs.shape
    reordered = mo_coeffs[ao_perm, :]  # AO permutation to molden ordering

    for j in range(n_mo):
        f.write(f" Sym=  A\n")
        f.write(f" Ene= {mo_energies[j]:18.10f}\n")
        f.write(f" Spin= {spin}\n")
        f.write(f" Occup= {occupations[j]:12.8f}\n")
        col = reordered[:, j]
        for i in range(n_ao):
            f.write(f" {i + 1:5d}   {col[i]:20.12E}\n")


def _rhf_occupations(n_mo: int, n_electrons: int) -> List[float]:
    n_occ = n_electrons // 2
    return [2.0 if i < n_occ else 0.0 for i in range(n_mo)]


def _uhf_occupations(
    n_mo: int, multiplicity: int, n_electrons: int
) -> Tuple[List[float], List[float]]:
    n_alpha = (n_electrons + multiplicity - 1) // 2
    n_beta = n_electrons - n_alpha
    alpha = [1.0 if i < n_alpha else 0.0 for i in range(n_mo)]
    beta = [1.0 if i < n_beta else 0.0 for i in range(n_mo)]
    return alpha, beta


[docs] def write_molden( path: os.PathLike | str, molecule: Molecule, basis: BasisSet, result, *, title: str = "", ) -> None: """Write geometry, basis, and molecular orbitals to a Molden-format file. Parameters ---------- path Output file path (``.molden`` is the conventional suffix). molecule The :class:`Molecule` used for the SCF run. basis The :class:`BasisSet` used for the SCF run. result RHFResult / UHFResult / RKSResult / UKSResult (or any object with matching ``mo_coeffs`` / ``mo_energies`` attributes, restricted or split into ``_alpha`` / ``_beta``). title Optional one-line title written to the ``[Title]`` block. Notes ----- The basis must use pure spherical harmonics (vibe-qc's default). MO coefficients are reordered from libint's native per-shell index to molden's ``(m=0, +1, -1, +2, -2, ...)`` convention on the fly. """ path = os.fspath(path) ao_perm = _build_ao_permutation(basis) n_atoms = len(molecule.atoms) n_electrons = molecule.n_electrons() multiplicity = molecule.multiplicity # Dispatch: restricted has `mo_coeffs`; unrestricted has `_alpha` / `_beta`. has_alpha = hasattr(result, "mo_coeffs_alpha") has_restricted = hasattr(result, "mo_coeffs") and not has_alpha if not (has_alpha or has_restricted): raise TypeError( "write_molden: result must expose either `mo_coeffs` " "(restricted) or `mo_coeffs_alpha` + `mo_coeffs_beta` " "(unrestricted); got " + type(result).__name__ ) with open(path, "w", encoding="ascii") as f: _write_header(f, title) _write_atoms(f, molecule) _write_gto(f, basis, n_atoms) _write_pure_flags(f, basis) f.write("[MO]\n") if has_restricted: coeffs = np.asarray(result.mo_coeffs) energies = np.asarray(result.mo_energies) occs = _rhf_occupations(coeffs.shape[1], n_electrons) _write_mo_block(f, energies, coeffs, occs, "Alpha", ao_perm) else: ca = np.asarray(result.mo_coeffs_alpha) cb = np.asarray(result.mo_coeffs_beta) ea = np.asarray(result.mo_energies_alpha) eb = np.asarray(result.mo_energies_beta) occ_a, occ_b = _uhf_occupations( ca.shape[1], multiplicity, n_electrons ) _write_mo_block(f, ea, ca, occ_a, "Alpha", ao_perm) _write_mo_block(f, eb, cb, occ_b, "Beta", ao_perm)