Source code for vibesketch.document.spectra

"""Calculated IR, Raman and UV-Vis spectra of a run's result (#87).

A verified result keeps the vibrations and excited states its output lists
(:class:`~vibesketch.document.model.ResultVibrations`,
:class:`~vibesketch.document.model.ResultExcitedStates`). This module reads
them from an ORCA 6 output or a vibe-qc QVF result, and exports a spectrum's
line list (CSV or plain text) and its broadened curve (plain text) with the
units, source and grid in the file's header.

Broadening follows vibe-view's ``vibeview.spectrum.broaden_lines`` so the two
programs draw the same curve: unit-area Gaussian and Cauchy (Lorentzian)
densities with scale FWHM / sqrt(8 ln 2) and FWHM / 2,
https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.norm.html
https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.cauchy.html
Lines are broadened in energy (cm⁻¹ for vibrations, eV for excited states)
and the curve is then sampled on the chosen axis, so a wavelength axis is
never broadened in nm. A finite range truncates the tails; the curve is not
renormalised to hide that.

ORCA's table vocabulary: ORCA 6.1 manual, Frequencies and Excited states,
https://www.faccts.de/docs/orca/6.1/manual/contents/spectroscopyproperties/ir.html
https://www.faccts.de/docs/orca/6.1/manual/contents/spectroscopyproperties/uvvis.html
"""

from __future__ import annotations

import csv
import io
import math
import re
from typing import Any, Literal

import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel, ConfigDict

from vibesketch.document.model import (
    Document,
    Job,
    ResultExcitedStates,
    ResultMode,
    ResultState,
    ResultVibrations,
    Run,
)
from vibesketch.units import EV_NM, EV_TO_CM1, HARTREE_TO_EV

__all__ = [
    "DEFAULT_FWHM",
    "MAX_POINTS",
    "MEDIA_TYPES",
    "SPECTRA",
    "SpectrumExport",
    "SpectrumRefused",
    "broadened_spectrum",
    "export_spectrum",
    "spectrum_lines",
]

Spectrum = Literal["ir", "raman", "uvvis"]
Lineshape = Literal["gaussian", "lorentzian"]
Axis = Literal["cm-1", "eV", "nm"]

SPECTRA: tuple[Spectrum, ...] = ("ir", "raman", "uvvis")
#: vibe-view's default widths: 10 cm⁻¹ for vibrations, 0.2 eV for excited states.
DEFAULT_FWHM = {"ir": 10.0, "raman": 10.0, "uvvis": 0.2}
#: The most points a broadened curve may have, as in vibe-view.
MAX_POINTS = 1_000_000
_NUMBER = r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[EeDd][+-]?\d+)?"
#: Printed precision: ORCA prints wavenumbers with two decimals.
_SAME_WAVENUMBER = 0.005001
_NATIVE = {"ir": "cm-1", "raman": "cm-1", "uvvis": "eV"}
_INTENSITY = {"ir": "km/mol", "raman": "angstrom^4/amu", "uvvis": "oscillator strength"}
_DEFAULT_PADDING = {"cm-1": 100.0, "eV": 0.5}


[docs] class SpectrumRefused(ValueError): """A spectrum the result does not contain, or an export that cannot be made as asked."""
# --- reading --------------------------------------------------------------- def _float(text: str) -> float: value = float(text.replace("D", "E").replace("d", "e")) if not math.isfinite(value): raise SpectrumRefused(f"non-finite value {text!r} in a spectrum table") return value def _line(text: str, index: int) -> int: return text[:index].count("\n") + 1 def _table(text: str, title: str) -> tuple[str, int] | None: """The body of ORCA's table under the heading ``title``, and where it starts. A major heading is framed by dashed lines; the table ends at the next one. A repeated table means several calculations, which are refused. """ matches = list(re.finditer(r"^[ \t]*" + re.escape(title) + r"[ \t]*$", text, re.M)) if not matches: return None if len(matches) != 1: raise SpectrumRefused(f"the ORCA output prints {title} more than once") rest = text[matches[0].end() :].lstrip("\n") rest = re.sub(r"^-+[ \t]*\n", "", rest, count=1) end = re.search(r"\n[ \t]*[A-Z][A-Z /().0-9-]{3,}[ \t]*\n[ \t]*-{4,}", rest) return (rest[: end.start()] if end else rest), matches[0].start() def orca_vibrations(text: str, atoms: int) -> ResultVibrations | None: """The vibrations of an ORCA 6 frequency output, with its IR and Raman lines. ORCA prints all 3N wavenumbers from 0, the projected translations and rotations as zero, and an imaginary mode as a negative wavenumber. """ found = _table(text, "VIBRATIONAL FREQUENCIES") if found is None: return None body, start = found rows = re.findall(rf"^\s*(\d+):\s*({_NUMBER})\s*cm\*\*-1", body, re.M) if [int(i) for i, _ in rows] != list(range(3 * atoms)): raise SpectrumRefused( f"the ORCA output lists {len(rows)} vibrational frequencies, not 3N = {3 * atoms}" ) wavenumbers = [_float(w) for _, w in rows] first = re.findall(r"The first frequency considered to be a vibration is\s+(\d+)", text) if len(set(first)) > 1: raise SpectrumRefused("the ORCA output names more than one first vibration") vibrations = ( range(int(first[0]), 3 * atoms) if first else [i for i, w in enumerate(wavenumbers) if w != 0.0] ) columns: dict[str, dict[int, float]] = {} for title, key, column, width in ( ("IR SPECTRUM", "ir", 2, None), ("RAMAN SPECTRUM", "raman", 1, 3), ): table = _table(text, title) if table is None: continue if key == "ir" and "km/mol" not in table[0]: raise SpectrumRefused("the ORCA IR table does not give intensities in km/mol") values: dict[int, float] = {} for match in re.finditer(r"^\s*(\d+):\s*(.*)$", table[0], re.M): mode = int(match[1]) numbers = [_float(v) for v in re.findall(_NUMBER, match[2])] if mode not in vibrations or mode in values or len(numbers) <= column: raise SpectrumRefused(f"the ORCA {title} has an unreadable row for mode {mode}") if width is not None and len(numbers) != width: raise SpectrumRefused(f"the ORCA {title} row for mode {mode} is not 3 values") if abs(numbers[0] - wavenumbers[mode]) > _SAME_WAVENUMBER: raise SpectrumRefused(f"the ORCA {title} disagrees with mode {mode}'s wavenumber") values[mode] = numbers[column] if not values: raise SpectrumRefused(f"the ORCA {title} is empty") columns[key] = values if not vibrations: return None return ResultVibrations( source=f"line {_line(text, start)}: vibrational frequencies", modes=[ ResultMode( mode=mode, wavenumber=wavenumbers[mode], ir_intensity=columns.get("ir", {}).get(mode), raman_activity=columns.get("raman", {}).get(mode), ) for mode in vibrations ], ) def orca_excited_states(text: str) -> ResultExcitedStates | None: """The absorption lines ORCA 6 prints from the transition electric dipoles.""" title = "ABSORPTION SPECTRUM VIA TRANSITION ELECTRIC DIPOLE MOMENTS" found = _table(text, title) if found is None: return None body, start = found if "(eV)" not in body or "(nm)" not in body: raise SpectrumRefused("the ORCA absorption table does not give energies in eV") states: list[ResultState] = [] for line in body.splitlines(): match = re.match(r"\s*\S+\s*->\s*(\S+)\s+(.*)", line) if not match: continue numbers = [_float(v) for v in match[2].split()] if len(numbers) < 4: raise SpectrumRefused(f"the ORCA absorption row for state {match[1]} is incomplete") energy, _, wavelength, strength = numbers[:4] if energy <= 0 or abs(EV_NM / energy - wavelength) > 0.051: raise SpectrumRefused(f"state {match[1]}'s energy and wavelength disagree") states.append(ResultState(state=match[1], energy=energy, oscillator_strength=strength)) if not states: raise SpectrumRefused("the ORCA absorption table lists no states in ORCA 6's layout") methods = re.findall( r"^\s*((?:TD-DFT|TDDFT|CIS|TDA|ROCIS|STEOM|EOM)[^\n]*?)\s*$", text[:start], re.M ) return ResultExcitedStates( source=f"line {_line(text, start)}: absorption via transition electric dipoles", method=methods[-1] if methods else "excited states as ORCA printed them", states=states, ) def _numbers(value: Any, what: str) -> list[float]: if not isinstance(value, list) or not all( isinstance(v, (int, float)) and not isinstance(v, bool) for v in value ): raise SpectrumRefused(f"QVF {what} must be a list of numbers") return [_float(str(v)) for v in value] def _cm1(payload: dict[str, Any]) -> bool: """Whether a QVF spectrum's x values are wavenumbers in cm⁻¹.""" units = payload.get("units") legacy = units.get("frequency") if isinstance(units, dict) else None return payload.get("x_units", "cm-1") == "cm-1" and legacy in (None, "cm^-1", "cm-1") def qvf_vibrations( metadata: dict[str, Any] | None, spectra: dict[str, dict[str, Any]], section: str ) -> ResultVibrations | None: """The vibrations of a QVF ``vibrations`` section, or of its IR/Raman spectra alone. Intensities come from the metadata's per-mode arrays, or from a ``spectra.ir``/``spectra.raman`` section listing the same wavenumbers. """ if metadata is None and not spectra: return None if metadata is not None: wavenumbers = _numbers(metadata.get("frequencies"), "vibration frequencies") else: first = spectra.get("ir") or spectra["raman"] wavenumbers = _numbers(first.get("frequencies"), "spectrum frequencies") columns: dict[str, list[float]] = {} for key, field in (("ir", "ir_intensities"), ("raman", "raman_activities")): if metadata is not None and field in metadata: columns[key] = _numbers(metadata[field], field) elif key in spectra: payload = spectra[key] if not _cm1(payload): raise SpectrumRefused(f"the QVF {key} spectrum is not in cm-1") lines = _numbers(payload.get("frequencies"), f"{key} frequencies") if len(lines) != len(wavenumbers) or any( abs(a - b) > 1e-6 for a, b in zip(lines, wavenumbers, strict=True) ): raise SpectrumRefused(f"the QVF {key} spectrum does not list the vibrations") columns[key] = _numbers(payload.get("intensities"), f"{key} intensities") if key in columns and len(columns[key]) != len(wavenumbers): raise SpectrumRefused(f"the QVF {key} values do not match the vibrations") modes = [ ResultMode( mode=i, wavenumber=w, ir_intensity=columns["ir"][i] if "ir" in columns else None, raman_activity=columns["raman"][i] if "raman" in columns else None, ) for i, w in enumerate(wavenumbers) if w != 0.0 ] return ResultVibrations(source=f"section {section}", modes=modes) if modes else None _EV_PER = {"ev": 1.0, "hartree": HARTREE_TO_EV, "cm-1": 1 / EV_TO_CM1} def qvf_excited_states(payload: dict[str, Any], section: str) -> ResultExcitedStates: """The states of a QVF ``excited_states`` section, energies converted to eV.""" meta = payload.get("meta") or {} declared = (meta.get("units") or {}).get("excitation_energy") states: list[ResultState] = [] for entry in payload.get("states") or []: energy = entry.get("excitation_energy") if isinstance(energy, dict): value, unit = energy.get("value"), energy.get("units") else: value, unit = energy, declared if not isinstance(value, (int, float)) or not isinstance(unit, str): raise SpectrumRefused(f"QVF state {entry.get('index')} has no energy with a unit") if unit.lower() not in _EV_PER: raise SpectrumRefused(f"QVF state {entry.get('index')}'s energy is in {unit!r}") name = f"S{entry['index']}" if entry.get("symmetry"): name += f" ({entry['symmetry']})" strength = entry.get("oscillator_strength") states.append( ResultState( state=name, energy=float(value) * _EV_PER[unit.lower()], oscillator_strength=None if strength is None else float(strength), ) ) if not states: raise SpectrumRefused("the QVF excited_states section lists no states") method = meta.get("method") return ResultExcitedStates( source=f"section {section}", method=method if isinstance(method, str) and method else "excited states", states=states, ) # --- export ---------------------------------------------------------------- def _result(document: Document, run_id: str) -> tuple[Job, Run]: for job in document.jobs: for run in job.runs or []: if run.id == run_id: if run.result is None: raise SpectrumRefused(f"run {run_id!r} has no imported result") return job, run raise SpectrumRefused(f"no run {run_id!r} in this document") def _lines(run: Run, spectrum: str) -> tuple[list[tuple[str, float, float | None]], str, str]: """``(label, position in the native unit, intensity or None)`` per line, method, source.""" result = run.result assert result is not None if spectrum not in SPECTRA: raise SpectrumRefused(f"spectrum must be one of {', '.join(SPECTRA)}, not {spectrum!r}") if spectrum == "uvvis": states = result.excited_states if states is None: raise SpectrumRefused(f"run {run.id!r}'s result lists no excited states") lines = [(s.state, s.energy, s.oscillator_strength) for s in states.states] return lines, states.method, states.source vibrations = result.vibrations if vibrations is None: raise SpectrumRefused(f"run {run.id!r}'s result lists no vibrations") field = "ir_intensity" if spectrum == "ir" else "raman_activity" rows = [(str(m.mode), m.wavenumber, getattr(m, field)) for m in vibrations.modes] if all(value is None for _, _, value in rows): what = "IR intensities" if spectrum == "ir" else "Raman activities" raise SpectrumRefused(f"run {run.id!r}'s result lists no {what}") return rows, f"harmonic vibrations ({vibrations.imaginary} imaginary)", vibrations.source def _number(value: float | None) -> str: return "" if value is None else repr(value) def _header( document: Document, job: Job, run: Run, spectrum: str, method: str, where: str ) -> list[str]: result = run.result assert result is not None species = document.species_by_id(job.species) return [ f"vibe-sketch {spectrum} spectrum of species {species.id}" + (f" ({species.label})" if species.label else ""), f"run: {run.id}; job: {job.id}; program: {job.program} {result.program_version}", f"source: {result.source_name}; sha256: {result.source_sha256}; {where}", f"method: {method}", ]
[docs] def spectrum_lines( document: Document, run_id: str, spectrum: str, *, format: Literal["csv", "text"] = "csv" ) -> str: """A spectrum's line list, with explicit units; an intensity not given is left blank. ``csv`` has one header row; ``text`` has ``#`` comment lines and whitespace-separated columns, a missing value written as ``-``. """ job, run = _result(document, run_id) lines, method, where = _lines(run, spectrum) native = _NATIVE[spectrum] if spectrum == "uvvis": columns = ["state", "energy_eV", "wavelength_nm", "oscillator_strength"] rows = [[label, repr(e), repr(EV_NM / e), _number(f)] for label, e, f in lines] else: unit = "km/mol" if spectrum == "ir" else "angstrom^4/amu" name = "intensity" if spectrum == "ir" else "activity" columns = ["mode", f"wavenumber_{native}", f"{name}_{unit}"] rows = [[label, repr(w), _number(v)] for label, w, v in lines] if format == "csv": out = io.StringIO() writer = csv.writer(out, lineterminator="\n") writer.writerow(columns) writer.writerows(rows) return out.getvalue() if format != "text": raise SpectrumRefused("a line list is csv or text") header = _header(document, job, run, spectrum, method, where) header.append("a missing value is written as -") body = [" ".join(cell or "-" for cell in row) for row in rows] return "".join(f"# {line}\n" for line in [*header, " ".join(columns)]) + "".join( f"{line}\n" for line in body )
def _to_native(value: float, unit: str, native: str) -> float: """An energy-like axis value in the spectrum's native unit.""" if unit == native: return value if unit == "nm": if value <= 0: raise SpectrumRefused("a wavelength must be positive") ev = EV_NM / value return ev if native == "eV" else ev * EV_TO_CM1 if unit == "eV" and native == "cm-1": return value * EV_TO_CM1 if unit == "cm-1" and native == "eV": return value / EV_TO_CM1 raise SpectrumRefused(f"cannot use {unit} with a {native} spectrum") def _broadened( centres: NDArray[np.float64], lines: list[tuple[float, float]], fwhm: float, lineshape: str ) -> NDArray[np.float64]: """The sum of unit-area lines of full width ``fwhm``, scaled by their intensities.""" total = np.zeros_like(centres) for position, strength in lines: x = (centres - position) / fwhm if lineshape == "gaussian": kernel = math.sqrt(4 * math.log(2) / math.pi) / fwhm * np.exp(-4 * math.log(2) * x * x) else: kernel = 2 / (math.pi * fwhm) / (1 + 4 * x * x) total += strength * kernel return total
[docs] def broadened_spectrum( document: Document, run_id: str, spectrum: str, *, lineshape: str = "lorentzian", fwhm: float | None = None, fwhm_unit: str | None = None, axis: str | None = None, start: float | None = None, stop: float | None = None, step: float | None = None, ) -> str: """A spectrum broadened with a Gaussian or Lorentzian line shape, as plain text. ``fwhm`` is in ``fwhm_unit`` (cm-1 or eV; the spectrum's own by default), ``axis`` is cm-1 for vibrations and eV, nm or cm-1 for excited states, and ``start``, ``stop`` and ``step`` give the grid in the axis unit. The curve is the sum of unit-area lines scaled by their intensities: its unit is the intensity's per native unit. Lines without an intensity are left out and counted in the header. """ job, run = _result(document, run_id) lines, method, where = _lines(run, spectrum) native = _NATIVE[spectrum] if lineshape not in ("gaussian", "lorentzian"): raise SpectrumRefused("the line shape is gaussian or lorentzian") axis = axis or native if axis not in (("cm-1",) if native == "cm-1" else ("eV", "nm", "cm-1")): raise SpectrumRefused(f"a {spectrum} spectrum's axis cannot be {axis}") width_unit = fwhm_unit or native if width_unit not in ("cm-1", "eV"): raise SpectrumRefused("the FWHM is given in cm-1 or eV, never in nm") width = DEFAULT_FWHM[spectrum] if fwhm is None else float(fwhm) if not math.isfinite(width) or width <= 0: raise SpectrumRefused("the FWHM must be finite and positive") native_width = _to_native(width, width_unit, native) used = [(position, value) for _, position, value in lines if value is not None] positions = [p for p, _ in used] if start is None or stop is None: padding = _DEFAULT_PADDING[native] low, high = min(positions) - padding, max(positions) + padding if native == "cm-1": low = max(0.0, low) if axis == "nm": low = max(low, 0.05) # a wavelength axis ends short of zero energy low, high = EV_NM / high, EV_NM / low elif axis != native: low, high = _to_native(low, native, axis), _to_native(high, native, axis) start = low if start is None else start stop = high if stop is None else stop if not all(math.isfinite(v) for v in (start, stop)) or stop <= start: raise SpectrumRefused("the grid's stop must exceed its start") if axis == "nm" and start <= 0: raise SpectrumRefused("a wavelength grid must start above 0 nm") if step is None: step = (stop - start) / 1999 if not math.isfinite(step) or step <= 0: raise SpectrumRefused("the grid step must be finite and positive") intervals = (stop - start) / step if intervals > MAX_POINTS - 1: raise SpectrumRefused(f"the grid would have more than {MAX_POINTS} points") grid = start + np.arange(int(math.floor(intervals + 1e-10)) + 1, dtype=np.float64) * step header = _header(document, job, run, spectrum, method, where) left_out = len(lines) - len(used) header += [ f"lines: {len(used)}" + (f" ({left_out} without an intensity left out)" if left_out else ""), f"line shape: {lineshape}; FWHM: {width!r} {width_unit}" f" ({native_width!r} {native}); broadened in {native}, unit-area lines", f"axis: {axis}; start: {start!r}; stop: {float(grid[-1])!r}; step: {step!r};" f" points: {len(grid)}", f"intensity: {_INTENSITY[spectrum]} per {native}", f"{'wavelength' if axis == 'nm' else 'energy' if axis == 'eV' else 'wavenumber'}_{axis}" " intensity", ] if axis == native: centres = grid elif axis == "nm": centres = (EV_NM / grid) * (1.0 if native == "eV" else EV_TO_CM1) else: centres = grid * (EV_TO_CM1 if native == "cm-1" else 1 / EV_TO_CM1) curve = _broadened(centres, used, native_width, lineshape) rows = (f"{x!r} {y!r}\n" for x, y in zip(grid.tolist(), curve.tolist(), strict=True)) return "".join(f"# {line}\n" for line in header) + "".join(rows)
[docs] class SpectrumExport(BaseModel): """What to export from a run's spectrum: its line list, or its broadened curve.""" model_config = ConfigDict(extra="forbid") run: str spectrum: Spectrum format: Literal["csv", "text", "broadened"] = "csv" lineshape: Lineshape = "lorentzian" fwhm: float | None = None fwhm_unit: Literal["cm-1", "eV"] | None = None axis: Axis | None = None start: float | None = None stop: float | None = None step: float | None = None
#: The media type of each export format. MEDIA_TYPES = { "csv": "text/csv; charset=utf-8", "text": "text/plain; charset=utf-8", "broadened": "text/plain; charset=utf-8", }
[docs] def export_spectrum(document: Document, request: SpectrumExport) -> str: """The export ``request`` asks for: a CSV or text line list, or a broadened curve.""" if request.format != "broadened": if request.model_fields_set & { "lineshape", "fwhm", "fwhm_unit", "axis", "start", "stop", "step", }: raise SpectrumRefused("line shape, width, axis and grid apply to a broadened spectrum") return spectrum_lines(document, request.run, request.spectrum, format=request.format) return broadened_spectrum( document, request.run, request.spectrum, lineshape=request.lineshape, fwhm=request.fwhm, fwhm_unit=request.fwhm_unit, axis=request.axis, start=request.start, stop=request.stop, step=request.step, )