From 958c76bd609a420918f45b8b181e485f15ce9776 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Holger-Dietrich=20Sa=C3=9Fnick?= Date: Wed, 9 Aug 2023 16:45:42 +0200 Subject: [PATCH] Add cell optimization to levelparser. --- .../blocks/{geo_opt.py => optimization.py} | 53 +++++++++++++++---- cp2k_output_tools/levelparser.py | 27 ++++++++-- 2 files changed, 65 insertions(+), 15 deletions(-) rename cp2k_output_tools/blocks/{geo_opt.py => optimization.py} (55%) diff --git a/cp2k_output_tools/blocks/geo_opt.py b/cp2k_output_tools/blocks/optimization.py similarity index 55% rename from cp2k_output_tools/blocks/geo_opt.py rename to cp2k_output_tools/blocks/optimization.py index 330902c..3a22188 100644 --- a/cp2k_output_tools/blocks/geo_opt.py +++ b/cp2k_output_tools/blocks/optimization.py @@ -1,6 +1,6 @@ import sys from dataclasses import dataclass -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import regex as re @@ -18,7 +18,17 @@ re.VERBOSE | re.MULTILINE, ) -GEO_OPT_STEP_RE = re.compile( +CELL_OPT_RE = re.compile( + r""" +^\ \*+\n +\ \*{3} \s+ STARTING\s{3}CELL\s{3}OPTIMIZATION .+\n +\ \*{3} \s+ (?P\S+) \s+ \*{3}\n +\ \*+\n +""", + re.VERBOSE | re.MULTILINE, +) + +OPT_STEP_RE = re.compile( r""" ^\ \-+\n \ OPTIMIZATION\ STEP: \s+ (?P\d+)\n @@ -27,7 +37,7 @@ re.VERBOSE | re.MULTILINE, ) -GEO_OPT_END_RE = re.compile( +OPT_END_RE = re.compile( r""" ^( \ \*{3} \s+ (?PMAXIMUM\ NUMBER .+?)\s+\*{3}\n @@ -52,8 +62,31 @@ class GeometryOptimizationStep(Level): messages: List[Message] -def match_geo_opt(content: str, start: int = 0, end: int = sys.maxsize) -> Optional[Tuple[GeometryOptimization, Tuple[int, int]]]: - start_match = GEO_OPT_RE.search(content, start, end) +@dataclass +class CellOptimization(Level): + converged: bool + + +@dataclass +class CellOptimizationStep(Level): + messages: List[Message] + + +def match_opt( + content: str, start: int = 0, end: int = sys.maxsize, opt_type: str = "geo" +) -> Optional[Tuple[Union[GeometryOptimization, CellOptimization], Tuple[int, int]]]: + if opt_type == "geo": + opt_re = GEO_OPT_RE + opt_step = GeometryOptimizationStep + opt = GeometryOptimization + elif opt_type == "cell": + opt_re = CELL_OPT_RE + opt_step = CellOptimizationStep + opt = CellOptimization + else: + print("Unsupported optimization type.") + + start_match = opt_re.search(content, start, end) if not start_match: return None, (start, end) @@ -61,13 +94,13 @@ def match_geo_opt(content: str, start: int = 0, end: int = sys.maxsize) -> Optio step_starts = [start_match.span()[1] + 1] step_ends = [] - for match in GEO_OPT_STEP_RE.finditer(content, start, end): + for match in OPT_STEP_RE.finditer(content, start, end): step_end, step_start = match.span() # the previous step ends where the next one starts step_starts.append(step_start) step_ends.append(step_end) - stop_match = GEO_OPT_END_RE.search(content, start, end) + stop_match = OPT_END_RE.search(content, start, end) converged = True if stop_match: step_ends.append(stop_match.span()[0]) @@ -86,7 +119,5 @@ def match_geo_opt(content: str, start: int = 0, end: int = sys.maxsize) -> Optio sublevels.append(scf) else: print("NO SCF FOUND in this", content[start:end]) - - steps.append(GeometryOptimizationStep(messages=list(match_messages(content, start, end)), sublevels=sublevels)) - - return GeometryOptimization(converged=converged, sublevels=steps), (start_match.span()[0], step_ends[-1]) + steps.append(opt_step(messages=list(match_messages(content, start, end)), sublevels=sublevels)) + return opt(converged=converged, sublevels=steps), (start_match.span()[0], step_ends[-1]) diff --git a/cp2k_output_tools/levelparser.py b/cp2k_output_tools/levelparser.py index b8dab20..90834fd 100644 --- a/cp2k_output_tools/levelparser.py +++ b/cp2k_output_tools/levelparser.py @@ -8,12 +8,14 @@ from .blocks.cell import CellInformation, match_cell from .blocks.common import Level, Tree -from .blocks.geo_opt import ( +from .blocks.linres import Linres, match_linres +from .blocks.optimization import ( + CellOptimization, + CellOptimizationStep, GeometryOptimization, GeometryOptimizationStep, - match_geo_opt, + match_opt, ) -from .blocks.linres import Linres, match_linres from .blocks.program_info import ProgramInfo, match_program_info from .blocks.scf import SCF, InnerSCF, OuterSCF, match_scf from .blocks.sirius import Sirius, SiriusSCF, match_sirius @@ -69,6 +71,19 @@ def _(level: GeometryOptimizationStep, indent=""): print(f"{indent} [{msg.type}]: {msg.message}") +@pretty_print.register +def _(level: CellOptimization, indent=""): + print(f"{indent}Cell Optimization:") + print(f"{indent} converged: {level.converged}") + + +@pretty_print.register +def _(level: CellOptimizationStep, indent=""): + print(f"{indent}Cell Optimization Step:") + for msg in level.messages: + print(f"{indent} [{msg.type}]: {msg.message}") + + @pretty_print.register def _(level: Linres, indent=""): print(f"{indent}Linres:") @@ -160,6 +175,10 @@ def parse_all(content: str, start: int = 0, end: int = sys.maxsize) -> Tree: program_info = match_program_info(content, start, end, as_tree_obj=True) + cell_opt, span = match_opt(content, start, end, "cell") + if cell_opt: + sublevels.append(cell_opt) + cell_infos = [] while True: cell_info, span = match_cell(content, start, end) @@ -169,7 +188,7 @@ def parse_all(content: str, start: int = 0, end: int = sys.maxsize) -> Tree: cell_infos.append(cell_info) start = span[1] # move the start ahead to the end of the cell section - geo_opt, span = match_geo_opt(content, start, end) + geo_opt, span = match_opt(content, start, end, "geo") if geo_opt: sublevels.append(geo_opt) start = span[1]