Merge pull request #365 from jungsdao/main

optimizer as an argument for optimize

gracefully handle optimizers that do and don't accept `smax`

Author: bernstei

.github/workflows/pytests.yml

@@ -209,9 +209,9 @@ jobs:
       - name: Install buildcell
         run: |
           sudo apt-get install gfortran
-          wget https://www.mtg.msm.cam.ac.uk/system/files/documents/airss-0.9.1.tgz
-          tar xzf airss-0.9.1.tgz
-          cd airss-0.9.1
+          wget https://www.mtg.msm.cam.ac.uk/files/airss-v0.9.4.tgz
+          tar xzf airss-v0.9.4.tgz
+          cd airss
           make buildcell
           mkdir -p $HOME/bin
           cp src/buildcell/src/buildcell $HOME/bin/

wfl/generate/optimize.py

@@ -34,7 +34,7 @@ def _new_log(self, forces=None):
 def _run_autopara_wrappable(atoms, calculator, fmax=1.0e-3, smax=None, steps=1000, pressure=None,
            stress_mask=None, keep_symmetry=True, traj_step_interval=1, traj_subselect=None,
            skip_failures=True, results_prefix='last_op__optimize_', verbose=False, update_config_type="append",
-           rng=None, _autopara_per_item_info=None,
+           optimizer=PreconLBFGS, rng=None, _autopara_per_item_info=None,
            **opt_kwargs):
     """runs a structure optimization. By default calculator properties will be stored in keys
     prefixed with "last_op__optimize_", which may be overwritten by next operation.
@@ -73,8 +73,10 @@ def _run_autopara_wrappable(atoms, calculator, fmax=1.0e-3, smax=None, steps=100
         optimisation logs are not printed unless this is True
     update_config_type: ["append" | "overwrite" | False], default "append"
         whether/how to add at.info['optimize_config_type'] to at.info['config_type']
+    optimizer : ASE optimizer, default PreconLBFGS
+        optimizer to use. 
     opt_kwargs
-        keyword arguments for PreconLBFGS
+        keyword arguments for optimizer
     rng: numpy.random.Generator, default None
         random number generator to use (needed for pressure sampling, initial temperature, or Langevin dynamics)
     _autopara_per_item_info: dict
@@ -93,9 +95,6 @@ def _run_autopara_wrappable(atoms, calculator, fmax=1.0e-3, smax=None, steps=100
 
     calculator = construct_calculator_picklesafe(calculator)
 
-    if smax is None:
-        smax = fmax
-
     if keep_symmetry:
         # noinspection PyUnresolvedReferences,PyUnresolvedReferences
         try:
@@ -134,7 +133,7 @@ def _run_autopara_wrappable(atoms, calculator, fmax=1.0e-3, smax=None, steps=100
         else:
             wrapped_at = at
 
-        opt = PreconLBFGS(wrapped_at, **opt_kwargs_to_use)
+        opt = optimizer(wrapped_at, **opt_kwargs_to_use)
 
         # default status, will be overwritten for first and last configs in traj
         at.info['optimize_config_type'] = 'optimize_mid'
@@ -162,7 +161,17 @@ def process_step():
         converged = False
 
         try:
-            converged = opt.run(fmax=fmax, smax=smax, steps=steps)
+            try:
+                converged = opt.run(fmax=fmax, smax=smax if smax is not None else fmax, steps=steps)
+            except TypeError as exc:
+                if "unexpected keyword argument" in str(exc):
+                    # opt.run doesn't accept smax
+                    if smax is not None:
+                        # passed in explicitly, fail
+                        raise
+                    converged = opt.run(fmax=fmax, steps=steps)
+                else: # some other error
+                    raise
         except Exception as exc:
             # label actual failed optimizations
             # when this happens, the atomic config somehow ends up with a 6-vector stress, which can't be