Update Optimizer (#82)

* 06102024 add optimizer/ with batch FIRE algorithm implemented with pytorch
also a hippynn wrapper under __init__.py

* 06282024 update optimizer

* 08052024@macbook:
update HippyNNBatchOptimizer for pull request

* update optimizer codebase to better match rest of library

* remove unneeded pt files

* update changes and authors

---------

Co-authored-by: Shuhao Zhang <[email protected]>
Co-authored-by: Nicholas Lubbers <[email protected]>

.gitignore

@@ -5,3 +5,6 @@ __pycache__/
 *.pyc
 build/
 hippynn.egg-info/*
+
+**/*.ipynb_checkpoints/
+**/__pycache__/

AUTHORS.txt

@@ -19,6 +19,7 @@ Emily Shinkle (LANL)
 Michael G. Taylor (LANL)
 Jan Janssen (LANL)
 Cagri Kaymak (LANL)
+Shuhao Zhang (CMU, LANL)
 
 Also thanks to testing and feedback from:
 

CHANGELOG.rst

@@ -7,6 +7,7 @@ New Features:
 -------------
 
 - Added a new custom cuda kernel implementation using triton. These are highly performant and now the default implementation.
+- Added Batch optimizer features in order to optimize geometries in parallel on the GPU. Algorithms include FIRE and BFGS.
 
 Improvements:
 -------------

hippynn/optimizer/__init__.py

@@ -1,24 +1,8 @@
-from ..graphs import Predictor
-from .BatchOptimizer import *
+"""
+Functionality for Batch optimization of configurations under a potential energy surface..
 
-class HippyNNBatchOptimizer():
-    def __init__(self, model, optimizer=BatchFIRE(), dump_traj=False):
-
-        # assume model is a hippynn model, and there will be 'Grad' in its output
-        self.predictor = Predictor.from_graph(model)
-        self.optimizer = optimizer
-        self.dump_traj = dump_traj
+Contributed by Shuhao Zhang (CMU, LANL)
+"""
 
-    def __call__(self, Z, R):
-        coords = R.clone()
-        self.optimizer.reset(coords)
-        while not self.optimizer.stop_signal:
-            ret = self.predictor(Z=Z, R=coords)
-            forces = -ret['Grad']
-            self.optimizer(forces)
-            coords = self.optimizer.coords
-        return coords
-
-    def dump_step(self, Z, R):
-        # need to figure out how to dump a batch of structures
-        pass
+from .algorithms import BFGSv1, BFGSv2, BFGSv3, FIRE, NewtonRaphson
+from .batch_optimizer import Optimizer