Add cell optimization to levelparser.

cp2k_output_tools/blocks/geo_opt.py → 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<type>\S+) \s+ \*{3}\n
+\ \*+\n
+""",
+    re.VERBOSE | re.MULTILINE,
+)
+
+OPT_STEP_RE = re.compile(
     r"""
 ^\ \-+\n
 \ OPTIMIZATION\ STEP: \s+ (?P<stepnr>\d+)\n
@@ -27,7 +37,7 @@
     re.VERBOSE | re.MULTILINE,
 )
 
-GEO_OPT_END_RE = re.compile(
+OPT_END_RE = re.compile(
     r"""
 ^(
 \ \*{3} \s+ (?P<msg>MAXIMUM\ NUMBER .+?)\s+\*{3}\n
@@ -52,22 +62,45 @@ 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)
 
     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])

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]