From 00336b9dae81a8241e637bebd4fe1c3776fbd253 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <ElliottKasoar@users.noreply.github.com>
Date: Wed, 20 Mar 2024 17:06:29 +0000
Subject: [PATCH] Add log and warning for max force

---
 janus_core/geom_opt.py     | 28 +++++++++++++++++++++-------
 janus_core/single_point.py |  7 ++++---
 janus_core/utils.py        | 22 ++++++++++++++++++++++
 3 files changed, 47 insertions(+), 10 deletions(-)
 create mode 100644 janus_core/utils.py

diff --git a/janus_core/geom_opt.py b/janus_core/geom_opt.py
index 6ae6bbc5..521c8643 100644
--- a/janus_core/geom_opt.py
+++ b/janus_core/geom_opt.py
@@ -13,11 +13,14 @@
 except ImportError:
     from ase.constraints import ExpCellFilter as DefaultFilter
 
+from numpy import linalg
+
 from janus_core.janus_types import ASEOptArgs, ASEWriteArgs
 from janus_core.log import config_logger
+from janus_core.utils import none_to_dict
 
 
-def optimize(  # pylint: disable=too-many-arguments,too-many-locals
+def optimize(  # pylint: disable=too-many-arguments,too-many-locals,too-many-branches
     atoms: Atoms,
     fmax: float = 0.1,
     steps: int = 1000,
@@ -67,11 +70,9 @@ def optimize(  # pylint: disable=too-many-arguments,too-many-locals
     atoms: Atoms
         Structure with geometry optimized.
     """
-    filter_kwargs = filter_kwargs if filter_kwargs else {}
-    opt_kwargs = opt_kwargs if opt_kwargs else {}
-    write_kwargs = write_kwargs if write_kwargs else {}
-    traj_kwargs = traj_kwargs if traj_kwargs else {}
-    log_kwargs = log_kwargs if log_kwargs else {}
+    [filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs, log_kwargs] = none_to_dict(
+        [filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs, log_kwargs]
+    )
 
     write_kwargs.setdefault(
         "filename",
@@ -106,8 +107,21 @@ def optimize(  # pylint: disable=too-many-arguments,too-many-locals
         logger.info("Starting geometry optimization")
 
     converged = dyn.run(fmax=fmax, steps=steps)
+
+    # Calculate current maximum force
+    if filter_func is not None:
+        max_force = linalg.norm(filtered_atoms.get_forces(), axis=1).max()
+    else:
+        max_force = linalg.norm(atoms.get_forces(), axis=1).max()
+
+    if logger:
+        logger.info("Max force: %.6f", max_force)
+
     if not converged:
-        warnings.warn(f"Optimization has not converged after {steps} steps")
+        warnings.warn(
+            f"Optimization has not converged after {steps} steps. "
+            f"Current max force {max_force} > target force {fmax}"
+        )
 
     # Write out optimized structure
     if write_results:
diff --git a/janus_core/single_point.py b/janus_core/single_point.py
index 0de65070..fa35bd4b 100644
--- a/janus_core/single_point.py
+++ b/janus_core/single_point.py
@@ -19,6 +19,7 @@
 )
 from janus_core.log import config_logger
 from janus_core.mlip_calculators import choose_calculator
+from janus_core.utils import none_to_dict
 
 
 class SinglePoint:
@@ -122,9 +123,9 @@ def __init__(
                 "or a path to the structure file (`struct_path`)"
             )
 
-        read_kwargs = read_kwargs if read_kwargs else {}
-        calc_kwargs = calc_kwargs if calc_kwargs else {}
-        log_kwargs = log_kwargs if log_kwargs else {}
+        [read_kwargs, calc_kwargs, log_kwargs] = none_to_dict(
+            [read_kwargs, calc_kwargs, log_kwargs]
+        )
 
         if log_kwargs and "filename" not in log_kwargs:
             raise ValueError("'filename' must be included in `log_kwargs`")
diff --git a/janus_core/utils.py b/janus_core/utils.py
new file mode 100644
index 00000000..afc85563
--- /dev/null
+++ b/janus_core/utils.py
@@ -0,0 +1,22 @@
+"""Utility functions for janus_core."""
+
+from typing import Optional
+
+
+def none_to_dict(dictionaries: list[Optional[dict]]) -> list[dict]:
+    """
+    Ensure dictionaries that may be None are dictionaires.
+
+    Parameters
+    ----------
+    dictionaries : list[dict]
+        List of dictionaries that be be None.
+
+    Returns
+    -------
+    list[dict]
+        Dictionaries set to {} if previously None.
+    """
+    for i, dictionary in enumerate(dictionaries):
+        dictionaries[i] = dictionary if dictionary else {}
+    return dictionaries