Add mock support for ml workflow

Author: elinscott

pyproject.toml

@@ -71,13 +71,15 @@ filterwarnings = [
     "ignore:Could not find 2nd order Makov-Payne energy; applying first order only:UserWarning",
     "ignore:Martyna-Tuckerman corrections not applied for an aperiodic calculation; do this with caution:UserWarning",
     "ignore:Makov-Payne corrections are not being used; do this with caution for periodic systems:UserWarning",
-    "ignore:eps_inf missing in input; it will default to 1.0. Proceed with caution for periodic systems:UserWarning",
+    "ignore:`eps_inf` missing in input; it will default to 1.0. Proceed with caution for periodic systems:UserWarning",
     "ignore:Some of the pseudopotentials do not have `PP_PSWFC` blocks, which means a projected DOS calculation is not possible. Skipping...:UserWarning",
     "ignore:Neither a pseudopotential library nor a list of pseudopotentials was provided; defaulting to `sg15_v1.2`:UserWarning",
     "ignore:The screening parameters for a KI calculation with no empty states will converge instantly; to save computational time set alpha_numsteps == 1:UserWarning",
     "ignore:The screening parameters have been calculated but are not necessarily self-consistent. You may want to increase `alpha_numsteps` to obtain a more accurate result.:UserWarning",
     "ignore:This system is not cubic and will therefore not have a uniform dielectric tensor. However, the image-correction schemes that are currently implemented assume a uniform dielectric. Proceed with caution:UserWarning",
     'ignore:Small box parameters `nrb` not provided in input:UserWarning',
+    'ignore:The computed screening parameter is greater than 1. Proceed with caution.:UserWarning',
+    'ignore:Predicting screening parameters with machine-learning is an experimental feature; proceed with caution:UserWarning'
 ]
 addopts = "--show-capture=no --capture=sys --strict-markers --tb=short -rfEs --basetemp=tests/tmp"
 markers = [

tests/benchmarks/.gitignore

@@ -0,0 +1,2 @@
+!*.pkl
+!??-*/

tests/calculators/test_koopmans_cp.py

@@ -8,7 +8,9 @@
 
 from koopmans import settings, utils
 from koopmans.calculators import KoopmansCPCalculator
-from koopmans.calculators._koopmans_cp import allowed, good_fft
+from koopmans.calculators._koopmans_cp import (allowed,
+                                               convert_flat_alphas_for_kcp,
+                                               good_fft)
 
 
 def test_convert_flat_alphas_for_kcp():

tests/helpers/patches/_benchmark.py

@@ -21,7 +21,7 @@
 from koopmans.processes.wannier import ExtendProcess, MergeProcess
 
 from ._utils import (benchmark_filename, find_subfiles_of_calc,
-                     metadata_filename)
+                     metadata_filename, recursively_find_files)
 
 
 def patch_calculator(c, monkeypatch):
@@ -84,6 +84,11 @@ def _run(self):
             filename.parent.mkdir(parents=True)
         write_pkl(self, filename)
 
+        # Copy over all files that are outputs of the process that need to be read
+        for filepointer in recursively_find_files([o for _, o in self.outputs]):
+            if filepointer.name in ['power_spectrum.npy']:
+                shutil.copy(filepointer, benchmark_filename(self).parent / filepointer.name)
+
     monkeypatch.setattr(p, '_run', _run)
 
 

tests/helpers/patches/_mock.py

@@ -11,7 +11,8 @@
 from koopmans.io import read_pkl
 from koopmans.utils import chdir, symlink
 
-from ._utils import benchmark_filename, metadata_filename
+from ._utils import (benchmark_filename, metadata_filename,
+                     recursively_find_files)
 
 
 def write_mock_file(filename: Union[Path, str], written_by: str):
@@ -125,28 +126,18 @@ def generate_dos(self):
     monkeypatch.setattr(calc_class, 'generate_dos', generate_dos)
 
 
-def recursively_find_files(obj):
-    if isinstance(obj, dict):
-        for k, v in obj.items():
-            yield from recursively_find_files(v)
-    elif isinstance(obj, list):
-        for v in obj:
-            yield from recursively_find_files(v)
-    elif isinstance(obj, FilePointer):
-        yield obj.aspath()
-    elif isinstance(obj, Path):
-        yield obj
-    return
-
-
 def mock_process_run(self):
     # Load the inputs from file
     bench_process = read_pkl(benchmark_filename(self))
     assert self.inputs == bench_process.inputs
     self.outputs = bench_process.outputs
 
     for f in recursively_find_files([o for _, o in self.outputs]):
-        write_mock_file(f, self.name)
+        if f.name in ['power_spectrum.npy']:
+            assert (benchmark_filename(self).parent / f.name).exists()
+            shutil.copy(benchmark_filename(self).parent / f.name, f)
+        else:
+            write_mock_file(f, self.name)
 
 
 def monkeypatch_mock(monkeypatch):

tests/helpers/patches/_utils.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 from typing import Set, Tuple
 
+from koopmans.files import FilePointer
+
 
 def benchmark_filename(obj) -> Path:
     base_directory = Path(__file__).parents[3]
@@ -37,3 +39,17 @@ def find_subfiles_of_dir(base_dir: Path) -> Set[Tuple[Path, float]]:
     else:
         files = set([])
     return files
+
+
+def recursively_find_files(obj):
+    if isinstance(obj, dict):
+        for k, v in obj.items():
+            yield from recursively_find_files(v)
+    elif isinstance(obj, list):
+        for v in obj:
+            yield from recursively_find_files(v)
+    elif isinstance(obj, FilePointer):
+        yield obj.aspath()
+    elif isinstance(obj, Path):
+        yield obj
+    return

tests/workflows/test_ml.py

@@ -12,7 +12,7 @@
 @pytest.mark.parametrize('descriptor', ['orbital_density', 'self_hartree'])
 @pytest.mark.parametrize('estimator', ['mean', 'ridge_regression'])
 @pytest.mark.parametrize('occ_and_emp_together', [True, False])
-def test_ml_train_water(tmp_path, water_snapshots, descriptor, estimator, occ_and_emp_together):
+def test_ml_train_water(tmp_path, workflow_patch, water_snapshots, descriptor, estimator, occ_and_emp_together):
     with chdir(tmp_path):
         projs = ProjectionBlocks.fromlist([[{'site': 'O', 'ang_mtm': 'sp3'}], [{'site': 'H', 'ang_mtm': 's'}]], spins=[
                                           None, None], atoms=water_snapshots['atoms'])