"""QVF (Quantum Visualization Format) writer for vibe-qc. Produces ``{stem}.qvf`` -- a zip archive with a JSON manifest (``manifest.json``) and typed binary or text payloads, one per output section. Each section is keyed by a ``kind`` string from the QVF registry. Consumers support whichever kinds they recognise and silently skip the rest. v1 scope (implemented) ---------------------- * ``structure`` -- atom positions, Z, labels, optional lattice * ``volume.density`` -- total electron density as raw float32 .dat * ``volume.orbital`` -- per-MO wavefunction on a grid * ``atom_properties`` -- Mulliken / Löwdin charges, spin populations * ``trajectory`` -- geometry-optimisation / IRC frames * ``vibrations`` -- normal-mode frequencies + displacements * ``spectra.ir`` -- IR spectrum from Hessian * ``bands`` -- band structure (eigenvalues + k-path) * ``provenance`` -- method, functional, basis, energy, convergence * ``citations`` -- BibTeX references Producer rules -------------- * Float32 default for volumetric arrays; float64 opt-in via ``volume_dtype="float64"``. * Deflate default compression; zstd if ``zipfile-zstd`` is importable. * ``manifest.json`` always stored uncompressed. * Every binary member carries a sha256 hex digest in the manifest. Public API ---------- ``write_qvf(stem, plan, **context) -> Path`` Write ``{stem}.qvf`` from an :class:`OutputPlan` and result data. Returns the written path. ``validate_qvf(path) -> dict`` Open a ``.qvf`` file and validate its manifest + binary payloads. Returns a validation report dict. """ from __future__ import annotations import copy import datetime as _dt import hashlib import io import json import math import os import struct import tempfile import zipfile from pathlib import Path from typing import Any, Optional, Sequence import numpy as np from .._text_safety import safe_json_bytes from ..plan import OutputPlan __all__ = [ "write_qvf", "qvf_bytes", "validate_qvf", "write_reaction_path_qvf", "write_scan_surface_qvf", "QVF_FORMAT_VERSION", "qvf_density_data", "qvf_mo_data", "qvf_wf_data", "qvf_bloch_wf_data", "scf_history_from_result", "qvf_ao_data", "QVF_BLOCH_WF_KIND", ] # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- QVF_FORMAT_VERSION = 1 QVF_FORMAT_VERSION_V2 = 2 _SCHEMA_URI = "https://vibe-qc.org/spec/qvf/1/manifest.schema.json" _SCHEMA_URI_V2 = "https://vibe-qc.org/spec/qvf/2/manifest.schema.json" _PRODUCER_NAME = "vibe-qc" QVF_BLOCH_WF_KIND = "x_vibeqc.bloch_wavefunction" # Bohr -> Ångström (CODATA 2018). Must match vibeqc.output.formats.xyz. _BOHR_TO_ANGSTROM = 0.529177210903 # Hartree -> eV (CODATA 2018). _HARTREE_TO_EV = 27.211386245988 # Single source of truth for archive size caps. The writer's voxel # guard and the validator's per-zip-member zip-bomb guard must agree: # a payload write_qvf will produce must not later be rejected by # validate_qvf as "too large". 1 Gvoxel covers ~8 GiB float64, which # is also the per-member uncompressed cap for binary blobs (large MO # coefficient matrices and trajectory coord arrays count against the # same limit, not just voxel grids). Compressed-on-disk bytes are # typically much smaller; the bomb guard runs against the # uncompressed `file_size` reported by zipfile. _MAX_VOXELS = 1024**3 _MAX_MEMBER_UNCOMPRESSED_BYTES = _MAX_VOXELS * 8 # float64 worst case _REAL_ARTIFACT_IMAG_ABS_TOL = 1.0e-10 _REAL_ARTIFACT_IMAG_REL_TOL = 1.0e-7 def _real_array_for_artifact( data: Any, *, dtype: np.dtype | type | str, label: str, abs_tol: float = _REAL_ARTIFACT_IMAG_ABS_TOL, rel_tol: float = _REAL_ARTIFACT_IMAG_REL_TOL, ) -> np.ndarray: """Return a real ndarray for real-only QVF members. NumPy's ``np.asarray(complex, dtype=float)`` emits ``ComplexWarning`` while discarding the imaginary component implicitly. QVF v1 scalar volumes and Γ-point wavefunction payloads are real-valued by schema, so do the projection here intentionally: tiny imaginary residuals from Bloch/time-reversal folds are accepted; genuinely complex payloads raise and let the optional writer be skipped instead of writing misleading data. """ arr = np.asarray(data) if np.iscomplexobj(arr): if arr.size: max_imag = float(np.max(np.abs(arr.imag))) max_real = float(np.max(np.abs(arr.real))) else: max_imag = 0.0 max_real = 0.0 limit = max(float(abs_tol), float(rel_tol) * max(max_real, 1.0)) if max_imag > limit: raise ValueError( f"{label} is real-only but carries a non-negligible " f"imaginary component (max|Im|={max_imag:.2e}, " f"max|Re|={max_real:.2e}, tolerance={limit:.2e})" ) arr = arr.real return np.asarray(arr, dtype=dtype) def _mo_coefficients_for_artifact( data: Any, *, label: str, abs_tol: float = _REAL_ARTIFACT_IMAG_ABS_TOL, rel_tol: float = _REAL_ARTIFACT_IMAG_REL_TOL, ) -> tuple[np.ndarray, str]: """Return a float64 MO-coefficient payload plus its QVF encoding. Real coefficients are stored as the historical 2-D float64 matrix. Genuine complex Bloch coefficients are stored losslessly as a float64 array with a trailing component axis ``[..., 0] = Re`` and ``[..., 1] = Im``. Tiny imaginary residuals keep the old real encoding. """ arr = np.asarray(data) if np.iscomplexobj(arr): if arr.size: max_imag = float(np.max(np.abs(arr.imag))) max_real = float(np.max(np.abs(arr.real))) else: max_imag = 0.0 max_real = 0.0 limit = max(float(abs_tol), float(rel_tol) * max(max_real, 1.0)) if max_imag > limit: packed = np.empty(arr.shape + (2,), dtype=np.float64) packed[..., 0] = arr.real packed[..., 1] = arr.imag return np.ascontiguousarray(packed), "complex_split_last_axis" arr = arr.real return np.ascontiguousarray(np.asarray(arr, dtype=np.float64)), "real" def _complex_split_array(data: Any) -> np.ndarray: """Pack a complex array as float64 ``[..., real/imag]`` components.""" arr = np.asarray(data, dtype=np.complex128) packed = np.empty(arr.shape + (2,), dtype=np.float64) packed[..., 0] = arr.real packed[..., 1] = arr.imag return np.ascontiguousarray(packed) # Element symbols, Z = 0..118. Mirrors _ELEMENT_SYMBOLS in trajectory.py. _ELEMENT_SYMBOLS = ( "X", # 0 -- placeholder / ghost atom "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", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds", "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og", ) def _symbol(z: int) -> str: if 0 <= z < len(_ELEMENT_SYMBOLS): return _ELEMENT_SYMBOLS[z] return "X" # Canonical section kinds that the writer can emit. This list must # stay in lock-step with the ``Section.oneOf`` branches in # qvf_manifest.schema.json -- tests/test_qvf_round_trip.py carries the # schema-drift guard. ``provenance`` and ``viewer_defaults`` are root # keys, not section kinds; they intentionally don't appear here. _IMPLEMENTED_KINDS = frozenset( { "structure", "volume.density", "volume.orbital", "volume.spin", "volume.elf", "volume.difference", "volume.generic", "volume.potential", "volume.rdg", "fermi_surface", "phonon_bands", "phonon_dos", "equation_of_state", "basis.ao", "wavefunction.gto", "atom_properties", "trajectory", "reaction.path", "reaction.waypoints", "scan.surface", "vibrations", "spectra.ir", "spectra.raman", "spectra.uvvis", "spectra.ecd", "spectra.vcd", "spectra.nmr", "spectra.epr", "spectra.generic", "bands", "structure.symmetry", "bonds", "scf_history", "citations", "dos.total", "dos.projected", "bond_orders", "topology.qtaim", "dos.coop", "dos.cohp", } ) # Kinds reserved in the design doc but not yet implemented in the # writer. The validator accepts them (so a vendor producer can ship # them ahead of the canonical writer) but the writer never emits one. # # `basis` was reserved before `wavefunction.gto` landed -- back then # we anticipated a separate kind that would carry just the AO basis # shells. `wavefunction.gto` now ships basis shells + MO coefficients # in a single section, so a standalone `basis` kind is dead weight # and has been removed from the registry. _RESERVED_KINDS = frozenset( { "volume.orbital_projection", "topology.elf_basins", "projections.lcao", } ) # --------------------------------------------------------------------------- # Grid-evaluation convenience helpers # --------------------------------------------------------------------------- # # These produce pre-packaged data dicts that can be passed directly to # write_qvf() as ``volume_data=`` and ``mo_data=``. They call the # existing grid evaluators in vibeqc.cube and repackage the results. def qvf_density_data( result: Any, basis: Any, molecule: Any, *, spacing: float = 0.25, padding: float = 4.0, label: str = "Electron density", ) -> dict[str, tuple]: """Evaluate total electron density on a uniform grid and return a dict suitable for ``write_qvf(..., volume_data=...)``. Parameters ---------- result Converged SCF result with a ``.density`` attribute (RHF/RKS) or ``.density_alpha`` + ``.density_beta`` (UHF/UKS). basis :class:`BasisSet` used in the calculation. molecule :class:`Molecule` defining the atomic positions. spacing Voxel spacing in bohr (default 0.25). padding Extra headroom around the molecular bounding box in bohr. label Human-readable label for the density section. Returns ------- dict ``{label: (data_3d, origin_3, span_3x3)}`` -- pass as ``volume_data=`` to :func:`write_qvf`. """ from vibeqc.cube import ( CubeGrid, _density_on_grid, make_uniform_grid, ) grid: CubeGrid = make_uniform_grid( molecule, spacing=spacing, padding=padding, ) # Build density matrix. if hasattr(result, "density_alpha"): D = np.asarray(result.density_alpha) + np.asarray(result.density_beta) else: D = np.asarray(result.density) try: from ...properties import _real_if_hermitian except ImportError: from vibeqc.properties import _real_if_hermitian # type: ignore[no-redef] D = _real_if_hermitian(D, what="QVF density matrix") rho = _density_on_grid(D, basis, grid) origin = np.asarray(grid.origin, dtype=np.float64) # Per-voxel step vectors (matches 'voxel_vectors' in the QVF schema). span = np.diag(np.asarray(grid.spacing, dtype=np.float64)) return {label: (rho, origin, span)} def qvf_mo_data( result: Any, basis: Any, molecule: Any, indices: list[int], *, spacing: float = 0.25, padding: float = 4.0, component: str = "real", ) -> list[dict[str, Any]]: """Evaluate MO wavefunctions on a uniform grid and return a list suitable for ``write_qvf(..., mo_data=...)``. Parameters ---------- result Converged SCF result with ``.mo_coefficients`` (RHF/RKS) or ``.mo_coefficients_alpha`` + ``.mo_coefficients_beta`` (UHF/UKS). basis :class:`BasisSet` used in the calculation. molecule :class:`Molecule` defining the atomic positions. indices 0-based MO indices to evaluate. Must be a list of plain ``int`` values; the tuple form returned by :func:`vibeqc.output.formats.cube.requested_mo_indices` (``[(index, name), ...]``) must be unpacked at the call site. spacing Voxel spacing in bohr (default 0.25). padding Extra headroom in bohr. component ``"real"`` (default), ``"imag"``, ``"abs"``, or ``"density"``. Returns ------- list[dict] Each dict has keys ``label``, ``data``, ``origin``, ``span``, ``band_index``, ``energy_eh``, ``occupation``, ``spin``, ``component``. Pass as ``mo_data=`` to :func:`write_qvf`. """ from vibeqc.cube import CubeGrid, _mo_on_grid, make_uniform_grid grid: CubeGrid = make_uniform_grid( molecule, spacing=spacing, padding=padding, ) origin = np.asarray(grid.origin, dtype=np.float64) # Per-voxel step vectors (matches 'voxel_vectors' in the QVF schema). span = np.diag(np.asarray(grid.spacing, dtype=np.float64)) # MO coefficients. if hasattr(result, "mo_coeffs"): C = np.asarray(result.mo_coeffs) spin = 0 elif hasattr(result, "mo_coefficients"): C = np.asarray(result.mo_coefficients) spin = 0 elif hasattr(result, "mo_coeffs_alpha"): C = np.asarray(result.mo_coeffs_alpha) spin = 0 else: raise ValueError( "qvf_mo_data: result has no mo_coeffs or mo_coefficients attribute" ) # MO energies. if hasattr(result, "mo_energies"): energies = np.asarray(result.mo_energies) elif hasattr(result, "mo_energies_alpha"): energies = np.asarray(result.mo_energies_alpha) else: energies = np.zeros(C.shape[1]) # Occupations. n_occ = getattr(result, "n_occ", None) if n_occ is None: n_elec = getattr(molecule, "n_electrons", None) if callable(n_elec): n_elec = n_elec() n_occ = int(n_elec // 2) if n_elec else 0 out: list[dict[str, Any]] = [] for idx in indices: C_col = C[:, idx] mo = _mo_on_grid(C_col, basis, grid) if component == "abs": mo = np.abs(mo) elif component == "density": mo = mo**2 occ = 2.0 if idx < n_occ else 0.0 out.append( { "label": f"MO_{idx}", "data": mo, "origin": origin, "span": span, "band_index": idx, "energy_eh": float(energies[idx]), "occupation": occ, "spin": spin, "component": component, } ) return out def qvf_ao_data( basis: Any, molecule: Any, *, grid: Any = None, spacing: float = 0.15, padding: float = 8.0, ao_indices: list[int] | None = None, primitive_indices: list[tuple[int, int]] | None = None, include_contracted: bool = True, include_primitives: bool = True, basis_label: str | None = None, ) -> list[dict[str, Any]]: """Build per-AO grid data for :func:`write_qvf` ``ao_data=``. Evaluates selected atomic orbitals (contracted shells and/or individual primitives) on a uniform grid and returns a list of dicts ready for the ``basis.ao`` section writer. Parameters ---------- basis :class:`BasisSet` carrying AO shells. molecule :class:`Molecule` defining atom centers. grid Pre-built :class:`CubeGrid`. When ``None``, one is built via :func:`make_uniform_grid` with the given ``spacing`` and ``padding``. spacing Voxel spacing in bohr (default 0.15 -- finer than the 0.25 density default to resolve compact primitives). padding Extra headroom around the bounding box in bohr (default 8.0). ao_indices List of global AO indices for contracted shells to evaluate. When ``None`` and ``primitive_indices`` is also ``None``, evaluates ALL contracted AOs. primitive_indices List of ``(shell_idx, prim_idx)`` tuples for individual primitives to evaluate. Each builds a temp single-primitive basis and evaluates it. include_contracted When ``True`` and ``ao_indices`` is ``None``, include every contracted AO. Ignored when ``ao_indices`` is explicit. include_primitives When ``True`` and ``primitive_indices`` is ``None``, include every primitive in every shell. Ignored when ``primitive_indices`` is explicit. basis_label Human-readable basis-set name (e.g. ``"pob-TZVP"``), written into each section's ``ao_metadata.basis_label``. Returns ------- list[dict] One dict per AO. Each dict has keys: ``label``, ``data``, ``origin``, ``span``, ``ao_metadata``, and ``viewer_hints``. Pass the list as ``ao_data=`` to :func:`write_qvf`. """ from vibeqc.cube import ( CubeGrid, _ao_ranges, evaluate_ao_on_grid, evaluate_single_primitive_on_grid, make_uniform_grid, ) # --- grid ----------------------------------------------------------- if grid is None: grid = make_uniform_grid(molecule, spacing=spacing, padding=padding) origin = np.asarray(grid.origin, dtype=np.float64) span = np.diag(np.asarray(grid.spacing, dtype=np.float64)) # --- shell metadata -------------------------------------------------- shells = list(basis.shells()) ao_ranges = _ao_ranges(basis) name = getattr(basis, "name", None) blabel = basis_label or (str(name) if name and name != "" else None) # Atom symbol lookup. atoms = list(molecule.atoms) _Z_to_sym: dict[int, str] = {} for ai, a in enumerate(atoms): _Z_to_sym[ai] = _symbol(int(a.Z)) results: list[dict[str, Any]] = [] # --- contracted AOs -------------------------------------------------- if include_contracted or ao_indices is not None: idxs: list[int] if ao_indices is not None: idxs = sorted(set(ao_indices)) else: idxs = list(range(int(basis.nbasis))) grids = evaluate_ao_on_grid(basis, grid, idxs) for ao_idx, data_3d in zip(idxs, grids): # Find which shell this AO belongs to. shell_idx = -1 for si, (lo, hi) in enumerate(ao_ranges): if lo <= ao_idx < hi: shell_idx = si break s = shells[shell_idx] ao_local = ao_idx - ao_ranges[shell_idx][0] l = int(s.l) m_val = ao_local - l # spherical: m = -l ... +l atom_idx = int(s.atom_index) sym = _Z_to_sym.get(atom_idx, "?") label = _ao_label(sym, shell_idx, l, m_val, prim_idx=None, contracted=True) section_id = _ao_section_id( sym, shell_idx, l, m_val, prim_idx=None, contracted=True ) meta: dict[str, Any] = { "atom_index": atom_idx, "atom_symbol": sym, "shell_index": shell_idx, "primitive_index": 0, "angular_momentum": [l, m_val], "shell_type": _l_to_shell_type(l), "exponent": float(s.exponents[0]), "coefficient": float(s.coefficients[0]), "is_primitive": False, "is_contracted": True, "ao_index": ao_idx, } if blabel: meta["basis_label"] = blabel results.append( { "label": label, "data": data_3d, "origin": origin, "span": span, "ao_metadata": meta, "section_id": section_id, } ) # --- individual primitives ------------------------------------------- if include_primitives or primitive_indices is not None: pairs: list[tuple[int, int]] if primitive_indices is not None: pairs = sorted(set(primitive_indices)) else: pairs = [ (si, pi) for si, s in enumerate(shells) for pi in range(len(s.exponents)) ] for shell_idx, prim_idx in pairs: s = shells[shell_idx] l = int(s.l) atom_idx = int(s.atom_index) sym = _Z_to_sym.get(atom_idx, "?") # Use m=0 label for primitives (representative). label = _ao_label(sym, shell_idx, l, 0, prim_idx=prim_idx, contracted=False) section_id = _ao_section_id( sym, shell_idx, l, 0, prim_idx=prim_idx, contracted=False ) data_3d = evaluate_single_primitive_on_grid( basis, molecule, shell_idx, prim_idx, grid, ) meta = { "atom_index": atom_idx, "atom_symbol": sym, "shell_index": shell_idx, "primitive_index": prim_idx, "angular_momentum": [l, 0], "shell_type": _l_to_shell_type(l), "exponent": float(s.exponents[prim_idx]), "coefficient": float(s.coefficients[prim_idx]), "is_primitive": True, "is_contracted": False, "ao_index": ao_ranges[shell_idx][0], } if blabel: meta["basis_label"] = blabel results.append( { "label": label, "data": data_3d, "origin": origin, "span": span, "ao_metadata": meta, "section_id": section_id, } ) return results # -- AO label / id helpers ------------------------------------------------- _L_LABELS: dict[int, str] = {0: "s", 1: "p", 2: "d", 3: "f", 4: "g", 5: "h"} def _l_to_shell_type(l: int) -> str: return _L_LABELS.get(l, f"l{l}") def _ao_label( sym: str, shell_idx: int, l: int, m: int, *, prim_idx: int | None = None, contracted: bool = False, ) -> str: """Human-readable AO label, e.g. 'O 2p_z exp=3.50'.""" lt = _L_LABELS.get(l, f"l={l}") if contracted: return f"{sym} {lt} (contracted) sh={shell_idx}" return f"{sym} {lt} sh={shell_idx} p{prim_idx}" def _ao_section_id( sym: str, shell_idx: int, l: int, m: int, *, prim_idx: int | None = None, contracted: bool = False, ) -> str: """Stable section id for an AO, e.g. 'ao_O_p_s2_p0'.""" lt = _L_LABELS.get(l, f"l{l}") if contracted: return f"ao_{sym}_{lt}_s{shell_idx}_contracted" return f"ao_{sym}_{lt}_s{shell_idx}_p{prim_idx}" def _ao_isovalue_default(exponent: float) -> float: """Heuristic isovalue based on Gaussian exponent a. Diffuse functions (a < 0.1) have low peak amplitude -- a high isovalue shows nothing. Tight functions (a >= 10) have high peak amplitude at the nucleus. """ if exponent < 0.1: return 0.005 if exponent < 1.0: return 0.02 if exponent < 10.0: return 0.05 return 0.10 def _primitive_norm(alpha: float, l: int) -> float: """libint2 primitive normalisation -- the factor by which libint scales contraction coefficients so they multiply *un-normalised* primitives. For a normalised Cartesian (lx, ly, lz) Gaussian primitive g(r) = (x-Ax)^lx (y-Ay)^ly (z-Az)^lz exp(-a r^2) the normalisation is: N = (2a/pi)^(3/4) . (4a)^(l/2) / √((2l-1)!!) where l = lx + ly + lz. libint's ``shell.coefficients`` already include this factor (they are "libint-normalised"). The QVF spec Sec. 4.6 requires coefficients that apply to *normalised* primitives, so the writer must divide libint's stored coefficients by this factor on the way out. """ radial = (2.0 * alpha / math.pi) ** 0.75 angular = (4.0 * alpha) ** (l / 2.0) df = 1.0 for k in range(1, 2 * l, 2): df *= k return radial * angular / math.sqrt(df) def _basis_shell_payload(basis: Any) -> tuple[list[dict[str, Any]], bool, int] | None: """Return QVF shell JSON, top-level purity, and AO count for ``basis``.""" try: shells_native = list(basis.shells()) except AttributeError: return None # libint stores per-shell `pure` flags; vibe-qc forces set_pure(true) # on every BasisSet (see symmetry_core.py), so `pure` is uniformly # True at the QVF top level. We still emit per-shell `pure` to # match the design exactly and stay correct if that ever changes. pure_top = all(bool(sh.pure) for sh in shells_native) if shells_native else True shell_list: list[dict[str, Any]] = [] n_ao = 0 for sh in shells_native: l = int(sh.l) shell_pure = bool(sh.pure) shell_list.append( { "center": int(sh.atom_index), "l": l, "exponents": [float(x) for x in sh.exponents], "coefficients": [ float(c) / _primitive_norm(float(a), l) for a, c in zip(sh.exponents, sh.coefficients) ], "pure": shell_pure, } ) n_ao += (2 * l + 1) if shell_pure else ((l + 1) * (l + 2) // 2) return shell_list, pure_top, n_ao def qvf_wf_data( result: Any, basis: Any, molecule: Any, *, structure_ref: str = "structure", orbital_kind: str = "canonical", k_point: "list[float] | None" = None, ) -> dict[str, Any] | None: """Package basis shells + MO coefficients + per-orbital metadata into a dict suitable for ``write_qvf(..., wf_data=...)``. Lets a re-sampling viewer (vibe-view, moltui) evaluate any orbital on a grid of its own choosing -- the Molden-style separation of concerns described in design Sec. 1.5. Parameters ---------- result Converged SCF result. RHF/RKS: ``.mo_coeffs`` + ``.mo_energies``. UHF/UKS: ``.mo_coeffs_alpha`` + ``_beta`` + matching energies. basis :class:`BasisSet` carrying the AO shells used in the calculation. molecule :class:`Molecule` defining atom centers (referenced by ``center`` indices in the basis shells). structure_ref ``id`` of the structure section the shell centers refer to (default ``"structure"`` -- matches what :func:`_write_structure_section` emits). orbital_kind ``"canonical"`` | ``"natural"`` | ``"localized"`` -- written verbatim into the mo_metadata. k_point Optional fractional reciprocal-space coordinate the MO coefficients describe. Pass ``[0.0, 0.0, 0.0]`` for the Γ-point of a periodic system; leave ``None`` for molecular wavefunctions. Recorded in ``mo_metadata["k_point"]`` so the archive is self-describing. Non-Gamma Bloch coefficients are stored as real/imag component pairs. Returns ------- dict | None A wf_data dict, or ``None`` if neither restricted nor unrestricted MO coefficients can be found on ``result`` (e.g. a DFTB result that exposes a different attribute layout). The shape matches what :func:`_write_wavefunction_gto_section` consumes. """ # --- shells ------------------------------------------------------- basis_payload = _basis_shell_payload(basis) if basis_payload is None: return None shell_list, pure_top, n_ao = basis_payload # --- MO coefficients + metadata ----------------------------------- # vibe-qc convention: mo_coeffs has shape [n_ao, n_mo] (columns are # MOs). The QVF format wants [n_mo, n_ao] (rows are MOs) so we # transpose on the way out. def _as_rowmo(arr: np.ndarray) -> tuple[np.ndarray, str]: row = np.asarray(arr).T coeff, encoding = _mo_coefficients_for_artifact( row, label="wavefunction.gto MO coefficients", ) return coeff, encoding # Restricted (RHF / RKS). if hasattr(result, "mo_coeffs") and not hasattr(result, "mo_coeffs_alpha"): C, encoding = _as_rowmo(result.mo_coeffs) n_mo = int(C.shape[0]) energies = ( np.asarray(result.mo_energies, dtype=np.float64).tolist() if hasattr(result, "mo_energies") else [0.0] * n_mo ) # Restricted occupations: 2.0 for the first n_elec/2 MOs. n_elec = molecule.n_electrons() n_doubly = int(n_elec // 2) occupations = [2.0 if i < n_doubly else 0.0 for i in range(n_mo)] mo_metadata: dict[str, Any] = { "n_mo": n_mo, "n_ao": n_ao, "spin": "restricted", "orbital_kind": orbital_kind, "energies": energies, "occupations": occupations, } if k_point is not None: mo_metadata["k_point"] = [float(x) for x in k_point] if encoding != "real": mo_metadata["coefficient_encoding"] = encoding mo_metadata["coefficient_components"] = ["real", "imag"] return { "basis": shell_list, "structure_ref": structure_ref, "pure": pure_top, "n_ao": n_ao, "mo_metadata": mo_metadata, "mo_coefficients": C, } # Unrestricted (UHF / UKS). if hasattr(result, "mo_coeffs_alpha") and hasattr(result, "mo_coeffs_beta"): Ca, enc_alpha = _as_rowmo(result.mo_coeffs_alpha) Cb, enc_beta = _as_rowmo(result.mo_coeffs_beta) n_elec = molecule.n_electrons() mult = int(getattr(molecule, "multiplicity", 1)) n_alpha = (n_elec + mult - 1) // 2 n_beta = (n_elec - mult + 1) // 2 e_alpha = ( np.asarray(result.mo_energies_alpha, dtype=np.float64).tolist() if hasattr(result, "mo_energies_alpha") else [0.0] * int(Ca.shape[0]) ) e_beta = ( np.asarray(result.mo_energies_beta, dtype=np.float64).tolist() if hasattr(result, "mo_energies_beta") else [0.0] * int(Cb.shape[0]) ) mo_metadata = { "n_ao": n_ao, "spin": "unrestricted", "orbital_kind": orbital_kind, "alpha": { "n_mo": int(Ca.shape[0]), "energies": e_alpha, "occupations": [ 1.0 if i < n_alpha else 0.0 for i in range(int(Ca.shape[0])) ], }, "beta": { "n_mo": int(Cb.shape[0]), "energies": e_beta, "occupations": [ 1.0 if i < n_beta else 0.0 for i in range(int(Cb.shape[0])) ], }, } if k_point is not None: mo_metadata["k_point"] = [float(x) for x in k_point] if enc_alpha != "real": mo_metadata["alpha"]["coefficient_encoding"] = enc_alpha mo_metadata["alpha"]["coefficient_components"] = ["real", "imag"] if enc_beta != "real": mo_metadata["beta"]["coefficient_encoding"] = enc_beta mo_metadata["beta"]["coefficient_components"] = ["real", "imag"] return { "basis": shell_list, "structure_ref": structure_ref, "pure": pure_top, "n_ao": n_ao, "mo_metadata": mo_metadata, "mo_coefficients_alpha": Ca, "mo_coefficients_beta": Cb, } return None def _multik_sequence(obj: Any, names: Sequence[str]) -> list[Any] | None: for name in names: value = getattr(obj, name, None) if isinstance(value, (list, tuple)) and len(value) > 0: return list(value) return None def _single_or_multik_coefficients(obj: Any) -> list[Any] | None: coeffs = _multik_sequence(obj, ("mo_coeffs_k", "mo_coeffs")) if coeffs is not None: return coeffs value = getattr(obj, "mo_coeffs", None) if value is not None and not hasattr(obj, "mo_coeffs_alpha"): return [value] return None def _single_or_multik_energies(obj: Any) -> list[Any] | None: energies = _multik_sequence(obj, ("mo_energies_k", "mo_energies")) if energies is not None: return energies value = getattr(obj, "mo_energies", None) if value is not None: return [value] return None def _single_or_multik_occupations( obj: Any, coeffs: Sequence[Any], molecule: Any | None, ) -> list[Any] | None: occs = _multik_sequence(obj, ("occupations_k", "occupations")) if occs is not None: return occs value = getattr(obj, "occupations", None) if value is not None: return [value] if len(coeffs) == 1 and molecule is not None and hasattr(molecule, "n_electrons"): C = np.asarray(coeffs[0]) if C.ndim != 2: return None n_mo = int(C.shape[1]) n_doubly = int(molecule.n_electrons() // 2) occ = np.zeros(n_mo, dtype=np.float64) occ[: min(n_doubly, n_mo)] = 2.0 return [occ] return None def qvf_bloch_wf_data( result: Any, basis: Any, molecule: Any | None = None, *, k_points: Sequence[Sequence[float]], structure_ref: str = "structure", orbital_kind: str = "canonical", ) -> dict[str, Any] | None: """Package all-k closed-shell Bloch orbitals for QVF READ restarts. This is a vibe-qc vendor extension, not the visualization-oriented ``wavefunction.gto`` contract. Coefficients are stored per k-point as MO-major arrays ``[n_mo, n_ao, 2]`` with the last axis equal to ``[real, imag]``. Per-k members are deliberately ragged so restart archives remain valid when canonical orthogonalisation prunes a different number of near-null directions at different k-points. """ basis_payload = _basis_shell_payload(basis) if basis_payload is None: return None shell_list, pure_top, n_ao = basis_payload coeffs = _single_or_multik_coefficients(result) occs = ( _single_or_multik_occupations(result, coeffs, molecule) if coeffs is not None else None ) if coeffs is None or occs is None: return None if len(coeffs) != len(occs): raise ValueError( "qvf_bloch_wf_data: coefficient and occupation k-block counts " f"differ ({len(coeffs)} vs {len(occs)})." ) kpts = np.asarray(k_points, dtype=np.float64).reshape(-1, 3) if kpts.shape[0] != len(coeffs): raise ValueError( "qvf_bloch_wf_data: k-point metadata count does not match " f"coefficient blocks ({kpts.shape[0]} vs {len(coeffs)})." ) energies = _single_or_multik_energies(result) if energies is not None and len(energies) != len(coeffs): raise ValueError( "qvf_bloch_wf_data: energy and coefficient k-block counts " f"differ ({len(energies)} vs {len(coeffs)})." ) coeff_blocks: list[np.ndarray] = [] occ_blocks: list[np.ndarray] = [] energy_blocks: list[np.ndarray] = [] block_meta: list[dict[str, Any]] = [] for ik, (C_raw, occ_raw) in enumerate(zip(coeffs, occs)): C = np.asarray(C_raw, dtype=np.complex128) occ = np.asarray(occ_raw, dtype=np.float64) if C.ndim != 2: raise ValueError( f"qvf_bloch_wf_data: coefficient block {ik} must be 2-D; " f"got shape {C.shape}." ) if C.shape[0] != n_ao: raise ValueError( f"qvf_bloch_wf_data: coefficient block {ik} has {C.shape[0]} " f"AO rows but the basis has {n_ao} functions." ) if occ.ndim != 1: raise ValueError( f"qvf_bloch_wf_data: occupation block {ik} must be 1-D; " f"got shape {occ.shape}." ) if C.shape[1] != occ.shape[0]: raise ValueError( f"qvf_bloch_wf_data: block {ik} has {C.shape[1]} orbitals " f"but {occ.shape[0]} occupations." ) if energies is None: eps = np.zeros(C.shape[1], dtype=np.float64) else: eps = np.asarray(energies[ik], dtype=np.float64) if eps.ndim != 1 or eps.shape[0] != C.shape[1]: raise ValueError( f"qvf_bloch_wf_data: energy block {ik} must have shape " f"({C.shape[1]},); got {eps.shape}." ) coeff_name = f"mo_coefficients_k{ik}" occ_name = f"occupations_k{ik}" energy_name = f"mo_energies_k{ik}" coeff_blocks.append(_complex_split_array(C.T)) occ_blocks.append(np.ascontiguousarray(occ)) energy_blocks.append(np.ascontiguousarray(eps)) block_meta.append( { "k_index": ik, "k_point": [float(x) for x in kpts[ik]], "n_mo": int(C.shape[1]), "n_ao": int(n_ao), "mo_coefficients": coeff_name, "occupations": occ_name, "mo_energies": energy_name, } ) mo_metadata: dict[str, Any] = { "schema_version": "1.0", "restart_kind": "closed_shell_bloch_kpoints", "spin": "restricted", "orbital_kind": orbital_kind, "n_kpoints": len(coeff_blocks), "n_ao": int(n_ao), "k_point_units": "fractional_reciprocal", "coefficient_encoding": "complex_split_last_axis", "coefficient_components": ["real", "imag"], "blocks": block_meta, } return { "basis": shell_list, "structure_ref": structure_ref, "pure": pure_top, "n_ao": n_ao, "kpoints": np.ascontiguousarray(kpts), "mo_metadata": mo_metadata, "mo_coefficients": coeff_blocks, "occupations": occ_blocks, "mo_energies": energy_blocks, } # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _sha256_hex(data: bytes) -> str: """Return the sha256 hex digest of ``data``.""" return hashlib.sha256(data).hexdigest() def _slug(label: str, *, fallback: str = "section") -> str: """Reduce a user-supplied label to a zip-path-safe slug. Keeps ASCII ``[A-Za-z0-9._-]``, replaces every other character with ``_``, collapses runs of ``_``, and trims leading/trailing ``_./-``. Returns ``fallback`` if the result is empty after trimming. Used for the label-derived components of zip paths (volume / orbital / spin / elf / difference) so that a label like ``"r(product) - r(reactant)"`` or a path-traversal attempt like ``"../etc/passwd"`` cannot reach the zip writer as-is. """ if not label: return fallback _SAFE = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-") out_chars: list[str] = [] prev_us = False for ch in label: if ch in _SAFE: out_chars.append(ch) prev_us = False elif not prev_us: out_chars.append("_") prev_us = True slug = "".join(out_chars).strip("._-") return slug or fallback def _require_3d_volume(data: np.ndarray, kind: str, label: str) -> None: """Raise ValueError if ``data`` is not a 3-D array. The volume writers previously skipped malformed sections silently (``if vol.ndim != 3: continue``), which let bad caller-supplied data ship as a QVF missing the section. We'd rather fail the write than produce an archive that quietly drops user data. """ if data.ndim != 3: raise ValueError( f"{kind}[{label!r}]: data must be 3-D (got ndim={data.ndim}, " f"shape={tuple(data.shape)})" ) def _binary_array_entry( path_in_zip: str, data: np.ndarray, ) -> tuple[bytes, dict[str, Any]]: """Serialize ``data`` as raw little-endian bytes and build the matching manifest member-entry dict. Returns ``(raw_bytes, member_spec)``. The caller writes ``raw_bytes`` into the zip and embeds ``member_spec`` in the appropriate section's ``members`` dict. The member spec has keys ``path``, ``format`` ("binary"), ``sha256``, and optionally ``dtype``, ``shape``. """ if not data.flags["C_CONTIGUOUS"]: data = np.ascontiguousarray(data) raw = data.tobytes() dtype_name = np.dtype(data.dtype).name member = { "path": path_in_zip, "format": "binary", "dtype": dtype_name, "shape": list(data.shape), "sha256": _sha256_hex(raw), } return raw, member def _write_binary_to_zip( zf: zipfile.ZipFile, path_in_zip: str, data: np.ndarray, ) -> dict[str, Any]: """Convenience: serialise ``data``, write it into ``zf``, and return the manifest member-spec dict.""" raw, member = _binary_array_entry(path_in_zip, data) zf.writestr(path_in_zip, raw) return member def _now_iso() -> str: return _dt.datetime.now().astimezone().isoformat(timespec="seconds") def _is_vendor_kind(kind: str) -> bool: return kind.startswith("x_") # --------------------------------------------------------------------------- # Main writer # --------------------------------------------------------------------------- def write_qvf( stem: os.PathLike | str, plan: OutputPlan, *, compression: int = zipfile.ZIP_DEFLATED, volume_dtype: str = "float32", atomic: bool = False, **context: Any, ) -> Path: """Write ``{stem}.qvf``. Parameters ---------- stem Path stem; ``.qvf`` suffix is appended. plan :class:`OutputPlan` declaring what artefacts are expected. compression ``zipfile.ZIP_DEFLATED`` (default), ``zipfile.ZIP_STORED``, or the zstd constant if ``zipfile-zstd`` is importable. volume_dtype ``"float32"`` (default) or ``"float64"`` for volumetric grids. atomic When ``True``, write the archive to a temporary sibling file and ``os.replace`` it onto ``{stem}.qvf`` only after write-time validation passes. A concurrent reader then sees either the previous complete archive or the new complete archive -- never a half-written zip. Used by the live-checkpoint path so vibe-view can hot-reload a running job's QVF safely. Default ``False`` preserves the historical in-place write. **context Data objects the section writers need. Typical keys: * ``run_status`` -- ``"running"`` | ``"converged"`` | ``"failed"``; emitted at ``provenance.run_status``. Lets a live consumer tell a mid-run checkpoint from the settled final archive. * ``checkpoint`` -- dict with ``seq`` (monotonic int), ``wall_time_s`` (float), ``written_at`` (ISO-8601 str), and optional ``scf_iteration`` / ``energy_eh`` running hints; emitted at ``provenance.checkpoint``. * ``partial_sections`` -- ``True`` to mark every section ``partial: true``, or an iterable of section ids / ``kind`` strings (e.g. ``{"trajectory"}``) to flag only the still-growing ones in a checkpoint snapshot. * ``biomolecule_data`` -- optional dict adding biomolecule metadata to the ``structure`` section for ribbon/cartoon rendering: ``chains`` (list[str]), ``residues`` (list of ``{name, seq, chain, atom_indices}``, 0-based atom indices), and ``secondary_structure`` (list of ``{type, chain, start_seq, end_seq}``). Additive peer keys on the section object; all three independently optional. See :func:`_normalize_biomolecule`. * ``molecule`` / ``system`` -- :class:`Molecule` or :class:`PeriodicSystem` * ``structure_lattice_bohr`` -- optional 3x3 lattice override for the ``structure`` section, using the same bohr column-vector convention as ``PeriodicSystem.lattice``. Periodic CCM callers use this to emit the full BvK torus cell while keeping the primitive input ``system`` intact. * ``result`` -- converged SCF result object * ``basis`` -- :class:`BasisSet` * ``population_summary`` -- :class:`PopulationSummary` * ``hessian_result`` -- :class:`HessianResult` * ``band_structure`` -- :class:`BandStructure` * ``trajectory_frames`` -- list of :class:`Molecule` * ``trajectory_energies`` -- list of float (Hartree) * ``trajectory_rms_grad`` -- list of float (optional) * ``bibtex_content`` -- str, the full BibTeX file body * ``bond_orders_data`` -- dict with keys ``method`` (str, e.g. ``"mayer"``) and ``pairs`` (list of dicts each with ``i``, ``j``, ``order`` and optional ``distance_ang``, ``symbol_i``, ``symbol_j``). Emitted as a ``bond_orders`` section. * ``volume_data`` -- dict of ``{label: (data_3d, origin, span)}`` * ``mo_data`` -- list of dicts with keys ``label``, ``data``, ``origin``, ``span``, ``band_index``, ``energy_eh``, ``occupation``, ``spin``, ``component`` * ``spin_data`` -- dict of ``{label: (data_3d, origin, span)}`` * ``elf_data`` -- dict of ``{label: (data_3d, origin, span)}`` * ``generic_volume_data`` -- dict of ``{label: (data_3d, origin, span)}`` for ``volume.generic`` (escape hatch for any scalar field that doesn't fit density/orbital/spin/elf/difference) * ``potential_data`` -- dict of ``{label: (data_3d, origin, span)}`` for ``volume.potential`` (electrostatic potential grid; same member structure as ``volume.density``, QVF spec Sec. 4.10) * ``rdg_data`` -- dict of ``{label: (data_3d, origin, span)}`` for ``volume.rdg`` (reduced density gradient for NCI analysis; same member structure, QVF spec Sec. 4.11) * ``diff_data`` -- dict of ``{label: spec}`` for difference density (e.g. r(product) - r(reactant)). ``spec`` is either a 3-tuple ``(data_3d, origin, span)`` for an unannotated difference, or a dict with keys ``data``, ``origin``, ``span``, and optionally ``operand_a`` (str, section id of minuend), ``operand_b`` (str, section id of subtrahend), ``description``. * ``reaction_path`` -- dict ``{frames, waypoints, energies?, reaction_coordinate?}`` for a self-contained ``reaction.path`` section. ``waypoints`` is a list of ``{frame_index, label, kind, energy_eh?}`` records where ``kind`` is one of ``"reactant" | "transition_state" | "intermediate" | "product" | "point"``. * ``reaction_waypoints`` -- dict ``{trajectory_ref, waypoints, reaction_coordinate?}`` for a lightweight ``reaction.waypoints`` annotation over an already-emitted ``trajectory`` section. ``trajectory_ref`` must name a trajectory section emitted in the same archive; the writer raises if it doesn't resolve. * ``viewer_defaults`` -- dict written verbatim to the manifest root. Recognised keys: ``auto_open`` (list of section ids), per-section render hints, and ``bookmarks`` (ordered list of ``{name, camera}`` records using the VTK camera model). * ``thermochemistry_data`` -- dict with keys ``zpve_eh``, ``enthalpy_eh``, ``entropy_cal_mol_k``, ``gibbs_free_energy_eh``, ``temperature_k``, ``pressure_atm`` for a root ``thermochemistry`` field (QVF spec Sec. 4.7). * ``dipole_moment_data`` -- dict with keys ``total_debye``, ``vector_debye`` (3-vector), ``origin`` (str) for a root ``dipole_moment`` field (QVF spec Sec. 4.7). * ``constraints_data`` -- dict with keys ``frozen_atoms`` (list of int), ``distance_constraints`` (list of ``{atoms, target_angstrom}``) for a root ``constraints`` field (QVF spec Sec. 4.7). * ``extensions`` -- dict of ``{vendor_ns: {version, schema_uri?, critical?}}`` for the root ``extensions`` governance block (QVF spec Sec. 5.4). * ``vendor_json_sections`` -- list of first- or third-party vendor JSON sections. Each entry is ``{id, kind, payload}`` with optional ``member`` (default ``"data"``), ``label``, and ``critical``. The kind must live in the ``x_.*`` namespace. * ``eos_data`` -- dict with keys ``volumes`` (float64 [n_points]), ``energies`` (float64 [n_points]), ``fit`` (dict with ``model``, ``V0``, ``E0``, ``B0``, ``B0_prime``, etc.) for an ``equation_of_state`` section (QVF spec Sec. 4.14). * ``fermi_surface_data`` -- dict with keys ``nk1``, ``nk2``, ``nk3`` (int), ``energies`` (float64 [nk1, nk2, nk3, n_bands]), ``band_indices`` (list of int), ``lattice_vectors`` (3x3), ``fermi_energy_ev`` (float), and optional ``n_spin`` (int, default 1) for a ``fermi_surface`` section (QVF spec Sec. 4.12). * ``wf_data`` -- dict with keys ``basis`` (list of shell dicts), ``mo_metadata`` (dict), ``mo_coefficients`` (2D || `[n_mo, n_ao]`), and optionally ``mo_coefficients_alpha`` / ``mo_coefficients_beta`` for unrestricted. Emitted as ``wavefunction.gto`` with id ``"wf"``. * ``bloch_wf_data`` -- dict from :func:`qvf_bloch_wf_data` carrying all-k closed-shell Bloch coefficients, occupations, and k-point metadata for QVF-backed periodic READ restarts. Emitted as the first-party vendor section ``x_vibeqc.bloch_wavefunction``. * ``wf_localized_data`` -- same shape as ``wf_data`` but emitted as a second ``wavefunction.gto`` section with id ``"wf_localized"``. ``mo_metadata["orbital_kind"]`` should be ``"localized"``. * ``wf_nto_hole_data`` -- same shape as ``wf_data``, emitted as ``wavefunction.gto`` with id ``"wf_nto_hole"`` and ``orbital_kind="natural"``. The "hole" side of a single NTO pair (simplest post-TD-DFT path). * ``wf_nto_electron_data`` -- same shape as ``wf_data``, emitted as ``wavefunction.gto`` with id ``"wf_nto_electron"`` and ``orbital_kind="natural"``. The "electron" side of a single NTO pair. * ``nto_data`` -- list of ``{"hole": wf_dict, "electron": wf_dict, "state_index": int, "excitation_energy_ev": float}`` records, one per excited state. Each emits two ``wavefunction.gto`` sections with ``orbital_kind="natural"`` and ids ``"wf_nto_S{n}_hole"`` / ``"wf_nto_S{n}_electron"``. Intended for Natural Transition Orbitals (post-TD-DFT). * ``ao_data`` -- list of dicts from :func:`qvf_ao_data`, each with keys ``label``, ``data`` (3-D array), ``origin``, ``span``, ``ao_metadata``, ``section_id``. Emitted as ``basis.ao`` sections. * ``coop_data`` -- dict with keys ``energies`` (float64 [n_points] in eV), ``projections`` (float64 [n_pairs, n_points] or [n_spin, n_pairs, n_points]), ``integrated`` (float64 [n_pairs]), ``energies_units``, ``n_spin``, ``fermi_energy_ev``, ``sigma_ev``, ``pairs`` for a ``dos.coop`` section (QVF spec Sec. 4.8b). * ``cohp_data`` -- dict with keys ``energies``, ``projections``, ``integrated``, ``energies_units``, ``n_spin``, ``fermi_energy_ev``, ``sigma_ev``, ``pairs`` for a ``dos.cohp`` section (QVF spec Sec. 4.8c). Returns ------- pathlib.Path The on-disk ``{stem}.qvf`` path. """ stem = Path(os.fspath(stem)) target = stem.with_suffix(".qvf") target.parent.mkdir(parents=True, exist_ok=True) # --- guardrails ----------------------------------------------------- if not isinstance(plan, OutputPlan): raise TypeError( f"write_qvf: 'plan' must be an OutputPlan, got {type(plan).__name__}" ) mol_or_sys = context.get("molecule") or context.get("system") # Volume size check. vol_data = context.get("volume_data") if vol_data: for label, (data, _o, _s) in vol_data.items(): nv = int(np.prod(data.shape)) if hasattr(data, "shape") else 0 if nv > _MAX_VOXELS: raise ValueError( f"volume_data[{label!r}]: {nv:_d} voxels exceeds " f"max {_MAX_VOXELS:_d}. Reduce grid or use a separate " f".cube/.xsf file." ) mo_data = context.get("mo_data") if mo_data: for mo in mo_data: data = mo.get("data") if data is not None and hasattr(data, "shape"): nv = int(np.prod(data.shape)) if nv > _MAX_VOXELS: raise ValueError( f"mo_data[{mo.get('label', '?')!r}]: {nv:_d} voxels " f"exceeds max {_MAX_VOXELS:_d}." ) # Warn (don't crash) on missing structure -- a QVF with no structure # section is valid but unusual. if mol_or_sys is None: import warnings warnings.warn( "write_qvf: no 'molecule' or 'system' in context -- the QVF " "will have no structure section. Pass molecule= or " "system= for a complete archive.", UserWarning, stacklevel=2, ) # Resolve compression. Default deflate; try zstd if available. _compression = compression if compression == zipfile.ZIP_DEFLATED: try: from zipfile_zstd import ZIP_ZSTANDARD # noqa: F811 _compression = ZIP_ZSTANDARD except ImportError: pass volume_dt = np.dtype(volume_dtype) # --- build manifest skeleton ---------------------------------------- _version = _resolve_version() manifest: dict[str, Any] = { "qvf_version": QVF_FORMAT_VERSION, "schema_uri": _SCHEMA_URI, "source": { "program": _PRODUCER_NAME, "version": _version, "calculation": ( f"{context.get('method', '?')}/{context.get('basis', '?')}" ), }, "sections": [], } sections: list[dict[str, Any]] = manifest["sections"] # --- provenance (manifest root) -------------------------------------- manifest["provenance"] = _build_provenance(context) # --- viewer_defaults (manifest root, optional) ----------------------- vd = context.get("viewer_defaults") if vd is not None: manifest["viewer_defaults"] = dict(vd) # --- thermochemistry (manifest root, optional) ----------------------- thermo = context.get("thermochemistry_data") if thermo is not None: manifest["thermochemistry"] = dict(thermo) # --- dipole_moment (manifest root, optional) ------------------------- dipole = context.get("dipole_moment_data") if dipole is not None: manifest["dipole_moment"] = dict(dipole) # --- constraints (manifest root, optional) --------------------------- constraints = context.get("constraints_data") if constraints is not None: manifest["constraints"] = dict(constraints) # --- extensions (manifest root, optional) ---------------------------- ext_block = context.get("extensions") if ext_block is not None: manifest["extensions"] = dict(ext_block) # Atomic mode writes to a temp sibling on the same filesystem and # ``os.replace``s it onto ``target`` only after validation passes, so # a concurrent reader never observes a partially-written zip. The # historical (default) path writes ``target`` in place. if atomic: fd, _tmp_name = tempfile.mkstemp( dir=str(target.parent), prefix=f".{target.name}.", suffix=".tmp", ) os.close(fd) write_target = Path(_tmp_name) else: write_target = target try: with zipfile.ZipFile(write_target, "w", _compression) as zf: _emit_qvf_into_zip( zf, manifest=manifest, sections=sections, context=context, mol_or_sys=mol_or_sys, volume_dt=volume_dt, ) # --- write-time validation gate --------------------------------- # Validate the freshly-written archive against the canonical # schema before returning. This is the producer-side enforcement # of the SSOT: write_qvf() never returns a path to an invalid # archive. Skipping the gate would let writer regressions ship # to consumers -- the exact failure mode the SSOT work fixes. report = validate_qvf(write_target) if not report["valid"]: raise ValueError( f"write_qvf produced an archive that fails canonical " f"validation -- this is a writer bug. Errors:\n - " + "\n - ".join(report["errors"][:8]) ) except BaseException: # Wipe the bad/partial output so a caller cannot mistake a stale # invalid file for a successful write. In atomic mode this only # touches the temp file, leaving any prior ``target`` intact. try: write_target.unlink() except OSError: pass raise if atomic: # POSIX ``os.replace`` is atomic within a filesystem; the temp # file and target share ``target.parent`` so the rename cannot # cross a filesystem boundary. os.replace(write_target, target) return target def _emit_qvf_into_zip( zf: zipfile.ZipFile, *, manifest: dict[str, Any], sections: list[dict[str, Any]], context: dict[str, Any], mol_or_sys: Any, volume_dt: np.dtype, ) -> None: """Write every QVF section into an open zipfile + finalize the manifest. Factored out of :func:`write_qvf` so the in-memory helper :func:`qvf_bytes` shares the same emission pipeline. """ # --- structure ---------------------------------------------- if mol_or_sys is not None: is_periodic = bool( context.get("system") is not None and context.get("molecule") is None ) _write_structure_section( zf, mol_or_sys, sections, periodic=is_periodic, biomolecule_data=context.get("biomolecule_data"), structure_lattice_bohr=context.get("structure_lattice_bohr"), ) # --- volume.density ----------------------------------------- vol_data = context.get("volume_data") if vol_data: _write_volume_density_section(zf, vol_data, sections, volume_dtype=volume_dt) # --- volume.orbital ----------------------------------------- mo_data = context.get("mo_data") if mo_data: _write_volume_orbital_section(zf, mo_data, sections, volume_dtype=volume_dt) # --- volume.spin -------------------------------------------- spin_data = context.get("spin_data") if spin_data: _write_volume_spin_section(zf, spin_data, sections, volume_dtype=volume_dt) # --- volume.elf --------------------------------------------- elf_data = context.get("elf_data") if elf_data: _write_volume_elf_section(zf, elf_data, sections, volume_dtype=volume_dt) # --- volume.difference -------------------------------------- diff_data = context.get("diff_data") if diff_data: _write_volume_difference_section( zf, diff_data, sections, volume_dtype=volume_dt ) # --- volume.generic ----------------------------------------- gen_vol_data = context.get("generic_volume_data") if gen_vol_data: _write_volume_generic_section( zf, gen_vol_data, sections, volume_dtype=volume_dt ) # --- volume.potential -------------------------------------- potential_data = context.get("potential_data") if potential_data: _write_volume_potential_section( zf, potential_data, sections, volume_dtype=volume_dt ) # --- volume.rdg ------------------------------------------- rdg_data = context.get("rdg_data") if rdg_data: _write_volume_rdg_section(zf, rdg_data, sections, volume_dtype=volume_dt) # --- wavefunction.gto (canonical) --------------------------- wf_data = context.get("wf_data") if wf_data: _write_wavefunction_gto_section(zf, wf_data, sections) # --- x_vibeqc.bloch_wavefunction (all-k READ restart) ------ bloch_wf_data = context.get("bloch_wf_data") if bloch_wf_data: _write_bloch_wavefunction_section(zf, bloch_wf_data, sections) # --- wavefunction.gto (localized) ---------------------------- wf_localized_data = context.get("wf_localized_data") if wf_localized_data: _write_wavefunction_gto_section( zf, wf_localized_data, sections, section_id="wf_localized" ) # --- wavefunction.gto (NTO hole / electron pairs) -------------- # Standalone keys: a single hole/electron pair (simplest path). # Emitted as sections "wf_nto_hole" and "wf_nto_electron". # The mo_metadata["orbital_kind"] should be "natural". wf_nto_hole_data = context.get("wf_nto_hole_data") if wf_nto_hole_data: wf_nto_hole_data.setdefault("mo_metadata", {}) wf_nto_hole_data["mo_metadata"].setdefault("orbital_kind", "natural") _write_wavefunction_gto_section( zf, wf_nto_hole_data, sections, section_id="wf_nto_hole" ) wf_nto_electron_data = context.get("wf_nto_electron_data") if wf_nto_electron_data: wf_nto_electron_data.setdefault("mo_metadata", {}) wf_nto_electron_data["mo_metadata"].setdefault("orbital_kind", "natural") _write_wavefunction_gto_section( zf, wf_nto_electron_data, sections, section_id="wf_nto_electron" ) # --- wavefunction.gto (NTO hole / electron pairs, multi-state) -- # Each entry is {"hole": wf_data_dict, "electron": wf_data_dict, # "state_index": int, "excitation_energy_ev": float, # optional "spin": "alpha"|"beta"}. # Emitted as two sections per state: wf_nto_S{n}_hole and # wf_nto_S{n}_electron (or wf_nto_S{n}_{spin}_hole etc. for UHF). # The mo_metadata["orbital_kind"] should be "natural". nto_states = context.get("nto_data") if nto_states: for state in nto_states: si = state.get("state_index", 0) e_ev = state.get("excitation_energy_ev") _spin = state.get("spin", "") _spin_tag = f"_{_spin}" if _spin else "" for role in ("hole", "electron"): wf = state.get(role) if wf is None: continue # Ensure orbital_kind is set. wf.setdefault("mo_metadata", {}) wf["mo_metadata"].setdefault("orbital_kind", "natural") if e_ev is not None: wf["mo_metadata"].setdefault("excitation_energy_ev", float(e_ev)) wf["mo_metadata"].setdefault("excitation_index", int(si)) sec_id = f"wf_nto_S{si}{_spin_tag}_{role}" _write_wavefunction_gto_section(zf, wf, sections, section_id=sec_id) # --- basis.ao ----------------------------------------------- ao_data = context.get("ao_data") if ao_data: _write_basis_ao_section(zf, ao_data, sections, volume_dtype=volume_dt) # --- atom_properties ---------------------------------------- pop = context.get("population_summary") if pop is not None: _write_atom_properties_section(zf, pop, sections) # --- trajectory --------------------------------------------- traj_frames = context.get("trajectory_frames") if traj_frames: _write_trajectory_section( zf, traj_frames, sections, energies=context.get("trajectory_energies"), rms_grad=context.get("trajectory_rms_grad"), trajectory_type=context.get( "trajectory_type", "geometry_optimization", ), ) # --- reaction.path (self-contained) ------------------------- rxn_path = context.get("reaction_path") if rxn_path: # A reaction path whose frames are PeriodicSystem instances carries the # per-frame `lattice` member + `dim` metadata that vibe-view needs to # render the cell and wrap atoms across periodic boundaries. Both are # *optional* additions, so the archive stays qvf_version=1: a consumer # detects a periodic path from the presence of the `lattice` member, # never from the manifest version. # # vibe-qc v0.10.0-v0.15.x bumped to qvf_version=2 here. That bump was # withdrawn by the governance ruling of 2026-07-10 (an optional member # is the additive case and must not bump the version, and the bump # obliged conforming v1-only consumers to refuse archives we routinely # wrote). See qvf-writer/GOVERNANCE.md, "Version history". _write_reaction_path_section( zf, rxn_path["frames"], rxn_path["waypoints"], sections, energies=rxn_path.get("energies"), reaction_coordinate=rxn_path.get("reaction_coordinate"), reaction_coordinate_label=rxn_path.get("reaction_coordinate_label"), reaction_coordinate_unit=rxn_path.get("reaction_coordinate_unit"), frame_volumes=rxn_path.get("frame_volumes"), volume_grid=rxn_path.get("volume_grid"), volume_frame_index=rxn_path.get("volume_frame_index"), volume_label=rxn_path.get("volume_label", "Electron density"), volume_isovalue=rxn_path.get("volume_isovalue"), ) # --- reaction.waypoints (annotation over a trajectory) ------ rxn_wps = context.get("reaction_waypoints") if rxn_wps: traj_ref = rxn_wps["trajectory_ref"] traj_section = next( (s for s in sections if s.get("id") == traj_ref), None, ) if traj_section is None or traj_section.get("kind") != "trajectory": raise ValueError( f"reaction.waypoints: trajectory_ref={traj_ref!r} does " "not name a trajectory section emitted in this archive. " "Producers MUST emit the referenced trajectory first." ) n_traj_frames = int(traj_section["members"]["coords"]["shape"][0]) _write_reaction_waypoints_section( zf, traj_ref, rxn_wps["waypoints"], n_traj_frames, sections, reaction_coordinate=rxn_wps.get("reaction_coordinate"), ) # --- scan.surface (2D relaxed-scan energy grid) ------------- scan_surf = context.get("scan_surface") if scan_surf: _write_scan_surface_section( zf, sections, axis_a=scan_surf["axis_a"], axis_b=scan_surf["axis_b"], energies=scan_surf["energies"], coordinate_a_label=scan_surf["coordinate_a_label"], coordinate_a_unit=scan_surf["coordinate_a_unit"], coordinate_b_label=scan_surf["coordinate_b_label"], coordinate_b_unit=scan_surf["coordinate_b_unit"], geometries=scan_surf.get("geometries"), atoms=scan_surf.get("atoms"), ) # --- vibrations --------------------------------------------- hess = context.get("hessian_result") if hess is not None: mol = context.get("molecule") syms = None if mol is not None: syms = [_symbol(int(a.Z)) for a in mol.atoms] _write_vibrations_section(zf, hess, sections, atom_symbols=syms, molecule=mol) # --- spectra.ir --------------------------------------------- if hess is not None: _write_spectra_ir_section(zf, hess, sections) # --- bands -------------------------------------------------- bs = context.get("band_structure") if bs is not None: _write_bands_section(zf, bs, sections) # --- citations ---------------------------------------------- bib = context.get("bibtex_content") if bib: _write_citations_section(zf, bib, sections) # --- dos.total ---------------------------------------------- dos_data = context.get("dos_data") if dos_data: _write_dos_total_section(zf, dos_data, sections) # --- dos.projected ------------------------------------------ pdos_data = context.get("pdos_data") if pdos_data: _write_dos_projected_section(zf, pdos_data, sections) # --- dos.coop ---------------------------------------------- coop_data = context.get("coop_data") if coop_data: _write_dos_coop_section(zf, coop_data, sections) # --- dos.cohp ---------------------------------------------- cohp_data = context.get("cohp_data") if cohp_data: _write_dos_cohp_section(zf, cohp_data, sections) # --- spectra.raman ------------------------------------------ raman = context.get("raman_data") if raman: _write_spectra_raman_section(zf, raman, sections) # --- spectra.uvvis ------------------------------------------ uvvis = context.get("uvvis_data") if uvvis: _write_spectra_uvvis_section(zf, uvvis, sections) # --- spectra.ecd -------------------------------------------- ecd = context.get("ecd_data") if ecd: _write_spectra_ecd_section(zf, ecd, sections) # --- spectra.vcd -------------------------------------------- vcd = context.get("vcd_data") if vcd: _write_spectra_vcd_section(zf, vcd, sections) # --- spectra.nmr -------------------------------------------- nmr = context.get("nmr_data") if nmr: _write_spectra_nmr_section(zf, nmr, sections) # --- spectra.epr -------------------------------------------- epr = context.get("epr_data") if epr: _write_spectra_epr_section(zf, epr, sections) # --- spectra.generic ---------------------------------------- generic_spec = context.get("generic_spectrum_data") if generic_spec: _write_spectra_generic_section(zf, generic_spec, sections) # --- structure.symmetry ------------------------------------- sym = context.get("symmetry_data") if sym: _write_symmetry_section(zf, sym, sections) # --- bonds -------------------------------------------------- bonds = context.get("bonds_data") if bonds: _write_bonds_section(zf, bonds, sections) # --- bond_orders --------------------------------------------- bo_data = context.get("bond_orders_data") if bo_data: _write_bond_orders_section(zf, bo_data["method"], bo_data["pairs"], sections) # --- scf_history -------------------------------------------- scf_hist = context.get("scf_history_data") if scf_hist: _write_scf_history_section(zf, scf_hist, sections) # --- equation_of_state -------------------------------------- eos_data = context.get("eos_data") if eos_data: _write_equation_of_state_section(zf, eos_data, sections) # --- fermi_surface ----------------------------------------- fermi_data = context.get("fermi_surface_data") if fermi_data: _write_fermi_surface_section(zf, fermi_data, sections) # --- phonon_bands ------------------------------------------ phonon_bands_data = context.get("phonon_bands_data") if phonon_bands_data: _write_phonon_bands_section(zf, phonon_bands_data, sections) # --- phonon_dos -------------------------------------------- phonon_dos_data = context.get("phonon_dos_data") if phonon_dos_data: _write_phonon_dos_section(zf, phonon_dos_data, sections) # --- topology.qtaim ----------------------------------------- qtaim = context.get("qtaim_data") if qtaim: _write_topology_qtaim_section( zf, qtaim["critical_points"], bond_paths=qtaim.get("bond_paths"), sections=sections, ) # --- x_ JSON sections ----------------------------------- vendor_json_sections = context.get("vendor_json_sections") if vendor_json_sections: _write_vendor_json_sections(zf, vendor_json_sections, sections) # --- partial-section marking (streaming checkpoints) -------- # A checkpoint snapshot flags every still-growing section (e.g. an # optimization ``trajectory`` or an SCF-history section not yet at # its final point) with ``partial: true`` so a live consumer knows # not to treat it as final. Sections carry no ``additionalProperties: # false`` in the schema, so the extra key validates under v1. _mark_partial_sections(sections, context.get("partial_sections")) # --- manifest.json (ZIP_STORED, always last by convention) --- manifest_bytes = safe_json_bytes(manifest, indent=2) zf.writestr( zipfile.ZipInfo("manifest.json"), manifest_bytes, compress_type=zipfile.ZIP_STORED, ) def _mark_partial_sections( sections: list[dict[str, Any]], partial_sections: Any, ) -> None: """Set ``partial=True`` on emitted sections named by ``partial_sections``. ``partial_sections`` may be ``None`` (no-op), ``True`` (mark every section partial -- the whole snapshot is mid-flight), or an iterable of section ids and/or ``kind`` strings. Matching is by section ``id`` first, then ``kind``, so a caller can say ``{"trajectory"}`` to flag every trajectory section without knowing the generated ids. """ if not partial_sections: return if partial_sections is True: for sec in sections: sec["partial"] = True return wanted = {str(x) for x in partial_sections} for sec in sections: if sec.get("id") in wanted or sec.get("kind") in wanted: sec["partial"] = True def _write_vendor_json_sections( zf: zipfile.ZipFile, vendor_sections: Sequence[dict[str, Any]], sections: list[dict[str, Any]], ) -> None: """Write caller-supplied JSON sections in an ``x_`` namespace.""" for idx, spec in enumerate(vendor_sections): kind = str(spec.get("kind", "")) if not _is_vendor_kind(kind): raise ValueError( "vendor_json_sections entries must use an x_ kind; " f"got {kind!r}" ) if "payload" not in spec: raise ValueError( "vendor_json_sections entries require a JSON 'payload' value" ) section_id = str(spec.get("id") or f"vendor_json_{idx}") member_name = str(spec.get("member", "data")) section_slug = _slug(section_id, fallback=f"vendor_json_{idx}") member_slug = _slug(member_name, fallback="data") path_in_zip = str( spec.get("path") or f"vendor/{section_slug}/{member_slug}.json" ) payload_json = safe_json_bytes(spec["payload"], indent=2) zf.writestr(path_in_zip, payload_json) section: dict[str, Any] = { "id": section_id, "kind": kind, "members": { member_name: { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(payload_json), }, }, } if "label" in spec: section["label"] = str(spec["label"]) if "critical" in spec: section["critical"] = bool(spec["critical"]) sections.append(section) def qvf_bytes( plan: "OutputPlan", # noqa: F821 -- same forward ref as write_qvf *, compression: int = zipfile.ZIP_DEFLATED, volume_dtype: str = "float32", **context: Any, ) -> bytes: """Build a QVF archive in memory and return the raw zip bytes. Mirrors :func:`write_qvf` -- same ``plan`` and ``**context`` surface, same emission pipeline, same canonical post-build validation gate -- but never touches the filesystem. Use this when vibe-qc wants to hand a freshly built QVF directly to ``vibe-view`` (see :func:`vibeview.launch_qvf`) without a temporary file. Returns ------- bytes The complete .qvf archive bytes. ``QVFReader()`` opens them directly. Raises ------ ValueError If the in-memory archive fails the SSOT validation gate. The bytes are *not* returned in that case -- same behaviour as :func:`write_qvf` unlinking the on-disk artefact. """ if not isinstance(plan, OutputPlan): raise TypeError( f"qvf_bytes: 'plan' must be an OutputPlan, got {type(plan).__name__}" ) mol_or_sys = context.get("molecule") or context.get("system") # Shared voxel size guard (single source = _MAX_VOXELS module const). vol_data = context.get("volume_data") if vol_data: for label, (data, _o, _s) in vol_data.items(): nv = int(np.prod(data.shape)) if hasattr(data, "shape") else 0 if nv > _MAX_VOXELS: raise ValueError( f"volume_data[{label!r}]: {nv:_d} voxels exceeds " f"max {_MAX_VOXELS:_d}." ) mo_data = context.get("mo_data") if mo_data: for mo in mo_data: data = mo.get("data") if data is not None and hasattr(data, "shape"): nv = int(np.prod(data.shape)) if nv > _MAX_VOXELS: raise ValueError( f"mo_data[{mo.get('label', '?')!r}]: {nv:_d} voxels " f"exceeds max {_MAX_VOXELS:_d}." ) if mol_or_sys is None: import warnings warnings.warn( "qvf_bytes: no 'molecule' or 'system' in context -- the QVF " "will have no structure section.", UserWarning, stacklevel=2, ) volume_dt = np.dtype(volume_dtype) _version = _resolve_version() manifest: dict[str, Any] = { "qvf_version": QVF_FORMAT_VERSION, "schema_uri": _SCHEMA_URI, "source": { "program": _PRODUCER_NAME, "version": _version, "calculation": ( f"{context.get('method', '?')}/{context.get('basis', '?')}" ), }, "sections": [], } sections: list[dict[str, Any]] = manifest["sections"] manifest["provenance"] = _build_provenance(context) vd = context.get("viewer_defaults") if vd is not None: manifest["viewer_defaults"] = dict(vd) buf = io.BytesIO() with zipfile.ZipFile(buf, "w", compression) as zf: _emit_qvf_into_zip( zf, manifest=manifest, sections=sections, context=context, mol_or_sys=mol_or_sys, volume_dt=volume_dt, ) # SSOT validation gate -- re-open the in-memory archive read-only # and run the canonical validator. Same blast radius as write_qvf: # raise rather than return bad bytes. buf.seek(0) with zipfile.ZipFile(buf, "r") as zf_ro: report = validate_qvf(zf_ro) if not report["valid"]: raise ValueError( "qvf_bytes produced an archive that fails canonical " "validation -- this is a writer bug. Errors:\n - " + "\n - ".join(report["errors"][:8]) ) return buf.getvalue() def _resolve_version() -> str: """Try to import vibe-qc's version string; fall back gracefully.""" try: from ...banner import VIBEQC_VERSION return str(VIBEQC_VERSION) except Exception: return "0.0.0" # --------------------------------------------------------------------------- # Section writers -- each adds its entry to the ``sections`` list # --------------------------------------------------------------------------- # -- structure ------------------------------------------------------------ def _normalize_biomolecule(bio: Any) -> dict[str, Any]: """Coerce a ``biomolecule_data`` dict to plain JSON-safe values. Returns a dict with only the fields that were supplied — all three are independently optional. numpy ints/arrays are coerced to plain ``int``/``list`` so :func:`safe_json_bytes` and the write-time :func:`validate_qvf` gate stay clean. Malformed entries are dropped rather than raising, mirroring the fail-soft writer policy. Fields (peer keys on the ``structure`` section object, additive — no ``qvf_version`` bump, riding the open Section schema): * ``chains`` -- list[str] of chain ids, e.g. ``["A", "B"]``. * ``residues`` -- list of ``{name, seq, chain, atom_indices}``; ``atom_indices`` are 0-based into the structure payload atoms. * ``secondary_structure`` -- list of ``{type, chain, start_seq, end_seq}`` with ``type`` in ``{"helix", "sheet", "coil"}``. """ if not isinstance(bio, dict): return {} out: dict[str, Any] = {} chains = bio.get("chains") if chains is not None: out["chains"] = [str(c) for c in chains] residues = bio.get("residues") if residues is not None: norm_res = [] for r in residues: if not isinstance(r, dict): continue # Explicit None check, not ``or []`` — a numpy-array # atom_indices would raise "truth value ambiguous". ai = r.get("atom_indices") ai = [] if ai is None else list(ai) norm_res.append( { "name": str(r.get("name", "")), "seq": int(r.get("seq", 0)), "chain": str(r.get("chain", "")), "atom_indices": [int(i) for i in ai], } ) out["residues"] = norm_res ss = bio.get("secondary_structure") if ss is not None: norm_ss = [] for s in ss: if not isinstance(s, dict): continue norm_ss.append( { "type": str(s.get("type", "coil")), "chain": str(s.get("chain", "")), "start_seq": int(s.get("start_seq", 0)), "end_seq": int(s.get("end_seq", 0)), } ) out["secondary_structure"] = norm_ss return out def _write_structure_section( zf: zipfile.ZipFile, mol_or_sys: Any, sections: list[dict[str, Any]], *, periodic: bool = False, biomolecule_data: Any = None, structure_lattice_bohr: Any = None, ) -> None: """Write the ``structure`` section. Molecules: ``Molecule.atoms``. Periodic systems: the bound ``PeriodicSystem`` exposes ``unit_cell`` (not ``atoms``) plus ``dim`` in {1,2,3} and ``lattice`` (Cartesian column vectors in bohr). Tries ``.atoms`` first for back-compat with molecular callers and any test stubs that still use the old name. """ # PeriodicSystem (the bound C++ type) carries atoms under # ``unit_cell``; Molecule carries them under ``atoms``. Pick the # populated one. Some stubs and the molecular path use ``atoms``; # the periodic path uses ``unit_cell``. atoms = list(getattr(mol_or_sys, "atoms", []) or []) if not atoms: atoms = list(getattr(mol_or_sys, "unit_cell", []) or []) if not atoms: import warnings warnings.warn( "_write_structure_section: molecule/system has no atoms " "(.atoms / .unit_cell both empty) -- skipping structure " "section.", UserWarning, stacklevel=2, ) return # Build atoms array for JSON. atom_list = [] for a in atoms: x, y, z = a.xyz atom_list.append( { "symbol": _symbol(int(a.Z)), "position": [ float(x) * _BOHR_TO_ANGSTROM, float(y) * _BOHR_TO_ANGSTROM, float(z) * _BOHR_TO_ANGSTROM, ], "atomic_number": int(a.Z), } ) payload: dict[str, Any] = { "atoms": atom_list, "pbc": [False, False, False], } if periodic: lattice_source = ( structure_lattice_bohr if structure_lattice_bohr is not None else mol_or_sys.lattice ) L_bohr = np.asarray(lattice_source, dtype=np.float64) if L_bohr.shape != (3, 3): raise ValueError( "structure_lattice_bohr must be a 3x3 bohr column-vector lattice" ) L_ang = (L_bohr * _BOHR_TO_ANGSTROM).T.tolist() # row vectors for JSON # PeriodicSystem exposes ``dim`` in {1,2,3}; molecular stubs may # use ``dimensionality``. Default to fully periodic. dim = int( getattr(mol_or_sys, "dim", None) or getattr(mol_or_sys, "dimensionality", 3) or 3 ) dim = max(1, min(3, dim)) payload["pbc"] = [i < dim for i in range(3)] payload["dimensionality"] = dim payload["lattice_vectors"] = [ [float(L_ang[0][0]), float(L_ang[0][1]), float(L_ang[0][2])], [float(L_ang[1][0]), float(L_ang[1][1]), float(L_ang[1][2])], [float(L_ang[2][0]), float(L_ang[2][1]), float(L_ang[2][2])], ] struct_json = safe_json_bytes(payload, indent=2) path_in_zip = "structure/structure.json" zf.writestr(path_in_zip, struct_json) section: dict[str, Any] = { "id": "structure", "kind": "structure", "members": { "structure": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(struct_json), }, }, } # Biomolecule metadata (Ask B) — optional, additive peer keys on the # structure Section object (NOT inside structure.json), riding the open # Section schema like the checkpoint fields; no qvf_version bump. if biomolecule_data is not None: section.update(_normalize_biomolecule(biomolecule_data)) sections.append(section) # -- volume.density / volume.orbital helpers ------------------------------ def _grid_descriptor( data: np.ndarray, origin: np.ndarray, span: np.ndarray, ) -> dict[str, Any]: """Build the ``grid`` member dict for a volumetric section.""" ox, oy, oz = float(origin[0]), float(origin[1]), float(origin[2]) nx, ny, nz = data.shape grid: dict[str, Any] = { "origin": [ox, oy, oz], "voxel_vectors": [ [float(span[0, 0]), float(span[0, 1]), float(span[0, 2])], [float(span[1, 0]), float(span[1, 1]), float(span[1, 2])], [float(span[2, 0]), float(span[2, 1]), float(span[2, 2])], ], "shape": [int(nx), int(ny), int(nz)], } return grid def _write_volume_density_section( zf: zipfile.ZipFile, vol_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.density`` sections. ``vol_data`` is ``{label: (data_3d, origin_bohr, span_bohr)}``. Origin and span are in bohr; we write the grid descriptor in bohr. """ for idx, (label, (data, origin, span)) in enumerate(vol_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.density[{label!r}]", ) _require_3d_volume(vol, "volume.density", label) section_id = f"vol_dens_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.density", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) def _write_volume_orbital_section( zf: zipfile.ZipFile, mo_data: list[dict[str, Any]], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.orbital`` sections. Each item in ``mo_data`` is a dict with keys: ``label``, ``data`` (3D array), ``origin`` (3-vector, bohr), ``span`` (3x3, bohr), ``band_index`` (int), ``energy_eh`` (float), optional ``occupation`` (float, default 2.0 for restricted), ``spin`` (int, default 0), ``component`` (str, "real" default). """ for idx, mo in enumerate(mo_data): label = mo.get("label", f"MO_{idx}") vol = _real_array_for_artifact( mo["data"], dtype=volume_dtype, label=f"volume.orbital[{label!r}]", ) _require_3d_volume(vol, "volume.orbital", label) section_id = f"vol_mo_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) comp = mo.get("component", "real") origin_arr = np.asarray(mo["origin"], dtype=np.float64) span_arr = np.asarray(mo["span"], dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.orbital", "label": label, "component": comp, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- volume.spin ----------------------------------------------------------- def _write_volume_spin_section( zf: zipfile.ZipFile, spin_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.spin`` sections. ``spin_data`` is ``{label: (data_3d, origin_bohr, span_bohr)}``. Same shape as ``volume.density``, different kind string. """ for idx, (label, (data, origin, span)) in enumerate(spin_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.spin[{label!r}]", ) _require_3d_volume(vol, "volume.spin", label) section_id = f"vol_spin_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}_spin.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_spin_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.spin", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- volume.elf ------------------------------------------------------------- def _write_volume_elf_section( zf: zipfile.ZipFile, elf_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.elf`` sections (electron localisation function).""" for idx, (label, (data, origin, span)) in enumerate(elf_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.elf[{label!r}]", ) _require_3d_volume(vol, "volume.elf", label) section_id = f"vol_elf_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}_elf.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_elf_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.elf", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- volume.generic ------------------------------------------------------- def _write_volume_generic_section( zf: zipfile.ZipFile, gen_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.generic`` sections. Escape hatch for any scalar field that doesn't fit one of the purpose-built kinds (density / orbital / spin / elf / difference). Producers should prefer a more specific kind when one applies; the viewer renders this with the same isosurface machinery but cannot apply kind-specific defaults (colormap, sign convention, etc.). ``gen_data`` is ``{label: (data_3d, origin_bohr, span_bohr)}`` -- same shape as ``volume.density``. """ for idx, (label, (data, origin, span)) in enumerate(gen_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.generic[{label!r}]", ) _require_3d_volume(vol, "volume.generic", label) section_id = f"vol_gen_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}_generic.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_generic_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.generic", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- volume.difference ---------------------------------------------------- def _write_volume_difference_section( zf: zipfile.ZipFile, diff_data: dict[str, Any], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.difference`` sections. ``diff_data`` is ``{label: spec}``, where ``spec`` is either: * ``(data_3d, origin_bohr, span_bohr)`` -- no operand metadata. * a dict with keys ``data``, ``origin``, ``span``, and optionally ``operand_a`` (str, section id of the minuend), ``operand_b`` (str, section id of the subtrahend), ``description`` (str). If one operand is given the other is required (schema's ``dependentRequired``). Sign convention: ``data = r(operand_a) - r(operand_b)``. """ for idx, (label, spec) in enumerate(diff_data.items()): if isinstance(spec, dict): data = spec["data"] origin = spec["origin"] span = spec["span"] operand_a = spec.get("operand_a") operand_b = spec.get("operand_b") description = spec.get("description") else: data, origin, span = spec operand_a = operand_b = description = None vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.difference[{label!r}]", ) _require_3d_volume(vol, "volume.difference", label) section_id = f"vol_diff_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}_diff.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_diff_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.difference", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } if operand_a is not None: section["operand_a"] = str(operand_a) if operand_b is not None: section["operand_b"] = str(operand_b) if description is not None: section["description"] = str(description) sections.append(section) # -- volume.potential ---------------------------------------------------- def _write_volume_potential_section( zf: zipfile.ZipFile, potential_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.potential`` sections (QVF spec 4.10). Member structure is identical to ``volume.density``: ``{label: (data_3d, origin_bohr, span_bohr)}``. Origin and span are in bohr; grid descriptor in bohr. """ for idx, (label, (data, origin, span)) in enumerate(potential_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.potential[{label!r}]", ) _require_3d_volume(vol, "volume.potential", label) section_id = f"vol_esp_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.potential", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- volume.rdg ---------------------------------------------------------- def _write_volume_rdg_section( zf: zipfile.ZipFile, rdg_data: dict[str, tuple], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``volume.rdg`` sections (QVF spec 4.11). Member structure is identical to ``volume.density``: ``{label: (data_3d, origin_bohr, span_bohr)}``. Data is dimensionless RDG = |grad rho| / (2*(3*pi^2)^(1/3) * rho^(4/3)). """ for idx, (label, (data, origin, span)) in enumerate(rdg_data.items()): vol = _real_array_for_artifact( data, dtype=volume_dtype, label=f"volume.rdg[{label!r}]", ) _require_3d_volume(vol, "volume.rdg", label) section_id = f"vol_rdg_{idx}" slug = _slug(label, fallback=section_id) path_in_zip = f"volumes/{slug}.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(origin, dtype=np.float64) span_arr = np.asarray(span, dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"volumes/{slug}_grid.json" zf.writestr(grid_path, grid_json) section: dict[str, Any] = { "id": section_id, "kind": "volume.rdg", "label": label, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } sections.append(section) # -- wavefunction.gto ----------------------------------------------------- def _write_wavefunction_gto_section( zf: zipfile.ZipFile, wf_data: dict[str, Any], sections: list[dict[str, Any]], *, section_id: str = "wf", ) -> None: """Write the ``wavefunction.gto`` section. ``wf_data`` is a dict with keys: * ``basis`` -- list of shell dicts, each with ``center`` (int, 0-based), ``l`` (int), ``exponents`` (list of float), ``coefficients`` (list of list of float, [n_prim, n_gen] for general contraction). * ``mo_metadata`` -- dict with ``spin`` ("restricted" or "unrestricted"), ``orbital_kind``, ``n_ao``, and either top-level ``energies``, ``occupations`` (restricted) or ``alpha``/``beta`` sub-dicts (unrestricted). * ``mo_coefficients`` -- 2D ||(`[n_mo, n_ao]`, float64) for restricted, or 3D ||(`[n_mo, n_ao, 2]`, float64) with last axis `[real, imag]`. * ``mo_coefficients_alpha``, ``mo_coefficients_beta`` -- for unrestricted. * ``structure_ref`` -- str, section id of the structure (default ``"structure"``). * ``pure`` -- bool, whether spherical harmonics are used (default True). ``section_id`` overrides the section id in the manifest (default ``"wf"``). Use ``"wf_localized"`` for a second wavefunction section carrying localized orbitals beside the canonical set. Covers molecular wavefunctions, Γ crystalline orbitals, and one selected complex Bloch k-point. The optional ``mo_metadata["k_point"]`` records the reciprocal fractional coordinate; complex coefficients are stored as real/imag component pairs. """ # Derive a subdirectory from the section id so two wavefunction # sections (canonical + localized) don't collide on zip paths. subdir = "wavefunction" if section_id == "wf" else f"wavefunction_{section_id}" # --- basis JSON ---------------------------------------------------- basis_shells = wf_data.get("basis", []) structure_ref = wf_data.get("structure_ref", "structure") pure = wf_data.get("pure", True) if "n_ao" in wf_data: n_ao = int(wf_data["n_ao"]) else: n_ao = sum( (2 * sh["l"] + 1) if pure else ((sh["l"] + 1) * (sh["l"] + 2) // 2) for sh in basis_shells ) basis_dict: dict[str, Any] = { "structure_ref": structure_ref, "pure": pure, "n_ao": n_ao, "shells": basis_shells, } basis_json = safe_json_bytes(basis_dict, indent=2) basis_path = f"{subdir}/basis.json" zf.writestr(basis_path, basis_json) # --- mo_metadata JSON ----------------------------------------------- mo_meta = dict(wf_data.get("mo_metadata", {})) mo_meta.setdefault("n_ao", n_ao) mo_meta_json = safe_json_bytes(mo_meta, indent=2) mo_meta_path = f"{subdir}/mo_metadata.json" zf.writestr(mo_meta_path, mo_meta_json) # --- mo_coefficients binary ----------------------------------------- members: dict[str, Any] = { "basis": { "path": basis_path, "format": "json", "sha256": _sha256_hex(basis_json), }, "mo_metadata": { "path": mo_meta_path, "format": "json", "sha256": _sha256_hex(mo_meta_json), }, } spin = mo_meta.get("spin", "restricted") if spin == "unrestricted": for tag in ("mo_coefficients_alpha", "mo_coefficients_beta"): coeff = wf_data.get(tag) if coeff is not None: arr, _encoding = _mo_coefficients_for_artifact( coeff, label=f"wavefunction.gto {tag}", ) if arr.ndim in (2, 3): path = f"{subdir}/{tag}.dat" members[tag] = _write_binary_to_zip(zf, path, arr) else: coeff = wf_data.get("mo_coefficients") if coeff is not None: arr, _encoding = _mo_coefficients_for_artifact( coeff, label="wavefunction.gto mo_coefficients", ) if arr.ndim in (2, 3): path = f"{subdir}/mo_coefficients.dat" members["mo_coefficients"] = _write_binary_to_zip(zf, path, arr) section: dict[str, Any] = { "id": section_id, "kind": "wavefunction.gto", "members": members, } sections.append(section) def _write_bloch_wavefunction_section( zf: zipfile.ZipFile, wf_data: dict[str, Any], sections: list[dict[str, Any]], *, section_id: str = "wf_bloch_kpoints", ) -> None: """Write the vibe-qc all-k Bloch wavefunction restart section.""" subdir = "wavefunction_bloch" basis_shells = wf_data.get("basis", []) structure_ref = wf_data.get("structure_ref", "structure") pure = wf_data.get("pure", True) n_ao = int( wf_data.get( "n_ao", sum( (2 * sh["l"] + 1) if bool(sh.get("pure", pure)) else ((sh["l"] + 1) * (sh["l"] + 2) // 2) for sh in basis_shells ), ) ) basis_dict: dict[str, Any] = { "structure_ref": structure_ref, "pure": pure, "n_ao": n_ao, "shells": basis_shells, } basis_json = safe_json_bytes(basis_dict, indent=2) basis_path = f"{subdir}/basis.json" zf.writestr(basis_path, basis_json) mo_meta = dict(wf_data.get("mo_metadata", {})) mo_meta.setdefault("n_ao", n_ao) mo_meta_json = safe_json_bytes(mo_meta, indent=2) mo_meta_path = f"{subdir}/mo_metadata.json" zf.writestr(mo_meta_path, mo_meta_json) members: dict[str, Any] = { "basis": { "path": basis_path, "format": "json", "sha256": _sha256_hex(basis_json), }, "mo_metadata": { "path": mo_meta_path, "format": "json", "sha256": _sha256_hex(mo_meta_json), }, } if "kpoints" in wf_data: kpts = np.ascontiguousarray(np.asarray(wf_data["kpoints"], dtype=np.float64)) members["kpoints"] = _write_binary_to_zip( zf, f"{subdir}/kpoints.dat", kpts, ) coeff_blocks = list(wf_data.get("mo_coefficients", [])) occ_blocks = list(wf_data.get("occupations", [])) energy_blocks = list(wf_data.get("mo_energies", [])) if not (len(coeff_blocks) == len(occ_blocks) == len(energy_blocks)): raise ValueError( "_write_bloch_wavefunction_section: coefficient, occupation, and " "energy block counts must match." ) blocks = mo_meta.get("blocks", []) if len(blocks) != len(coeff_blocks): raise ValueError( "_write_bloch_wavefunction_section: metadata block count does not " f"match payload blocks ({len(blocks)} vs {len(coeff_blocks)})." ) for ik, block in enumerate(blocks): coeff_name = str(block.get("mo_coefficients", f"mo_coefficients_k{ik}")) occ_name = str(block.get("occupations", f"occupations_k{ik}")) energy_name = str(block.get("mo_energies", f"mo_energies_k{ik}")) coeff = np.ascontiguousarray( np.asarray(coeff_blocks[ik], dtype=np.float64) ) if coeff.ndim != 3 or coeff.shape[-1] != 2: raise ValueError( "_write_bloch_wavefunction_section: coefficient blocks must " f"have shape [n_mo, n_ao, 2]; block {ik} has {coeff.shape}." ) members[coeff_name] = _write_binary_to_zip( zf, f"{subdir}/{coeff_name}.dat", coeff, ) members[occ_name] = _write_binary_to_zip( zf, f"{subdir}/{occ_name}.dat", np.ascontiguousarray(np.asarray(occ_blocks[ik], dtype=np.float64)), ) members[energy_name] = _write_binary_to_zip( zf, f"{subdir}/{energy_name}.dat", np.ascontiguousarray(np.asarray(energy_blocks[ik], dtype=np.float64)), ) sections.append( { "id": section_id, "kind": QVF_BLOCH_WF_KIND, "members": members, } ) # -- spectra.raman --------------------------------------------------------- def _write_spectra_raman_section( zf: zipfile.ZipFile, raman: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.raman`` section. ``raman`` is a dict with keys like ``frequencies_cm1``, ``intensities``, ``broadening``, ``units_x``, ``units_y``, etc. Mirrors the IR spectrum JSON format. """ spec = dict(raman) spec.setdefault("kind", "spectra.raman") spec.setdefault("version", "1.0") # Consumer expects {frequencies, intensities}. if "frequencies" not in spec: spec["frequencies"] = spec.pop("frequencies_cm1", []) spec.pop("kind", None) spec.pop("version", None) spec.pop("units_x", None) spec.pop("units_y", None) spec.pop("label_x", None) spec.pop("label_y", None) spec.pop("broadening", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/raman.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "raman_spec", "kind": "spectra.raman", "label": "Raman spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.uvvis --------------------------------------------------------- def _write_spectra_uvvis_section( zf: zipfile.ZipFile, uvvis_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.uvvis`` section. ``uvvis_data`` is a dict with keys: * ``energies_ev`` -- list of float, excitation energies (eV) * ``intensities`` -- list of float, oscillator strengths * optionally ``wavelength_nm``, ``broadening``, ``units_x``, ``units_y`` """ spec = dict(uvvis_data) spec.setdefault("kind", "spectra.uvvis") spec.setdefault("version", "1.0") # Consumer expects {frequencies, intensities} -- we store energies as # "frequencies" in eV (the consumer can convert to nm). if "frequencies" not in spec and "energies_ev" in spec: spec["frequencies"] = spec.pop("energies_ev") if "frequencies" not in spec and "wavelength_nm" in spec: import numpy as _np wl = _np.asarray(spec.pop("wavelength_nm"), dtype=float) # eV ≈ 1240 / l(nm) spec["frequencies"] = (1240.0 / wl).tolist() spec.pop("kind", None) spec.pop("version", None) spec.pop("units_x", None) spec.pop("units_y", None) spec.pop("broadening", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/uvvis.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "uvvis_spec", "kind": "spectra.uvvis", "label": "UV-Vis spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.ecd ----------------------------------------------------------- def _write_spectra_ecd_section( zf: zipfile.ZipFile, ecd: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.ecd`` section (electronic circular dichroism).""" spec = dict(ecd) spec.setdefault("kind", "spectra.ecd") spec.setdefault("version", "1.0") if "frequencies" not in spec and "energies_ev" in spec: spec["frequencies"] = spec.pop("energies_ev") spec.pop("kind", None) spec.pop("version", None) spec.pop("units_x", None) spec.pop("units_y", None) spec.pop("broadening", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/ecd.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "ecd_spec", "kind": "spectra.ecd", "label": "ECD spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.vcd ----------------------------------------------------------- def _write_spectra_vcd_section( zf: zipfile.ZipFile, vcd: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.vcd`` section (vibrational circular dichroism).""" spec = dict(vcd) spec.setdefault("kind", "spectra.vcd") spec.setdefault("version", "1.0") if "frequencies" not in spec and "frequencies_cm1" in spec: spec["frequencies"] = spec.pop("frequencies_cm1") spec.pop("kind", None) spec.pop("version", None) spec.pop("units_x", None) spec.pop("units_y", None) spec.pop("broadening", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/vcd.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "vcd_spec", "kind": "spectra.vcd", "label": "VCD spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.nmr ----------------------------------------------------------- def _write_spectra_nmr_section( zf: zipfile.ZipFile, nmr: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.nmr`` section. ``nmr`` is a dict with keys like ``chemical_shifts``, ``shielding_tensors``, ``j_couplings``, ``isotope``, ``reference``, ``solvent``. """ spec = dict(nmr) spec.setdefault("kind", "spectra.nmr") spec.setdefault("version", "1.0") spec.pop("kind", None) spec.pop("version", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/nmr.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "nmr_spec", "kind": "spectra.nmr", "label": "NMR spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.epr ----------------------------------------------------------- def _write_spectra_epr_section( zf: zipfile.ZipFile, epr: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.epr`` section (QVF spec § 4.4). ``epr`` is a dict carrying EPR parameters, e.g. ``g_tensor`` (with ``matrix`` / ``isotropic`` / ``principal``), ``hyperfine`` (list of per-nucleus ``a_tensor_mhz`` / ``a_iso_mhz``), and ``zero_field_splitting`` (``d_mhz`` / ``e_mhz``). Payload shape is intentionally loose in v1, mirroring ``spectra.nmr``. vibe-qc does not yet compute EPR; this is a dormant writer (like the phonon / EOS writers) driven by the ``epr_data`` context key, present so an external producer's format and vibe-view's renderer share one contract. """ spec = dict(epr) spec.pop("kind", None) spec.pop("version", None) spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/epr.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "epr_spec", "kind": "spectra.epr", "label": "EPR spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- spectra.generic ------------------------------------------------------- def _write_spectra_generic_section( zf: zipfile.ZipFile, generic: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write a ``spectra.generic`` section. ``generic`` is a dict with at least ``frequencies`` and ``intensities``, plus any user-defined metadata keys. """ spec = dict(generic) spec.setdefault("kind", "spectra.generic") spec.setdefault("version", "1.0") label = spec.pop("label", "Generic spectrum") spec_id = spec.pop("section_id", "gen_spec") spec.pop("kind", None) spec.pop("version", None) spec_json = safe_json_bytes(spec, indent=2) # Slug the caller-supplied section_id before it becomes a zip member # name -- like the volume labels above -- so a bidi/control char or a # path-traversal attempt cannot reach the archive path. See _slug. path_in_zip = f"spectra/{_slug(spec_id, fallback='gen_spec')}.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": spec_id, "kind": "spectra.generic", "label": label, "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- structure.symmetry ---------------------------------------------------- def _write_symmetry_section( zf: zipfile.ZipFile, sym: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``structure.symmetry`` section. ``sym`` is a dict with keys like ``space_group_number``, ``space_group_symbol``, ``point_group``, etc. (spglib output). """ payload = dict(sym) payload.setdefault("kind", "structure.symmetry") payload.setdefault("version", "1.0") sym_json = safe_json_bytes(payload, indent=2) path_in_zip = "structure/symmetry.json" zf.writestr(path_in_zip, sym_json) section: dict[str, Any] = { "id": "sym0", "kind": "structure.symmetry", "members": { "data": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(sym_json), }, }, } sections.append(section) # -- bonds ------------------------------------------------------------------ def _write_bonds_section( zf: zipfile.ZipFile, bonds: list[tuple[int, int, float]], sections: list[dict[str, Any]], ) -> None: """Write the ``bonds`` section. ``bonds`` is a list of ``(i, j, order)`` tuples. Emitted as one JSON member ``bonds`` carrying ``{"pairs": [{"i", "j", "order"}, ...]}`` -- a small table that doesn't justify a binary blob, and that round-trips through every consumer without a custom struct format. """ payload = { "pairs": [ {"i": int(i), "j": int(j), "order": float(order)} for (i, j, order) in bonds ], } bonds_json = safe_json_bytes(payload) path_in_zip = "bonds/connectivity.json" zf.writestr(path_in_zip, bonds_json) section: dict[str, Any] = { "id": "bonds0", "kind": "bonds", "label": "Bond connectivity", "members": { "bonds": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(bonds_json), }, }, } sections.append(section) # -- bond_orders ------------------------------------------------------------ def _write_bond_orders_section( zf: zipfile.ZipFile, method: str, pairs: list[dict[str, Any]], sections: list[dict[str, Any]], ) -> None: """Write the ``bond_orders`` section. ``method`` is the bond-order definition (``"mayer"``, ``"wiberg"``, etc.). ``pairs`` is a list of dicts each with ``i``, ``j``, ``order`` and optional ``distance_ang``, ``symbol_i``, ``symbol_j``. """ payload: dict[str, Any] = { "method": method, "pairs": [{k: v for k, v in p.items()} for p in pairs], } bo_json = safe_json_bytes(payload, indent=2) path_in_zip = "bonds/orders.json" zf.writestr(path_in_zip, bo_json) section: dict[str, Any] = { "id": "bond_orders", "kind": "bond_orders", "members": { "bond_orders": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(bo_json), }, }, } sections.append(section) # -- topology.qtaim --------------------------------------------------------- def _write_topology_qtaim_section( zf: zipfile.ZipFile, critical_points: list[dict[str, Any]], *, bond_paths: list[dict[str, Any]] | None = None, sections: list[dict[str, Any]], ) -> None: """Write the ``topology.qtaim`` section. ``critical_points`` is a list of dicts each with at minimum ``type`` (``"bcp"``, ``"rcp"``, ``"ccp"``), ``position`` (3-vector in Angstrom), ``rho`` (e/ų), and ``laplacian`` (e/Å⁵). Optional fields: ``ellipticity``, ``atom_pair`` (for BCPs). ``bond_paths`` is an optional list of ``{atoms: [i, j], path: [[x,y,z],...]}`` records for gradient-path polylines. The compute module (critical-point search, Hessian evaluation, gradient-path tracing) is not yet in vibe-qc; the writer accepts pre-computed data so it is usable as soon as the compute side lands. """ payload: dict[str, Any] = {"points": list(critical_points)} if bond_paths: payload["bond_paths"] = list(bond_paths) cp_json = safe_json_bytes(payload, indent=2) path_in_zip = "topology/critical_points.json" zf.writestr(path_in_zip, cp_json) section: dict[str, Any] = { "id": "qtaim", "kind": "topology.qtaim", "members": { "critical_points": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(cp_json), }, }, } sections.append(section) # -- scf_history ------------------------------------------------------------ def _write_scf_history_section( zf: zipfile.ZipFile, history: list[dict[str, Any]], sections: list[dict[str, Any]], ) -> None: """Write the ``scf_history`` section. ``history`` is a list of per-iteration records with keys like ``iter``, ``energy_eh``, ``delta_e``, ``diis_error``. Emitted as a JSON document carrying ``{"iterations": [...]}`` (not JSONL -- the canonical schema declares this member's format as ``json`` and a JSONL payload isn't a valid JSON document). """ payload = {"iterations": list(history)} hist_json = safe_json_bytes(payload) path_in_zip = "scf_history/iterations.json" zf.writestr(path_in_zip, hist_json) section: dict[str, Any] = { "id": "scf_hist0", "kind": "scf_history", "members": { "iterations": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(hist_json), }, }, } sections.append(section) # -- atom_properties ------------------------------------------------------ def _write_atom_properties_section( zf: zipfile.ZipFile, pop: Any, sections: list[dict[str, Any]], ) -> None: """Write the ``atom_properties`` section. ``pop`` is a :class:`PopulationSummary` with fields ``mulliken_atoms``, ``loewdin_atoms``, ``mayer_bonds``, ``dipole``, ``errors``. """ section: dict[str, Any] = { "id": "props0", "kind": "atom_properties", "members": {}, } if pop.mulliken_atoms: charges = np.array([row[3] for row in pop.mulliken_atoms], dtype=np.float64) member = _write_binary_to_zip( zf, "atom_properties/mulliken_charge.bin", charges, ) section["members"]["mulliken_charge"] = member if pop.loewdin_atoms: charges = np.array([row[3] for row in pop.loewdin_atoms], dtype=np.float64) member = _write_binary_to_zip( zf, "atom_properties/loewdin_charge.bin", charges, ) section["members"]["loewdin_charge"] = member if section["members"]: sections.append(section) # -- trajectory ----------------------------------------------------------- def _write_trajectory_section( zf: zipfile.ZipFile, frames: Sequence[Any], sections: list[dict[str, Any]], *, energies: Optional[Sequence[float]] = None, rms_grad: Optional[Sequence[float]] = None, trajectory_type: str = "geometry_optimization", ) -> None: """Write the ``trajectory`` section.""" n_steps = len(frames) # Metadata JSON first_atoms = list(frames[0].atoms) meta_atoms = [ { "symbol": _symbol(int(a.Z)), "position": [0.0, 0.0, 0.0], "atomic_number": int(a.Z), } for a in first_atoms ] meta = {"atoms": meta_atoms} if energies is not None: meta["energies"] = [float(e) for e in energies] meta_json = safe_json_bytes(meta) meta_path = "trajectory/metadata.json" zf.writestr(meta_path, meta_json) # Coords binary: (n_steps, n_atoms, 3) float64 in Å. n_atoms_f = len(first_atoms) coords = np.zeros((n_steps, n_atoms_f, 3), dtype=np.float64) for i, mol in enumerate(frames): for j, a in enumerate(mol.atoms): coords[i, j, 0] = float(a.xyz[0]) * _BOHR_TO_ANGSTROM coords[i, j, 1] = float(a.xyz[1]) * _BOHR_TO_ANGSTROM coords[i, j, 2] = float(a.xyz[2]) * _BOHR_TO_ANGSTROM coords_member = _write_binary_to_zip( zf, "trajectory/coords.bin", coords, ) section: dict[str, Any] = { "id": "traj0", "kind": "trajectory", "members": { "metadata": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, "coords": coords_member, }, } sections.append(section) # -- reaction.path / reaction.waypoints ----------------------------------- _WAYPOINT_KINDS = frozenset( {"reactant", "transition_state", "intermediate", "product", "point"} ) def _validate_waypoints( waypoints: Sequence[dict[str, Any]], n_frames: int, *, context_label: str, ) -> list[dict[str, Any]]: """Sanity-check a waypoint list and return a normalised copy. Raises ``ValueError`` if a waypoint is missing required fields, the ``kind`` is outside the registry, or ``frame_index`` is out of ``[0, n_frames)``. """ if not waypoints: raise ValueError( f"{context_label}: at least one waypoint is required " "(reaction.path / reaction.waypoints both need them)" ) out: list[dict[str, Any]] = [] for i, wp in enumerate(waypoints): if "frame_index" not in wp or "label" not in wp or "kind" not in wp: raise ValueError( f"{context_label}: waypoint #{i} must carry " "'frame_index', 'label', and 'kind'" ) fi = int(wp["frame_index"]) if not 0 <= fi < n_frames: raise ValueError( f"{context_label}: waypoint #{i} frame_index={fi} " f"is outside [0, {n_frames})" ) kind = str(wp["kind"]) if kind not in _WAYPOINT_KINDS: raise ValueError( f"{context_label}: waypoint #{i} kind={kind!r} is not " f"one of {sorted(_WAYPOINT_KINDS)}" ) rec: dict[str, Any] = { "frame_index": fi, "label": str(wp["label"]), "kind": kind, } if "energy_eh" in wp: rec["energy_eh"] = float(wp["energy_eh"]) out.append(rec) return out def _frame_atoms(frame: Any) -> list[Any]: """Pick the populated atom collection on ``frame``. Molecule exposes atoms under ``.atoms``; PeriodicSystem under ``.unit_cell``. Mirrors the dispatch in the structure-section writer so reaction.path frames can be either type. """ atoms = list(getattr(frame, "atoms", []) or []) if not atoms: atoms = list(getattr(frame, "unit_cell", []) or []) return atoms def _frame_is_periodic(frame: Any) -> bool: """A reaction.path frame is periodic iff it carries a lattice.""" return getattr(frame, "lattice", None) is not None and bool( list(getattr(frame, "unit_cell", []) or []) ) def _reaction_path_is_periodic(frames: Sequence[Any]) -> bool: """True iff every frame in ``frames`` is a periodic system. Raises ValueError on a mixed reaction path (molecular + periodic interleaved) -- that's malformed input, not a silent fall-through. """ flags = [_frame_is_periodic(f) for f in frames] if all(flags): return True if not any(flags): return False raise ValueError( "reaction.path: mixed molecular/periodic frames are not " "supported -- every frame must be the same system type" ) def _write_reaction_path_section( zf: zipfile.ZipFile, frames: Sequence[Any], waypoints: Sequence[dict[str, Any]], sections: list[dict[str, Any]], *, energies: Optional[Sequence[float]] = None, reaction_coordinate: Optional[Sequence[float]] = None, reaction_coordinate_label: Optional[str] = None, reaction_coordinate_unit: Optional[str] = None, frame_volumes: Optional[np.ndarray] = None, volume_grid: Optional[dict[str, Any]] = None, volume_frame_index: Optional[Sequence[int]] = None, volume_label: str = "Electron density", volume_isovalue: Optional[float] = None, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write a self-contained ``reaction.path`` section. Binary layout matches ``trajectory`` (coords float64 [n_frames, n_atoms, 3] in Å) so the same readers decode coords. The waypoint annotations live in the JSON metadata member. Frames may be ``Molecule`` (atoms under ``.atoms``) or ``PeriodicSystem`` (atoms under ``.unit_cell``; lattice + dim available). Periodic reaction paths add a binary ``lattice`` member (float64, columns = a, b, c, in bohr -- matching ``PeriodicSystem.lattice``) and a ``dim`` integer in the metadata JSON. A shared lattice across all frames is stored once as shape [3, 3]; per-frame lattices (forward-compat with variable-cell) are stored as shape [n_frames, 3, 3]. Mixed molecular/periodic reaction paths are rejected at write time. """ n_frames = len(frames) if n_frames == 0: raise ValueError("reaction.path: frames is empty") is_periodic = _reaction_path_is_periodic(frames) norm_wps = _validate_waypoints(waypoints, n_frames, context_label="reaction.path") first_atoms = _frame_atoms(frames[0]) if not first_atoms: raise ValueError("reaction.path: first frame has no atoms") meta_atoms = [ { "symbol": _symbol(int(a.Z)), "position": [0.0, 0.0, 0.0], "atomic_number": int(a.Z), } for a in first_atoms ] meta: dict[str, Any] = {"atoms": meta_atoms, "waypoints": norm_wps} if energies is not None: meta["energies"] = [float(e) for e in energies] if reaction_coordinate is not None: meta["reaction_coordinate"] = [float(x) for x in reaction_coordinate] # Optional human-readable axis annotations for the energy plot. # Additive + optional: v1/v2 readers that don't know them ignore # them, so this needs no qvf_version bump. if reaction_coordinate_label is not None: meta["reaction_coordinate_label"] = str(reaction_coordinate_label) if reaction_coordinate_unit is not None: meta["reaction_coordinate_unit"] = str(reaction_coordinate_unit) # v2: pull lattice + dim off periodic frames before writing the # metadata JSON so its sha256 includes them. lattice_array: Optional[np.ndarray] = None if is_periodic: lattices = [np.asarray(f.lattice, dtype=np.float64) for f in frames] for L in lattices: if L.shape != (3, 3): raise ValueError( "reaction.path: PeriodicSystem.lattice must be " f"3x3; got shape {L.shape}" ) all_equal = all(np.allclose(L, lattices[0]) for L in lattices[1:]) if all_equal: lattice_array = lattices[0] else: lattice_array = np.stack(lattices, axis=0) dims = [int(getattr(f, "dim", 3) or 3) for f in frames] if len(set(dims)) == 1: meta["dim"] = dims[0] else: meta["dim_per_frame"] = dims # Optional per-frame volumetric data (W1). A 4D array # [n_emitted, nx, ny, nz] of scalar field values morphing along the # path, plus a shared grid descriptor and an index map saying which # path frame each emitted slab corresponds to (so callers can # decimate -- emit every Nth frame's cube without claiming a cube for # every frame). Validated here so the sha256 includes the metadata. vol_array: Optional[np.ndarray] = None if frame_volumes is not None: vol_array = _real_array_for_artifact( frame_volumes, dtype=volume_dtype, label="reaction.path frame_volumes", ) if vol_array.ndim != 4: raise ValueError( "reaction.path: frame_volumes must be 4D " f"[n_emitted, nx, ny, nz]; got shape {vol_array.shape}" ) if volume_grid is None: raise ValueError( "reaction.path: frame_volumes requires volume_grid " "(an {origin, voxel_vectors, shape} descriptor)" ) n_emit = vol_array.shape[0] if volume_frame_index is None: vfi = list(range(n_emit)) else: vfi = [int(i) for i in volume_frame_index] if len(vfi) != n_emit: raise ValueError( "reaction.path: volume_frame_index length " f"({len(vfi)}) must equal n_emitted ({n_emit})" ) for i in vfi: if not 0 <= i < n_frames: raise ValueError( f"reaction.path: volume_frame_index entry {i} out of " f"range [0, {n_frames})" ) grid_shape = list(int(x) for x in volume_grid["shape"]) if list(vol_array.shape[1:]) != grid_shape: raise ValueError( "reaction.path: frame_volumes grid dims " f"{list(vol_array.shape[1:])} disagree with volume_grid " f"shape {grid_shape}" ) meta["volume_frame_index"] = vfi meta["volume_label"] = str(volume_label) if volume_isovalue is not None: meta["volume_isovalue"] = float(volume_isovalue) meta_json = safe_json_bytes(meta) meta_path = "reaction/metadata.json" zf.writestr(meta_path, meta_json) n_atoms_f = len(first_atoms) coords = np.zeros((n_frames, n_atoms_f, 3), dtype=np.float64) for i, frame in enumerate(frames): frame_atoms = _frame_atoms(frame) if len(frame_atoms) != n_atoms_f: raise ValueError( f"reaction.path: frame {i} has {len(frame_atoms)} " f"atoms; expected {n_atoms_f} (matching frame 0)" ) for j, a in enumerate(frame_atoms): coords[i, j, 0] = float(a.xyz[0]) * _BOHR_TO_ANGSTROM coords[i, j, 1] = float(a.xyz[1]) * _BOHR_TO_ANGSTROM coords[i, j, 2] = float(a.xyz[2]) * _BOHR_TO_ANGSTROM coords_member = _write_binary_to_zip(zf, "reaction/coords.bin", coords) members: dict[str, Any] = { "metadata": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, "coords": coords_member, } if lattice_array is not None: members["lattice"] = _write_binary_to_zip( zf, "reaction/lattice.bin", lattice_array ) if vol_array is not None: assert volume_grid is not None # guarded above members["frame_volumes"] = _write_binary_to_zip( zf, "reaction/frame_volumes.bin", vol_array ) grid_json = safe_json_bytes(volume_grid) grid_path = "reaction/volume_grid.json" zf.writestr(grid_path, grid_json) members["volume_grid"] = { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), } section: dict[str, Any] = { "id": "rxn0", "kind": "reaction.path", "members": members, } sections.append(section) def write_reaction_path_qvf( stem: "os.PathLike | str", *, frames: Sequence[Any], energies: Sequence[float], waypoints: Sequence[dict[str, Any]], reaction_coordinate: Optional[Sequence[float]] = None, reaction_coordinate_label: Optional[str] = None, reaction_coordinate_unit: Optional[str] = None, frame_volumes: Optional[np.ndarray] = None, volume_grid: Optional[dict[str, Any]] = None, volume_frame_index: Optional[Sequence[int]] = None, volume_label: str = "Electron density", volume_isovalue: Optional[float] = None, method: str, basis: str, functional: Optional[str] = None, extra_assemble_kwargs: Optional[dict[str, Any]] = None, compression: Optional[int] = None, ) -> Path: """High-level helper: emit a vibe-view reaction.path archive. Wraps :func:`write_qvf` for the common pattern shared by :meth:`vibeqc.NEBResult.write_qvf` and :meth:`vibeqc.ScanResult.write_qvf`: a structure section (reactant geometry -- first frame), a reaction.path section (frames + waypoints + energies + reaction coordinate), and a citations section (BibTeX assembled with the caller's flags). Periodic vs molecular dispatch is automatic: if ``frames[0]`` is a :class:`PeriodicSystem` the writer detects it (via ``_reaction_path_is_periodic``) and bumps the manifest to QVF v2 with the per-frame lattice + dim -- no extra knobs needed. Parameters ---------- stem Path stem; ``.qvf`` is appended. frames Sequence of :class:`Molecule` or :class:`PeriodicSystem`, one per image. All frames must be the same type; mixed molecular/periodic raises in the lower-level writer. energies Per-frame energies in Hartree, ``len == len(frames)``. waypoints Iterable of ``{frame_index, label, kind, energy_eh?}`` dicts. ``kind`` is one of ``"reactant" | "transition_state" | "intermediate" | "product" | "point"``. reaction_coordinate Per-frame coordinate values; whatever the caller wants the x-axis of the energy plot to show -- arc length for NEB (normalised 0-1), bond length / angle for a relaxed scan, etc. ``None`` => the plot uses frame indices. reaction_coordinate_label, reaction_coordinate_unit Optional human-readable name + unit for the reaction coordinate (e.g. ``"O-H distance"`` / ``"bohr"``). The viewer labels the energy-plot x-axis ``"{label} ({unit})"`` when present, else falls back to ``"Reaction coordinate"``. Additive optional metadata -- no ``qvf_version`` bump. method, basis, functional SCF flavour for the OutputPlan + citation routing. ``functional`` is None for HF methods. extra_assemble_kwargs Forwarded to ``CitationDatabase.assemble`` -- e.g. ``{"uses_neb": True, "uses_ci_neb": True}`` for a climbing-image NEB run. Use this to fire driver-specific citation routes; the per-image SCF citations (method / basis / functional) fire automatically from the args above. compression Optional ``zipfile`` compression constant; ``None`` => the writer's default. Returns ------- pathlib.Path The on-disk ``{stem}.qvf`` path. """ from ..citations.bibtex import format_bibtex from ..citations.registry import load_default_database from ..plan import OutputPlan as _OutputPlan frames_list = list(frames) if not frames_list: raise ValueError("write_reaction_path_qvf: frames is empty") is_periodic = _reaction_path_is_periodic(frames_list) first = frames_list[0] rc: Optional[list[float]] = None if reaction_coordinate is not None: rc = [float(v) for v in reaction_coordinate] if len(rc) != len(frames_list): raise ValueError( f"write_reaction_path_qvf: reaction_coordinate length " f"({len(rc)}) must equal n_frames ({len(frames_list)})" ) energy_list = [float(e) for e in energies] if len(energy_list) != len(frames_list): raise ValueError( f"write_reaction_path_qvf: energies length " f"({len(energy_list)}) must equal n_frames " f"({len(frames_list)})" ) # OutputPlan needs a string method/basis. Caller is responsible # for not passing None there -- citations + provenance get # garbled otherwise. plan = _OutputPlan.from_run_job_kwargs( output=stem, method=method, basis=basis, functional=functional, job_kind="periodic_scf" if is_periodic else "molecular_scf", output_qvf=True, citations=True, write_xyz=False, write_molden_file=False, write_population=False, ) # Stub SCF result carrying the first frame's energy -- enough # for the structure-section writer; the reaction.path section # carries its own per-image energies. stub_result = type( "_ReactionPathStubResult", (), { "converged": True, "energy": energy_list[0], }, )() # Assemble citations. The method / basis / functional routes # always fire; extra_assemble_kwargs fires driver-level routes # (e.g. uses_neb / uses_ci_neb). db = load_default_database() asm_kwargs: dict[str, Any] = dict(extra_assemble_kwargs or {}) asm_kwargs.setdefault("method", method.lower()) asm_kwargs.setdefault("basis", basis) if functional is not None: asm_kwargs.setdefault("functional", functional) asm_kwargs.setdefault("periodic", is_periodic) assembled = db.assemble(**asm_kwargs) bibtex_content = format_bibtex(assembled) context: dict[str, Any] = { "method": method, "basis": basis, "functional": functional, "result": stub_result, "reaction_path": { "frames": frames_list, "waypoints": list(waypoints), "energies": energy_list, }, "bibtex_content": bibtex_content, } if rc is not None: context["reaction_path"]["reaction_coordinate"] = rc if reaction_coordinate_label is not None: context["reaction_path"]["reaction_coordinate_label"] = str( reaction_coordinate_label ) if reaction_coordinate_unit is not None: context["reaction_path"]["reaction_coordinate_unit"] = str( reaction_coordinate_unit ) if frame_volumes is not None: context["reaction_path"]["frame_volumes"] = frame_volumes context["reaction_path"]["volume_grid"] = volume_grid context["reaction_path"]["volume_frame_index"] = volume_frame_index context["reaction_path"]["volume_label"] = volume_label if volume_isovalue is not None: context["reaction_path"]["volume_isovalue"] = volume_isovalue if is_periodic: context["system"] = first else: context["molecule"] = first kwargs: dict[str, Any] = {} if compression is not None: kwargs["compression"] = compression return write_qvf(stem, plan, **context, **kwargs) def _write_scan_surface_section( zf: zipfile.ZipFile, sections: list[dict[str, Any]], *, axis_a: Sequence[float], axis_b: Sequence[float], energies: np.ndarray, coordinate_a_label: str, coordinate_a_unit: str, coordinate_b_label: str, coordinate_b_unit: str, geometries: Optional[np.ndarray] = None, atoms: Optional[Sequence[dict[str, Any]]] = None, ) -> None: """Write a self-contained ``scan.surface`` section (2D energy grid). ``energies`` is ``[nA, nB]`` (Hartree); ``axis_a`` / ``axis_b`` are the 1D driven-coordinate values. Optional ``geometries`` is the relaxed structure at each node, flattened ``[nA*nB, n_atoms, 3]`` in Å, row-major over ``(a, b)``. """ e = np.asarray(energies, dtype=np.float64) if e.ndim != 2: raise ValueError(f"scan.surface: energies must be 2D; got {e.shape}") n_a, n_b = e.shape a_arr = np.asarray(list(axis_a), dtype=np.float64) b_arr = np.asarray(list(axis_b), dtype=np.float64) if a_arr.shape != (n_a,): raise ValueError(f"scan.surface: axis_a length {a_arr.shape} != nA={n_a}") if b_arr.shape != (n_b,): raise ValueError(f"scan.surface: axis_b length {b_arr.shape} != nB={n_b}") meta: dict[str, Any] = { "shape": [int(n_a), int(n_b)], "coordinate_a_label": str(coordinate_a_label), "coordinate_a_unit": str(coordinate_a_unit), "coordinate_b_label": str(coordinate_b_label), "coordinate_b_unit": str(coordinate_b_unit), } if geometries is not None and atoms is not None: meta["atoms"] = list(atoms) meta_json = safe_json_bytes(meta) meta_path = "scan_surface/metadata.json" zf.writestr(meta_path, meta_json) members: dict[str, Any] = { "metadata": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, "axis_a": _write_binary_to_zip(zf, "scan_surface/axis_a.bin", a_arr), "axis_b": _write_binary_to_zip(zf, "scan_surface/axis_b.bin", b_arr), "energies": _write_binary_to_zip(zf, "scan_surface/energies.bin", e), } if geometries is not None and atoms is not None: g = np.asarray(geometries, dtype=np.float64) if g.ndim != 3 or g.shape[0] != n_a * n_b: raise ValueError( "scan.surface: geometries must be " f"[nA*nB, n_atoms, 3]; got {g.shape} (nA*nB={n_a * n_b})" ) members["geometries"] = _write_binary_to_zip( zf, "scan_surface/geometries.bin", g ) sections.append({"id": "scan0", "kind": "scan.surface", "members": members}) def write_scan_surface_qvf( stem: "os.PathLike | str", *, axis_a: Sequence[float], axis_b: Sequence[float], energies: np.ndarray, coordinate_a_label: str, coordinate_a_unit: str, coordinate_b_label: str, coordinate_b_unit: str, structure: Any, geometries: Optional[np.ndarray] = None, method: str, basis: str, functional: Optional[str] = None, compression: Optional[int] = None, ) -> Path: """High-level helper: emit a vibe-view ``scan.surface`` archive for a 2D relaxed scan -- a structure section (the reference geometry), a scan.surface section (the energy grid + axes), and a citations section. Mirrors :func:`write_reaction_path_qvf`. """ from ..citations.bibtex import format_bibtex from ..citations.registry import load_default_database from ..plan import OutputPlan as _OutputPlan is_periodic = _frame_is_periodic(structure) ref_mol = structure.unit_cell_molecule() if is_periodic else structure atoms_meta = None if geometries is not None: atoms_meta = [ {"symbol": _symbol(int(a.Z)), "atomic_number": int(a.Z)} for a in _frame_atoms(structure) ] plan = _OutputPlan.from_run_job_kwargs( output=stem, method=method, basis=basis, functional=functional, job_kind="periodic_scf" if is_periodic else "molecular_scf", output_qvf=True, citations=True, write_xyz=False, write_molden_file=False, write_population=False, ) e0 = float(np.asarray(energies, dtype=np.float64).flat[0]) stub_result = type( "_ScanSurfaceStubResult", (), {"converged": True, "energy": e0}, )() db = load_default_database() asm_kwargs: dict[str, Any] = { "method": method.lower(), "basis": basis, "periodic": is_periodic, } if functional is not None: asm_kwargs["functional"] = functional bibtex_content = format_bibtex(db.assemble(**asm_kwargs)) context: dict[str, Any] = { "method": method, "basis": basis, "functional": functional, "result": stub_result, "bibtex_content": bibtex_content, "scan_surface": { "axis_a": list(axis_a), "axis_b": list(axis_b), "energies": energies, "coordinate_a_label": coordinate_a_label, "coordinate_a_unit": coordinate_a_unit, "coordinate_b_label": coordinate_b_label, "coordinate_b_unit": coordinate_b_unit, "geometries": geometries, "atoms": atoms_meta, }, } if is_periodic: context["system"] = structure else: context["molecule"] = ref_mol kwargs: dict[str, Any] = {} if compression is not None: kwargs["compression"] = compression return write_qvf(stem, plan, **context, **kwargs) def _write_reaction_waypoints_section( zf: zipfile.ZipFile, trajectory_ref: str, waypoints: Sequence[dict[str, Any]], n_trajectory_frames: int, sections: list[dict[str, Any]], *, reaction_coordinate: Optional[Sequence[float]] = None, ) -> None: """Write a ``reaction.waypoints`` section pointing at an existing ``trajectory`` section. No coords are emitted -- they live in the referenced trajectory. The producer is responsible for ensuring ``trajectory_ref`` names a section actually present in this archive; the validator checks that cross-reference at write time. """ norm_wps = _validate_waypoints( waypoints, n_trajectory_frames, context_label=f"reaction.waypoints(trajectory_ref={trajectory_ref!r})", ) payload: dict[str, Any] = {"waypoints": norm_wps} if reaction_coordinate is not None: if len(reaction_coordinate) != n_trajectory_frames: raise ValueError( "reaction.waypoints: reaction_coordinate length " f"({len(reaction_coordinate)}) must equal " f"n_trajectory_frames ({n_trajectory_frames})" ) payload["reaction_coordinate"] = [float(x) for x in reaction_coordinate] wps_json = safe_json_bytes(payload) path_in_zip = "reaction/waypoints.json" zf.writestr(path_in_zip, wps_json) section: dict[str, Any] = { "id": "rxn_wp0", "kind": "reaction.waypoints", "trajectory_ref": str(trajectory_ref), "members": { "waypoints": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(wps_json), }, }, } sections.append(section) # -- vibrations ----------------------------------------------------------- def _write_vibrations_section( zf: zipfile.ZipFile, hess: Any, sections: list[dict[str, Any]], *, atom_symbols: list[str] | None = None, molecule: Any | None = None, ) -> None: """Write the ``vibrations`` section. ``hess`` is a :class:`HessianResult` with ``frequencies_cm1``, ``normal_modes`` (and optionally ``masses_amu``). ``atom_symbols`` provides element symbols; ``molecule`` provides the equilibrium geometry + atomic numbers for the metadata atoms array. Two correctness fixes vs. the original implementation (audit findings A5-01/A5-02/A5-05): * The metadata atoms array now carries the **real equilibrium positions** (Å) and atomic numbers, not ``[0,0,0]`` / ``Z=0``. Without this the viewer animated every atom collapsed at the origin. * The displacements are now **un-mass-weighted Cartesian** patterns. ``HessianResult.normal_modes`` are orthonormal *mass-weighted* eigenvectors; the Cartesian motion of atom ``i`` is ``modes[3i:3i+3, k] / sqrt(M_i)``. Shipping the raw mass-weighted vector made heavy atoms over-move and light atoms under-move (a C-H stretch showed the C moving as much as the H). Each mode is then normalized to unit maximum atomic displacement so the viewer's amplitude slider gives a consistent, visible scale while preserving the (now correct) relative per-atom amplitudes. """ freqs = np.asarray(hess.frequencies_cm1, dtype=np.float64) n_modes = freqs.shape[0] n_atoms = n_modes // 3 modes = np.asarray(hess.normal_modes, dtype=np.float64) # Equilibrium geometry (bohr -> Å) + atomic numbers from the molecule. positions_ang: list[list[float]] | None = None atomic_numbers: list[int] | None = None if molecule is not None: try: positions_ang = [ [ float(a.xyz[0]) * _BOHR_TO_ANGSTROM, float(a.xyz[1]) * _BOHR_TO_ANGSTROM, float(a.xyz[2]) * _BOHR_TO_ANGSTROM, ] for a in molecule.atoms ] atomic_numbers = [int(a.Z) for a in molecule.atoms] except Exception: # noqa: BLE001 -- degrade to zeros if molecule API differs positions_ang = None atomic_numbers = None syms = atom_symbols or ["?"] * n_atoms atoms_list = [] for a_idx in range(n_atoms): sym = syms[a_idx] if a_idx < len(syms) else "?" pos = ( positions_ang[a_idx] if positions_ang is not None and a_idx < len(positions_ang) else [0.0, 0.0, 0.0] ) z = ( atomic_numbers[a_idx] if atomic_numbers is not None and a_idx < len(atomic_numbers) else 0 ) atoms_list.append({"symbol": sym, "position": pos, "atomic_number": z}) meta = { "frequencies": [float(freqs[p]) for p in range(n_modes)], "atoms": atoms_list, } meta_json = safe_json_bytes(meta) meta_path = "vibrations/metadata.json" zf.writestr(meta_path, meta_json) # Un-mass-weighting factor 1/sqrt(M_i) per atom (electron-mass units, to # match HessianResult.normal_modes). Falls back to 1 when masses are # unavailable (e.g. minimal test stubs) so we never crash. masses_amu = getattr(hess, "masses_amu", None) inv_sqrt_m: np.ndarray | None = None if masses_amu is not None: m = np.asarray(masses_amu, dtype=np.float64) if m.shape == (n_atoms,) and np.all(m > 0): _AMU_TO_E = 1822.8884862 inv_sqrt_m = 1.0 / np.sqrt(m * _AMU_TO_E) # Displacements (n_modes, n_atoms, 3): un-mass-weighted Cartesian # patterns, each mode normalized to unit max atomic displacement. disp = np.zeros((n_modes, n_atoms, 3), dtype=np.float64) for p in range(n_modes): block = modes[:, p].reshape(n_atoms, 3) if inv_sqrt_m is not None: block = block * inv_sqrt_m[:, None] max_norm = float(np.max(np.linalg.norm(block, axis=1))) if n_atoms else 0.0 if max_norm > 1e-12: block = block / max_norm disp[p] = block disp_member = _write_binary_to_zip( zf, "vibrations/displacements.bin", disp, ) section: dict[str, Any] = { "id": "vib0", "kind": "vibrations", "members": { "metadata": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, "displacements": disp_member, }, } sections.append(section) # -- spectra.ir ----------------------------------------------------------- def _write_spectra_ir_section( zf: zipfile.ZipFile, hess: Any, sections: list[dict[str, Any]], ) -> None: """Write the ``spectra.ir`` section.""" try: from ...hessian import ir_intensities intensities = np.asarray(ir_intensities(hess), dtype=np.float64) except Exception: return freqs = np.asarray(hess.frequencies_cm1, dtype=np.float64) # Only positive (real) frequencies -- skip imaginary and trans/rot. mask = freqs > 1.0 spec = { "frequencies": freqs[mask].tolist(), "intensities": intensities[mask].tolist(), } spec_json = safe_json_bytes(spec, indent=2) path_in_zip = "spectra/ir.json" zf.writestr(path_in_zip, spec_json) section: dict[str, Any] = { "id": "ir_spec", "kind": "spectra.ir", "label": "IR spectrum", "members": { "spectrum": { "path": path_in_zip, "format": "json", "sha256": _sha256_hex(spec_json), }, }, } sections.append(section) # -- basis.ao ------------------------------------------------------------- def _write_basis_ao_section( zf: zipfile.ZipFile, ao_data: list[dict[str, Any]], sections: list[dict[str, Any]], *, volume_dtype: np.dtype = np.dtype("float32"), ) -> None: """Write ``basis.ao`` sections -- one per atomic orbital. ``ao_data`` is a list of dicts from :func:`qvf_ao_data`. Each dict has keys: ``label``, ``data`` (3-D array), ``origin`` (bohr), ``span`` (3x3 bohr), ``ao_metadata``, ``section_id``. The grid+data member layout is identical to ``volume.orbital``; the ``ao_metadata`` is embedded at the section root. """ for idx, ao in enumerate(ao_data): label = ao.get("label", f"AO_{idx}") vol = _real_array_for_artifact( ao["data"], dtype=volume_dtype, label=f"basis.ao[{label!r}]", ) _require_3d_volume(vol, "basis.ao", label) section_id = ao.get("section_id", f"ao_{idx}") slug = _slug(label, fallback=section_id) path_in_zip = f"basis_ao/{slug}.dat" file_member = _write_binary_to_zip(zf, path_in_zip, vol) origin_arr = np.asarray(ao["origin"], dtype=np.float64) span_arr = np.asarray(ao["span"], dtype=np.float64).reshape(3, 3) grid = _grid_descriptor(vol, origin_arr, span_arr) grid_json = safe_json_bytes(grid) grid_path = f"basis_ao/{slug}_grid.json" zf.writestr(grid_path, grid_json) meta = ao.get("ao_metadata", {}) section: dict[str, Any] = { "id": section_id, "kind": "basis.ao", "label": label, "ao_metadata": { "atom_index": int(meta.get("atom_index", 0)), "atom_symbol": str(meta.get("atom_symbol", "?")), "shell_index": int(meta.get("shell_index", 0)), "primitive_index": int(meta.get("primitive_index", 0)), "angular_momentum": [ int(meta["angular_momentum"][0]), int(meta["angular_momentum"][1]), ], "shell_type": str(meta.get("shell_type", "s")), "exponent": float(meta.get("exponent", 0.0)), "coefficient": float(meta.get("coefficient", 0.0)), "is_primitive": bool(meta.get("is_primitive", False)), "is_contracted": bool(meta.get("is_contracted", True)), "ao_index": int(meta.get("ao_index", 0)), }, "members": { "data": file_member, "grid": { "path": grid_path, "format": "json", "sha256": _sha256_hex(grid_json), }, }, } bl = meta.get("basis_label") if bl: section["ao_metadata"]["basis_label"] = str(bl) sections.append(section) # -- bands ---------------------------------------------------------------- def _write_bands_section( zf: zipfile.ZipFile, bs: Any, sections: list[dict[str, Any]], ) -> None: """Write the ``bands`` section. ``bs`` is a :class:`BandStructure` with ``kpath``, ``energies`` (n_points, n_bands, Hartree), and optional ``e_fermi``, ``projections``, ``channels``. """ kp = bs.kpath energies = np.asarray(bs.energies, dtype=np.float64) n_points, n_bands = energies.shape e_fermi_eh = bs.e_fermi if bs.e_fermi is not None else 0.0 e_fermi_ev = float(e_fermi_eh) * _HARTREE_TO_EV # Eigenvalues as (1, n_k, n_bands) float64, eV. energies_ev = energies * _HARTREE_TO_EV data = energies_ev.reshape(1, n_points, n_bands) data_member = _write_binary_to_zip( zf, "bands/eigenvalues.bin", data, ) # k-path JSON segs: list[dict[str, Any]] = [] kpath_json_data: dict[str, Any] = { "kind": "bands", "version": "1.0", "n_kpoints": int(n_points), "n_bands": int(n_bands), "n_spin": 1, "fermi": float(e_fermi_ev), "fermi_energy_ev": float(e_fermi_ev), "reciprocal_space": True, "segments": [], } # Build segments from label pairs. labels = kp.labels if hasattr(kp, "labels") else [] if len(labels) >= 2: for i in range(len(labels) - 1): d_start, name_start = labels[i] d_end, name_end = labels[i + 1] seg_mask = (kp.distances >= d_start) & (kp.distances <= d_end) n_pts = int(np.sum(seg_mask)) n_pts = max(n_pts, 2) segs.append( { "label_start": name_start, "label_end": name_end, "k_start": kp.kpoints_frac[ int(np.argmin(np.abs(kp.distances - d_start))) ].tolist(), "k_end": kp.kpoints_frac[ int(np.argmin(np.abs(kp.distances - d_end))) ].tolist(), "n_points": n_pts, } ) kpath_json_data["segments"] = segs # --- band-character channels (fat bands) ------------------------------ projections = getattr(bs, "projections", None) channels = getattr(bs, "channels", None) if projections is not None: proj = np.asarray(projections, dtype=np.float64) if proj.ndim != 3: raise ValueError( f"_write_bands_section: projections must be rank-3, got " f"ndim={proj.ndim}" ) if proj.shape[0] != n_points or proj.shape[1] != n_bands: raise ValueError( f"_write_bands_section: projections shape {proj.shape} " f"mismatches eigenvalues ({n_points}, {n_bands})" ) n_channels = proj.shape[2] if channels is not None and len(channels) != n_channels: raise ValueError( f"_write_bands_section: channels length {len(channels)} " f"mismatches projections n_channels={n_channels}" ) # Write channels list into the kpath JSON. kpath_json_data["channels"] = [str(c) for c in channels] if channels else [] kpath_json = safe_json_bytes(kpath_json_data, indent=2) kpath_path = "bands/kpath.json" zf.writestr(kpath_path, kpath_json) members: dict[str, Any] = { "kpath": { "path": kpath_path, "format": "json", "sha256": _sha256_hex(kpath_json), }, "eigenvalues": data_member, } if projections is not None: proj_member = _write_binary_to_zip(zf, "bands/projections.bin", proj) members["projections"] = proj_member section: dict[str, Any] = { "id": "bands0", "kind": "bands", "members": members, } sections.append(section) # -- citations ------------------------------------------------------------ def _write_citations_section( zf: zipfile.ZipFile, bibtex_content: str, sections: list[dict[str, Any]], ) -> None: """Write the ``citations`` section (embedded BibTeX). BibTeX is utf-8 bytes; the manifest format is ``binary`` (not ``json``) and carries a sha256 like every other binary member. Consumers decode the bytes as utf-8 when they want the text. """ bib_bytes = bibtex_content.encode("utf-8") path_in_zip = "citations/references.bib" zf.writestr(path_in_zip, bib_bytes) section: dict[str, Any] = { "id": "citations0", "kind": "citations", "members": { "references": { "path": path_in_zip, "format": "binary", "sha256": _sha256_hex(bib_bytes), }, }, } sections.append(section) # -- dos.total -------------------------------------------------------------- def _write_dos_total_section( zf: zipfile.ZipFile, dos_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``dos.total`` section (QVF spec Sec.4.8). ``dos_data`` is a dict with keys: * ``energies`` -- float64 `[n_points]` in eV, Fermi = 0 * ``dos`` -- float64 `[n_points]` (restricted) or `[2, n_points]` (spin-polarized), states / eV / cell * ``smearing`` -- float, broadening width in eV * ``smearing_type`` -- str, e.g. ``"gaussian"`` * ``fermi_energy_ev`` -- float, absolute Fermi level in eV * ``n_electrons`` -- float, integrated electron count * ``n_spin`` -- int, 1 (restricted) or 2 (spin-polarized) """ energies = np.asarray(dos_data["energies"], dtype=np.float64) dos_arr = np.asarray(dos_data["dos"], dtype=np.float64) n_spin = int(dos_data.get("n_spin", 1)) if n_spin == 1: if dos_arr.ndim != 1: raise ValueError( f"dos.total: for n_spin=1, dos must be 1-D; got shape {dos_arr.shape}" ) dos_shape = list(dos_arr.shape) else: if dos_arr.ndim != 2 or dos_arr.shape[0] != 2: raise ValueError( "dos.total: for n_spin=2, dos must be shape [2, n_points]; " f"got shape {dos_arr.shape}" ) dos_shape = list(dos_arr.shape) # Binary payloads. energies_member = _write_binary_to_zip( zf, "dos/energies.bin", energies, ) dos_member = _write_binary_to_zip( zf, "dos/total.bin", dos_arr, ) # Metadata JSON (per-spec "role meta"). smearing_ev = float(dos_data.get("smearing", 0.05)) smearing_type = str(dos_data.get("smearing_type", "gaussian")) fermi_ev = float(dos_data.get("fermi_energy_ev", 0.0)) n_elec = float(dos_data.get("n_electrons", 0.0)) section: dict[str, Any] = { "id": "dos_total", "kind": "dos.total", "members": { "energies": energies_member, "dos": dos_member, }, # Per the spec, metadata is embedded at the section level as # optional JSON keys (not as a separate member JSON). "smearing": smearing_ev, "smearing_type": smearing_type, "fermi_energy_ev": fermi_ev, "n_electrons": n_elec, "n_spin": n_spin, } sections.append(section) # -- dos.projected ---------------------------------------------------------- def _write_dos_projected_section( zf: zipfile.ZipFile, pdos_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``dos.projected`` section (QVF spec Sec.4.9). ``pdos_data`` is a dict with keys: * ``energies`` -- float64 `[n_points]` in eV, Fermi = 0 * ``projections`` -- float64 `[n_channels, n_points]` (restricted) or `[n_spin, n_channels, n_points]` (spin-polarized), states / eV / cell * ``energies_units`` -- str, ``"eV"`` * ``n_spin`` -- int, 1 or 2 * ``fermi_energy_ev`` -- float, absolute Fermi level in eV * ``channels`` -- list of dicts, each with ``atom_index``, ``symbol``, ``l``, ``label`` """ energies = np.asarray(pdos_data["energies"], dtype=np.float64) projections = np.asarray(pdos_data["projections"], dtype=np.float64) n_spin = int(pdos_data.get("n_spin", 1)) if n_spin == 1: if projections.ndim != 2: raise ValueError( "dos.projected: for n_spin=1, projections must be " f"[n_channels, n_points]; got shape {projections.shape}" ) else: if projections.ndim != 3 or projections.shape[0] != 2: raise ValueError( "dos.projected: for n_spin=2, projections must be " f"[2, n_channels, n_points]; got shape {projections.shape}" ) # Binary payloads. Use a projected-specific energies path so a # dos.projected section can coexist with a dos.total section in the # same archive -- both used to write "dos/energies.bin", which # produced a duplicate zip member (the energy grid member path is # per-section, so the consumer resolves each independently). energies_member = _write_binary_to_zip( zf, "dos/projected_energies.bin", energies, ) proj_member = _write_binary_to_zip( zf, "dos/projections.bin", projections, ) # Channel metadata. channels: list[dict[str, Any]] = list(pdos_data.get("channels", [])) energies_units = str(pdos_data.get("energies_units", "eV")) fermi_ev = float(pdos_data.get("fermi_energy_ev", 0.0)) section: dict[str, Any] = { "id": "dos_pdos", "kind": "dos.projected", "members": { "energies": energies_member, "projections": proj_member, }, # Metadata at section level per spec. Note: energies_units is # a separate meta key to distinguish from dos.total's units layout. "energies_units": energies_units, "n_spin": n_spin, "fermi_energy_ev": fermi_ev, "channels": channels, } sections.append(section) # -- dos.coop -------------------------------------------------------------- def _write_dos_coop_section( zf: zipfile.ZipFile, coop_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``dos.coop`` section (QVF spec Sec.4.8b). ``coop_data`` is a dict with keys: * ``energies`` -- float64 ``[n_points]`` in eV, Fermi = 0 * ``projections`` -- float64 ``[n_pairs, n_points]`` (restricted) or ``[n_spin, n_pairs, n_points]`` (spin-polarized), COOP(E) * ``integrated`` -- float64 ``[n_pairs]``, ICOOP integrated to E_F * ``energies_units`` -- str, ``"eV"`` * ``n_spin`` -- int, 1 or 2 * ``fermi_energy_ev`` -- float, absolute Fermi level in eV * ``sigma_ev`` -- float, Gaussian broadening in eV * ``pairs`` -- list of dicts, each with ``i``, ``j``, ``symbol_i``, ``symbol_j``, ``distance_ang`` """ energies = np.asarray(coop_data["energies"], dtype=np.float64) projections = np.asarray(coop_data["projections"], dtype=np.float64) integrated = np.asarray(coop_data["integrated"], dtype=np.float64) n_spin = int(coop_data.get("n_spin", 1)) if n_spin == 1: if projections.ndim != 2: raise ValueError( "dos.coop: for n_spin=1, projections must be " f"[n_pairs, n_points]; got shape {projections.shape}" ) if integrated.ndim != 1: raise ValueError( "dos.coop: for n_spin=1, integrated must be " f"[n_pairs]; got shape {integrated.shape}" ) else: if projections.ndim != 3 or projections.shape[0] != 2: raise ValueError( "dos.coop: for n_spin=2, projections must be " f"[2, n_pairs, n_points]; got shape {projections.shape}" ) if integrated.ndim != 2 or integrated.shape[0] != 2: raise ValueError( "dos.coop: for n_spin=2, integrated must be " f"[2, n_pairs]; got shape {integrated.shape}" ) energies_member = _write_binary_to_zip(zf, "dos/coop_energies.bin", energies) proj_member = _write_binary_to_zip(zf, "dos/coop_projections.bin", projections) integ_member = _write_binary_to_zip(zf, "dos/coop_integrated.bin", integrated) # Pair metadata as JSON pairs: list[dict[str, Any]] = list(coop_data.get("pairs", [])) meta = { "method": "coop", "fermi_energy_ev": float(coop_data.get("fermi_energy_ev", 0.0)), "sigma_ev": float(coop_data.get("sigma_ev", 0.27)), "n_spin": n_spin, "pairs": pairs, } meta_json = safe_json_bytes(meta) meta_path = "dos/coop_meta.json" zf.writestr(meta_path, meta_json) section: dict[str, Any] = { "id": "coop0", "kind": "dos.coop", "members": { "energies": energies_member, "projections": proj_member, "integrated": integ_member, "meta": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, }, "energies_units": str(coop_data.get("energies_units", "eV")), "n_spin": n_spin, "fermi_energy_ev": float(coop_data.get("fermi_energy_ev", 0.0)), "sigma_ev": float(coop_data.get("sigma_ev", 0.27)), "pairs": pairs, } sections.append(section) # -- dos.cohp -------------------------------------------------------------- def _write_dos_cohp_section( zf: zipfile.ZipFile, cohp_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``dos.cohp`` section (QVF spec Sec.4.8c). ``cohp_data`` is a dict with keys: * ``energies`` -- float64 ``[n_points]`` in eV, Fermi = 0 * ``projections`` -- float64 ``[n_pairs, n_points]`` (restricted) or ``[n_spin, n_pairs, n_points]`` (spin-polarized), -COHP(E) with bonding positive * ``integrated`` -- float64 ``[n_pairs]``, -ICOHP * ``energies_units`` -- str, ``"eV"`` * ``n_spin`` -- int, 1 or 2 * ``fermi_energy_ev`` -- float * ``sigma_ev`` -- float * ``pairs`` -- list of dicts """ energies = np.asarray(cohp_data["energies"], dtype=np.float64) projections = np.asarray(cohp_data["projections"], dtype=np.float64) integrated = np.asarray(cohp_data["integrated"], dtype=np.float64) n_spin = int(cohp_data.get("n_spin", 1)) if n_spin == 1: if projections.ndim != 2: raise ValueError( "dos.cohp: for n_spin=1, projections must be " f"[n_pairs, n_points]; got shape {projections.shape}" ) if integrated.ndim != 1: raise ValueError( "dos.cohp: for n_spin=1, integrated must be " f"[n_pairs]; got shape {integrated.shape}" ) else: if projections.ndim != 3 or projections.shape[0] != 2: raise ValueError( "dos.cohp: for n_spin=2, projections must be " f"[2, n_pairs, n_points]; got shape {projections.shape}" ) if integrated.ndim != 2 or integrated.shape[0] != 2: raise ValueError( "dos.cohp: for n_spin=2, integrated must be " f"[2, n_pairs]; got shape {integrated.shape}" ) energies_member = _write_binary_to_zip(zf, "dos/cohp_energies.bin", energies) proj_member = _write_binary_to_zip(zf, "dos/cohp_projections.bin", projections) integ_member = _write_binary_to_zip(zf, "dos/cohp_integrated.bin", integrated) pairs: list[dict[str, Any]] = list(cohp_data.get("pairs", [])) meta = { "method": "cohp", "fermi_energy_ev": float(cohp_data.get("fermi_energy_ev", 0.0)), "sigma_ev": float(cohp_data.get("sigma_ev", 0.27)), "n_spin": n_spin, "pairs": pairs, } meta_json = safe_json_bytes(meta) meta_path = "dos/cohp_meta.json" zf.writestr(meta_path, meta_json) section: dict[str, Any] = { "id": "cohp0", "kind": "dos.cohp", "members": { "energies": energies_member, "projections": proj_member, "integrated": integ_member, "meta": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, }, "energies_units": str(cohp_data.get("energies_units", "eV")), "n_spin": n_spin, "fermi_energy_ev": float(cohp_data.get("fermi_energy_ev", 0.0)), "sigma_ev": float(cohp_data.get("sigma_ev", 0.27)), "pairs": pairs, } sections.append(section) # -- equation_of_state ----------------------------------------------------- def _write_equation_of_state_section( zf: zipfile.ZipFile, eos_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``equation_of_state`` section (QVF spec Sec. 4.14). ``eos_data`` is a dict with keys: * ``volumes`` -- float64 `[n_points]`, unit-cell volumes * ``energies`` -- float64 `[n_points]`, total energies * ``fit`` -- dict with ``model``, ``V0``, ``E0``, ``B0``, ``B0_prime``, ``energy_unit``, ``volume_unit``, ``pressure_unit``, ``residual_rms``, and optional ``pressures_gpa``. """ volumes = np.asarray(eos_data["volumes"], dtype=np.float64) energies = np.asarray(eos_data["energies"], dtype=np.float64) if volumes.shape != energies.shape: raise ValueError( "equation_of_state: volumes and energies must have the same shape; " f"got {volumes.shape} vs {energies.shape}" ) vols_member = _write_binary_to_zip(zf, "eos/volumes.bin", volumes) energies_member = _write_binary_to_zip(zf, "eos/energies.bin", energies) fit = dict(eos_data.get("fit", {})) fit_json = safe_json_bytes(fit) fit_path = "eos/fit.json" zf.writestr(fit_path, fit_json) section: dict[str, Any] = { "id": "eos", "kind": "equation_of_state", "members": { "volumes": vols_member, "energies": energies_member, "fit": { "path": fit_path, "format": "json", "sha256": _sha256_hex(fit_json), }, }, } sections.append(section) # -- fermi_surface -------------------------------------------------------- def _write_fermi_surface_section( zf: zipfile.ZipFile, fermi_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``fermi_surface`` section (QVF spec Sec. 4.12). ``fermi_data`` is a dict with keys: * ``nk1``, ``nk2``, ``nk3`` -- int, Monkhorst-Pack mesh dimensions * ``energies`` -- float64 `[nk1, nk2, nk3, n_bands]`, signed distance from E_F in eV for bands near the Fermi level * ``fermi_energy_ev`` -- float, absolute Fermi level in eV * ``band_indices`` -- list of int, which bands are included * ``lattice_vectors`` -- 3x3 array, real-space lattice in Angstrom * ``n_spin`` -- int, 1 (default) or 2 """ nk1 = int(fermi_data["nk1"]) nk2 = int(fermi_data["nk2"]) nk3 = int(fermi_data["nk3"]) energies = np.asarray(fermi_data["energies"], dtype=np.float64) if energies.ndim != 4: raise ValueError( "fermi_surface: energies must be 4-D [nk1, nk2, nk3, n_bands]; " f"got shape {energies.shape}" ) fermi_ev = float(fermi_data.get("fermi_energy_ev", 0.0)) band_indices = [int(b) for b in fermi_data.get("band_indices", [])] lattice = np.asarray(fermi_data["lattice_vectors"], dtype=np.float64) n_spin = int(fermi_data.get("n_spin", 1)) # Mesh descriptor (JSON). mesh = { "nk1": nk1, "nk2": nk2, "nk3": nk3, "n_spin": n_spin, "fermi_energy_ev": fermi_ev, "band_indices": band_indices, "lattice_vectors": lattice.tolist(), } mesh_json = safe_json_bytes(mesh) mesh_path = "fermi/mesh.json" zf.writestr(mesh_path, mesh_json) # Binary payload. energies_member = _write_binary_to_zip(zf, "fermi/energies.bin", energies) section: dict[str, Any] = { "id": "fermi0", "kind": "fermi_surface", "members": { "mesh": { "path": mesh_path, "format": "json", "sha256": _sha256_hex(mesh_json), }, "energies": energies_member, }, } sections.append(section) # -- phonon_bands --------------------------------------------------------- def _write_phonon_bands_section( zf: zipfile.ZipFile, bands_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``phonon_bands`` section (QVF spec Sec. 4.13). ``bands_data`` is a dict with keys: * ``qpath`` -- dict with ``n_atoms``, ``n_modes``, ``has_eigenvectors`` (bool), ``segments`` (list of dicts with ``label_start``, ``label_end``, ``k_start``, ``k_end``, ``n_points``) * ``frequencies`` -- float64 `[n_qpts, n_modes]`, in cm^-1 * ``eigenvectors`` -- optional float64 `[n_qpts, n_modes, n_atoms, 3]` """ qpath = dict(bands_data["qpath"]) qpath_json = safe_json_bytes(qpath) qpath_path = "phonons/qpath.json" zf.writestr(qpath_path, qpath_json) freq = np.asarray(bands_data["frequencies"], dtype=np.float64) freq_member = _write_binary_to_zip(zf, "phonons/frequencies.bin", freq) members: dict[str, Any] = { "qpath": { "path": qpath_path, "format": "json", "sha256": _sha256_hex(qpath_json), }, "frequencies": freq_member, } eig = bands_data.get("eigenvectors") if eig is not None: eig_arr = np.asarray(eig, dtype=np.float64) members["eigenvectors"] = _write_binary_to_zip( zf, "phonons/eigenvectors.bin", eig_arr ) section: dict[str, Any] = { "id": "phonon_bands", "kind": "phonon_bands", "members": members, } sections.append(section) # -- phonon_dos ----------------------------------------------------------- def _write_phonon_dos_section( zf: zipfile.ZipFile, dos_data: dict[str, Any], sections: list[dict[str, Any]], ) -> None: """Write the ``phonon_dos`` section (QVF spec Sec. 4.13). ``dos_data`` is a dict with keys: * ``frequencies`` -- float64 `[n_points]`, frequency grid in cm^-1 * ``dos`` -- float64 `[n_points]`, phonon DOS in states / cm^-1 * ``meta`` -- optional dict with ``smearing``, ``smearing_type``, ``n_atoms``, ``n_modes`` * ``projected`` -- optional float64 `[n_atoms, n_points]` """ freq = np.asarray(dos_data["frequencies"], dtype=np.float64) dos_arr = np.asarray(dos_data["dos"], dtype=np.float64) if freq.shape != dos_arr.shape: raise ValueError( "phonon_dos: frequencies and dos must have the same shape; " f"got {freq.shape} vs {dos_arr.shape}" ) meta = dict(dos_data.get("meta", {})) meta_json = safe_json_bytes(meta) meta_path = "phonons/dos_meta.json" zf.writestr(meta_path, meta_json) freq_member = _write_binary_to_zip(zf, "phonons/dos_freq.bin", freq) dos_member = _write_binary_to_zip(zf, "phonons/dos_total.bin", dos_arr) members: dict[str, Any] = { "meta": { "path": meta_path, "format": "json", "sha256": _sha256_hex(meta_json), }, "frequencies": freq_member, "dos": dos_member, } proj = dos_data.get("projected") if proj is not None: proj_arr = np.asarray(proj, dtype=np.float64) members["projected"] = _write_binary_to_zip( zf, "phonons/dos_projected.bin", proj_arr ) section: dict[str, Any] = { "id": "phonon_dos", "kind": "phonon_dos", "members": members, } sections.append(section) # -- scf history extraction ----------------------------------------------- def scf_history_from_result(result: Any) -> list[dict[str, Any]] | None: """Extract a QVF ``scf_history`` payload from an SCF result object. Reads ``result.scf_trace`` -- the canonical per-iteration record both the molecular (C++ ``SolverResult``) and periodic (e.g. ``GpwScfResult``) solvers carry. Molecular trace steps are objects with ``.iter`` / ``.energy`` / ``.delta_e`` attributes; periodic steps are dicts with the same keys. Returns a list of records shaped for :func:`_write_scf_history_section` (``iter`` / ``energy_eh`` / ``delta_e``), or ``None`` when no per-iteration trace is available. ``delta_e`` is omitted for the first iteration, mirroring the SCF log / perf-tracker convention (no meaningful energy delta exists before the second cycle). """ trace = getattr(result, "scf_trace", None) if not trace: return None history: list[dict[str, Any]] = [] for step in trace: if isinstance(step, dict): it = step.get("iter") energy = step.get("energy") delta_e = step.get("delta_e") else: it = getattr(step, "iter", None) energy = getattr(step, "energy", None) delta_e = getattr(step, "delta_e", None) if it is None or energy is None: continue it = int(it) record: dict[str, Any] = {"iter": it, "energy_eh": float(energy)} if delta_e is not None and it > 1: record["delta_e"] = float(delta_e) history.append(record) return history or None # -- provenance ----------------------------------------------------------- def _build_provenance(context: dict[str, Any]) -> dict[str, Any]: """Build the ``provenance`` block for the manifest root.""" result = context.get("result") provenance: dict[str, Any] = {} if context.get("method"): provenance["method"] = str(context["method"]) if context.get("functional"): provenance["functional"] = str(context["functional"]) if context.get("basis"): provenance["basis"] = str(context["basis"]) if context.get("jk_method"): provenance["jk_method"] = str(context["jk_method"]) if context.get("jk_method_resolved"): provenance["jk_method_resolved"] = str(context["jk_method_resolved"]) if context.get("jk_method_executed"): provenance["jk_method_executed"] = str(context["jk_method_executed"]) if "dft_plus_u" in context: provenance["dft_plus_u"] = bool(context["dft_plus_u"]) if context.get("dft_plus_u_route"): provenance["dft_plus_u_route"] = str(context["dft_plus_u_route"]) mol_or_sys = context.get("molecule") or context.get("system") if mol_or_sys is not None: provenance["charge"] = int(getattr(mol_or_sys, "charge", 0)) provenance["multiplicity"] = int( getattr(mol_or_sys, "multiplicity", 1), ) n_elec = getattr(mol_or_sys, "n_electrons", None) if n_elec is not None: if callable(n_elec): n_elec = n_elec() provenance["n_electrons"] = int(n_elec) if result is not None: provenance["scf_converged"] = bool( getattr(result, "converged", False), ) # SCF iteration count. Both the molecular ``SolverResult`` and # the periodic solver results carry ``n_iter``; fall back to the # length of the per-iteration trace when the count isn't a field. n_iter = getattr(result, "n_iter", None) if n_iter is None: _trace = getattr(result, "scf_trace", None) if _trace: n_iter = len(_trace) if n_iter is not None: provenance["n_scf_iterations"] = int(n_iter) energy = getattr(result, "energy", None) if energy is not None: provenance["scf_energy"] = { "value": float(energy), "units": "Eh", } fermi = getattr(result, "fermi_energy", None) if fermi is not None: provenance["fermi_energy"] = { "value": float(fermi), "units": "Eh", } if context.get("wall_seconds") is not None: provenance["wall_seconds"] = float(context["wall_seconds"]) try: import socket provenance["hostname"] = socket.gethostname() except Exception: provenance["hostname"] = "unknown" if mol_or_sys is not None: # The bound PeriodicSystem exposes ``dim`` in {1,2,3}; molecular # stubs may carry ``dimensionality``; a bare Molecule has neither # (0-D). Reading only ``dimensionality`` used to record 0 for # every real periodic system. dim = getattr(mol_or_sys, "dim", None) if dim is None: dim = getattr(mol_or_sys, "dimensionality", None) if dim is None: dim = 3 if context.get("system") is not None else 0 provenance["dimensionality"] = int(dim) # --- streaming / checkpoint fields (Ask-1 / Ask-2) ------------------- # ``run_status`` and ``checkpoint`` let a live consumer (vibe-view) # tell a running snapshot from the settled final one, and order a # sequence of snapshots by ``checkpoint.seq`` without diffing bytes. # Both live under the ``provenance`` block, which the schema declares # ``additionalProperties: true`` -- so these validate against the # current QVF v1 manifest schema with no v3 bump required. run_status = context.get("run_status") if run_status is not None: rs = str(run_status) if rs not in ("running", "converged", "failed"): raise ValueError( "write_qvf: run_status must be one of " f"'running' | 'converged' | 'failed', got {rs!r}." ) provenance["run_status"] = rs checkpoint = context.get("checkpoint") if checkpoint is not None: # Copy defensively; the caller's dict must not be mutated and any # non-JSON-native values (e.g. numpy floats) are coerced here. ckpt: dict[str, Any] = {} if "seq" in checkpoint: ckpt["seq"] = int(checkpoint["seq"]) if "wall_time_s" in checkpoint: ckpt["wall_time_s"] = float(checkpoint["wall_time_s"]) if checkpoint.get("written_at") is not None: ckpt["written_at"] = str(checkpoint["written_at"]) # Optional running-state hints (additive; provenance is open). if checkpoint.get("scf_iteration") is not None: ckpt["scf_iteration"] = int(checkpoint["scf_iteration"]) if checkpoint.get("energy_eh") is not None: ckpt["energy_eh"] = float(checkpoint["energy_eh"]) provenance["checkpoint"] = ckpt return provenance # --------------------------------------------------------------------------- # Validation tool # --------------------------------------------------------------------------- # # validate_qvf() drives off the canonical schema at # qvf_manifest.schema.json -- the SSOT. The hand-rolled per-kind checks # that used to live in this function are gone; what remains is the # cross-cutting work that JSON Schema can't express: sha256 of every # member matches the bytes on disk, every referenced zip path exists, # binary dtype/shape add up to the byte count on disk, operand_a / # operand_b / trajectory_ref cross-references resolve to existing # sections. _SCHEMA_PATH_V1 = Path(__file__).parent / "qvf_manifest.schema.json" _SCHEMA_PATH_V2 = Path(__file__).parent / "qvf_manifest_v2.schema.json" # Back-compat alias for any caller still reaching for ``_SCHEMA_PATH``; # defaults to the v1 path (the only version that existed before v2). _SCHEMA_PATH = _SCHEMA_PATH_V1 _SCHEMA_CACHE: dict[int, dict[str, Any]] = {} # --------------------------------------------------------------------------- # v2 = v1 + the periodic-reaction-path delta, derived — never forked. # # v2 began life as a hand-copied snapshot of v1 and then fell behind it: # every section kind added to v1 afterwards (bond_orders, spectra.epr, # volume.rdg, volume.potential, basis.ao, fermi_surface, phonon.bands, # phonon.dos, equation_of_state, topology.qtaim), the root # thermochemistry / dipole_moment / constraints / extensions blocks, the # Section ``critical`` flag, and the ``bands`` fat-band ``projections`` # member were all absent from the fork. Because the manifest root sets # ``additionalProperties: false`` and ``Section`` is a closed ``oneOf``, # a periodic reaction.path -- which forces qvf_version=2 -- carrying any # of them tripped the write-time validation gate in :func:`write_qvf`. # # Deriving v2 from v1 makes the delta below the *only* way the two can # differ, so a section added to v1 tomorrow lands in v2 for free. # ``scripts/gen_qvf_v2_schema.py`` renders this derivation to # ``qvf_manifest_v2.schema.json`` for external validators; # ``tests/test_qvf_v2_schema_drift.py`` pins that file to the derivation. # --------------------------------------------------------------------------- _V2_TITLE = "QVF Manifest (v2)" _V2_DESCRIPTION = ( "DEPRECATED. Canonical JSON Schema for manifest.json inside a .qvf " "archive that declares qvf_version=2. GENERATED FILE -- do not " "hand-edit; regenerate with scripts/gen_qvf_v2_schema.py. " "qvf_version=2 was WITHDRAWN by the governance ruling of 2026-07-10 " "(qvf-writer/GOVERNANCE.md, Version history): its only delta over v1 " "was the optional SectionReactionPath `lattice` member, which is the " "additive case and must not bump the version. `lattice` is now an " "optional member of v1, and a periodic reaction path is detected by " "its presence. This schema is therefore v1 with a different " "qvf_version const, retained only so that archives written by " "vibe-qc v0.10.0-v0.15.x (which are v1-compatible) still validate and " "so anything resolving its $id keeps working. Producers MUST NOT emit " "qvf_version=2; consumers SHOULD accept it and treat it as 1." ) _V2_QVF_VERSION_DESCRIPTION = ( "QVF format version. DEPRECATED: 2 was withdrawn (see the schema " "description); it is v1 in every respect but this const. Producers " "must emit 1." ) def _derive_v2_schema(v1_schema: dict[str, Any]) -> dict[str, Any]: """Return the (deprecated) v2 manifest schema derived from ``v1_schema``. ``qvf_version: 2`` was withdrawn by the governance ruling of 2026-07-10 (``qvf-writer/GOVERNANCE.md``, "Version history"): its only delta over v1 was the optional ``SectionReactionPath.lattice`` member, which is the *additive* case and must not bump the version. ``lattice`` now lives in the v1 schema, so the v2 delta has collapsed to pure identity metadata: ``$id``, ``title``, ``description``, and ``properties.qvf_version`` (const + description). Everything else is v1 verbatim. v2 is therefore v1 with a different ``qvf_version`` const. It is retained only so archives written by vibe-qc v0.10.0-v0.15.x still validate, and so anything resolving the v2 ``$id`` keeps working. Producers must emit 1. """ schema = copy.deepcopy(v1_schema) schema["$id"] = _SCHEMA_URI_V2 schema["title"] = _V2_TITLE schema["description"] = _V2_DESCRIPTION version_prop = schema["properties"]["qvf_version"] version_prop["const"] = QVF_FORMAT_VERSION_V2 version_prop["description"] = _V2_QVF_VERSION_DESCRIPTION return schema def _load_canonical_schema(qvf_version: int = 1) -> dict[str, Any]: """Return the canonical QVF manifest schema for ``qvf_version``. v1 is loaded from ``qvf_manifest.schema.json``, the single source of truth. v2 is *derived* from it by :func:`_derive_v2_schema` rather than read from disk, so the on-disk v2 copy can never drift out from under the validator. v2 extends only ``SectionReactionPath``, with an optional ``lattice`` binary member + ``dim`` carried in metadata, so that periodic reaction paths (slabs, surfaces, NEB) round-trip cleanly through vibe-view. Every other section is v1's. """ cached = _SCHEMA_CACHE.get(qvf_version) if cached is not None: return cached if qvf_version not in (1, 2): raise ValueError(f"unknown qvf_version {qvf_version!r}; supported: 1, 2") with open(_SCHEMA_PATH_V1, encoding="utf-8") as f: schema = json.load(f) _SCHEMA_CACHE[1] = schema if qvf_version == 2: schema = _derive_v2_schema(schema) _SCHEMA_CACHE[2] = schema return schema def validate_qvf( source: "os.PathLike | str | zipfile.ZipFile", ) -> dict[str, Any]: """Validate a QVF against the canonical SSOT schema. ``source`` may be either a filesystem path to a ``.qvf`` file or an already-open :class:`zipfile.ZipFile`. The latter form lets :func:`qvf_bytes` validate an in-memory archive without round-tripping through disk. Returns a dict with keys ``valid`` (bool), ``summary`` (list of per-section result strs), and ``errors`` (list of error strs). Checks performed: * The archive is a valid zip, no member exceeds the zip-bomb cap. * ``manifest.json`` exists and parses as JSON. * The manifest validates against :data:`_SCHEMA_PATH` (the canonical schema) -- this catches per-kind member shape, dtype, format, and unknown kinds. * Every member's declared zip path exists in the archive. * Every member's declared sha256 matches the bytes on disk. * Every binary member's ``len(bytes) == np.dtype(dtype).itemsize * product(shape)`` (no silent under/over-sized buffer). * On ``volume.difference``: both ``operand_a`` and ``operand_b`` (if present) resolve to section ids that exist in the archive. * On ``reaction.waypoints``: ``trajectory_ref`` resolves to a section in the archive whose kind is ``trajectory``. Sections whose kind is in :data:`_RESERVED_KINDS` are not shape-validated (no schema branch yet); their file refs are still checked. Vendor (``x_*``) sections must conform to the ``SectionVendor`` schema branch (members must be valid Json/Binary members) but the shape of those members is unconstrained. """ report: dict[str, Any] = {"valid": True, "summary": [], "errors": []} summary: list[str] = [] errors: list[str] = [] # Accept either a path-like (open + close here) or an already-open # ZipFile (do not close -- caller owns the handle). own_zf = True if isinstance(source, zipfile.ZipFile): zf = source own_zf = False else: path = Path(os.fspath(source)) if not path.is_file(): report.update(valid=False, errors=[f"file not found: {path}"]) return report try: zf = zipfile.ZipFile(path, "r") except zipfile.BadZipFile as exc: report.update(valid=False, errors=[f"not a valid zip file: {exc}"]) return report # --- zip-bomb guard: per-member uncompressed-size ceiling ---------- # Aligned with the writer's _MAX_VOXELS payload guard (see module # constants) so a valid write_qvf output cannot subsequently fail # validation as "too large". Compressed bytes on disk are typically # far smaller; this guard reads `file_size` which is the # uncompressed length encoded in the zip central directory. try: for _info in zf.infolist(): if _info.file_size > _MAX_MEMBER_UNCOMPRESSED_BYTES: errors.append( f"member {_info.filename!r}: uncompressed size " f"{_info.file_size:_d} bytes exceeds max " f"{_MAX_MEMBER_UNCOMPRESSED_BYTES:_d}; possible zip bomb" ) report["valid"] = False except Exception as exc: if own_zf: zf.close() report.update(valid=False, errors=[f"cannot read zip directory: {exc}"]) return report try: # --- manifest.json ------------------------------------------------- try: manifest_bytes = zf.read("manifest.json") except KeyError: report.update( valid=False, errors=errors + ["manifest.json missing from archive"] ) return report try: manifest = json.loads(manifest_bytes.decode("utf-8")) except (json.JSONDecodeError, UnicodeDecodeError) as exc: report.update( valid=False, errors=errors + [f"manifest.json is not valid JSON: {exc}"] ) return report # --- jsonschema validation against the canonical SSOT -------------- # # Sections with reserved kinds have no schema branch yet -- temporarily # remove them from the to-be-validated manifest and shape-skip them. # Their file refs are still verified below. try: import jsonschema # noqa: F401 (soft import; required dep) except ImportError: errors.append( "validate_qvf requires the 'jsonschema' package " "(`pip install jsonschema`) -- the canonical schema " "cannot be enforced without it." ) report.update(valid=False, errors=errors) return report # Dispatch by the manifest's qvf_version: v1 archives validate # against the v1 schema; v2 archives (periodic reaction.path) # against the v2 schema. An unknown version falls through to # the v1 schema's `const: 1` rule, which produces the natural # "qvf_version was X, expected 1" validation error. try: manifest_qvf_version = int(manifest.get("qvf_version", 1)) except (TypeError, ValueError): manifest_qvf_version = 1 schema_version = manifest_qvf_version if manifest_qvf_version in (1, 2) else 1 schema = _load_canonical_schema(schema_version) # Build a "check manifest" that mirrors the input but with # reserved-kind sections pulled out (they have no schema branch # yet; the validator still file-ref-checks them below). If # `sections` was absent from the input we leave it absent so # the schema's required-field rule fires. check_manifest = dict(manifest) raw_sections = manifest.get("sections") if isinstance(raw_sections, list): shape_check_secs = [ s for s in raw_sections if not (isinstance(s, dict) and s.get("kind") in _RESERVED_KINDS) ] check_manifest["sections"] = shape_check_secs # else: `sections` is missing or not a list -- leave the manifest # as-is so the schema validator reports it. validator = jsonschema.Draft202012Validator(schema) schema_errors = sorted( validator.iter_errors(check_manifest), key=lambda e: list(e.absolute_path) ) for e in schema_errors: loc = "/".join(str(p) for p in e.absolute_path) or "" errors.append(f"schema: {loc}: {e.message}") report["valid"] = False fmt_ver = manifest.get("qvf_version", "?") n_secs = len(raw_sections) if isinstance(raw_sections, list) else 0 summary.append( f"manifest.json: qvf_version={fmt_ver}, " f"{n_secs} section(s); " f"{'schema OK' if not schema_errors else f'{len(schema_errors)} schema error(s)'}" ) # If sections is missing/malformed, the per-section loop is a # no-op -- the schema error already covered it. if not isinstance(raw_sections, list): return report # --- critical-flag enforcement (QVF spec 5.4-5.5) ----------- # Per-section critical: warn if a section is marked critical # but its kind is reserved (no schema branch) or an unknown # non-vendor kind -- no consumer can render it, so a critical # flag is a contract violation. if isinstance(raw_sections, list): for sec in raw_sections: if not isinstance(sec, dict): continue if sec.get("critical") is True: kind = sec.get("kind", "") sec_id = sec.get("id", "?") if kind in _RESERVED_KINDS: errors.append( f"section {sec_id} ({kind}): critical=true but " f"kind is reserved (no consumer renders it)" ) report["valid"] = False elif not _is_vendor_kind(kind) and kind not in _IMPLEMENTED_KINDS: errors.append( f"section {sec_id} ({kind}): critical=true but " f"kind is not a known canonical or vendor kind" ) report["valid"] = False # Root extensions block: warn if a critical extension is declared # but no sections use its vendor namespace. ext_block = manifest.get("extensions") if isinstance(ext_block, dict) and isinstance(raw_sections, list): section_kinds = { s["kind"] for s in raw_sections if isinstance(s, dict) and isinstance(s.get("kind"), str) } for ns, ext_info in ext_block.items(): if not isinstance(ext_info, dict): continue if ext_info.get("critical") is True: prefix = ns if ns.endswith(".") else ns + "." used = any(k.startswith(prefix) for k in section_kinds) if not used: summary.append( f"extensions: {ns} declares critical=true but " f"no section uses its namespace" ) # --- per-section cross-checks -------------------------------------- zip_names = set(zf.namelist()) section_ids = {s.get("id") for s in raw_sections if isinstance(s, dict)} seen_section_ids: set[str] = set() for sec in raw_sections: if not isinstance(sec, dict): continue sec_id = sec.get("id") if not isinstance(sec_id, str): continue if sec_id in seen_section_ids: errors.append(f"manifest: duplicate section id {sec_id!r}") report["valid"] = False seen_section_ids.add(sec_id) for sec in raw_sections: if not isinstance(sec, dict): continue sec_id = sec.get("id", "?") kind = sec.get("kind", "") if _is_vendor_kind(kind): file_refs = _collect_file_refs(sec) payload_ok = _validate_file_refs( zf, file_refs, zip_names, sec_id, errors ) summary.append( f"section {sec_id} ({kind}): vendor " f"({'refs OK' if payload_ok else 'refs FAILED'}); " "member shape not validated" ) continue if kind in _RESERVED_KINDS: file_refs = _collect_file_refs(sec) payload_ok = _validate_file_refs( zf, file_refs, zip_names, sec_id, errors ) summary.append( f"section {sec_id} ({kind}): reserved, not yet " f"implemented; refs {'OK' if payload_ok else 'FAILED'}" ) continue # Implemented (canonical) kind: file refs + dtype/shape match. file_refs = _collect_file_refs(sec) refs_ok = _validate_file_refs(zf, file_refs, zip_names, sec_id, errors) sizes_ok = _validate_binary_shapes(zf, file_refs, sec_id, errors) # Cross-reference checks for the kinds that carry them. xref_ok = True if kind == "volume.difference": for key in ("operand_a", "operand_b"): ref = sec.get(key) if ref is not None and ref not in section_ids: errors.append( f"section {sec_id} (volume.difference): {key}=" f"{ref!r} does not name a section in this archive" ) xref_ok = False if kind == "reaction.waypoints": ref = sec.get("trajectory_ref") target_sec = next( ( s for s in raw_sections if isinstance(s, dict) and s.get("id") == ref ), None, ) if target_sec is None: errors.append( f"section {sec_id} (reaction.waypoints): " f"trajectory_ref={ref!r} does not name a section " "in this archive" ) xref_ok = False elif target_sec.get("kind") != "trajectory": errors.append( f"section {sec_id} (reaction.waypoints): " f"trajectory_ref={ref!r} resolves to a section " f"of kind {target_sec.get('kind')!r}, expected " "'trajectory'" ) xref_ok = False status_bits = [] if refs_ok and sizes_ok and xref_ok: status_bits.append("OK") else: if not refs_ok: status_bits.append("refs FAILED") if not sizes_ok: status_bits.append("sizes FAILED") if not xref_ok: status_bits.append("xref FAILED") summary.append(f"section {sec_id} ({kind}): {', '.join(status_bits)}") if not (refs_ok and sizes_ok and xref_ok): report["valid"] = False finally: if own_zf: zf.close() report["summary"] = summary report["errors"] = errors if errors: report["valid"] = False return report def _collect_file_refs(section: dict[str, Any]) -> list[dict[str, Any]]: """Walk a section dict and return every file-like sub-dict that has a ``path`` key and (optionally) a ``sha256`` key.""" refs: list[dict[str, Any]] = [] def _recurse(obj: Any) -> None: if isinstance(obj, dict): if "path" in obj and isinstance(obj["path"], str): refs.append(obj) for v in obj.values(): _recurse(v) elif isinstance(obj, list): for item in obj: _recurse(item) _recurse(section) return refs def _validate_binary_shapes( zf: zipfile.ZipFile, file_refs: list[dict[str, Any]], sec_id: str, errors: list[str], ) -> bool: """Check that every binary member's dtype x shape matches the byte length of the member on disk. Catches dtype-shape lies the schema cannot express. Returns True if all pass, False otherwise. """ ok = True for ref in file_refs: if ref.get("format") != "binary": continue dtype = ref.get("dtype") shape = ref.get("shape") if dtype is None or shape is None: # schema-required fields were missing; the schema error # already covers this. continue try: itemsize = int(np.dtype(dtype).itemsize) except TypeError: errors.append( f"section {sec_id}: member {ref['path']!r} declares " f"dtype={dtype!r} which numpy cannot resolve" ) ok = False continue expected = itemsize for d in shape: expected *= int(d) try: actual = zf.getinfo(ref["path"]).file_size except KeyError: # missing-file error already reported by _validate_file_refs. continue if actual != expected: errors.append( f"section {sec_id}: member {ref['path']!r} byte length " f"{actual} != dtype({dtype}).itemsize x prod({shape}) " f"= {expected}" ) ok = False return ok def _validate_file_refs( zf: zipfile.ZipFile, file_refs: list[dict[str, Any]], zip_names: set[str], sec_id: str, errors: list[str], ) -> bool: """Validate sha256 and existence for a list of file reference dicts. Returns True if all pass, False otherwise.""" ok = True for ref in file_refs: path_in_zip = ref["path"] fmt = ref.get("format", "binary") if fmt not in ("json", "binary"): errors.append( f"section {sec_id}: member {path_in_zip!r} has invalid " f"format {fmt!r} (expected 'json' or 'binary')", ) ok = False continue if path_in_zip not in zip_names: errors.append( f"section {sec_id}: file {path_in_zip!r} " "declared in manifest but missing from zip", ) ok = False continue expected_sha = ref.get("sha256") data: bytes | None = None if expected_sha is not None or fmt == "json": try: data = zf.read(path_in_zip) except Exception as exc: errors.append( f"section {sec_id}: cannot read {path_in_zip!r}: {exc}", ) ok = False continue actual = _sha256_hex(data) if actual != expected_sha: errors.append( f"section {sec_id}: sha256 mismatch for {path_in_zip!r} " f"(expected {expected_sha[:12]}..., got {actual[:12]}...)", ) ok = False if fmt == "json" and data is not None: try: json.loads(data.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: errors.append( f"section {sec_id}: JSON member {path_in_zip!r} " f"does not parse: {exc}" ) ok = False return ok def _print_validate_report(report: dict[str, Any]) -> str: """Print-friendly summary of a :func:`validate_qvf` report.""" lines: list[str] = [] for s in report["summary"]: if "vendor" in s.lower(): prefix = "⚠" elif "OK" in s or "format " in s: prefix = "✓" else: prefix = "✗" lines.append(f"{prefix} {s}") for e in report["errors"]: lines.append(f"✗ ERROR: {e}") if report["valid"]: lines.append("\n✓ QVF file is valid.") else: lines.append("\n✗ QVF file has validation errors.") return "\n".join(lines) # --------------------------------------------------------------------------- # CLI entry point for qvf-validate # --------------------------------------------------------------------------- def _qvf_validate_cli() -> None: """CLI: ``python -m vibeqc.output.formats.qvf ``.""" import sys if len(sys.argv) < 2: print("usage: python -m vibeqc.output.formats.qvf ", file=sys.stderr) sys.exit(2) report = validate_qvf(sys.argv[1]) print(_print_validate_report(report)) sys.exit(0 if report["valid"] else 1) if __name__ == "__main__": _qvf_validate_cli()