Skip to content

Commit

Permalink
Added encoders and parsers for CREST to support energy, gradient, hes…
Browse files Browse the repository at this point in the history
…sian, and optimization calculations.
  • Loading branch information
coltonbh committed Oct 2, 2024
1 parent e123ca9 commit ec788d0
Show file tree
Hide file tree
Showing 10 changed files with 729 additions and 8 deletions.
2 changes: 2 additions & 0 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
"qcparse",
"rotamer",
"rotamers",
"runtypes",
"singlepoint",
"spinmult",
"tcin",
"tcout",
Expand Down
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),

## [unreleased]

### Added

- Encoders and parsers for CREST to support `energy`, `gradient`, `hessian`, and `optimization` calculations.

## [0.6.3] - 2024-09-12

### Added
Expand Down
60 changes: 56 additions & 4 deletions qcparse/encoders/crest.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,13 @@
from qcparse.exceptions import EncoderError
from qcparse.models import NativeInput

SUPPORTED_CALCTYPES = {CalcType.conformer_search}
SUPPORTED_CALCTYPES = {
CalcType.conformer_search,
CalcType.optimization,
CalcType.energy,
CalcType.gradient,
CalcType.hessian,
}


def encode(inp_obj: ProgramInput) -> NativeInput:
Expand Down Expand Up @@ -42,13 +48,47 @@ def validate_input(inp_obj: ProgramInput):
"""
# These values come from other parts of the ProgramInput and should not be set
# in the keywords.
non_allowed_keywords = ["charge", "uhf", "runtype"]
non_allowed_keywords = ["charge", "uhf"]
for keyword in non_allowed_keywords:
if keyword in inp_obj.keywords:
raise EncoderError(
f"{keyword} should not be set in keywords for CREST. It is already set "
"on the Structure or ProgramInput elsewhere.",
)
if "runtype" in inp_obj.keywords:
_validate_runtype_calctype(inp_obj.keywords["runtype"], inp_obj.calctype)


def _validate_runtype_calctype(runtype: str, calctype: CalcType):
"""Validate that the runtype is supported for the calctype."""
invalid_runtype = False
valid_runtypes = set()

if calctype == CalcType.conformer_search:
valid_runtypes = {"imtd-gc", "imtd-smtd", "entropy", "nci", "nci-mtd"}
if runtype not in valid_runtypes:
invalid_runtype = True

elif calctype == CalcType.optimization:
valid_runtypes = {"optimize", "ancopt"}
if runtype not in valid_runtypes:
invalid_runtype = True

elif calctype in {CalcType.energy, CalcType.gradient}:
valid_runtypes = {"singlepoint"}
if runtype not in valid_runtypes:
invalid_runtype = True

elif calctype == CalcType.hessian:
valid_runtypes = {"numhess"}
if runtype not in valid_runtypes:
invalid_runtype = True

if invalid_runtype:
raise EncoderError(
f"Unsupported runtype {runtype} for calctype {calctype}. Valid runtypes "
f"are: {valid_runtypes}.",
)


def _to_toml_dict(inp_obj: ProgramInput, struct_filename: str) -> Dict[str, Any]:
Expand All @@ -64,8 +104,20 @@ def _to_toml_dict(inp_obj: ProgramInput, struct_filename: str) -> Dict[str, Any]
toml_dict.setdefault("threads", os.cpu_count())
toml_dict["input"] = struct_filename

# TODO: May need to deal with non-covalent mode at some point
toml_dict["runtype"] = "imtd-gc"
# Set default runtype if not already set
if "runtype" not in inp_obj.keywords:
if inp_obj.calctype == CalcType.conformer_search:
toml_dict["runtype"] = "imtd-gc"
elif inp_obj.calctype == CalcType.optimization:
toml_dict["runtype"] = "optimize"
elif inp_obj.calctype in {CalcType.energy, CalcType.gradient}:
toml_dict["runtype"] = "singlepoint"
elif inp_obj.calctype == CalcType.hessian:
toml_dict["runtype"] = "numhess"
else:
raise EncoderError(
f"Unsupported calctype {inp_obj.calctype} for CREST encoder.",
)

# Calculation level keywords
calculation = toml_dict.pop("calculation", {})
Expand Down
153 changes: 151 additions & 2 deletions qcparse/parsers/crest.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,18 @@
import re
from pathlib import Path
from typing import List, Optional, Union
from typing import Any, Dict, List, Optional, Union

import numpy as np
from qcio import ConformerSearchResults, Structure
from qcio import (
CalcType,
ConformerSearchResults,
OptimizationResults,
ProgramInput,
ProgramOutput,
Provenance,
SinglePointResults,
Structure,
)

from .utils import regex_search

Expand Down Expand Up @@ -87,3 +97,142 @@ def parse_conformer_search_dir(
rotamers=rotamers,
rotamer_energies=np.array(rotamer_energies),
)


def parse_energy_grad(text: str) -> SinglePointResults:
"""Parse the output of a CREST energy and gradient calculation.
Args:
text: The text of the output file.
Returns:
The parsed energy and gradient as a SinglePointResults object.
"""
# Parse the energy
energy_regex = r"# Energy \( Eh \)\n#*\n\s*([-\d.]+)"
gradient_regex = r"# Gradient \( Eh/a0 \)\n#\s*\n((?:\s*[-\d.]+\n)+)"

energy = float(regex_search(energy_regex, text).group(1))
gradient = np.array(
[float(x) for x in regex_search(gradient_regex, text).group(1).split()]
)
return SinglePointResults(
energy=energy,
gradient=gradient,
)


def parse_singlepoint_dir(
directory: Union[Path, str], filename: str = "crest.engrad"
) -> SinglePointResults:
"""Parse the output directory of a CREST single point calculation.
Args:
directory: Path to the directory containing the CREST output files.
filename: The name of the file containing the single point results.
Default is 'crest.engrad'.
Returns:
The parsed single point results as a SinglePointResults object.
"""
directory = Path(directory)
text = (directory / filename).read_text()

return parse_energy_grad(text)


def parse_numhess_dir(
directory: Union[Path, str],
filename: str = "numhess1",
stdout: Optional[str] = None,
) -> SinglePointResults:
"""Parse the output directory of a CREST numerical Hessian calculation.
Args:
directory: Path to the directory containing the CREST output files.
filename: The name of the file containing the numerical Hessian results.
Default is 'numhess1'.
Returns:
The parsed numerical Hessian results as a SinglePointResults object.
"""
data = (Path(directory) / filename).read_text()
float_regex = r"[-+]?\d*\.\d+|\d+"
numbers = re.findall(float_regex, data)
array = np.array(numbers, dtype=float)
spr_dict: Dict[str, Any] = {"hessian": array}
if stdout:
energy_regex = r"Energy\s*=\s*([-+]?\d+\.\d+)\s*Eh"
energy = float(regex_search(energy_regex, stdout).group(1))
spr_dict["energy"] = energy
return SinglePointResults(**spr_dict)


def parse_optimization_dir(
directory: Union[Path, str],
*,
inp_obj: ProgramInput,
stdout: str,
) -> OptimizationResults:
"""Parse the output directory of a CREST optimization calculation.
Args:
directory: Path to the directory containing the CREST output files.
inp_obj: The qcio ProgramInput object for the optimization.
stdout: The stdout from CREST.
Returns:
The parsed optimization results as a OptimizationResults object.
"""
# Read in the xyz file containing the trajectory
directory = Path(directory)
xyz_text = (directory / "crestopt.log").read_text()

# Parse structures and energies from the xyz file
structures = Structure.from_xyz_multi(
xyz_text,
charge=inp_obj.structure.charge,
multiplicity=inp_obj.structure.multiplicity,
)
energies = [
float(struct.extras[Structure._xyz_comment_key][1]) for struct in structures
]

# Fake gradient for each step because CREST does not output it
fake_gradient = np.zeros(len(inp_obj.structure.symbols) * 3)

# Parse program version
program_version = parse_version_string(stdout)

# Collect final gradient if calculation succeeded
try:
final_spr = parse_singlepoint_dir(directory)
except FileNotFoundError:
# Calculation failed, so we don't have the final energy or gradient
final_spr = SinglePointResults(gradient=fake_gradient)

# Create the optimization trajectory
trajectory: List[ProgramOutput] = [
ProgramOutput(
input_data=ProgramInput(
calctype=CalcType.gradient,
structure=struct,
model=inp_obj.model,
),
success=True,
results=SinglePointResults(energy=energy, gradient=fake_gradient),
provenance=Provenance(
program="crest",
program_version=program_version,
),
)
for struct, energy in zip(structures, energies)
]

# Fill in final gradient
# https://github.com/crest-lab/crest/issues/354
trajectory[-1].results.gradient[:] = final_spr.gradient

return OptimizationResults(
trajectory=trajectory,
)
20 changes: 20 additions & 0 deletions tests/data/crest_output/crest.engrad
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#
# Atoms
#
3
#
# Energy ( Eh )
#
-0.335557824179335
#
# Gradient ( Eh/a0 )
#
-0.005962071557911
-0.004419818102026
0.003139227894649
0.003048425211480
0.001982394235964
-0.001779667371498
0.002913646346432
0.002437423866062
-0.001359560523152
20 changes: 20 additions & 0 deletions tests/data/crest_output/numhess1
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
$hessian
0.02040569 -0.00018059 0.02080099 -0.02081319 0.01511689
0.00867078 0.00037976 -0.01495837 -0.02946283
-0.00018059 -0.01341723 -0.03209513 0.01368595 0.03374600
0.01874084 -0.01351862 -0.02035995 0.01336374
0.02080099 -0.03209513 0.00327178 0.00784908 0.01737681
-0.01812512 -0.02863169 0.01472059 0.01485103
-0.02081319 0.01368595 0.00784908 0.01933555 -0.01625843
-0.00694960 0.00149263 0.00258608 -0.00090575
0.01511689 0.03374600 0.01737681 -0.01625843 -0.03409225
-0.01710500 0.00114214 0.00035657 -0.00027546
0.00867078 0.01874084 -0.01812512 -0.00694960 -0.01710500
0.01843539 -0.00173455 -0.00164242 -0.00030677
0.00037976 -0.01351862 -0.02863169 0.00149263 0.00114214
-0.00173455 -0.00185964 0.01238496 0.03036359
-0.01495837 -0.02035995 0.01472059 0.00258608 0.00035657
-0.00164242 0.01238496 0.02002423 -0.01308397
-0.02946283 0.01336374 0.01485103 -0.00090575 -0.00027546
-0.00030677 0.03036359 -0.01308397 -0.01454546
$end
Loading

0 comments on commit ec788d0

Please sign in to comment.