Add output files

examples/fsm_example.py

@@ -23,8 +23,10 @@
 
 import numpy as np
 
+from mlfsm import __version__
 from mlfsm.cos import FreezingString
 from mlfsm.opt import CartesianOptimizer, InternalsOptimizer, Optimizer
+from mlfsm.output import FSMOutput
 from mlfsm.utils import load_xyz, load_xyz_fixed
 
 HERE = Path(__file__).parent
@@ -74,6 +76,10 @@ def run_fsm(
 
     outdir.mkdir(parents=True, exist_ok=True)
 
+    fsm_output: FSMOutput | None = None
+    if not interpolate:
+        fsm_output = FSMOutput(outdir)
+
     # get fixed atom indices
     def parse_indices(text):
         if text is None or text.strip() == "":
@@ -166,8 +172,18 @@ def parse_indices(text):
     else:
         raise ValueError(f"Unknown calculator {calculator}")
 
+    if fsm_output is not None:
+        fsm_output.write_header(__version__)
+        fsm_output.write_parameters(optcoords, interp, method, maxiter, maxls, dmax, nnodes_min, ninterp, stepsize)
+        fsm_output.write_system_info(reactant, product, chg, mult, fixed_atoms if len(fixed_atoms) > 0 else None)
+        fsm_output.write_calculator_info(calc)
+        fsm_output.write_initial_structures(reactant, product)
+
     # Initialize FSM string
-    string = FreezingString(reactant, product, nnodes_min, interp, ninterp, stepsize)
+    string = FreezingString(reactant, product, nnodes_min, interp, ninterp, stepsize, output=fsm_output)
+    if fsm_output is not None:
+        fsm_output.write_path_init(string.dist, string.stepsize, string.nnodes_min, string.init_coordsobj)
+
     if interpolate:
         string.interpolate(outdir)
         return
@@ -187,6 +203,10 @@ def parse_indices(text):
         string.optimize(optimizer)
         string.write(outdir)
 
+    if fsm_output is not None:
+        fsm_output.write_final_summary(string)
+        fsm_output.close()
+
     print(f"Gradient calls: {string.ngrad}")
 
 

src/mlfsm/__init__.py

@@ -1,3 +1,3 @@
 """mlfsm package."""
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"

src/mlfsm/cos.py

@@ -11,7 +11,9 @@
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
-from mlfsm.coords import Cartesian
+    from mlfsm.output import FSMOutput
+
+from mlfsm.coords import Cartesian, Redundant
 from mlfsm.geom import (
     calculate_arc_length,
     distance,
@@ -80,7 +82,9 @@ def __init__(
         interp_method: str = "ric",
         ninterp: int = 100,
         stepsize: float = 0.0,
+        output: Optional["FSMOutput"] = None,
     ) -> None:
+        self.output = output
         self.interp: Any
         self.interp_method = interp_method
         self.nnodes_min = int(nnodes_min)
@@ -100,17 +104,25 @@ def __init__(
         self.natoms = len(self.atoms.numbers)
 
         if not self.use_cartesian_distance:
-            interp = self.interp(reactant, product, ninterp=self.ninterp)
-            s = calculate_arc_length(interp())
+            _interp_init = self.interp(reactant, product, ninterp=self.ninterp)
+            s = calculate_arc_length(_interp_init())
             self.dist = s[-1]
             self.stepsize = self.dist / self.nnodes_min
         else:
-            interp = Linear(reactant, product, ninterp=self.ninterp)
-            s = calculate_arc_length(interp())
+            _interp_init = Linear(reactant, product, ninterp=self.ninterp)
+            s = calculate_arc_length(_interp_init())
             self.dist = s[-1]
             self.stepsize = float(stepsize)
             self.nnodes_min = int(self.dist / self.stepsize)
 
+        self.init_coordsobj: Optional[Redundant] = None
+        if interp_method == "ric":
+            self.init_coordsobj = (
+                _interp_init.coords if isinstance(_interp_init, RIC) else RIC(reactant, product, ninterp=2).coords
+            )
+        else:
+            self.init_coordsobj = None
+
         logger.info(f"NNODES_MIN: {self.nnodes_min}")
         logger.info(f"DIST: {self.dist:.3f} STEPSIZE: {self.stepsize:.3f}")
 
@@ -177,6 +189,9 @@ def grow(self) -> None:
         ``self.growing = False`` when the two frontiers are within one
         step-size of each other.
         """
+        if self.output is not None:
+            self.output._ensure_iteration_header(self.iteration + 1, self.dist)
+
         r_atoms = self.r_string[-1]
         p_atoms = self.p_string[-1]
 
@@ -200,8 +215,12 @@ def grow(self) -> None:
                 self.growing = False
                 return
 
+            if self.output is not None:
+                self.output.write_current_frontier_node("r", r_atoms)
+
             r_prev = r_xyz.copy().reshape(-1, 3)
             r_idx = 1
+            r_s = 0.0
             for qtarget in string[1:-1]:
                 r_next = interp.coords.x(r_prev, qtarget)
                 _, r_next = project_trans_rot(r_xyz.reshape(-1, 3), r_next)
@@ -228,13 +247,19 @@ def grow(self) -> None:
             self.r_energy += [None]
             self.r_tangent += [normalize(dxds)]
             self.r_nnodes = len(self.r_string)
+            if self.output is not None:
+                self.output.write_frontier_node("r", r_frontier, r_s)
 
             if self.dist <= 2 * self.stepsize:
                 self.growing = False
                 return
 
+            if self.output is not None:
+                self.output.write_current_frontier_node("p", p_atoms)
+
             p_prev = p_xyz.copy().reshape(-1, 3)
             p_idx = 1
+            p_s = 0.0
             for qtarget in string[1:-1][::-1]:
                 p_next = interp.coords.x(p_prev, qtarget)
                 _, p_next = project_trans_rot(p_xyz.reshape(-1, 3), p_next)
@@ -261,6 +286,8 @@ def grow(self) -> None:
             self.p_energy += [None]
             self.p_tangent += [normalize(dxds)]
             self.p_nnodes = len(self.p_string)
+            if self.output is not None:
+                self.output.write_frontier_node("p", p_frontier, p_s)
 
         else:
             string = interp()
@@ -274,6 +301,10 @@ def grow(self) -> None:
 
             r_idx = np.abs(s - self.stepsize).argmin()
             p_idx = np.abs(s - (s[-1] - self.stepsize)).argmin()
+
+            if self.output is not None:
+                self.output.write_current_frontier_node("r", r_atoms)
+
             r_frontier = self.atoms.copy()
             r_frontier.set_positions(string[r_idx].reshape(-1, 3))
 
@@ -282,11 +313,16 @@ def grow(self) -> None:
             self.r_energy += [None]
             self.r_tangent += [normalize(cs(s[r_idx], 1))]
             self.r_nnodes = len(self.r_string)
+            if self.output is not None:
+                self.output.write_frontier_node("r", r_frontier, float(s[r_idx]))
 
             if self.dist <= 2 * self.stepsize:
                 self.growing = False
                 return
 
+            if self.output is not None:
+                self.output.write_current_frontier_node("p", p_atoms)
+
             p_frontier = self.atoms.copy()
             p_frontier.set_positions(string[p_idx].reshape(-1, 3))
 
@@ -295,6 +331,8 @@ def grow(self) -> None:
             self.p_energy += [None]
             self.p_tangent += [normalize(cs(s[p_idx], 1))]
             self.p_nnodes = len(self.p_string)
+            if self.output is not None:
+                self.output.write_frontier_node("p", p_frontier, float(s[-1] - s[p_idx]))
 
     def optimize(self, optimizer: Any) -> None:
         """Relax all unfixed frontier nodes perpendicular to the local tangent.
@@ -317,37 +355,45 @@ def optimize(self, optimizer: Any) -> None:
             if self.r_energy[i] is None and self.r_fix[i]:
                 energy = optimizer.calc.get_potential_energy(self.r_string[i])
                 self.r_energy[i] = float_check(energy)
+                if self.output is not None:
+                    self.output.write_optimized_node("r", i, self.r_string[i], self.r_energy[i], 0, 0)
             elif not self.r_fix[i]:
                 assert self.r_tangent[i] is not None
                 atoms = self.r_string[i]
                 try:
-                    atoms, energy, ngrad = optimizer.optimize(atoms, self.r_tangent[i])
+                    atoms, energy, nfev, nit = optimizer.optimize(atoms, self.r_tangent[i])
                     self.r_string[i] = atoms
                     self.r_energy[i] = float_check(energy)
                 except Exception:
                     energy = optimizer.calc.get_potential_energy(atoms)
                     self.r_energy[i] = float_check(energy)
-                    ngrad = 0
+                    nfev, nit = 0, 0
                 self.r_fix[i] = True
-                self.ngrad += ngrad
+                self.ngrad += nfev
+                if self.output is not None:
+                    self.output.write_optimized_node("r", i, self.r_string[i], self.r_energy[i], nfev, nit)
 
         for i in range(self.p_nnodes):
             if self.p_energy[i] is None and self.p_fix[i]:
                 energy = optimizer.calc.get_potential_energy(self.p_string[i])
                 self.p_energy[i] = float_check(energy)
+                if self.output is not None:
+                    self.output.write_optimized_node("p", i, self.p_string[i], self.p_energy[i], 0, 0)
             elif not self.p_fix[i]:
                 assert self.p_tangent[i] is not None
                 atoms = self.p_string[i]
                 try:
-                    atoms, energy, ngrad = optimizer.optimize(atoms, self.p_tangent[i])
+                    atoms, energy, nfev, nit = optimizer.optimize(atoms, self.p_tangent[i])
                     self.p_string[i] = atoms
                     self.p_energy[i] = float_check(energy)
                 except Exception:
                     energy = optimizer.calc.get_potential_energy(atoms)
                     self.p_energy[i] = float_check(energy)
-                    ngrad = 0
+                    nfev, nit = 0, 0
                 self.p_fix[i] = True
-                self.ngrad += ngrad
+                self.ngrad += nfev
+                if self.output is not None:
+                    self.output.write_optimized_node("p", i, self.p_string[i], self.p_energy[i], nfev, nit)
 
         self.dist = distance(self.r_string[-1].get_positions().flatten(), self.p_string[-1].get_positions().flatten())
 
@@ -399,6 +445,9 @@ def write(self, outdir: Path | str) -> None:
         energy_str = np.array2string(energy, precision=1, floatmode="fixed")
         logging.info(f"ITERATION: {self.iteration} DIST: {self.dist:.2f} ENERGY: {energy_str}")
 
+        if self.output is not None:
+            self.output.write_iteration_summary(self.iteration, self.r_energy, self.p_energy, self.dist)
+
         if not self.growing:
             with gradfile.open("w") as f:
                 f.write(f"{self.ngrad}\n")

src/mlfsm/opt.py

@@ -91,7 +91,7 @@ def obj(self, xyz: NDArray[Any], tangent: NDArray[Any], atoms: Atoms) -> tuple[f
         pgrads = proj @ grads
         return energy, pgrads
 
-    def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, int]:
+    def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, int, int]:
         """Run optimization in Cartesian coordinates using user specified method.
 
         Args:
@@ -100,8 +100,8 @@ def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, i
 
         Returns
         -------
-            tuple[ASE.Atoms,float,int]: ASE.Atoms with final positions, energy of final structure, and number
-            of gradient calculations used by optimization.
+            tuple[ASE.Atoms,float,int,int]: ASE.Atoms with final positions, energy, total function
+            evaluations (nfev), and number of optimizer iterations (nit).
         """
         xyz = atoms.get_positions().flatten()
         config = {
@@ -119,7 +119,7 @@ def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, i
         res = minimize(**config)
         atomsf = atoms.copy()
         atomsf.set_positions(res.x.reshape(-1, 3))
-        return atomsf, res.fun, res.njev
+        return atomsf, res.fun, res.nfev, res.nit
 
 
 @dataclass
@@ -211,7 +211,7 @@ def obj(
 
         return energy, pgrads
 
-    def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, int]:
+    def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, int, int]:
         """Run optimization in internal coordinates using user specified method.
 
         Args:
@@ -220,8 +220,8 @@ def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, i
 
         Returns
         -------
-            tuple[ASE.Atoms,float,int]: ASE.Atoms with final positions, energy of final structure, and number
-            of gradient calculations used by optimization.
+            tuple[ASE.Atoms,float,int,int]: ASE.Atoms with final positions, energy, total function
+            evaluations (nfev), and number of optimizer iterations (nit).
         """
         assert self.coordsobj is not None, "Coordsobj must be initialized"
 
@@ -241,4 +241,4 @@ def optimize(self, atoms: Atoms, tangent: NDArray[Any]) -> tuple[Atoms, float, i
         atomsf = atoms.copy()
         atomsf.set_positions(xf)
 
-        return atomsf, res.fun, res.njev
+        return atomsf, res.fun, res.nfev, res.nit