Fix convergence check in optimizers, and skip phono3py tests if not available

Author: bernstei

docs/source/overview.parallelisation.rst

@@ -69,7 +69,7 @@ The number of threads is controlled by an integer, passed in to the
 function as an optional ``num_python_subprocesses`` argument, or stored
 in the env var ``WFL_NUM_PYTHON_SUBPROCESSES``.  The script should be
 started with a normal run of the python executable. Setting
-the ``WFL_TORCH_DEVICE_IS_SUBPROCESS_ID`` env var (to any value)
+the ``WFL_TORCH_N_GPUS`` env var to the number of GPUs
 causes ``wfl`` to call ``torch.cuda.set_device()`` for each subprocess
 ensuring that it gets a unique GPU from the other subprocesses.
 

tests/local_scripts/complete_pytest.tin

@@ -2,7 +2,7 @@
 
 module purge
 # module load compiler/gnu python/system python_extras/quippy lapack/mkl
-module load compiler/gnu python python_extras/quippy lapack/mkl
+module load compiler/gnu python python_extras/structure python_extras/quippy lapack/mkl
 # for wfl dependencies
 module load python_extras/wif
 module load python_extras/torch/cpu
@@ -14,7 +14,7 @@ WFL_PYTEST_EXPYRE_INFO=$(
 cat << EOF | python3
 import json, os
 i = {"pre_cmds": ["module purge",
-                  "module load compiler/gnu lapack/mkl python python_extras/quippy python_extras/wif dft/vasp dft/pwscf",
+                  "module load compiler/gnu lapack/mkl python python_extras/structure python_extras/quippy python_extras/wif dft/vasp dft/pwscf",
                   "module list"]}
 ienv = json.loads(os.environ.get("WFL_PYTEST_EXPYRE_INFO", "{}"))
 i.update(ienv)

tests/test_phonopy.py

@@ -7,6 +7,10 @@
 
 from wfl.configset import ConfigSet, OutputSpec
 from wfl.generate.phonopy import phonopy
+try:
+    import phono3py
+except ImportError:
+    phono3py = None
 
 
 def test_phonopy(tmp_path):
@@ -33,6 +37,7 @@ def test_phonopy(tmp_path):
         for v in at.positions[1:]:
             assert min(np.linalg.norm(sc.positions[1:] - v, axis=1)) < 1.0e-7
 
+@pytest.mark.skipif(phono3py is None, reason="No phono3py module")
 def test_phono3py(tmp_path):
     at0 = Atoms(numbers=[29], cell = [[0, 2, 2], [2, 0, 2], [2, 2, 0]], positions = [[0, 0, 0]], pbc = [True]*3)
     at1 = Atoms(numbers=[29], cell = [[0, 1.9, 1.9], [1.9, 0, 1.9], [1.9, 1.9, 0]], positions = [[0, 0, 0]], pbc = [True]*3)
@@ -62,6 +67,7 @@ def test_phono3py(tmp_path):
     assert sum([at.info["config_type"] == "phonon_cubic_1" for at in pert]) == 13*2
 
 
+@pytest.mark.skipif(phono3py is None, reason="No phono3py module")
 def test_phono3py_same_supercell(tmp_path):
     at0 = Atoms(numbers=[29], cell = [[0, 2, 2], [2, 0, 2], [2, 2, 0]], positions = [[0, 0, 0]], pbc = [True]*3)
     at1 = Atoms(numbers=[29], cell = [[0, 1.9, 1.9], [1.9, 0, 1.9], [1.9, 1.9, 0]], positions = [[0, 0, 0]], pbc = [True]*3)

tests/test_remote_run.py

@@ -69,7 +69,7 @@ def test_vasp_fail(tmp_path, expyre_systems, monkeypatch, remoteinfo_env):
 
 def do_vasp_fail(tmp_path, sys_name, monkeypatch, remoteinfo_env):
     ri = {'sys_name': sys_name, 'job_name': 'pytest_vasp_'+sys_name,
-          'env_vars' : ['ASE_VASP_COMMAND=NONE', 'ASE_VASP_COMMAND_GAMMA=NONE'],
+          'env_vars' : ['ASE_VASP_COMMAND=NO_VASP_FAIL', 'ASE_VASP_COMMAND_GAMMA=NO_VASP_FAIL'],
           'input_files' : ['POTCARs'],
           'resources': {'max_time': '5m', 'num_nodes': 1},
           'num_inputs_per_queued_job': 1, 'check_interval': 10}

wfl/generate/neb.py

@@ -89,9 +89,11 @@ def process_step():
 
         # preliminary value
         final_status = 'unconverged'
+        is_converged = False
 
         try:
-            opt.run(fmax=fmax, steps=steps)
+            for is_converged in opt.irun(fmax=fmax, steps=steps):
+                pass
         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
@@ -109,7 +111,9 @@ def process_step():
         for at in traj[0]:
             at.info['neb_config_type'] = 'neb_initial'
 
-        if opt.converged():
+        # from value returned by Optimizer.irun(), since Optimizer.converged requires
+        # a gradient as of ASE 3.26
+        if is_converged:
             final_status = 'converged'
 
         for intermed_images in traj[1:-1]:

wfl/generate/optimize.py

@@ -159,9 +159,11 @@ def process_step():
 
         # preliminary value
         final_status = 'unconverged'
+        is_converged = False
 
         try:
-            opt.run(fmax=fmax, smax=smax, steps=steps)
+            for is_converged in opt.irun(fmax=fmax, smax=smax, steps=steps):
+                pass
         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
@@ -184,7 +186,8 @@ def process_step():
 
         # as of 3.26 converged() requires a gradient, but for PreconLBFGS it's not used
         # See https://gitlab.com/ase/ase/-/issues/1744
-        if opt.converged(None):
+        # use Optimizer.irun instead, since it returns converged status
+        if is_converged:
             final_status = 'converged'
 
         traj[-1].info['optimize_config_type'] = f'optimize_last_{final_status}'