diff --git a/pyproject.toml b/pyproject.toml index b97f2ddd..b38c0977 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -236,17 +236,4 @@ module = [ "q2mm.backends.mm.jax_md_engine", "q2mm.backends.mm.openmm", "q2mm.backends.mm.tinker", - "q2mm.io._helpers", - "q2mm.io.amber", - "q2mm.io.cmap", - "q2mm.io.gaussian", - "q2mm.io.jaguar", - "q2mm.io.macromodel", - "q2mm.io.mm3", - "q2mm.io.openmm", - "q2mm.io.tinker", - "q2mm.models.forcefield", - "q2mm.models.hessian", - "q2mm.models.molecule", - "q2mm.models.seminario", ] diff --git a/q2mm/io/_helpers.py b/q2mm/io/_helpers.py index 41f8e4cf..a19be4d5 100644 --- a/q2mm/io/_helpers.py +++ b/q2mm/io/_helpers.py @@ -5,8 +5,10 @@ from __future__ import annotations +import os from collections.abc import Callable, Sequence -from typing import TYPE_CHECKING, Any +from pathlib import Path +from typing import TYPE_CHECKING if TYPE_CHECKING: from q2mm.models.forcefield import ( @@ -23,8 +25,36 @@ canonicalize_bond_env_id, ) + +class _TextFileBase: + """Shared path and lazy-line-reading base for legacy text-file parsers.""" + + __slots__ = ("_lines", "path", "directory", "filename") + + def __init__(self, path: str) -> None: + self._lines: list[str] | None = None + self.path = os.path.abspath(path) + self.directory = os.path.dirname(self.path) + self.filename = os.path.basename(self.path) + + @property + def lines(self) -> list[str]: + """Return the lines of the file.""" + if self._lines is None: + with open(self.path) as f: + self._lines = f.readlines() + return self._lines + + def write(self, path: str, lines: list[str] | None = None) -> None: + """Write lines to file at *path* (defaults to ``self.lines``).""" + if lines is None: + lines = self.lines + with open(path, "w") as f: + f.writelines(lines) + + # --------------------------------------------------------------------------- -# Format compatibility +# Format compatibility & template resolution # --------------------------------------------------------------------------- _FORMAT_COMPATIBLE_FORMS: dict[str, set[str]] = { @@ -46,8 +76,21 @@ def _validate_form_for_format(ff: ForceField, target_format: str) -> None: ) +def _resolve_template_path( + ff: ForceField, + template_path: str | Path | None, + source_format: str, +) -> Path | None: + """Resolve an explicit *template_path* or fall back to ``ff.source_path`` when formats match.""" + if template_path is not None: + return Path(template_path) + if ff.source_format == source_format and ff.source_path is not None: + return ff.source_path + return None + + # --------------------------------------------------------------------------- -# Env-id splitting +# Env-id splitting & atom-type cleaning # --------------------------------------------------------------------------- @@ -58,12 +101,7 @@ def _split_env_id(env_id: str, expected_len: int) -> list[str]: return [] -# --------------------------------------------------------------------------- -# Atom-type cleaning -# --------------------------------------------------------------------------- - - -def _clean_atom_types(atom_types: list[str] | tuple[str, ...] | None, expected_len: int) -> list[str]: +def _clean_atom_types(atom_types: Sequence[str] | None, expected_len: int) -> list[str]: if atom_types is None: return [] if isinstance(atom_types, str): @@ -79,28 +117,30 @@ def _clean_atom_types(atom_types: list[str] | tuple[str, ...] | None, expected_l # --------------------------------------------------------------------------- -def _build_param_maps(params: list, secondary_key: str) -> tuple[dict, dict]: - """Build ff_row and secondary-key lookup dicts for a list of parameters.""" +def _build_param_maps[P: (BondParam, AngleParam, StretchBendParam, VdwParam)]( + params: Sequence[P], secondary_key: str +) -> tuple[dict[int, P], dict[str, P]]: + """Build ff_row and secondary-key lookup dicts for a sequence of parameters.""" by_row = {p.ff_row: p for p in params if p.ff_row is not None} by_key = {getattr(p, secondary_key): p for p in params if getattr(p, secondary_key, None)} return by_row, by_key -def _build_bond_maps(bonds: list[BondParam]) -> tuple[dict[int, BondParam], dict[str, BondParam]]: +def _build_bond_maps(bonds: Sequence[BondParam]) -> tuple[dict[int, BondParam], dict[str, BondParam]]: return _build_param_maps(bonds, "env_id") -def _build_angle_maps(angles: list[AngleParam]) -> tuple[dict[int, AngleParam], dict[str, AngleParam]]: +def _build_angle_maps(angles: Sequence[AngleParam]) -> tuple[dict[int, AngleParam], dict[str, AngleParam]]: return _build_param_maps(angles, "env_id") def _build_sb_maps( - stretch_bends: list[StretchBendParam], + stretch_bends: Sequence[StretchBendParam], ) -> tuple[dict[int, StretchBendParam], dict[str, StretchBendParam]]: return _build_param_maps(stretch_bends, "env_id") -def _build_vdw_maps(vdws: list[VdwParam]) -> tuple[dict[int, VdwParam], dict[str, VdwParam]]: +def _build_vdw_maps(vdws: Sequence[VdwParam]) -> tuple[dict[int, VdwParam], dict[str, VdwParam]]: return _build_param_maps(vdws, "atom_type") @@ -109,14 +149,14 @@ def _build_vdw_maps(vdws: list[VdwParam]) -> tuple[dict[int, VdwParam], dict[str # --------------------------------------------------------------------------- -def _match_for_export( +def _match_for_export[P: (BondParam, AngleParam, StretchBendParam)]( ff_row: int | None, - atom_types: list[str] | tuple[str, ...] | None, - by_row: dict, - by_env: dict, + atom_types: Sequence[str] | None, + by_row: dict[int, P], + by_env: dict[str, P], expected_len: int, - canonicalize_fn: Callable, -) -> Any: + canonicalize_fn: Callable[[Sequence[str]], str], +) -> P | None: """Match a parsed file row to an internal param by ``ff_row`` or ``env_id``. Pure: takes the row's primitive identity data (*ff_row*, *atom_types*) @@ -133,7 +173,7 @@ def _match_for_export( def _match_bond_for_export( ff_row: int | None, - atom_types: list[str] | tuple[str, ...] | None, + atom_types: Sequence[str] | None, bond_by_row: dict[int, BondParam], bond_by_env: dict[str, BondParam], ) -> BondParam | None: @@ -142,7 +182,7 @@ def _match_bond_for_export( def _match_angle_for_export( ff_row: int | None, - atom_types: list[str] | tuple[str, ...] | None, + atom_types: Sequence[str] | None, angle_by_row: dict[int, AngleParam], angle_by_env: dict[str, AngleParam], ) -> AngleParam | None: @@ -151,7 +191,7 @@ def _match_angle_for_export( def _match_sb_for_export( ff_row: int | None, - atom_types: list[str] | tuple[str, ...] | None, + atom_types: Sequence[str] | None, sb_by_row: dict[int, StretchBendParam], sb_by_env: dict[str, StretchBendParam], ) -> StretchBendParam | None: diff --git a/q2mm/io/amber.py b/q2mm/io/amber.py index 75d554b0..c24d883c 100644 --- a/q2mm/io/amber.py +++ b/q2mm/io/amber.py @@ -9,7 +9,7 @@ from decimal import Decimal from pathlib import Path -from q2mm.io._helpers import _validate_form_for_format +from q2mm.io._helpers import _resolve_template_path, _validate_form_for_format from q2mm.models.forcefield import ( AngleParam, BondParam, @@ -308,10 +308,10 @@ def load_amber_frcmod(path: str | Path) -> ForceField: types, rest = _parse_amber_types(line, 2) vals = _parse_floats(rest) if len(types) == 2 and all(types) and len(vals) >= 2: - elems = tuple(_amber_type_to_element(t, mass_map) for t in types) + t0, t1 = types bonds.append( BondParam( - elements=elems, + elements=(_amber_type_to_element(t0, mass_map), _amber_type_to_element(t1, mass_map)), equilibrium=vals[1], force_constant=vals[0], env_id=canonicalize_bond_env_id(types), @@ -324,10 +324,14 @@ def load_amber_frcmod(path: str | Path) -> ForceField: types, rest = _parse_amber_types(line, 3) vals = _parse_floats(rest) if len(types) == 3 and all(types) and len(vals) >= 2: - elems = tuple(_amber_type_to_element(t, mass_map) for t in types) + t0, t1, t2 = types angles.append( AngleParam( - elements=elems, + elements=( + _amber_type_to_element(t0, mass_map), + _amber_type_to_element(t1, mass_map), + _amber_type_to_element(t2, mass_map), + ), equilibrium=vals[1], force_constant=vals[0], env_id=canonicalize_angle_env_id(types), @@ -338,10 +342,15 @@ def load_amber_frcmod(path: str | Path) -> ForceField: elif section == "DIHE": record = dihe_rows[row] - elems = tuple(_amber_type_to_element(t, mass_map) for t in record.atom_types) + t0, t1, t2, t3 = record.atom_types torsions.append( TorsionParam( - elements=elems, + elements=( + _amber_type_to_element(t0, mass_map), + _amber_type_to_element(t1, mass_map), + _amber_type_to_element(t2, mass_map), + _amber_type_to_element(t3, mass_map), + ), periodicity=abs(record.pn), force_constant=record.barrier / record.effective_idivf, phase=record.phase, @@ -356,7 +365,7 @@ def load_amber_frcmod(path: str | Path) -> ForceField: vals = _parse_floats(rest) # vals: barrier, phase, periodicity (no IDIVF) if len(types) == 4 and all(types) and len(vals) >= 3: - elems = tuple(_amber_type_to_element(t, mass_map) for t in types) + t0, t1, t2, t3 = types pn = vals[2] # IMPROPER has neither DIHE's continuation sign nor a zero-to-one convention. if not math.isfinite(pn) or pn <= 0 or pn != int(pn): @@ -364,7 +373,12 @@ def load_amber_frcmod(path: str | Path) -> ForceField: periodicity = int(pn) torsions.append( TorsionParam( - elements=elems, + elements=( + _amber_type_to_element(t0, mass_map), + _amber_type_to_element(t1, mass_map), + _amber_type_to_element(t2, mass_map), + _amber_type_to_element(t3, mass_map), + ), periodicity=periodicity, force_constant=vals[0], phase=vals[1], @@ -397,10 +411,10 @@ def load_amber_frcmod(path: str | Path) -> ForceField: return ForceField( name=f"AMBER from {path.name}", - bonds=bonds, - angles=angles, - torsions=torsions, - vdws=vdws, + bonds=tuple(bonds), + angles=tuple(angles), + torsions=tuple(torsions), + vdws=tuple(vdws), source_path=path, source_format="amber_frcmod", functional_form=FunctionalForm.HARMONIC, @@ -453,9 +467,7 @@ def save_amber_frcmod( """ _validate_form_for_format(ff, "amber_frcmod") output_path = Path(path) - template = Path(template_path) if template_path is not None else None - if template is None and ff.source_format == "amber_frcmod" and ff.source_path is not None: - template = ff.source_path + template = _resolve_template_path(ff, template_path, "amber_frcmod") if template is not None: return _save_amber_frcmod_template(ff, output_path, template) diff --git a/q2mm/io/cmap.py b/q2mm/io/cmap.py index 9e36b32a..d6bb2958 100644 --- a/q2mm/io/cmap.py +++ b/q2mm/io/cmap.py @@ -40,25 +40,36 @@ def parse_cmap_section(text: str) -> list[CmapGrid]: lines = text.splitlines() grids: list[CmapGrid] = [] - def append_grid(header: tuple[tuple[str, ...], tuple[str, ...], int], values: list[float]) -> None: - phi_types, psi_types, resolution = header - try: - grids.append( - CmapGrid( - atom_types_phi=phi_types, - atom_types_psi=psi_types, - resolution=resolution, - energy=values, - ) - ) - except ValueError as exc: - logger.warning("Skipping malformed CMAP grid: %s", exc) - in_cmap_section = False - current_header: tuple[tuple[str, ...], tuple[str, ...], int] | None = None + current_header: tuple[tuple[str, str, str, str], tuple[str, str, str, str], int] | None = None current_values: list[float] = [] expected_count = 0 + def finalize_current(*, where: str = "") -> None: + if current_header is None: + return + if len(current_values) == expected_count: + phi_types, psi_types, resolution = current_header + try: + grids.append( + CmapGrid( + atom_types_phi=phi_types, + atom_types_psi=psi_types, + resolution=resolution, + energy=tuple(current_values), + ) + ) + except ValueError as exc: + logger.warning("Skipping malformed CMAP grid: %s", exc) + else: + suffix = f" {where}" if where else "" + logger.warning( + "Incomplete CMAP grid%s: expected %d values, got %d. Skipping.", + suffix, + expected_count, + len(current_values), + ) + for raw_line in lines: line = raw_line.strip() @@ -81,16 +92,7 @@ def append_grid(header: tuple[tuple[str, ...], tuple[str, ...], int], values: li upper, ): if upper != "CMAP": - # Finalize any in-progress grid - if current_header is not None: - if len(current_values) == expected_count: - append_grid(current_header, current_values) - else: - logger.warning( - "Incomplete CMAP grid at section boundary: expected %d values, got %d. Skipping.", - expected_count, - len(current_values), - ) + finalize_current(where="at section boundary") in_cmap_section = False current_header = None current_values = [] @@ -102,17 +104,7 @@ def append_grid(header: tuple[tuple[str, ...], tuple[str, ...], int], values: li # Try to parse as a CMAP header line: 8 atom types + resolution tokens = line.split() if len(tokens) == 9 and _is_header_line(tokens): - # Finalize previous grid if any - if current_header is not None: - if len(current_values) == expected_count: - append_grid(current_header, current_values) - else: - logger.warning( - "Incomplete CMAP grid: expected %d values, got %d. Skipping.", - expected_count, - len(current_values), - ) - + finalize_current() resolution = int(tokens[8]) current_header = ( (tokens[0], tokens[1], tokens[2], tokens[3]), @@ -131,16 +123,7 @@ def append_grid(header: tuple[tuple[str, ...], tuple[str, ...], int], values: li except ValueError: logger.warning("Non-numeric token in CMAP grid data: %r", token) - # Finalize last grid - if current_header is not None and len(current_values) == expected_count: - append_grid(current_header, current_values) - elif current_header is not None: - logger.warning( - "Incomplete CMAP grid at end of file: expected %d values, got %d.", - expected_count, - len(current_values), - ) - + finalize_current(where="at end of file") logger.info("Parsed %d CMAP grid(s).", len(grids)) return grids diff --git a/q2mm/io/gaussian.py b/q2mm/io/gaussian.py index 0d900bc3..95e4bbf1 100644 --- a/q2mm/io/gaussian.py +++ b/q2mm/io/gaussian.py @@ -7,7 +7,6 @@ from __future__ import annotations import logging -import os import re from dataclasses import dataclass, field from pathlib import Path @@ -16,12 +15,15 @@ from q2mm import constants as co from q2mm.constants import DEFAULT_BOND_TOLERANCE +from q2mm.io._helpers import _TextFileBase from q2mm.models.hessian import HessianProvenance, HessianUnits, hessian_to_atomic_units from q2mm.models.identifiers import _extract_element from q2mm.models.molecule import Molecule from q2mm.models.observations import ObservationSet logger = logging.getLogger(__name__) +_MASS_SYMBOLS: tuple[str, ...] = tuple(co.MASSES.keys()) +_MASS_VALUES: tuple[float, ...] = tuple(co.MASSES.values()) @dataclass @@ -57,7 +59,7 @@ def symbol(self) -> str: if self.element: return _extract_element(self.element) if self.atomic_num is not None and self.atomic_num >= 1: - return list(co.MASSES.keys())[self.atomic_num - 1] + return _MASS_SYMBOLS[self.atomic_num - 1] return "" @@ -177,25 +179,22 @@ def _gaussian_archive_records(path: str, filename: str, *, au_hessian: bool) -> return records -class GaussLog: +class GaussLog(_TextFileBase): """Retrieves data from Gaussian log files. If you are extracting frequencies/Hessian data from this file, use the keyword NoSymmetry when running the Gaussian calculation. """ - __slots__ = [ - "_lines", - "path", - "directory", - "filename", + __slots__ = ( + "_body_parsed", "_evals", "_evecs", "_frequencies_cm", "_records", "_esp_rms", "_au_hessian", - ] + ) def __init__(self, path: str, au_hessian: bool = False) -> None: """Instantiate a file object for the file at the location path passed. @@ -209,44 +208,15 @@ def __init__(self, path: str, au_hessian: bool = False) -> None: (Hartree/Bohr²). Defaults to False. """ - self._lines = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) - self._evals = None - self._evecs = None - self._frequencies_cm = None - self._records = None - self._esp_rms = None + super().__init__(path) + self._body_parsed = False + self._evals: np.ndarray | None = None + self._evecs: np.ndarray | None = None + self._frequencies_cm: np.ndarray | None = None + self._records: list[_GaussianRecord] | None = None + self._esp_rms: float | None = None self._au_hessian = au_hessian - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) - @property def evecs(self) -> np.ndarray | None: """Returns eigenvectors of frequency analysis if applicable. @@ -258,7 +228,7 @@ def evecs(self) -> np.ndarray | None: Gaussian frequency analysis, or None if not a frequency job. """ - if self._evecs is None: + if not self._body_parsed: self.read_out() return self._evecs @@ -277,7 +247,7 @@ def evals(self) -> np.ndarray | None: or None if not a frequency job. """ - if self._evals is None: + if not self._body_parsed: self.read_out() return self._evals @@ -291,7 +261,7 @@ def frequencies(self) -> np.ndarray | None: (np.ndarray | None): Frequencies in cm⁻¹, or None if not a frequency job. """ - if self._frequencies_cm is None: + if not self._body_parsed: self.read_out() return self._frequencies_cm @@ -303,6 +273,7 @@ def _records_parsed(self) -> list[_GaussianRecord]: """ if self._records is None: self.read_archive() + assert self._records is not None return self._records @property @@ -323,23 +294,19 @@ def read_out(self) -> None: and ``_esp_rms`` from the Gaussian log file body. """ logger.log(5, f"READING: {self.filename}") - self._evals = [] - self._evecs = [] - self._frequencies_cm = [] - self._records = [] - force_constants = [] - evecs = [] + self._evals = None + self._evecs = None + self._frequencies_cm = None + evals_list: list[float] = [] + evecs_list: list[list[float]] = [] + frequencies_cm_list: list[float] = [] + records: list[_GaussianRecord] = [] + force_constants: list[float] = [] + evecs: list[list[float]] = [] with open(self.path) as f: - # The keyword "harmonic" shows up before the section we're - # interested in. It can show up multiple times depending on the - # options in the Gaussian .com file. past_first_harm = False - # High precision mode, turned on by including "freq=hpmodes" in the - # Gaussian .com file. hpmodes = False file_iterator = iter(f) - # This while loop breaks when the end of the file is reached, or - # if the high quality modes have been read already. while True: try: line = next(file_iterator) @@ -348,10 +315,10 @@ def read_out(self) -> None: if "Charges from ESP fit" in line: pattern = re.compile(rf"RMS=\s+({co.RE_FLOAT})") match = pattern.search(line) - self._esp_rms = float(match.group(1)) - # Gathering some geometric information. + if match is not None: + self._esp_rms = float(match.group(1)) elif "Standard orientation:" in line: - self._records.append(_GaussianRecord(self.filename, self.path)) + records.append(_GaussianRecord(self.filename, self.path)) next(file_iterator) next(file_iterator) next(file_iterator) @@ -359,7 +326,7 @@ def read_out(self) -> None: line = next(file_iterator) while "---" not in line: cols = line.split() - self._records[-1].atoms.append( + records[-1].atoms.append( _GaussianAtom( index=int(cols[0]), atomic_num=int(cols[1]), @@ -371,87 +338,38 @@ def read_out(self) -> None: line = next(file_iterator) logger.log( 5, - f" -- Found {len(self._records[-1].atoms)} atoms.", + f" -- Found {len(records[-1].atoms)} atoms.", ) elif "Harmonic" in line: - # The high quality eigenvectors come before the low quality - # ones. If you see "Harmonic" again, it means you're at the - # low quality ones now, so break. if past_first_harm: break - else: - past_first_harm = True + past_first_harm = True elif "Frequencies" in line: - # We're going to keep reusing these. - # We accumulate sets of eigevectors and eigenvalues, add - # them to self._evecs and self._evals, and then reuse this - # for the next set. del force_constants[:] del evecs[:] - # Values inside line look like: - # "Frequencies --- xxxx.xxxx xxxx.xxxx" - # That's why we remove the 1st two columns. This is - # consistent with and without "hpmodes". - # For "hpmodes" option, there are 5 of these frequencies. - # Without "hpmodes", there are 3. - # Thus the eigenvectors and eigenvalues will come in sets of - # either 5 or 3. cols = line.split() for frequency in map(float, cols[2:]): - # Store actual frequency in cm⁻¹ - self._frequencies_cm.append(frequency) - # Has 1. or -1. depending on the sign of the frequency. - if frequency < 0.0: - force_constants.append(-1.0) - else: - force_constants.append(1.0) - # For now this is empty, but we will add to it soon. + frequencies_cm_list.append(frequency) + force_constants.append(-1.0 if frequency < 0.0 else 1.0) evecs.append([]) - # Moving on to the reduced masses. line = next(file_iterator) cols = line.split() - # Again, trim the "Reduced masses ---". - # It's "Red. masses --" for without "hpmodes". for i, mass in enumerate(map(float, cols[3:])): - # +/- 1 / reduced mass - force_constants[i] = force_constants[i] / mass + force_constants[i] /= mass - # Now we are on the line with the force constants. line = next(file_iterator) cols = line.split() - # Trim "Force constants ---". It's "Frc consts --" without - # "hpmodes". for i, force_constant in enumerate(map(float, cols[3:])): - # AU_TO_MDYNA conversion (15.569141) force_constants[i] *= force_constant / co.AU_TO_MDYNA - # Force constants were calculated above as follows: - # a = +/- 1 depending on the sign of the frequency - # b = a / reduced mass (obtained from the Gaussian log) - # c = b * force constant / conversion factor (force - # (constant obtained from Gaussian log) (conversion - # factor is inside constants module) - - # Skip the IR intensities. next(file_iterator) - # This is different depending on whether you use "hpmodes". line = next(file_iterator) - # "Coord" seems to only appear when the "hpmodes" is used. if "Coord" in line: hpmodes = True - # This is different depending on whether you use - # "freq=projected". line = next(file_iterator) - # The "projected" keyword seems to add "IRC Coupling". if "IRC Coupling" in line: line = next(file_iterator) - # We're on to the eigenvectors. - # Until the end of this section containing the eigenvectors, - # the number of columns remains constant. When that changes, - # we know we're to the next set of frequencies, force - # constants and eigenvectors. - # Actually check that we've moved on, sometimes a "Depolar" entry is if "Depolar" in line: line = next(file_iterator) if "Atom" in line: @@ -460,89 +378,41 @@ def read_out(self) -> None: cols_len = len(cols) while len(cols) == cols_len: - # This will come after all the eigenvectors have been - # read. We can break out then. if "Harmonic" in line: break - # If "hpmodes" is used, you have an extra column here - # that is simply an index. if hpmodes: cols = cols[1:] - # cols corresponds to line(s) (maybe only 1st line) - # under section "Coord Atom Element:" (at least for - # "hpmodes"). - - # Just the square root of the mass from co.MASSES. - # co.MASSES currently has the average mass. - # Gaussian may use the mass of the most abundant - # isotope. This may be a problem. - mass_sqrt = np.sqrt(list(co.MASSES.items())[int(cols[1]) - 1][1]) - + mass_sqrt = np.sqrt(_MASS_VALUES[int(cols[1]) - 1]) cols = cols[2:] - # This corresponds to the same line still, but without - # the atom elements. - - # This loop expands the LoL, evecs, as so. - # Iteration 1: each sub-list has 1 element - # Iteration 2: each sub-list has 2 elements - # ... etc. until the length of the sublist is equal to - # the number of atoms. Remember, for low precision - # eigenvectors it only adds in sets of 3, not 5. - - # Elements of evecs are simply the data under - # "Coord Atom Element" multiplied by the square root - # of the weight. for i in range(len(evecs)): if hpmodes: - # evecs is a LoL. Length of sublist is - # equal to # of columns in section "Coord Atom - # Element" minus 3, for the 1st 3 columns - # (index, atom index, atomic number). evecs[i].append(float(cols[i]) * mass_sqrt) else: - # This is fow low precision eigenvectors. It's a - # funny way to go in sets of 3. Take a look at - # your low precision Gaussian log and it will - # make more sense. for _ in range(3): x = float(cols.pop(0)) evecs[i].append(x * mass_sqrt) line = next(file_iterator) cols = line.split() - # Here the overall number of eigenvalues and eigenvectors is - # increased by 5 (high precision) or 3 (low precision). The - # total number goes to 3N - 6 for non-linear and 3N - 5 for - # linear. Same goes for self._evecs. for i in range(len(evecs)): - self._evals.append(force_constants[i]) - self._evecs.append(evecs[i]) - # We know we're done if this is in the line. + evals_list.append(force_constants[i]) + evecs_list.append(evecs[i]) if "Harmonic" in line: break - if self._evals and self._evecs: - for evec in self._evecs: - # Each evec is a single eigenvector. - # Add up the sum of squares over an eigenvector. - sum_of_squares = 0.0 - # Appropriately named, element is an element of that single - # eigenvector. - for element in evec: - sum_of_squares += element * element - # Now x is the inverse of the square root of the sum of squares - # for an individual eigenvector. - element = 1 / np.sqrt(sum_of_squares) + self._records = records + if evals_list and evecs_list: + for evec in evecs_list: + inv_norm = 1.0 / float(np.sqrt(sum(element * element for element in evec))) for i in range(len(evec)): - evec[i] *= element - self._evals = np.array(self._evals) - self._evecs = np.array(self._evecs) - self._frequencies_cm = np.array(self._frequencies_cm) + evec[i] *= inv_norm + self._evals = np.array(evals_list) + self._evecs = np.array(evecs_list) + self._frequencies_cm = np.array(frequencies_cm_list) logger.log(1, f">>> self._evals: {self._evals}") logger.log(1, f">>> self._evecs: {self._evecs}") logger.log(5, f" -- {len(self._records)} structures found.") + self._body_parsed = True - # May want to move some attributes assigned to the structure class onto - # this filetype class. def read_archive(self) -> None: """Read the last complete archive section, preserving legacy behavior.""" logger.log(5, f"READING: {self.filename}") diff --git a/q2mm/io/jaguar.py b/q2mm/io/jaguar.py index fdeb46d5..02848f2d 100644 --- a/q2mm/io/jaguar.py +++ b/q2mm/io/jaguar.py @@ -8,16 +8,16 @@ from __future__ import annotations import logging -import os import re from dataclasses import dataclass, field from importlib.util import find_spec from string import digits -from typing import Any +from typing import Any, cast import numpy as np from q2mm import constants as co +from q2mm.io._helpers import _TextFileBase from q2mm.models.hessian import HessianProvenance, HessianUnits from q2mm.models.identifiers import _extract_element from q2mm.models.molecule import Molecule @@ -110,7 +110,7 @@ def _molecule_from_jaguar_record(record: _JaguarRecord) -> Molecule: ) -class JaguarIn: +class JaguarIn(_TextFileBase): """Retrieve data from Jaguar ``.in`` files. The Hessian is **not** mass-weighted and is returned in atomic units @@ -128,39 +128,9 @@ def __init__(self, path: str) -> None: path (str): Path to the Jaguar ``.in`` file. """ - self._lines = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) - self._hessian = None - self._molecules = None - - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) + super().__init__(path) + self._hessian: np.ndarray | None = None + self._molecules: list[Molecule] | None = None def get_hessian(self, num_atoms: int, *, tag_units: bool = False) -> np.ndarray: """Read the Hessian matrix from a Jaguar ``.in`` file. @@ -243,7 +213,7 @@ def get_hessian(self, num_atoms: int, *, tag_units: bool = False) -> np.ndarray: if tag_units: ureg = _get_pint_ureg() if ureg is not None: - return ureg.Quantity(self._hessian, "hartree/bohr**2") + return cast(np.ndarray, ureg.Quantity(self._hessian, "hartree/bohr**2")) return self._hessian @property @@ -316,7 +286,7 @@ def attach_hessian(self, molecule: Molecule) -> Molecule: ) -class JaguarOut: +class JaguarOut(_TextFileBase): """Retrieve data from Schrödinger Jaguar ``.out`` files. Eigenvalues and eigenvectors are **not** mass-weighted. @@ -329,42 +299,12 @@ def __init__(self, path: str) -> None: path (str): Path to the Jaguar ``.out`` file. """ - self._lines = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) - self._records = None - self._eigenvalues = None - self._eigenvectors = None - self._frequencies = None - self._dummy_atom_eigenvector_indices = None - - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) + super().__init__(path) + self._records: list[_JaguarRecord] | None = None + self._eigenvalues: np.ndarray | None = None + self._eigenvectors: np.ndarray | None = None + self._frequencies: np.ndarray | None = None + self._dummy_atom_eigenvector_indices: list[int] | None = None @property def _records_parsed(self) -> list[_JaguarRecord]: @@ -374,6 +314,7 @@ def _records_parsed(self) -> list[_JaguarRecord]: """ if self._records is None: self.import_file() + assert self._records is not None return self._records @property @@ -386,6 +327,7 @@ def eigenvalues(self) -> np.ndarray: """numpy.ndarray: Eigenvalues derived from force constants and frequencies.""" if self._eigenvalues is None: self.import_file() + assert self._eigenvalues is not None return self._eigenvalues @property @@ -393,6 +335,7 @@ def eigenvectors(self) -> np.ndarray: """numpy.ndarray: Cartesian eigenvectors with dummy-atom rows removed.""" if self._eigenvectors is None: self.import_file() + assert self._eigenvectors is not None return self._eigenvectors @property @@ -400,6 +343,7 @@ def frequencies(self) -> np.ndarray: """numpy.ndarray: Vibrational frequencies in cm⁻¹.""" if self._frequencies is None: self.import_file() + assert self._frequencies is not None return self._frequencies @property @@ -407,6 +351,7 @@ def dummy_atom_eigenvector_indices(self) -> list[int]: """list[int]: Row indices in the eigenvector matrix that correspond to dummy atoms.""" if self._dummy_atom_eigenvector_indices is None: self.import_file() + assert self._dummy_atom_eigenvector_indices is not None return self._dummy_atom_eigenvector_indices def import_file(self) -> None: @@ -417,9 +362,9 @@ def import_file(self) -> None: eigenvectors. """ logger.log(10, f"READING: {self.filename}") - frequencies = [] - force_constants = [] - eigenvectors = [] + frequencies: list[float] = [] + force_constants: list[float] = [] + eigenvectors: list[list[float]] = [] records: list[_JaguarRecord] = [] with open(self.path) as f: section_geometry = False @@ -428,7 +373,7 @@ def import_file(self) -> None: current_record: _JaguarRecord | None = None charge = 0 multiplicity = 1 - temp_eigenvectors = [] + temp_eigenvectors: list[list[float]] = [] for i, line in enumerate(f): charge_match = re.search(r"net molecular charge:\s*(-?\d+)", line, re.IGNORECASE) if charge_match is not None: @@ -440,6 +385,7 @@ def import_file(self) -> None: cols = line.split() if len(cols) == 0: section_geometry = False + assert current_record is not None records.append(current_record) continue elif len(cols) == 1: @@ -447,11 +393,13 @@ def import_file(self) -> None: else: match = re.match(rf"\s+([\d\w]+)\s+({co.RE_FLOAT})\s+({co.RE_FLOAT})\s+({co.RE_FLOAT})", line) if match is not None: - current_atom = _JaguarAtom() - current_atom.element = match.group(1).translate(str.maketrans("", "", digits)) - current_atom.x = float(match.group(2)) - current_atom.y = float(match.group(3)) - current_atom.z = float(match.group(4)) + assert current_record is not None + current_atom = _JaguarAtom( + element=match.group(1).translate(str.maketrans("", "", digits)), + x=float(match.group(2)), + y=float(match.group(3)), + z=float(match.group(4)), + ) current_record.atoms.append(current_atom) logger.log( 0, @@ -478,11 +426,10 @@ def import_file(self) -> None: section_eigenvectors = False eigenvectors.extend(temp_eigenvectors) continue - else: - for i, x in enumerate(cols[2:]): - if not len(temp_eigenvectors) > i: - temp_eigenvectors.append([]) - temp_eigenvectors[i].append(float(x)) + for j, x in enumerate(cols[2:]): + if len(temp_eigenvectors) <= j: + temp_eigenvectors.append([]) + temp_eigenvectors[j].append(float(x)) if section_eigenvalues is True and section_eigenvectors is False: if "frequencies" in line: cols = line.split() @@ -496,47 +443,21 @@ def import_file(self) -> None: section_eigenvalues = True logger.log(1, f">>> len(frequencies): {len(frequencies)}") logger.log(1, f">>> frequencies:\n{frequencies}") - # logger.log(1, '>>> frequencies:\n{}'.format( - # [x / co.FORCE_CONVERSION for x in frequencies])) - # logger.log(1, '>>> frequencies:\n{}'.format( - # [x * 4.55633e-6 for x in frequencies])) - # logger.log(1, '>>> frequencies:\n{}'.format( - # [x * 1.23981e-4 for x in frequencies])) - # logger.log(1, '>>> frequencies:\n{}'.format( - # [x / 219474.6305 for x in frequencies])) eigenvalues = [ -fc / co.FORCE_CONVERSION if f < 0 else fc / co.FORCE_CONVERSION for fc, f in zip(force_constants, frequencies) ] logger.log(1, f">>> eigenvalues:\n{eigenvalues}") # Remove eigenvector components related to dummy atoms. - # Find the index of the atoms that are dummies. - dummy_atom_indices = [] - for i, atom in enumerate(records[-1].atoms): - if atom.is_dummy: - dummy_atom_indices.append(i) + dummy_atom_indices = [i for i, atom in enumerate(records[-1].atoms) if atom.is_dummy] logger.log(10, f" -- Located {len(dummy_atom_indices)} dummy atoms.") - # Correlate those indices to the rows in the cartesian eigenvector. - dummy_atom_eigenvector_indices = [] - for dummy_atom_index in dummy_atom_indices: - start = dummy_atom_index * 3 - dummy_atom_eigenvector_indices.append(start) - dummy_atom_eigenvector_indices.append(start + 1) - dummy_atom_eigenvector_indices.append(start + 2) - new_eigenvectors = [] - # Create new eigenvectors without the rows corresponding to the - # dummy atoms. - for eigenvector in eigenvectors: - new_eigenvectors.append([]) - for i, eigenvector_row in enumerate(eigenvector): - if i not in dummy_atom_eigenvector_indices: - new_eigenvectors[-1].append(eigenvector_row) - # Replace old eigenvectors with new where dummy atoms aren't included. - eigenvectors = np.array(new_eigenvectors) + dummy_atom_eigenvector_indices = [idx * 3 + offset for idx in dummy_atom_indices for offset in range(3)] + dummy_set = set(dummy_atom_eigenvector_indices) + filtered_eigenvectors = [[row for i, row in enumerate(evec) if i not in dummy_set] for evec in eigenvectors] self._dummy_atom_eigenvector_indices = dummy_atom_eigenvector_indices self._records = records self._eigenvalues = np.array(eigenvalues) - self._eigenvectors = np.array(eigenvectors) + self._eigenvectors = np.array(filtered_eigenvectors) self._frequencies = np.array(frequencies) logger.log(5, f" -- Read {len(self._records)} structures") logger.log(5, f" -- Read {len(self.frequencies)} frequencies.") diff --git a/q2mm/io/macromodel.py b/q2mm/io/macromodel.py index a50b7d2a..2d2de9ef 100644 --- a/q2mm/io/macromodel.py +++ b/q2mm/io/macromodel.py @@ -8,13 +8,13 @@ from __future__ import annotations import logging -import os import re from dataclasses import dataclass, field import numpy as np from q2mm import constants as co +from q2mm.io._helpers import _TextFileBase from q2mm.models.identifiers import ( _extract_element, canonicalize_angle_env_id, @@ -51,28 +51,8 @@ def symbol(self) -> str: @dataclass -class _MacroModelBond: - """One bond parsed from a MacroModel ``.mmo`` bond-energy line.""" - - atom_nums: list[int] - value: float - ff_row: int - comment: str = "" - - -@dataclass -class _MacroModelAngle: - """One angle parsed from a MacroModel ``.mmo`` angle-energy line.""" - - atom_nums: list[int] - value: float - ff_row: int - comment: str = "" - - -@dataclass -class _MacroModelTorsion: - """One torsion parsed from a MacroModel ``.mmo`` torsion-energy line.""" +class _MacroModelTerm: + """One bond, angle, or torsion parsed from a MacroModel ``.mmo`` energy line.""" atom_nums: list[int] value: float @@ -92,73 +72,39 @@ class _MacroModelRecord: origin_name: str atoms: list[_MacroModelAtom] = field(default_factory=list) - bonds: list[_MacroModelBond] = field(default_factory=list) - angles: list[_MacroModelAngle] = field(default_factory=list) - torsions: list[_MacroModelTorsion] = field(default_factory=list) - - -def _read_bond_line(line: str) -> _MacroModelBond | None: - """Parse a single line for bond data. + bonds: list[_MacroModelTerm] = field(default_factory=list) + angles: list[_MacroModelTerm] = field(default_factory=list) + torsions: list[_MacroModelTerm] = field(default_factory=list) - Args: - line: A line from the bond section of the ``.mmo`` file. - Returns: - A :class:`_MacroModelBond` if the line matches the bond pattern, - otherwise ``None``. - - """ +def _read_bond_line(line: str) -> _MacroModelTerm | None: + """Parse a single line for bond data.""" match = co.RE_BOND.match(line) - # atom_nums are 1-based atom indices (not atomic numbers) if match: atom_nums = [int(x) for x in [match.group(1), match.group(2)]] value = float(match.group(3)) comment = match.group(4).strip() ff_row = int(match.group(5)) - return _MacroModelBond(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) + return _MacroModelTerm(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) return None -def _read_angle_line(line: str) -> _MacroModelAngle | None: - """Parse a single line for angle data. - - Terminal atoms are reordered so that the lower index comes first. - - Args: - line: A line from the angle section of the ``.mmo`` file. - - Returns: - A :class:`_MacroModelAngle` if the line matches the angle pattern, - otherwise ``None``. - - """ +def _read_angle_line(line: str) -> _MacroModelTerm | None: + """Parse a single line for angle data, ordering terminal atoms ascending.""" match = co.RE_ANGLE.match(line) if match: atom_nums = [int(x) for x in [match.group(1), match.group(2), match.group(3)]] - # Reorder the terminal atoms so that the lower index atom is first. if atom_nums[0] > atom_nums[2]: atom_nums = [atom_nums[2], atom_nums[1], atom_nums[0]] value = float(match.group(4)) comment = match.group(5).strip() ff_row = int(match.group(6)) - return _MacroModelAngle(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) + return _MacroModelTerm(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) return None -def _read_torsion_line(line: str) -> _MacroModelTorsion | None: - """Parse a single line for torsion data. - - Atom indices are reordered so that the lower central-atom index comes - first. - - Args: - line: A line from the torsion section of the ``.mmo`` file. - - Returns: - A :class:`_MacroModelTorsion` if the line matches the torsion - pattern, otherwise ``None``. - - """ +def _read_torsion_line(line: str) -> _MacroModelTerm | None: + """Parse a single line for torsion data, ordering central atoms ascending.""" match = co.RE_TORSION.match(line) if match: atom_nums = [int(x) for x in [match.group(1), match.group(2), match.group(3), match.group(4)]] @@ -167,7 +113,7 @@ def _read_torsion_line(line: str) -> _MacroModelTorsion | None: value = float(match.group(5)) comment = match.group(6).strip() ff_row = int(match.group(7)) - return _MacroModelTorsion(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) + return _MacroModelTerm(atom_nums=atom_nums, comment=comment, value=value, ff_row=ff_row) return None @@ -237,7 +183,7 @@ def _atom(num: int) -> _MacroModelAtom: ) -class MacroModel: +class MacroModel(_TextFileBase): """Extract structural data from MacroModel ``.mmo`` files. Reads bond lengths, angles, and torsions for each structure entry @@ -251,38 +197,8 @@ def __init__(self, path: str) -> None: path (str): Path to the MacroModel ``.mmo`` file. """ - self._lines = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) - self._records = None - - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) + super().__init__(path) + self._records: list[_MacroModelRecord] | None = None @property def _records_parsed(self) -> list[_MacroModelRecord]: @@ -290,35 +206,28 @@ def _records_parsed(self) -> list[_MacroModelRecord]: Not part of the public API — see :attr:`molecules`. """ - # Atom reading not yet implemented; would be needed for - # Hessian extraction (requires atom count for matrix shape). - if self._records is None or self._records == []: + if not self._records: logger.log(10, f"READING: {self.filename}") - self._records = [] + records: list[_MacroModelRecord] = [] with open(self.path) as f: count_current = 0 count_input = 0 count_structure = 0 count_previous = 0 - bonds = [] - angles = [] - torsions = [] - atoms = [] - current_record = None - section = None + bonds: list[_MacroModelTerm] = [] + angles: list[_MacroModelTerm] = [] + torsions: list[_MacroModelTerm] = [] + atoms: list[_MacroModelAtom] = [] + current_record: _MacroModelRecord | None = None + section: str | None = None for line in f: - # The MMO file lists the bonds, angles, and torsions in - # some order that I am unsure of. It seems consistent - # with the same filename but with two files with the - # exact same structure the ordering is off. This - # reorders the lists before being added to the record. if "Atomic Charges, Coordinates and Connectivity" in line: section = "atoms" continue if section == "atoms": if "(" in line: split = [item.strip() for item in line.split()] - atom_num = split[2][:-1] # same as index + atom_num = split[2][:-1] ele_name = re.sub(r"[0-9]", "", split[0]) atom = _MacroModelAtom( atom_type_name=split[0], @@ -331,13 +240,12 @@ def _records_parsed(self) -> list[_MacroModelRecord]: atoms.append(atom) if "Total" in line: section = None - # Sort the bonds, angles, and torsions before the start - # of a new structure bonds.sort(key=lambda x: (x.atom_nums[0], x.atom_nums[1])) angles.sort(key=lambda x: (x.atom_nums[1], x.atom_nums[0], x.atom_nums[2])) torsions.sort( key=lambda x: (x.atom_nums[1], x.atom_nums[2], x.atom_nums[0], x.atom_nums[3]) ) + assert current_record is not None current_record.bonds = bonds current_record.angles = angles current_record.torsions = torsions @@ -349,24 +257,14 @@ def _records_parsed(self) -> list[_MacroModelRecord]: if "Input Structure Name" in line: count_structure += 1 count_previous = count_current - # Sometimes only one of the above ("Input filename" and - # "Input Structure Name") is used, sometimes both are used. - # count_current will make sure you catch both. count_current = max(count_input, count_structure) - # If these don't match, then we reached the end of a - # structure. if count_current != count_previous: bonds = [] angles = [] torsions = [] atoms = [] current_record = _MacroModelRecord(self.filename) - self._records.append(current_record) - # For each structure we come across, look for sections that - # we are interested in: those pertaining to bonds, angles, - # and torsions. Of course more could be added. We set the - # section to None to mark the end of a section, and we leave - # it None for parts of the file we don't care about. + records.append(current_record) if "BOND LENGTHS AND STRETCH ENERGIES" in line: section = "bond" if "ANGLES, BEND AND STRETCH BEND ENERGIES" in line: @@ -389,6 +287,7 @@ def _records_parsed(self) -> list[_MacroModelRecord]: torsion = _read_torsion_line(line) if torsion is not None: torsions.append(torsion) + self._records = records return self._records @property @@ -397,7 +296,7 @@ def molecules(self) -> list[Molecule]: return [_molecule_from_macromodel_record(record) for record in self._records_parsed] -class MacroModelLog: +class MacroModelLog(_TextFileBase): """Retrieve data from MacroModel log files. The Hessian matrix read from these files is mass-weighted. @@ -410,39 +309,9 @@ def __init__(self, path: str) -> None: path (str): Path to the MacroModel log file. """ - self._lines = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) - self._hessian = None - self._records = None - - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) + super().__init__(path) + self._hessian: np.ndarray | None = None + self._records: list[_MacroModelRecord] | None = None @property def hessian(self) -> np.ndarray: @@ -457,7 +326,10 @@ def hessian(self) -> np.ndarray: logger.log(10, f"READING: {self.filename}") with open(self.path) as f: lines = f.read() - num_atoms = int(re.search(r"Read\s+(\d+)\s+atoms.", lines).group(1)) + atoms_match = re.search(r"Read\s+(\d+)\s+atoms.", lines) + if atoms_match is None: + raise ValueError(f"Could not find atom count in MacroModel log {self.filename}.") + num_atoms = int(atoms_match.group(1)) logger.log(5, f" -- Read {num_atoms} atoms.") hessian = np.zeros([num_atoms * 3, num_atoms * 3], dtype=float) @@ -467,20 +339,17 @@ def hessian(self) -> np.ndarray: start_row = False start_col = False row_num = 0 - col_nums = [] - elements = [] + col_nums: list[int] = [] + elements: list[float] = [] for i, word in enumerate(words): - # 1. Start of Hessian section. if word == "Mass-weighted": section_hessian = True continue - # 5. End of Hessian. Add last row of Hessian and break. if word == "Eigenvalues:": for col_num, element in zip(col_nums, elements): hessian[row_num - 1, col_num - 1] = element section_hessian = False break - # 4. End of a Hessian row. Add to matrix and reset. if section_hessian and start_col and word == "Element": for col_num, element in zip(col_nums, elements): hessian[row_num - 1, col_num - 1] = element @@ -490,15 +359,12 @@ def hessian(self) -> np.ndarray: col_nums = [] elements = [] continue - # 2. Start of a Hessian row. if section_hessian and word == "Element": row_num = int(words[i + 1]) col_nums = [] elements = [] start_row = True continue - # 3. Okay, made it through the row number. Now look for columns - # and elements. if section_hessian and start_row and word == ":": start_row = False start_col = True @@ -521,61 +387,42 @@ def _records_parsed(self) -> list[_MacroModelRecord]: """ if self._records is None: logger.log(10, f"READING: {self.filename}") - self._records = [] + records: list[_MacroModelRecord] = [] with open(self.path) as f: count_current = 0 count_input = 0 count_structure = 0 count_previous = 0 - bonds = [] - angles = [] - torsions = [] - current_record = None - section = None + bonds: list[_MacroModelTerm] = [] + angles: list[_MacroModelTerm] = [] + torsions: list[_MacroModelTerm] = [] + current_record: _MacroModelRecord | None = None + section: str | None = None for line in f: if "m_atom" in line: section = "atom" elif "m_bond" in line: section = "bond" - elif ":::" in line and "ready" not in section: + elif ":::" in line and section is not None and "ready" not in section: section = section + "ready" - elif ":::" in line and "ready" in section: + elif ":::" in line and section is not None and "ready" in section: section = None - elif section == "atom ready": - # read in atoms to list - continue - elif section == "bond ready": - # read in bond atom numbers, populate later with atoms + elif section in ("atom ready", "bond ready"): continue else: continue - # The MMO file lists the bonds, angles, and torsions in - # some order that I am unsure of. It seems consistent - # with the same filename but with two files with the - # exact same structure the ordering is off. This - # reorders the lists before being added to the record. if "Input filename" in line: count_input += 1 if "Input Structure Name" in line: count_structure += 1 count_previous = count_current - # Sometimes only one of the above ("Input filename" and - # "Input Structure Name") is used, sometimes both are used. - # count_current will make sure you catch both. count_current = max(count_input, count_structure) - # If these don't match, then we reached the end of a - # structure. if count_current != count_previous: bonds = [] angles = [] torsions = [] current_record = _MacroModelRecord(self.filename) - self._records.append(current_record) - # For each structure we come across, look for sections that - # we are interested in: those pertaining to bonds, angles, - # and torsions. Of course more could be added. We set the - # section to None to mark the end of a section, and we leave - # it None for parts of the file we don't care about. + records.append(current_record) if "BOND LENGTHS AND STRETCH ENERGIES" in line: section = "bond" if "ANGLES, BEND AND STRETCH BEND ENERGIES" in line: @@ -598,15 +445,14 @@ def _records_parsed(self) -> list[_MacroModelRecord]: torsion = _read_torsion_line(line) if torsion is not None: torsions.append(torsion) - if "Connection Table" in line: - # Sort the bonds, angles, and torsions before the start - # of a new structure + if "Connection Table" in line and current_record is not None: bonds.sort(key=lambda x: (x.atom_nums[0], x.atom_nums[1])) angles.sort(key=lambda x: (x.atom_nums[1], x.atom_nums[0], x.atom_nums[2])) torsions.sort(key=lambda x: (x.atom_nums[1], x.atom_nums[2], x.atom_nums[0], x.atom_nums[3])) current_record.bonds = bonds current_record.angles = angles current_record.torsions = torsions + self._records = records logger.log(5, f" -- Imported {len(self._records)} structure(s).") return self._records diff --git a/q2mm/io/mm3.py b/q2mm/io/mm3.py index 9e9b31f4..4738d1fc 100644 --- a/q2mm/io/mm3.py +++ b/q2mm/io/mm3.py @@ -21,6 +21,7 @@ _match_bond_for_export, _match_sb_for_export, _normalize_equilibrium_angle, + _resolve_template_path, _split_env_id, _torsion_file_value, _validate_form_for_format, @@ -1005,12 +1006,12 @@ def load_mm3_fld(path: str | Path, *, include_standard: bool = True) -> ForceFie atom_types = [t.strip() for t in row.atom_types if t.strip() and t.strip() != "-"] if row.ptype == "bf" and len(atom_types) >= 2: - elems = tuple(_extract_element(t) for t in atom_types[:2]) + t0, t1 = atom_types[:2] env_id = canonicalize_bond_env_id(atom_types[:2]) eq_val = eq_lookup.get(("be", row.ff_row), 0.0) bonds.append( BondParam( - elements=elems, + elements=(_extract_element(t0), _extract_element(t1)), equilibrium=eq_val, force_constant=mm3_bond_k_to_canonical(row.value), label=f"MM3 row {row.ff_row}", @@ -1023,12 +1024,12 @@ def load_mm3_fld(path: str | Path, *, include_standard: bool = True) -> ForceFie ) elif row.ptype == "af" and len(atom_types) >= 3: - elems = tuple(_extract_element(t) for t in atom_types[:3]) + t0, t1, t2 = atom_types[:3] env_id = canonicalize_angle_env_id(atom_types[:3]) eq_val = eq_lookup.get(("ae", row.ff_row), 0.0) angles.append( AngleParam( - elements=elems, + elements=(_extract_element(t0), _extract_element(t1), _extract_element(t2)), equilibrium=eq_val, force_constant=mm3_angle_k_to_canonical(row.value), label=f"MM3 row {row.ff_row}", @@ -1037,10 +1038,11 @@ def load_mm3_fld(path: str | Path, *, include_standard: bool = True) -> ForceFie ) ) - elif row.ptype == "df" and len(atom_types) >= 4: - elems = tuple(_extract_element(t) for t in atom_types[:4]) + elif row.ptype in ("df", "imp1", "imp2") and len(atom_types) >= 4: + t0, t1, t2, t3 = atom_types[:4] env_id = "-".join(t.strip() for t in atom_types[:4]) - periodicity = row.ff_col + is_improper = row.ptype != "df" + periodicity = row.ff_col if not is_improper else (1 if row.ptype == "imp1" else 2) # MM3 torsion alternates signs by order: # (V1/2)(1+cos ω) + (V2/2)(1−cos 2ω) + (V3/2)(1+cos 3ω) # + (V4/2)(1−cos 4ω) + (V5/2)(1+cos 5ω) + (V6/2)(1−cos 6ω) @@ -1048,42 +1050,31 @@ def load_mm3_fld(path: str | Path, *, include_standard: bool = True) -> ForceFie # k*(1+cos(nφ−γ)) with k = V_n/2. Even orders need γ=180° for # the minus sign: (1+cos(nω−π)) = (1−cos nω). phase = 180.0 if periodicity % 2 == 0 else 0.0 + label_kind = " imp" if is_improper else "" torsions.append( TorsionParam( - elements=elems, + elements=( + _extract_element(t0), + _extract_element(t1), + _extract_element(t2), + _extract_element(t3), + ), periodicity=periodicity, force_constant=row.value / 2.0, phase=phase, - label=f"MM3 row {row.ff_row} V{periodicity}", + label=f"MM3 row {row.ff_row}{label_kind} V{periodicity}", env_id=env_id, ff_row=row.ff_row, - ) - ) - - elif row.ptype in ("imp1", "imp2") and len(atom_types) >= 4: - elems = tuple(_extract_element(t) for t in atom_types[:4]) - env_id = "-".join(t.strip() for t in atom_types[:4]) - periodicity = 1 if row.ptype == "imp1" else 2 - phase = 180.0 if periodicity == 2 else 0.0 - torsions.append( - TorsionParam( - elements=elems, - periodicity=periodicity, - force_constant=row.value / 2.0, - phase=phase, - label=f"MM3 row {row.ff_row} imp V{periodicity}", - env_id=env_id, - ff_row=row.ff_row, - is_improper=True, + is_improper=is_improper, ) ) elif row.ptype == "sb" and len(atom_types) >= 3: - elems = tuple(_extract_element(t) for t in atom_types[:3]) + t0, t1, t2 = atom_types[:3] env_id = canonicalize_angle_env_id(atom_types[:3]) stretch_bends.append( StretchBendParam( - elements=elems, + elements=(_extract_element(t0), _extract_element(t1), _extract_element(t2)), force_constant=mm3_sb_k_to_canonical(row.value), label=f"MM3 row {row.ff_row} SB", env_id=env_id, @@ -1093,11 +1084,11 @@ def load_mm3_fld(path: str | Path, *, include_standard: bool = True) -> ForceFie ff = ForceField( name=f"MM3 from {Path(path).name}", - bonds=bonds, - angles=angles, - stretch_bends=stretch_bends, - torsions=torsions, - vdws=vdws, + bonds=tuple(bonds), + angles=tuple(angles), + stretch_bends=tuple(stretch_bends), + torsions=tuple(torsions), + vdws=tuple(vdws), source_path=Path(path), source_format="mm3_fld", functional_form=FunctionalForm.MM3, @@ -1148,9 +1139,7 @@ def save_mm3_fld( """ _validate_form_for_format(ff, "mm3_fld") output_path = Path(path) - template = Path(template_path) if template_path is not None else None - if template is None and ff.source_format == "mm3_fld" and ff.source_path is not None: - template = ff.source_path + template = _resolve_template_path(ff, template_path, "mm3_fld") if template is not None: template_rows, template_lines = _mm3_import_ff(template) diff --git a/q2mm/io/mol2.py b/q2mm/io/mol2.py index d83673df..5ca63b6b 100644 --- a/q2mm/io/mol2.py +++ b/q2mm/io/mol2.py @@ -7,13 +7,13 @@ from __future__ import annotations import logging -import os from dataclasses import dataclass, field import numpy as np from q2mm.elements import MASSES, TWO_LETTER_ELEMENTS from q2mm.geometry import bond_length +from q2mm.io._helpers import _TextFileBase from q2mm.models.identifiers import canonicalize_bond_env_id from q2mm.models.molecule import Bond, Molecule @@ -148,7 +148,7 @@ def _atom(num: int) -> _Mol2Atom: ) -class Mol2: +class Mol2(_TextFileBase): """Retrieve structural data from Tripos Mol2 files. Please ensure that mol2 atom types match the atom types specified @@ -165,7 +165,7 @@ class Mol2: ATOM_FLAG = "ATOM" BOND_FLAG = "BOND" - __slots__ = ["_lines", "path", "directory", "filename", "_records"] + __slots__ = ("_records",) def __init__(self, path: str) -> None: """Initialize a Mol2 instance. @@ -174,39 +174,9 @@ def __init__(self, path: str) -> None: path (str): Absolute path of the mol2 file. """ - self._lines: list[str] | None = None - self.path = os.path.abspath(path) - self.directory = os.path.dirname(self.path) - self.filename = os.path.basename(self.path) + super().__init__(path) self._records: list[_Mol2Record] | None = None - @property - def lines(self) -> list[str]: - """Return the lines of the file. - - Returns: - (list[str]): Lines of the file. - - """ - if self._lines is None: - with open(self.path) as f: - self._lines = f.readlines() - return self._lines - - def write(self, path: str, lines: list[str] | None = None) -> None: - """Write lines to file at path. - - Args: - path (str): Location of file to write. - lines (list[str] | None): Lines to write. Defaults to - ``self.lines``. - - """ - if lines is None: - lines = self.lines - with open(path, "w") as f: - f.writelines(lines) - @property def _records_parsed(self) -> list[_Mol2Record]: """Records extracted from the mol2 file (private staging). diff --git a/q2mm/io/openmm.py b/q2mm/io/openmm.py index bb71d64f..f0cb5d05 100644 --- a/q2mm/io/openmm.py +++ b/q2mm/io/openmm.py @@ -257,7 +257,7 @@ def save_openmm_xml( # Collect unique (element, atom_type) pairs across all molecules seen_types: dict[str, tuple[str, float]] = {} for mol in molecules: - for symbol, atype in zip(mol.symbols, mol.atom_types): + for symbol, atype in zip(mol.symbols, mol.atom_types or mol.symbols): type_name = atype if atype else symbol if type_name not in seen_types: if symbol not in MASSES: @@ -284,7 +284,7 @@ def save_openmm_xml( for mol_idx, mol in enumerate(molecules): res_name = f"Q2MM{mol_idx}" if len(molecules) > 1 else "Q2MM" res_el = ET.SubElement(residues_el, "Residue", name=res_name) - for i, (symbol, atype) in enumerate(zip(mol.symbols, mol.atom_types)): + for i, (symbol, atype) in enumerate(zip(mol.symbols, mol.atom_types or mol.symbols)): type_name = atype if atype else symbol ET.SubElement( res_el, @@ -292,7 +292,7 @@ def save_openmm_xml( name=f"{symbol}{i + 1}", type=f"q2mm-{type_name}", ) - for bond in mol.bonds: + for bond in mol.bonds or (): ET.SubElement( res_el, "Bond", @@ -311,15 +311,15 @@ def save_openmm_xml( ET.SubElement(bond_force_el, "PerBondParameter", name="k") ET.SubElement(bond_force_el, "PerBondParameter", name="r0") - for bond in ff.bonds: - k_openmm = canonical_to_openmm_bond_k(bond.force_constant) - r0_nm = ang_to_nm(bond.equilibrium) + for bond_param in ff.bonds: + k_bond_openmm = canonical_to_openmm_bond_k(bond_param.force_constant) + r0_nm = ang_to_nm(bond_param.equilibrium) - class1, class2 = bond.env_id.split("-") if bond.env_id else bond.elements + class1, class2 = bond_param.env_id.split("-") if bond_param.env_id else bond_param.elements bond_el = ET.SubElement(bond_force_el, "Bond") bond_el.set("class1", class1) bond_el.set("class2", class2) - bond_el.set("k", f"{k_openmm:.6f}") + bond_el.set("k", f"{k_bond_openmm:.6f}") bond_el.set("r0", f"{r0_nm:.6f}") # ---- Custom angle force (MM3 sextic bend) ---- @@ -341,7 +341,7 @@ def save_openmm_xml( ET.SubElement(angle_force_el, "PerAngleParameter", name="theta0") for angle in ff.angles: - k_openmm = canonical_to_openmm_angle_k(angle.force_constant) + k_angle_openmm = canonical_to_openmm_angle_k(angle.force_constant) theta0_rad = deg_to_rad(angle.equilibrium) class1, class2, class3 = angle.env_id.split("-") if angle.env_id else angle.elements @@ -349,7 +349,7 @@ def save_openmm_xml( angle_el.set("class1", class1) angle_el.set("class2", class2) angle_el.set("class3", class3) - angle_el.set("k", f"{k_openmm:.6f}") + angle_el.set("k", f"{k_angle_openmm:.6f}") angle_el.set("theta0", f"{theta0_rad:.6f}") # ---- Custom torsion force ---- diff --git a/q2mm/io/reference.py b/q2mm/io/reference.py index 1c0fdc61..7fbf0c9e 100644 --- a/q2mm/io/reference.py +++ b/q2mm/io/reference.py @@ -191,8 +191,8 @@ def _parse_datum( ] # ---- Scalar kinds ----------------------------------------------------- - if kind in ("energy",): - value = _as_float(_require_key(datum, "value", context), "value", context) + value = _as_float(_require_key(datum, "value", context), "value", context) + if kind == "energy": return [ Observation( kind=kind, @@ -204,7 +204,6 @@ def _parse_datum( ] if kind in ("bond_length", "bond_angle", "torsion_angle"): - value = _as_float(_require_key(datum, "value", context), "value", context) atoms_raw = datum.get("atoms") data_idx_raw = datum.get("data_idx") @@ -246,7 +245,6 @@ def _parse_datum( raise ReferenceYAMLError(f"'{kind}' requires either 'atoms' or 'data_idx' in {context}.") if kind == "eig_diagonal": - value = _as_float(_require_key(datum, "value", context), "value", context) mode_idx = _as_int(_require_key(datum, "mode_idx", context), "mode_idx", context) if mode_idx < 0: raise ReferenceYAMLError(f"'mode_idx' must be non-negative in {context}, got {mode_idx}.") @@ -261,35 +259,18 @@ def _parse_datum( ) ] - if kind == "eig_offdiagonal": - value = _as_float(_require_key(datum, "value", context), "value", context) + if kind in ("eig_offdiagonal", "hessian_element"): row = _as_int(_require_key(datum, "row", context), "row", context) col = _as_int(_require_key(datum, "col", context), "col", context) if row < 0 or col < 0: raise ReferenceYAMLError(f"'row' and 'col' must be non-negative in {context}, got row={row}, col={col}.") + prefix = "eig" if kind == "eig_offdiagonal" else "hess" return [ Observation( - kind="eig_offdiagonal", - value=value, - weight=weight, - label=label if "label" in datum else f"eig[{row},{col}]", - case_id=case_id, - atom_indices=(row, col), - ) - ] - - if kind == "hessian_element": - value = _as_float(_require_key(datum, "value", context), "value", context) - row = _as_int(_require_key(datum, "row", context), "row", context) - col = _as_int(_require_key(datum, "col", context), "col", context) - if row < 0 or col < 0: - raise ReferenceYAMLError(f"'row' and 'col' must be non-negative in {context}, got row={row}, col={col}.") - return [ - Observation( - kind="hessian_element", + kind=kind, value=value, weight=weight, - label=label if "label" in datum else f"hess[{row},{col}]", + label=label if "label" in datum else f"{prefix}[{row},{col}]", case_id=case_id, atom_indices=(row, col), ) @@ -297,7 +278,6 @@ def _parse_datum( if kind == "frequency": # Single frequency entry (no bulk 'values') - value = _as_float(_require_key(datum, "value", context), "value", context) data_idx = _as_int(datum.get("data_idx", datum.get("index", 0)), "data_idx", context) if data_idx < 0: raise ReferenceYAMLError(f"'data_idx' must be non-negative in {context}, got {data_idx}.") diff --git a/q2mm/io/tinker.py b/q2mm/io/tinker.py index 8c457d8c..51e71fa9 100644 --- a/q2mm/io/tinker.py +++ b/q2mm/io/tinker.py @@ -9,9 +9,12 @@ from collections.abc import Callable, Hashable, Sequence from pathlib import Path +from typing import Any, cast + from q2mm.io._helpers import ( _clean_atom_types, _normalize_equilibrium_angle, + _resolve_template_path, _split_env_id, _validate_form_for_format, ) @@ -283,6 +286,9 @@ def _tinker_import_ff(path: str | Path) -> tuple[list[_TinkerParameterRow], list for slot, (amplitude, _phase, _fold) in enumerate(_tinker_torsion_terms(parts, row), start=1): rows.append(_TinkerParameterRow("df", amplitude, row, slot, parts[1:5])) continue + atom_count: int + ptypes: tuple[str, ...] + allowed_lengths: tuple[int, ...] if record in _BONDS: atom_count, ptypes, allowed_lengths = 2, ("bf", "be"), (5,) elif record in _ANGLES: @@ -430,12 +436,12 @@ def _elem(atom_type: str) -> str: atom_types = _clean_atom_types(row.atom_types, 4) if row.ptype == "bf" and len(atom_types) >= 2: - elems = tuple(_elem(t) for t in atom_types[:2]) + t0, t1 = atom_types[:2] env_id = canonicalize_bond_env_id(atom_types[:2]) eq_val = eq_lookup.get(("be", row.ff_row), 0.0) bonds.append( BondParam( - elements=elems, + elements=(_elem(t0), _elem(t1)), equilibrium=eq_val, force_constant=mm3_bond_k_to_canonical(row.value), label=f"Tinker row {row.ff_row}", @@ -444,12 +450,12 @@ def _elem(atom_type: str) -> str: ) ) elif row.ptype == "af" and len(atom_types) >= 3: - elems = tuple(_elem(t) for t in atom_types[:3]) + t0, t1, t2 = atom_types[:3] env_id = canonicalize_angle_env_id(atom_types[:3]) eq_val = eq_lookup.get(("ae", row.ff_row), 0.0) angles.append( AngleParam( - elements=elems, + elements=(_elem(t0), _elem(t1), _elem(t2)), equilibrium=eq_val, force_constant=mm3_angle_k_to_canonical(row.value), label=f"Tinker row {row.ff_row}", @@ -458,7 +464,7 @@ def _elem(atom_type: str) -> str: ) ) elif row.ptype == "df" and len(atom_types) >= 4: - elems = tuple(_elem(t) for t in atom_types[:4]) + t0, t1, t2, t3 = atom_types[:4] env_id = canonicalize_torsion_env_id(atom_types[:4]) amplitude, phase, periodicity = _tinker_torsion_terms(_tinker_tokens(lines[row.ff_row - 1]), row.ff_row)[ row.ff_col - 1 @@ -466,7 +472,7 @@ def _elem(atom_type: str) -> str: coefficient = _tinker_scaled_torsion(amplitude, torsion_unit, row.ff_row) torsions.append( TorsionParam( - elements=elems, + elements=(_elem(t0), _elem(t1), _elem(t2), _elem(t3)), periodicity=periodicity, force_constant=coefficient, phase=phase, @@ -491,10 +497,10 @@ def _elem(atom_type: str) -> str: return ForceField( name=f"Tinker from {Path(path).name}", - bonds=bonds, - angles=angles, - torsions=torsions, - vdws=vdws, + bonds=tuple(bonds), + angles=tuple(angles), + torsions=tuple(torsions), + vdws=tuple(vdws), source_path=Path(path), source_format="tinker_prm", functional_form=FunctionalForm.MM3, @@ -530,7 +536,7 @@ def _tinker_template_pairs[TemplateParam: (BondParam, AngleParam, TorsionParam, if isinstance(comparison, BondParam) and isinstance(expected, BondParam): if comparison.context in ("", "0000 0000") and expected.context in ("", "0000 0000"): # Compare generic aliases without changing either bound parameter. - comparison = dataclasses.replace(comparison, context=expected.context) + comparison = cast(TemplateParam, dataclasses.replace(cast(Any, comparison), context=expected.context)) if comparison != expected: raise ValueError(f"Tinker row {original.ff_row}: template cannot represent non-scalar parameter edits") for field in editable: @@ -559,11 +565,12 @@ def _tinker_replace_token(lines: list[str], row: int, column: int, value: float) def _tinker_torsion_for_binding(param: TorsionParam) -> TorsionParam: """Use proper-torsion reversal equivalence without rewriting source columns.""" atom_types = _split_env_id(param.env_id, 4) - elements = tuple(param.elements) + e0, e1, e2, e3 = param.elements + elements = min((e0, e1, e2, e3), (e3, e2, e1, e0)) return dataclasses.replace( param, env_id=canonicalize_torsion_env_id(atom_types) if atom_types else param.env_id, - elements=min(elements, elements[::-1]), + elements=elements, ) @@ -573,47 +580,47 @@ def _tinker_template_lines(ff: ForceField, template: Path) -> list[str]: lines = f.readlines() unit = _tinker_torsion_unit(lines) - for before, after in _tinker_template_pairs( + for bond_before, bond_after in _tinker_template_pairs( original.bonds, ff.bonds, lambda p: p.env_id, ("force_constant", "equilibrium") ): - assert before.ff_row is not None - if before.force_constant != after.force_constant: - _tinker_replace_token(lines, before.ff_row, 3, canonical_to_mm3_bond_k(after.force_constant)) - if before.equilibrium != after.equilibrium: - _tinker_replace_token(lines, before.ff_row, 4, after.equilibrium) - for before, after in _tinker_template_pairs( + assert bond_before.ff_row is not None + if bond_before.force_constant != bond_after.force_constant: + _tinker_replace_token(lines, bond_before.ff_row, 3, canonical_to_mm3_bond_k(bond_after.force_constant)) + if bond_before.equilibrium != bond_after.equilibrium: + _tinker_replace_token(lines, bond_before.ff_row, 4, bond_after.equilibrium) + for angle_before, angle_after in _tinker_template_pairs( original.angles, ff.angles, lambda p: p.env_id, ("force_constant", "equilibrium") ): - assert before.ff_row is not None - if before.force_constant != after.force_constant: - _tinker_replace_token(lines, before.ff_row, 4, canonical_to_mm3_angle_k(after.force_constant)) - if before.equilibrium != after.equilibrium: - _tinker_replace_token(lines, before.ff_row, 5, _normalize_equilibrium_angle(after.equilibrium)) - for before, after in _tinker_template_pairs( + assert angle_before.ff_row is not None + if angle_before.force_constant != angle_after.force_constant: + _tinker_replace_token(lines, angle_before.ff_row, 4, canonical_to_mm3_angle_k(angle_after.force_constant)) + if angle_before.equilibrium != angle_after.equilibrium: + _tinker_replace_token(lines, angle_before.ff_row, 5, _normalize_equilibrium_angle(angle_after.equilibrium)) + for tor_before, tor_after in _tinker_template_pairs( tuple(_tinker_torsion_for_binding(param) for param in original.torsions), tuple(_tinker_torsion_for_binding(param) for param in ff.torsions), lambda p: (p.env_id, p.periodicity), ("force_constant", "phase"), ): - assert before.ff_row is not None - terms = _tinker_torsion_terms(_tinker_tokens(lines[before.ff_row - 1]), before.ff_row) - slot = next(index for index, term in enumerate(terms) if term[2] == before.periodicity) - if before.force_constant != after.force_constant: - amplitude = after.force_constant / unit - if after.force_constant != 0.0 and amplitude == 0.0: - raise ValueError(f"Tinker row {before.ff_row}: torsion scaling underflow") - _tinker_replace_token(lines, before.ff_row, 5 + 3 * slot, amplitude) - serialized = _tinker_tokens(lines[before.ff_row - 1])[5 + 3 * slot] - _tinker_scaled_torsion(_tinker_float(serialized, before.ff_row), unit, before.ff_row) - if before.phase != after.phase: - _tinker_replace_token(lines, before.ff_row, 6 + 3 * slot, after.phase) - for before, after in _tinker_template_pairs( + assert tor_before.ff_row is not None + terms = _tinker_torsion_terms(_tinker_tokens(lines[tor_before.ff_row - 1]), tor_before.ff_row) + slot = next(index for index, term in enumerate(terms) if term[2] == tor_before.periodicity) + if tor_before.force_constant != tor_after.force_constant: + amplitude = tor_after.force_constant / unit + if tor_after.force_constant != 0.0 and amplitude == 0.0: + raise ValueError(f"Tinker row {tor_before.ff_row}: torsion scaling underflow") + _tinker_replace_token(lines, tor_before.ff_row, 5 + 3 * slot, amplitude) + serialized = _tinker_tokens(lines[tor_before.ff_row - 1])[5 + 3 * slot] + _tinker_scaled_torsion(_tinker_float(serialized, tor_before.ff_row), unit, tor_before.ff_row) + if tor_before.phase != tor_after.phase: + _tinker_replace_token(lines, tor_before.ff_row, 6 + 3 * slot, tor_after.phase) + for vdw_before, vdw_after in _tinker_template_pairs( original.vdws, ff.vdws, lambda p: p.atom_type, ("radius", "epsilon", "reduction") ): - assert before.ff_row is not None + assert vdw_before.ff_row is not None for column, field in enumerate(("radius", "epsilon", "reduction"), start=2): - if getattr(before, field) != getattr(after, field): - _tinker_replace_token(lines, before.ff_row, column, getattr(after, field)) + if getattr(vdw_before, field) != getattr(vdw_after, field): + _tinker_replace_token(lines, vdw_before.ff_row, column, getattr(vdw_after, field)) return lines @@ -668,9 +675,7 @@ def save_tinker_prm( _validate_form_for_format(ff, "tinker_prm") _validate_tinker_export_terms(ff, supports_reduction=True) output_path = Path(path) - template = Path(template_path) if template_path is not None else None - if template is None and ff.source_format == "tinker_prm" and ff.source_path is not None: - template = ff.source_path + template = _resolve_template_path(ff, template_path, "tinker_prm") if template is not None: lines = _tinker_template_lines(ff, template) diff --git a/q2mm/models/forcefield.py b/q2mm/models/forcefield.py index aa9dc83f..5f232e66 100644 --- a/q2mm/models/forcefield.py +++ b/q2mm/models/forcefield.py @@ -80,7 +80,8 @@ class BondParam: @property def key(self) -> tuple[str, str]: """Sorted element pair for canonical matching (e.g., ``('C', 'F')``).""" - return tuple(sorted(self.elements)) + a, b = self.elements + return (a, b) if a <= b else (b, a) @dataclass(frozen=True) @@ -107,8 +108,8 @@ class AngleParam: @property def key(self) -> tuple[str, str, str]: """Canonical key: center fixed, outers sorted.""" - outer = tuple(sorted([self.elements[0], self.elements[2]])) - return (outer[0], self.elements[1], outer[1]) + a, b, c = self.elements + return (a, b, c) if a <= c else (c, b, a) @dataclass(frozen=True) @@ -132,8 +133,8 @@ class StretchBendParam: @property def key(self) -> tuple[str, str, str]: """Canonical key: center fixed, outers sorted.""" - outer = tuple(sorted([self.elements[0], self.elements[2]])) - return (outer[0], self.elements[1], outer[1]) + a, b, c = self.elements + return (a, b, c) if a <= c else (c, b, a) @dataclass(frozen=True) @@ -622,43 +623,33 @@ def create_for_molecule( """ # Unique bond types bond_types: dict[tuple[str, str], list[float]] = {} - for bond in molecule.bonds: - key = bond.element_pair - if key not in bond_types: - bond_types[key] = [] - bond_types[key].append(bond.length) - - bonds = [] - for key, lengths in bond_types.items(): - avg_len = np.mean(lengths) - bonds.append( - BondParam( - elements=key, - equilibrium=avg_len, - force_constant=default_bond_k, - label=f"{key[0]}-{key[1]} (auto)", - ) + for bond in molecule.bonds or (): + bond_types.setdefault(bond.element_pair, []).append(bond.length) + + bonds = [ + BondParam( + elements=key, + equilibrium=float(np.mean(lengths)), + force_constant=default_bond_k, + label=f"{key[0]}-{key[1]} (auto)", ) + for key, lengths in bond_types.items() + ] # Unique angle types angle_types: dict[tuple[str, str, str], list[float]] = {} - for angle in molecule.angles: - key = angle.element_triple - if key not in angle_types: - angle_types[key] = [] - angle_types[key].append(angle.value) - - angles = [] - for key, values in angle_types.items(): - avg_val = np.mean(values) - angles.append( - AngleParam( - elements=key, - equilibrium=avg_val, - force_constant=default_angle_k, - label=f"{key[0]}-{key[1]}-{key[2]} (auto)", - ) + for angle in molecule.angles or (): + angle_types.setdefault(angle.element_triple, []).append(angle.value) + + angles = [ + AngleParam( + elements=key, + equilibrium=float(np.mean(values)), + force_constant=default_angle_k, + label=f"{key[0]}-{key[1]}-{key[2]} (auto)", ) + for key, values in angle_types.items() + ] return cls( name=name or f"Auto FF for {molecule.name}", diff --git a/q2mm/models/hessian.py b/q2mm/models/hessian.py index be41ab22..c33afcda 100644 --- a/q2mm/models/hessian.py +++ b/q2mm/models/hessian.py @@ -20,10 +20,10 @@ import copy import logging import warnings -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum -from typing import Any, Literal +from typing import Any, Literal, cast import numpy as np @@ -33,6 +33,10 @@ logger = logging.getLogger(__name__) +_BOHR_TO_M: float = co.BOHR_TO_ANG * 1e-10 +_EIGVAL_AU_TO_SI: float = co.HARTREE_TO_J / (co.AMU_TO_KG * _BOHR_TO_M**2) +_RAD_S_TO_CM1_DENOM: float = 2.0 * np.pi * co.SPEED_OF_LIGHT_MS * 100.0 + def _resolve_symbols(atoms_or_symbols: Sequence[str] | object) -> list[str]: """Normalise the *atoms* argument to a plain list of element symbols. @@ -47,15 +51,20 @@ def _resolve_symbols(atoms_or_symbols: Sequence[str] | object) -> list[str]: """ # Molecule (has .symbols attribute) if hasattr(atoms_or_symbols, "symbols"): - return list(atoms_or_symbols.symbols) + return [str(s) for s in getattr(atoms_or_symbols, "symbols")] + + if not isinstance(atoms_or_symbols, Iterable) or isinstance(atoms_or_symbols, (str, bytes)): + raise TypeError(f"Expected element symbols (list[str]) or a Molecule; got {type(atoms_or_symbols).__name__}.") - items = list(atoms_or_symbols) - if not items: + raw_items = list(atoms_or_symbols) + if not raw_items: return [] - if not isinstance(items[0], str): - raise TypeError(f"Expected element symbols (list[str]) or a Molecule; got list of {type(items[0]).__name__}.") + if not all(isinstance(s, str) for s in raw_items): + first_bad = next(s for s in raw_items if not isinstance(s, str)) + raise TypeError(f"Expected element symbols (list[str]) or a Molecule; got list of {type(first_bad).__name__}.") + items = [str(s) for s in raw_items] unknown = [s for s in items if s not in co.MASSES] if unknown: raise ValueError( @@ -167,7 +176,7 @@ def hessian_to_atomic_units(hessian: Any, units: HessianUnits | str | None) -> n stripped = _strip_pint(hessian) if stripped is not hessian: # Was pint-like — its own units are authoritative, already AU. - return stripped + return cast(np.ndarray, stripped) array = np.asarray(hessian, dtype=float) unit_value = getattr(units, "value", units) if unit_value == HessianUnits.ATOMIC.value: @@ -310,16 +319,13 @@ def hessian_to_frequencies( raise # 5. Convert eigenvalues [Hartree / (amu · Bohr²)] → s⁻² - bohr_to_m = co.BOHR_TO_ANG * 1e-10 - factor = co.HARTREE_TO_J / (co.AMU_TO_KG * bohr_to_m**2) - vals_si = eigenvalues * factor + vals_si = eigenvalues * _EIGVAL_AU_TO_SI # 6. eigenvalue → angular frequency → cm⁻¹ # Negative eigenvalue → imaginary frequency (negative cm⁻¹) - freqs = np.sign(vals_si) * np.sqrt(np.abs(vals_si)) - freqs /= 2.0 * np.pi * co.SPEED_OF_LIGHT_MS * 100.0 + freqs = np.sign(vals_si) * np.sqrt(np.abs(vals_si)) / _RAD_S_TO_CM1_DENOM - result = freqs.tolist() + result: list[float] = freqs.tolist() if sort: result.sort() return result @@ -366,21 +372,14 @@ def frequency_param_jacobian( if hessian_au.shape != (n3, n3): raise ValueError(f"hessian_au must have shape ({n3}, {n3}), got {hessian_au.shape}") - n_params = dH_dp_au.shape[2] - - # Mass-weighting scale: 1/sqrt(m_i * m_j) - inv_sqrt = np.array([1.0 / np.sqrt(co.MASSES[s]) for s in symbols_list for _ in range(3)]) - scale = np.outer(inv_sqrt, inv_sqrt) + scale = mass_weight_scale_3n(symbols_list) # Mass-weight the Hessian - hess = hessian_au.copy() - hess *= scale + hess = hessian_au * scale hess = 0.5 * (hess + hess.T) # Mass-weight each dH/dp slice - mw_dH_dp = np.empty((n3, n3, n_params)) - for j in range(n_params): - mw_dH_dp[:, :, j] = dH_dp_au[:, :, j] * scale + mw_dH_dp = dH_dp_au * scale[:, :, np.newaxis] # Eigendecompose the mass-weighted Hessian eigenvalues, eigenvectors = np.linalg.eigh(hess) @@ -389,22 +388,19 @@ def frequency_param_jacobian( d_eig_dp = np.einsum("ik,ijp,jk->kp", eigenvectors, mw_dH_dp, eigenvectors) # Chain through eigenvalue → frequency conversion - bohr_to_m = co.BOHR_TO_ANG * 1e-10 - factor = co.HARTREE_TO_J / (co.AMU_TO_KG * bohr_to_m**2) - vals_si = eigenvalues * factor - denom = 2.0 * np.pi * co.SPEED_OF_LIGHT_MS * 100.0 + vals_si = eigenvalues * _EIGVAL_AU_TO_SI # freq = sign(λ) * sqrt(|λ|*factor) / denom # d(freq)/d(λ) = factor / (2 * sqrt(|λ|*factor) * denom) # = sqrt(factor) / (2 * sqrt(|λ|) * denom) [regularised] abs_vals_si = np.maximum(np.abs(vals_si), epsilon) - d_freq_d_eig = factor / (2.0 * np.sqrt(abs_vals_si) * denom) + d_freq_d_eig = _EIGVAL_AU_TO_SI / (2.0 * np.sqrt(abs_vals_si) * _RAD_S_TO_CM1_DENOM) # d(freq_k)/dp_j = d(freq_k)/d(λ_k) * d(λ_k)/dp_j d_freq_dp = d_freq_d_eig[:, np.newaxis] * d_eig_dp # Compute frequencies (same as hessian_to_frequencies) - freqs_arr = np.sign(vals_si) * np.sqrt(np.abs(vals_si)) / denom + freqs_arr = np.sign(vals_si) * np.sqrt(np.abs(vals_si)) / _RAD_S_TO_CM1_DENOM if sort: order = np.argsort(freqs_arr) @@ -469,21 +465,17 @@ def _jax_frequency_param_jacobian( d_eig_dp = jnp.einsum("ik,ijp,jk->kp", eigenvectors, mw_dH_dp, eigenvectors) # Chain through eigenvalue → frequency conversion - bohr_to_m = co.BOHR_TO_ANG * 1e-10 - factor = co.HARTREE_TO_J / (co.AMU_TO_KG * bohr_to_m**2) - denom = 2.0 * jnp.pi * co.SPEED_OF_LIGHT_MS * 100.0 - - vals_si = eigenvalues * factor + vals_si = eigenvalues * _EIGVAL_AU_TO_SI # d(freq)/d(λ) with regularisation for near-zero eigenvalues abs_vals_si = jnp.maximum(jnp.abs(vals_si), epsilon) - d_freq_d_eig = factor / (2.0 * jnp.sqrt(abs_vals_si) * denom) + d_freq_d_eig = _EIGVAL_AU_TO_SI / (2.0 * jnp.sqrt(abs_vals_si) * _RAD_S_TO_CM1_DENOM) # d(freq_k)/dp_j = d(freq_k)/d(λ_k) * d(λ_k)/dp_j d_freq_dp = d_freq_d_eig[:, None] * d_eig_dp # Compute frequencies - freqs_arr = jnp.sign(vals_si) * jnp.sqrt(jnp.abs(vals_si)) / denom + freqs_arr = jnp.sign(vals_si) * jnp.sqrt(jnp.abs(vals_si)) / _RAD_S_TO_CM1_DENOM # Sort ascending (argsort is fine here — used only for output ordering, # not differentiated through) @@ -526,13 +518,9 @@ def _jax_frequencies_from_hessian( eigenvalues = jnp.linalg.eigvalsh(hess) - bohr_to_m = co.BOHR_TO_ANG * 1e-10 - factor = co.HARTREE_TO_J / (co.AMU_TO_KG * bohr_to_m**2) - denom = 2.0 * jnp.pi * co.SPEED_OF_LIGHT_MS * 100.0 - - vals_si = eigenvalues * factor - freqs_arr = jnp.sign(vals_si) * jnp.sqrt(jnp.abs(vals_si)) / denom - return jnp.sort(freqs_arr) + vals_si = eigenvalues * _EIGVAL_AU_TO_SI + freqs_arr = jnp.sign(vals_si) * jnp.sqrt(jnp.abs(vals_si)) / _RAD_S_TO_CM1_DENOM + return cast(np.ndarray, jnp.sort(freqs_arr)) def symbols_to_masses_3n(symbols: list[str] | Sequence[str]) -> list[float]: @@ -654,8 +642,7 @@ def reform_hessian(eigenvalues: np.ndarray, eigenvectors: np.ndarray) -> np.ndar np.ndarray: Hessian matrix """ - reformed_hessian = eigenvectors.dot(np.diag(eigenvalues).dot(eigenvectors.T)) - return reformed_hessian + return cast(np.ndarray, eigenvectors.dot(np.diag(eigenvalues).dot(eigenvectors.T))) # ---- Eigenmatrix operations ---- @@ -697,7 +684,7 @@ def transform_to_eigenmatrix( :func:`mass_weight_hessian`). """ - return eigenvectors.T @ hessian @ eigenvectors + return cast(np.ndarray, eigenvectors.T @ hessian @ eigenvectors) def extract_eigenmatrix_data( @@ -833,7 +820,7 @@ def invert_ts_curvature_jax( rest = jnp.where(evals[1:] < 0, 0.0, evals[1:]) new_evals = jnp.concatenate([new_smallest[jnp.newaxis], rest]) - return (evecs * new_evals) @ evecs.T + return cast(np.ndarray, (evecs * new_evals) @ evecs.T) def invert_ts_curvature( diff --git a/q2mm/models/identifiers.py b/q2mm/models/identifiers.py index eebcb3c2..643211ef 100644 --- a/q2mm/models/identifiers.py +++ b/q2mm/models/identifiers.py @@ -23,9 +23,14 @@ def _extract_element(atom_type: str) -> str: return letters[0].upper() +def _clean_env_tokens(atom_types: Sequence[str]) -> list[str]: + """Strip and filter non-empty, non-separator atom-type tokens.""" + return [item.strip() for item in atom_types if item and item.strip() and item.strip() != "-"] + + def canonicalize_bond_env_id(atom_types: Sequence[str]) -> str: """Canonicalize bond atom-type labels so bond direction does not matter.""" - cleaned = [item.strip() for item in atom_types if item and item.strip() and item.strip() != "-"] + cleaned = _clean_env_tokens(atom_types) if len(cleaned) < 2: return "-".join(cleaned) return "-".join(sorted(cleaned[:2])) @@ -33,7 +38,7 @@ def canonicalize_bond_env_id(atom_types: Sequence[str]) -> str: def canonicalize_angle_env_id(atom_types: Sequence[str]) -> str: """Canonicalize angle atom-type labels with fixed center and sorted outers.""" - cleaned = [item.strip() for item in atom_types if item and item.strip() and item.strip() != "-"] + cleaned = _clean_env_tokens(atom_types) if len(cleaned) < 3: return "-".join(cleaned) outer = sorted([cleaned[0], cleaned[2]]) @@ -42,7 +47,7 @@ def canonicalize_angle_env_id(atom_types: Sequence[str]) -> str: def canonicalize_torsion_env_id(atom_types: Sequence[str]) -> str: """Canonicalize torsion atom-type labels so A-B-C-D == D-C-B-A.""" - cleaned = [item.strip() for item in atom_types if item and item.strip() and item.strip() != "-"] + cleaned = _clean_env_tokens(atom_types) if len(cleaned) < 4: return "-".join(cleaned) forward = cleaned[:4] diff --git a/q2mm/models/molecule.py b/q2mm/models/molecule.py index 93c9b069..c9da8926 100644 --- a/q2mm/models/molecule.py +++ b/q2mm/models/molecule.py @@ -21,7 +21,7 @@ import numpy as np -from q2mm.geometry import dihedral_angle +from q2mm.geometry import bond_angle, bond_length, dihedral_angle from q2mm.models.identifiers import ( canonicalize_angle_env_id, canonicalize_bond_env_id, @@ -39,13 +39,22 @@ from q2mm.elements import COVALENT_RADII -def _immutable_array(values: Any, *, dtype: type = float) -> np.ndarray: +def _immutable_array(values: Any, *, dtype: Any = float) -> np.ndarray: """Return a defensive, read-only copy of *values* as a NumPy array.""" - array = np.array(values, dtype=dtype, copy=True) + array: np.ndarray = np.array(values, dtype=dtype, copy=True) array.setflags(write=False) return array +def _build_adjacency(n_atoms: int, bonds: tuple[Bond, ...]) -> dict[int, list[int]]: + """Build an atom-index adjacency map from *bonds*.""" + adj: dict[int, list[int]] = {i: [] for i in range(n_atoms)} + for bond in bonds: + adj[bond.atom_i].append(bond.atom_j) + adj[bond.atom_j].append(bond.atom_i) + return adj + + @dataclass(frozen=True) class Bond: """An immutable bond between two atoms in a :class:`Molecule`. @@ -67,7 +76,8 @@ class Bond: @property def element_pair(self) -> tuple[str, str]: """Sorted element pair for matching (e.g., ('C', 'F')).""" - return tuple(sorted(self.elements)) + a, b = self.elements + return (a, b) if a <= b else (b, a) @dataclass(frozen=True) @@ -85,8 +95,8 @@ class Angle: @property def element_triple(self) -> tuple[str, str, str]: """Canonical element triple: (outer, center, outer) sorted by outer elements.""" - outer = tuple(sorted([self.elements[0], self.elements[2]])) - return (outer[0], self.elements[1], outer[1]) + a, b, c = self.elements + return (a, b, c) if a <= c else (c, b, a) @dataclass(frozen=True) @@ -202,21 +212,23 @@ def __post_init__(self) -> None: raise ValueError("hessian_provenance requires a Hessian.") bonds_explicit = self.bonds is not None - bonds = tuple(self.bonds) if bonds_explicit else self._detect_bonds() + bonds = tuple(self.bonds) if self.bonds is not None else self._detect_bonds(atom_types) object.__setattr__(self, "bonds", bonds) object.__setattr__(self, "bonds_explicit", bonds_explicit) + adj = _build_adjacency(len(symbols), bonds) + angles_explicit = self.angles is not None - angles = tuple(self.angles) if angles_explicit else self._detect_angles(bonds) + angles = tuple(self.angles) if self.angles is not None else self._detect_angles(adj, atom_types) object.__setattr__(self, "angles", angles) object.__setattr__(self, "angles_explicit", angles_explicit) torsions_explicit = self.torsions is not None - torsions = tuple(self.torsions) if torsions_explicit else self._detect_torsions(bonds) + torsions = tuple(self.torsions) if self.torsions is not None else self._detect_torsions(bonds, adj, atom_types) object.__setattr__(self, "torsions", torsions) object.__setattr__(self, "torsions_explicit", torsions_explicit) - object.__setattr__(self, "improper_torsions", self._detect_improper_torsions(bonds)) + object.__setattr__(self, "improper_torsions", self._detect_improper_torsions(adj, atom_types)) @property def n_atoms(self) -> int: @@ -225,7 +237,7 @@ def n_atoms(self) -> int: # ---- Topology detection (geometry-derived) ---- - def _detect_bonds(self) -> tuple[Bond, ...]: + def _detect_bonds(self, atom_types: tuple[str, ...]) -> tuple[Bond, ...]: """Detect bonds based on covalent radii with tolerance factor.""" tolerance = self.bond_tolerance bonds = [] @@ -233,7 +245,7 @@ def _detect_bonds(self) -> tuple[Bond, ...]: for j in range(i + 1, self.n_atoms): ri = COVALENT_RADII.get(self.symbols[i], 0.76) rj = COVALENT_RADII.get(self.symbols[j], 0.76) - dist = float(np.linalg.norm(self.geometry[i] - self.geometry[j])) + dist = bond_length(self.geometry[i], self.geometry[j]) if dist < tolerance * (ri + rj): bonds.append( Bond( @@ -241,54 +253,40 @@ def _detect_bonds(self) -> tuple[Bond, ...]: atom_j=j, elements=(self.symbols[i], self.symbols[j]), length=dist, - env_id=canonicalize_bond_env_id([self.atom_types[i], self.atom_types[j]]), + env_id=canonicalize_bond_env_id([atom_types[i], atom_types[j]]), ) ) return tuple(bonds) - def _detect_angles(self, bonds: tuple[Bond, ...]) -> tuple[Angle, ...]: + def _detect_angles(self, adj: dict[int, list[int]], atom_types: tuple[str, ...]) -> tuple[Angle, ...]: """Detect angles from resolved bonds.""" - adj: dict[int, list[int]] = {i: [] for i in range(self.n_atoms)} - for bond in bonds: - adj[bond.atom_i].append(bond.atom_j) - adj[bond.atom_j].append(bond.atom_i) - angles = [] for center in range(self.n_atoms): neighbors = adj[center] for ii in range(len(neighbors)): for jj in range(ii + 1, len(neighbors)): a, b = neighbors[ii], neighbors[jj] - v1 = self.geometry[a] - self.geometry[center] - v2 = self.geometry[b] - self.geometry[center] - cos_a = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)) - angle_val = np.degrees(np.arccos(np.clip(cos_a, -1, 1))) angles.append( Angle( atom_i=a, atom_j=center, atom_k=b, elements=(self.symbols[a], self.symbols[center], self.symbols[b]), - value=float(angle_val), - env_id=canonicalize_angle_env_id( - [self.atom_types[a], self.atom_types[center], self.atom_types[b]] - ), + value=bond_angle(self.geometry[a], self.geometry[center], self.geometry[b]), + env_id=canonicalize_angle_env_id([atom_types[a], atom_types[center], atom_types[b]]), ) ) return tuple(angles) - def _detect_torsions(self, bonds: tuple[Bond, ...]) -> tuple[Torsion, ...]: + def _detect_torsions( + self, bonds: tuple[Bond, ...], adj: dict[int, list[int]], atom_types: tuple[str, ...] + ) -> tuple[Torsion, ...]: """Detect torsion/dihedral angles from resolved bonds. For each bond B-C, finds all atoms A bonded to B (A != C) and all atoms D bonded to C (D != B) to form torsions A-B-C-D. Deduplicates so that A-B-C-D and D-C-B-A are not both stored. """ - adj: dict[int, list[int]] = {i: [] for i in range(self.n_atoms)} - for bond in bonds: - adj[bond.atom_i].append(bond.atom_j) - adj[bond.atom_j].append(bond.atom_i) - seen: set[tuple[int, int, int, int]] = set() torsions: list[Torsion] = [] for bond in bonds: @@ -319,13 +317,13 @@ def _detect_torsions(self, bonds: tuple[Bond, ...]) -> tuple[Torsion, ...]: ), value=value, env_id=canonicalize_torsion_env_id( - [self.atom_types[a], self.atom_types[b], self.atom_types[c], self.atom_types[d]] + [atom_types[a], atom_types[b], atom_types[c], atom_types[d]] ), ) ) return tuple(torsions) - def _detect_improper_torsions(self, bonds: tuple[Bond, ...]) -> tuple[Torsion, ...]: + def _detect_improper_torsions(self, adj: dict[int, list[int]], atom_types: tuple[str, ...]) -> tuple[Torsion, ...]: """Detect improper torsions at trigonal (sp2) centres. For each atom with exactly 3 bonded neighbours, generates an @@ -335,11 +333,6 @@ def _detect_improper_torsions(self, bonds: tuple[Bond, ...]) -> tuple[Torsion, . Neighbours are sorted by index for deterministic ordering. """ - adj: dict[int, list[int]] = {i: [] for i in range(self.n_atoms)} - for bond in bonds: - adj[bond.atom_i].append(bond.atom_j) - adj[bond.atom_j].append(bond.atom_i) - impropers: list[Torsion] = [] for centre in range(self.n_atoms): nbrs = sorted(adj[centre]) @@ -366,7 +359,7 @@ def _detect_improper_torsions(self, bonds: tuple[Bond, ...]) -> tuple[Torsion, . ), value=value, env_id=canonicalize_torsion_env_id( - [self.atom_types[a], self.atom_types[centre], self.atom_types[b], self.atom_types[d]] + [atom_types[a], atom_types[centre], atom_types[b], atom_types[d]] ), ) ) @@ -402,7 +395,7 @@ def _replace(self, *, reinfer_topology: bool = False, **overrides: Any) -> Molec def _replace_preserving_topology(self, **overrides: Any) -> Molecule: """Replace fields while preserving connectivity and refreshing geometry metadata.""" geometry = np.asarray(overrides.get("geometry", self.geometry), dtype=float) - atom_types = tuple(overrides.get("atom_types", self.atom_types)) + atom_types = tuple(overrides.get("atom_types", self.atom_types or self.symbols)) if geometry.shape != (self.n_atoms, 3): raise ValueError(f"geometry must have shape ({self.n_atoms}, 3), got {geometry.shape}.") if len(atom_types) != self.n_atoms: @@ -411,26 +404,22 @@ def _replace_preserving_topology(self, **overrides: Any) -> Molecule: bonds = tuple( replace( bond, - length=float(np.linalg.norm(geometry[bond.atom_i] - geometry[bond.atom_j])), + length=bond_length(geometry[bond.atom_i], geometry[bond.atom_j]), env_id=canonicalize_bond_env_id([atom_types[bond.atom_i], atom_types[bond.atom_j]]), ) - for bond in self.bonds + for bond in self.bonds or () ) - angles = [] - for angle in self.angles: - v1 = geometry[angle.atom_i] - geometry[angle.atom_j] - v2 = geometry[angle.atom_k] - geometry[angle.atom_j] - cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)) - angles.append( - replace( - angle, - value=float(np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))), - env_id=canonicalize_angle_env_id( - [atom_types[angle.atom_i], atom_types[angle.atom_j], atom_types[angle.atom_k]] - ), - ) + angles = tuple( + replace( + angle, + value=bond_angle(geometry[angle.atom_i], geometry[angle.atom_j], geometry[angle.atom_k]), + env_id=canonicalize_angle_env_id( + [atom_types[angle.atom_i], atom_types[angle.atom_j], atom_types[angle.atom_k]] + ), ) + for angle in self.angles or () + ) torsions = tuple( replace( @@ -450,14 +439,14 @@ def _replace_preserving_topology(self, **overrides: Any) -> Molecule: ] ), ) - for torsion in self.torsions + for torsion in self.torsions or () ) kwargs: dict[str, Any] = { "geometry": geometry, "atom_types": atom_types, "bonds": bonds, - "angles": tuple(angles), + "angles": angles, "torsions": torsions, } kwargs.update(overrides) @@ -547,4 +536,4 @@ def with_overrides( def __repr__(self) -> str: formula = "".join(f"{s}{self.symbols.count(s)}" for s in dict.fromkeys(self.symbols)) hess_str = f", hessian={self.hessian.shape}" if self.hessian is not None else "" - return f"Molecule({formula}, {self.n_atoms} atoms, {len(self.bonds)} bonds{hess_str})" + return f"Molecule({formula}, {self.n_atoms} atoms, {len(self.bonds or ())} bonds{hess_str})" diff --git a/q2mm/models/observations.py b/q2mm/models/observations.py index 855b745d..b83f8c26 100644 --- a/q2mm/models/observations.py +++ b/q2mm/models/observations.py @@ -115,6 +115,16 @@ def _validate_scalar_observation( return normalized_value, normalized_weight, normalized_case_id, str(label) +def _coerce_enum[E: Enum](value: E | str, enum_cls: type[E], label: str) -> E: + """Coerce *value* to *enum_cls*, raising a uniform ValueError on invalid members.""" + if isinstance(value, enum_cls): + return value + try: + return enum_cls(value) + except ValueError as exc: + raise ValueError(f"Unknown {label} {value!r}.") from exc + + @dataclass(frozen=True) class ScanCoordinate: """One explicit coordinate value defining a constrained scan point.""" @@ -125,11 +135,7 @@ class ScanCoordinate: unit: str def __post_init__(self) -> None: - if not isinstance(self.kind, ScanCoordinateKind): - try: - object.__setattr__(self, "kind", ScanCoordinateKind(self.kind)) - except ValueError as exc: - raise ValueError(f"Unknown scan coordinate kind {self.kind!r}.") from exc + object.__setattr__(self, "kind", _coerce_enum(self.kind, ScanCoordinateKind, "scan coordinate kind")) atom_indices = tuple(int(index) for index in self.atom_indices) expected = { ScanCoordinateKind.DISTANCE: (2, "angstrom"), @@ -169,11 +175,7 @@ def __post_init__(self) -> None: ) if not isinstance(self.atom_index, int) or isinstance(self.atom_index, bool) or self.atom_index < 0: raise ValueError("AtomicPartialChargeObservation.atom_index must be a non-negative integer.") - if not isinstance(self.unit, ChargeUnit): - try: - object.__setattr__(self, "unit", ChargeUnit(self.unit)) - except ValueError as exc: - raise ValueError(f"Unknown atomic-charge unit {self.unit!r}.") from exc + object.__setattr__(self, "unit", _coerce_enum(self.unit, ChargeUnit, "atomic-charge unit")) object.__setattr__(self, "value", value) object.__setattr__(self, "weight", weight) object.__setattr__(self, "case_id", case_id) @@ -205,11 +207,9 @@ def __post_init__(self) -> None: raise ValueError("Direct ESP grid points must contain three finite Cartesian coordinates.") if self.point_unit != "angstrom": raise ValueError("Direct ESP grid points must use angstrom.") - if not isinstance(self.unit, ElectrostaticPotentialUnit): - try: - object.__setattr__(self, "unit", ElectrostaticPotentialUnit(self.unit)) - except ValueError as exc: - raise ValueError(f"Unknown electrostatic-potential unit {self.unit!r}.") from exc + object.__setattr__( + self, "unit", _coerce_enum(self.unit, ElectrostaticPotentialUnit, "electrostatic-potential unit") + ) object.__setattr__(self, "value", value) object.__setattr__(self, "weight", weight) object.__setattr__(self, "case_id", case_id) @@ -242,16 +242,10 @@ def __post_init__(self) -> None: reference_case_id = str(self.reference_case_id).strip() if not group_id or not reference_case_id: raise ValueError("Relative-energy group and reference case IDs must be non-empty.") - if not isinstance(self.unit, ObservationEnergyUnit): - try: - object.__setattr__(self, "unit", ObservationEnergyUnit(self.unit)) - except ValueError as exc: - raise ValueError(f"Unknown relative-energy unit {self.unit!r}.") from exc - if not isinstance(self.quantity, ThermodynamicQuantity): - try: - object.__setattr__(self, "quantity", ThermodynamicQuantity(self.quantity)) - except ValueError as exc: - raise ValueError(f"Unknown thermodynamic quantity {self.quantity!r}.") from exc + object.__setattr__(self, "unit", _coerce_enum(self.unit, ObservationEnergyUnit, "relative-energy unit")) + object.__setattr__( + self, "quantity", _coerce_enum(self.quantity, ThermodynamicQuantity, "thermodynamic quantity") + ) object.__setattr__(self, "value", value) object.__setattr__(self, "weight", weight) object.__setattr__(self, "case_id", case_id) @@ -287,11 +281,7 @@ def __post_init__(self) -> None: raise ValueError("Scan-energy group and reference case IDs must be non-empty.") if not isinstance(self.coordinate, ScanCoordinate): raise TypeError("ScanEnergyObservation.coordinate must be a ScanCoordinate.") - if not isinstance(self.unit, ObservationEnergyUnit): - try: - object.__setattr__(self, "unit", ObservationEnergyUnit(self.unit)) - except ValueError as exc: - raise ValueError(f"Unknown scan-energy unit {self.unit!r}.") from exc + object.__setattr__(self, "unit", _coerce_enum(self.unit, ObservationEnergyUnit, "scan-energy unit")) object.__setattr__(self, "value", value) object.__setattr__(self, "weight", weight) object.__setattr__(self, "case_id", case_id) @@ -321,11 +311,7 @@ def __post_init__(self) -> None: ) if not isinstance(self.parameter_id, ParameterId): raise TypeError("ParameterTetherObservation.parameter_id must be a ParameterId.") - if not isinstance(self.unit, ParameterUnit): - try: - object.__setattr__(self, "unit", ParameterUnit(self.unit)) - except ValueError as exc: - raise ValueError(f"Unknown parameter-tether unit {self.unit!r}.") from exc + object.__setattr__(self, "unit", _coerce_enum(self.unit, ParameterUnit, "parameter-tether unit")) object.__setattr__(self, "value", value) object.__setattr__(self, "weight", weight) object.__setattr__(self, "case_id", case_id) diff --git a/q2mm/models/seminario.py b/q2mm/models/seminario.py index 5f565ca8..c44077cc 100644 --- a/q2mm/models/seminario.py +++ b/q2mm/models/seminario.py @@ -16,7 +16,7 @@ from collections.abc import Iterable import logging import numpy as np -from typing import Literal +from typing import Any, Literal from q2mm.constants import BOHR_TO_ANG from q2mm.models.units import ( @@ -79,7 +79,7 @@ def _collect_matching( param: BondParam | AngleParam, items_attr: str, element_key_attr: str, -) -> list[tuple[Molecule, object]]: +) -> list[tuple[Molecule, Any]]: """Collect all items (bonds or angles) across molecules that match a parameter.""" all_items = [(mol, item) for mol in molecules for item in getattr(mol, items_attr)] match_mode = _match_mode(param, [item for _, item in all_items]) @@ -99,38 +99,32 @@ def _should_keep_force_constant(value: float, invalid_policy: Literal["keep", "s return True +def _project_subblock_along(hessian: np.ndarray, atom_a: int, atom_b: int, direction: np.ndarray) -> float: + """Project the -(atom_a, atom_b) 3x3 Hessian sub-block onto unit vector *direction*.""" + a3, b3 = 3 * atom_a, 3 * atom_b + h_sub = -hessian[a3 : a3 + 3, b3 : b3 + 3] + # General eigenvalue decomposition (NOT eigh — sub-block is NOT symmetric). + # Keep complex eigenpairs through projection — only take real at the end + # (upstream seminario_sum uses np.abs on complex dot products). + eigenvalues, eigenvectors = np.linalg.eig(h_sub) + k = 0.0 + for n in range(3): + k += eigenvalues[n] * np.abs(np.dot(eigenvectors[:, n], direction)) + return float(np.real(k)) + + def _project_hessian_block(hessian: np.ndarray, atom_i: int, atom_j: int, coords: np.ndarray, au_units: bool) -> float: """Project a single Hessian sub-block onto the bond vector. Returns the projected force constant in atomic units (Hartree/Bohr^2) or input units if au_units=False. """ - if au_units: - coords_work = coords / BOHR_TO_ANG - else: - coords_work = coords.copy() - + coords_work = coords / BOHR_TO_ANG if au_units else coords.copy() r_vec = coords_work[atom_j] - coords_work[atom_i] - r_len = np.linalg.norm(r_vec) + r_len = float(np.linalg.norm(r_vec)) if r_len < 1e-10: return 0.0 - r_hat = r_vec / r_len - - i3, j3 = 3 * atom_i, 3 * atom_j - h_sub = -hessian[i3 : i3 + 3, j3 : j3 + 3] - - # General eigenvalue decomposition (NOT eigh — sub-block is NOT symmetric) - eigenvalues, eigenvectors = np.linalg.eig(h_sub) - # Keep complex eigenpairs through projection — only take real at the end - # (upstream seminario_sum uses np.abs on complex dot products) - - # Seminario projection: k = sum_n lambda_n * |e_n · r_hat| - # np.abs handles complex dot products correctly (returns magnitude) - k = 0.0 - for n in range(3): - k += eigenvalues[n] * np.abs(np.dot(eigenvectors[:, n], r_hat)) - # Result should be real (imaginary parts cancel in conjugate pairs) - return k.real + return _project_subblock_along(hessian, atom_i, atom_j, r_vec / r_len) def seminario_bond_fc( @@ -235,24 +229,8 @@ def seminario_angle_fc( u_kj = np.cross(n_hat, r_kj_hat) u_kj /= np.linalg.norm(u_kj) - # Sub-block Hessians — keep complex eigenpairs through projection - i3, j3, k3 = 3 * atom_i, 3 * atom_j, 3 * atom_k - - # For i-j interaction - h_ij = -hessian[i3 : i3 + 3, j3 : j3 + 3] - evals_ij, evecs_ij = np.linalg.eig(h_ij) - k_ij = 0.0 - for n in range(3): - k_ij += evals_ij[n] * np.abs(np.dot(evecs_ij[:, n], u_ij)) - k_ij = k_ij.real - - # For k-j interaction - h_kj = -hessian[k3 : k3 + 3, j3 : j3 + 3] - evals_kj, evecs_kj = np.linalg.eig(h_kj) - k_kj = 0.0 - for n in range(3): - k_kj += evals_kj[n] * np.abs(np.dot(evecs_kj[:, n], u_kj)) - k_kj = k_kj.real + k_ij = _project_subblock_along(hessian, atom_i, atom_j, u_ij) + k_kj = _project_subblock_along(hessian, atom_k, atom_j, u_kj) # Combine via reciprocal sum: 1/k = 1/(k_ij·r_ij²) + 1/(k_kj·r_kj²) # Standard Seminario angle formula. Note: FUERZA still overestimates @@ -273,7 +251,7 @@ def seminario_angle_fc( if au_units: k_angle *= HARTREE_RAD2_TO_KCALMOLRAD2 - return k_angle + return float(k_angle) def qfuerza_fresh( @@ -443,12 +421,12 @@ def qfuerza_into( raise ValueError("Molecule must have a Hessian attached. Use molecule.with_hessian(hess)") # Invert TS curvature if this is a transition-state Hessian + processed_hessians: dict[int, np.ndarray] | None = None if invert_ts_curvature: - processed_hessians: dict[int, np.ndarray] = {} + processed_hessians = {} for mol in molecules: + assert mol.hessian is not None processed_hessians[id(mol)] = _invert_ts_curvature(mol.hessian, replace_with=replace_with) - else: - processed_hessians = None return _estimate_into_ff( ff, @@ -482,6 +460,7 @@ def _estimate_bond( equilibria = [bond.length for _, bond in matching_bonds] for molecule_item, bond in matching_bonds: hess = processed_hessians[id(molecule_item)] if processed_hessians else molecule_item.hessian + assert hess is not None k = seminario_bond_fc( bond.atom_i, bond.atom_j, @@ -499,18 +478,14 @@ def _estimate_bond( else: logger.warning(f" Bond {bond.elements} ({bond.atom_i}-{bond.atom_j}): invalid FC = {k} — skipped") - updates: dict[str, float] = {} - if equilibria: - updates["equilibrium"] = float(np.mean(equilibria)) + new_eq = float(np.mean(equilibria)) if equilibria else bond_param.equilibrium if force_constants: - updates["force_constant"] = float(np.mean(force_constants)) - logger.info( - f" Bond {bond_param.key}: k={updates['force_constant']:.4f} kcal/(mol·Å²), " - f"r0={updates.get('equilibrium', bond_param.equilibrium):.4f} Å" - ) + new_fc = float(np.mean(force_constants)) + logger.info(f" Bond {bond_param.key}: k={new_fc:.4f} kcal/(mol·Å²), r0={new_eq:.4f} Å") else: + new_fc = bond_param.force_constant logger.warning(f" Bond {bond_param.key}: no valid force constants found, keeping existing force constant") - return dataclasses.replace(bond_param, **updates) if updates else bond_param + return dataclasses.replace(bond_param, equilibrium=new_eq, force_constant=new_fc) def _estimate_angle( @@ -532,6 +507,7 @@ def _estimate_angle( equilibria = [angle.value for _, angle in matching_angles] for molecule_item, angle in matching_angles: hess = processed_hessians[id(molecule_item)] if processed_hessians else molecule_item.hessian + assert hess is not None k = seminario_angle_fc( angle.atom_i, angle.atom_j, @@ -547,27 +523,23 @@ def _estimate_angle( else: logger.warning(f" Angle {angle.elements}: invalid FC = {k} — skipped") - updates: dict[str, float] = {} - if equilibria: - updates["equilibrium"] = float(np.mean(equilibria)) + new_eq = float(np.mean(equilibria)) if equilibria else angle_param.equilibrium if force_constants: fuerza_value = float(np.mean(force_constants)) if strategy == "qfuerza" and _is_hydrogen_angle(angle_param.elements): - updates["force_constant"] = QFUERZA_H_ANGLE_DEFAULT_CANONICAL + new_fc = QFUERZA_H_ANGLE_DEFAULT_CANONICAL logger.info( f" Angle {angle_param.key}: QFUERZA H-angle substitution — " f"{QFUERZA_H_ANGLE_DEFAULT_MDYNA} mdyn·Å/rad² " f"(FUERZA was {fuerza_value:.4f} kcal/(mol·rad²))" ) else: - updates["force_constant"] = fuerza_value - logger.info( - f" Angle {angle_param.key}: k={updates['force_constant']:.4f}, " - f"theta0={updates.get('equilibrium', angle_param.equilibrium):.1f} deg" - ) + new_fc = fuerza_value + logger.info(f" Angle {angle_param.key}: k={new_fc:.4f}, theta0={new_eq:.1f} deg") else: + new_fc = angle_param.force_constant logger.warning(f" Angle {angle_param.key}: no valid force constants found, keeping existing force constant") - return dataclasses.replace(angle_param, **updates) if updates else angle_param + return dataclasses.replace(angle_param, equilibrium=new_eq, force_constant=new_fc) def _estimate_into_ff( diff --git a/test/test_io_helpers.py b/test/test_io_helpers.py index ca5b9b5f..e8e39d5e 100644 --- a/test/test_io_helpers.py +++ b/test/test_io_helpers.py @@ -10,11 +10,12 @@ """ import dataclasses +from pathlib import Path import pytest -from q2mm.io._helpers import _normalize_equilibrium_angle, _torsion_file_value -from q2mm.models.forcefield import TorsionParam +from q2mm.io._helpers import _TextFileBase, _normalize_equilibrium_angle, _resolve_template_path, _torsion_file_value +from q2mm.models.forcefield import ForceField, FunctionalForm, TorsionParam class TestNormalizeEquilibriumAngle: @@ -102,6 +103,34 @@ def test_torsion_file_value_not_a_package_attribute(self) -> None: assert not hasattr(q2mm.io, "_torsion_file_value") + def test_text_file_base_and_resolve_template_not_package_attributes(self) -> None: + import q2mm.io + + assert not hasattr(q2mm.io, "_TextFileBase") + assert not hasattr(q2mm.io, "_resolve_template_path") + + +class TestSharedTextFileBaseAndTemplateResolution: + """Legacy file-wrapper parsers share ``_TextFileBase`` and template path resolution.""" + + def test_legacy_file_parsers_subclass_text_file_base(self) -> None: + from q2mm.io.gaussian import GaussLog + from q2mm.io.jaguar import JaguarIn, JaguarOut + from q2mm.io.macromodel import MacroModel, MacroModelLog + from q2mm.io.mol2 import Mol2 + + for cls in (GaussLog, JaguarIn, JaguarOut, MacroModel, MacroModelLog, Mol2): + assert issubclass(cls, _TextFileBase) + + def test_resolve_template_path_prefers_explicit_then_matching_source_format(self, tmp_path: Path) -> None: + source = tmp_path / "source.fld" + override = tmp_path / "override.fld" + ff = ForceField(functional_form=FunctionalForm.MM3, source_path=source, source_format="mm3_fld") + + assert _resolve_template_path(ff, override, "mm3_fld") == override + assert _resolve_template_path(ff, None, "mm3_fld") == source + assert _resolve_template_path(ff, None, "tinker_prm") is None + class TestFormatStagingRowsArePrivate: """Per-format parameter-row staging records are private and non-overlapping. diff --git a/test/test_seminario.py b/test/test_seminario.py index 12ba93c9..8ff9aaab 100644 --- a/test/test_seminario.py +++ b/test/test_seminario.py @@ -97,6 +97,50 @@ def test_archive_hessian_reproduces_gaussian_frequencies(self) -> None: "returned.", ) + def test_no_frequency_log_caches_none_without_rereading(self) -> None: + """Repeated access to frequencies/evals/evecs on a no-frequency log reads the body once.""" + from unittest.mock import patch + + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "single_point.log" + path.write_text( + " Standard orientation:\n" + " ---------------------------------------------------------------------\n" + " Center Atomic Atomic Coordinates (Angstroms)\n" + " Number Number Type X Y Z\n" + " ---------------------------------------------------------------------\n" + " 1 1 0 0.000000 0.000000 0.000000\n" + " ---------------------------------------------------------------------\n", + encoding="utf-8", + ) + log = GaussLog(str(path)) + with patch.object(GaussLog, "read_out", autospec=True, side_effect=GaussLog.read_out) as spy: + self.assertIsNone(log.frequencies) + path.unlink() + self.assertIsNone(log.evals) + self.assertIsNone(log.evecs) + self.assertIsNone(log.frequencies) + self.assertEqual(spy.call_count, 1) + + +class TestSeminarioAngleCancellingProjections(unittest.TestCase): + """Regression test for equal-and-opposite angle projections in seminario_angle_fc.""" + + def test_cancelling_projections_return_inf_without_zero_division(self) -> None: + from unittest.mock import patch + + from q2mm.models.seminario import seminario_angle_fc + + coords = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=float) + hessian = np.zeros((9, 9), dtype=float) + with ( + np.errstate(divide="ignore"), + patch("q2mm.models.seminario._project_subblock_along", side_effect=[1.0, -1.0]), + ): + result = seminario_angle_fc(0, 1, 2, coords, hessian, au_units=False) + self.assertIsInstance(result, float) + self.assertTrue(np.isinf(result)) + @unittest.skipUnless((ETHANE_DIR / "GS.mol2").exists(), "Ethane mol2 fixture not found") class TestMol2Parsing(unittest.TestCase):