fix optimizer initialization routine, it was badly clever

lanl · Aug 14, 2024 · c027a10 · c027a10
1 parent dcd78d6
commit c027a10
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/hippynn/optimizer/__init__.py b/hippynn/optimizer/__init__.py
@@ -1,5 +1,5 @@
 """
-Functionality for Batch optimization of configurations under a potential energy surface..
+Functionality for Batch geometry optimization of configurations under a potential energy surface..
 
 Contributed by Shuhao Zhang (CMU, LANL)
 """

diff --git a/hippynn/optimizer/batch_optimizer.py b/hippynn/optimizer/batch_optimizer.py
@@ -50,10 +50,12 @@ def __init__(
                     force_node = find_unique_relative(model.nodes_to_compute, lambda node: 'force' in node.name,)
                 except Exception as ee:
                     raise ValueError("No automatic force node could be found for optimizer.") from ee
-                # TODO even find an energy node if we can?
+                # TODO even find an energy node if we can and add forces to it?
+        else:
+            force_node = find_unique_relative(model.nodes_to_compute, lambda node: node.db_name == force_key)
 
         additional_outputs = None
-        if force_node not in model.nodes_to_compute:
+        if force_node is not None and force_node not in model.nodes_to_compute:
             additional_outputs = [force_node]
 
         self.predictor = Predictor.from_graph(