Adding link() method to workflows; WIP

src/koopmans/calculators/_pw.py

@@ -84,7 +84,7 @@ def vbm_energy(self) -> float:
 
         # Fetch the total number of electrons in the system
         nelec = nelec_from_pseudos(self.atoms, self.parameters.pseudopotentials,
-                                   self.parameters.pseudo_dir) + self.parameters.get('tot_charge', 0)
+                                   self.directory / self.parameters.pseudo_dir) + self.parameters.get('tot_charge', 0)
 
         # Determine the number of occupied bands in each spin channel
         if self.parameters.nspin == 1:

src/koopmans/calculators/_utils.py

@@ -24,7 +24,8 @@ class ExtendedCalc(CalculatorExt, ASECalc, CalculatorABC):
 import os
 from abc import ABC, abstractmethod, abstractproperty
 from pathlib import Path
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
+from typing import (Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar,
+                    Union)
 
 import ase.io as ase_io
 import numpy as np
@@ -84,6 +85,9 @@ def __init__(self, skip_qc: bool = False, **kwargs: Any):
         # skip_qc = True
         self.skip_qc = skip_qc
 
+        # Prepare a dictionary to store a record of linked files
+        self._linked_files: Dict[str, Tuple[CalculatorExt | None, Path]] = {}
+
     def __repr__(self):
         entries = []
 
@@ -152,12 +156,33 @@ def _post_calculate(self):
 
         return
 
+    def _fetch_linked_files(self):
+        """Link all files provided in self._linked_files
+
+        This function is called in _pre_calculate() i.e. immediately before a calculation is run.
+        """
+        for dest_filename, (src_calc, src_filename) in self._linked_files.items():
+            if src_calc is None:
+                src_filename = src_filename.resolve()
+            else:
+                src_filename = (src_calc.directory / src_filename).resolve()
+            if not src_filename.exists():
+                raise FileNotFoundError(
+                    f'Tried to link {src_filename} with the {self.prefix} calculator but it does not exist')
+            dest_filename = (self.directory / dest_filename).resolve()
+            if not dest_filename.exists():
+                dest_filename.parent.mkdir(parents=True, exist_ok=True)
+                utils.symlink(src_filename, dest_filename)
+
     def _pre_calculate(self):
         """Perform any necessary pre-calculation steps before running the calculation"""
 
         # By default, check the corresponding program is installed
         self.check_code_is_installed()
 
+        # Copy over all files linked to this calculation
+        self._fetch_linked_files()
+
         return
 
     def _calculate(self):
@@ -236,6 +261,14 @@ def fromdict(cls: Type[TCalc], dct: Any) -> TCalc:
             setattr(calc, k.lstrip('_'), v)
         return calc
 
+    def link_file(self, src_calc: CalculatorExt | None, src_filename: Path, dest_filename: Path):
+        if src_filename.is_absolute() and src_calc is not None:
+            raise ValueError(f'"src_filename" in {self.__class__.__name__}.link_file() must be a relative path if a '
+                             f'"src_calc" is provided')
+        if dest_filename.is_absolute():
+            raise ValueError(f'"dest_filename" in {self.__class__.__name__}.link_file() must be a relative path')
+        self._linked_files[str(dest_filename)] = (src_calc, src_filename)
+
 
 class CalculatorABC(ABC, Generic[TCalc]):
 

src/koopmans/settings/_utils.py

@@ -130,14 +130,14 @@ def __setitem__(self, key: str, value: Any):
             self.pop(key, None)
             return
 
-        # Insisting that all values corresponding to paths are absolute and are Path objects
+        # Insisting that all values corresponding to paths are relative and are Path objects
         if key in self.are_paths:
             if isinstance(value, str):
                 value = Path(value)
             elif not isinstance(value, Path):
                 raise ValueError(f'{key} must be either a string or a Path')
-            if not value.is_absolute():
-                value = (self.directory / value).resolve()
+            if value.is_absolute() and key not in ['pseudo_dir', 'pseudo_directory']:
+                raise ValueError(f'{key} must be a relative path')
 
         # Parse any units provided
         if key in self.physicals:

src/koopmans/utils/_xml.py

@@ -1,6 +1,6 @@
 import xml.etree.ElementTree as ET
 from pathlib import Path
-from typing import List, Tuple
+from typing import List
 
 import numpy as np
 
@@ -60,15 +60,15 @@ def read_xml_array(
     # Extract the array
     array_xml = np.zeros((nr_xml[2], nr_xml[1], nr_xml[0]), dtype=float)
     for k in range(nr_xml[2]):
-        current_name = 'z.' + str(k % (nr_xml[2]-1)+1)
+        current_name = 'z.' + str(k % (nr_xml[2] - 1) + 1)
         entry = branch.find(current_name)
         assert isinstance(entry, ET.Element)
         text = entry.text
         assert isinstance(text, str)
         rho_tmp = np.array(text.split('\n')[1:-1], dtype=float)
         for j in range(nr_xml[1]):
             for i in range(nr_xml[0]):
-                array_xml[k, j, i] = rho_tmp[(j % (nr_xml[1]-1))*(nr_xml[0]-1)+(i % (nr_xml[0]-1))]
+                array_xml[k, j, i] = rho_tmp[(j % (nr_xml[1] - 1))*(nr_xml[0] - 1) + (i % (nr_xml[0] - 1))]
     array_xml *= norm_const
 
     if retain_final_element:

src/koopmans/utils/_xsf.py

@@ -17,7 +17,7 @@ def write_xsf(filename: Path, atoms: Atoms, arrays: List[np.ndarray], nr_xml: Tu
         for k in range(nr_xml[2]):
             for j in range(nr_xml[1]):
                 for i in range(nr_xml[0]):
-                    array_xsf[k, j, i] = array[k % (nr_xml[2]-1), j % (nr_xml[1]-1), i % (nr_xml[0]-1)]
+                    array_xsf[k, j, i] = array[k % (nr_xml[2] - 1), j % (nr_xml[1] - 1), i % (nr_xml[0] - 1)]
         arrays_xsf.append(array_xsf)
 
     cell_parameters = atoms.get_cell()

src/koopmans/workflows/_convergence_ml.py

@@ -42,7 +42,7 @@ def _fromjsondct(cls, bigdct: Dict[str, Any], override: Dict[str, Any] = {}):
         """
         try:
             snapshots_file = bigdct['atoms']['atomic_positions'].pop('snapshots')
-        except:
+        except KeyError:
             raise ValueError(f'To calculate a trajectory, please provide a xyz-file containing the atomic positions '
                              'of the snapshots in the setup-block of the json-input file.')
 
@@ -98,10 +98,11 @@ def inititalize_directories(self, grid_search_mode: bool):
             if not (grid_search_mode and convergence_point != self.convergence_points[-1]):
                 self.dirs[f'convergence_{convergence_point}'] = self.dirs[f'convergence_pred'] / \
                     f"predicted_after_{convergence_point+1}"
-            # in case of the grid search we won't produce all the plots, hence we also don't need the corresponding folders
+            # in case of the grid search we won't produce all the plots, hence we also don't need the corresponding
+            # folders
             if not grid_search_mode:
-                self.dirs[f'convergence_final_results_{convergence_point}'] = self.dirs[f'convergence_final_results'] / \
-                    f"predicted_after_{convergence_point+1}"
+                self.dirs[f'convergence_final_results_{convergence_point}'] = \
+                    self.dirs[f'convergence_final_results'] / f"predicted_after_{convergence_point+1}"
 
         # create all directories
         for dir in self.dirs.values():
@@ -184,7 +185,7 @@ def convert_to_list(param, type):
             try:
                 precompute_parameters_of_radial_basis(
                     self.ml.n_max, self.ml.l_max, self.ml.r_min, self.ml.r_max)
-            except:
+            except ValueError:
                 # skip this set of parameters if it is not possible to find the coefficients of the radial basis
                 # function
                 utils.warn(

src/koopmans/workflows/_koopmans_dfpt.py

@@ -167,13 +167,14 @@ def _run(self):
         base_outdir = self.calculator_parameters['pw'].outdir
         base_outdir.mkdir(exist_ok=True)
         scf_calcs = [c for c in self.calculations if isinstance(c, PWCalculator) and c.parameters.calculation == 'scf']
-        init_outdir = scf_calcs[-1].parameters.outdir
-        if self.parameters.from_scratch and init_outdir != base_outdir:
-            utils.symlink(f'{init_outdir}/*', base_outdir)
+        init_scf = scf_calcs[-1]
+        # if self.parameters.from_scratch and init_outdir != base_outdir:
+        #     utils.symlink(f'{init_outdir}/*', base_outdir)
 
         # Convert from wannier to KC
         self.print('Conversion to Koopmans format', style='subheading')
         wann2kc_calc = self.new_calculator('wann2kc')
+        self.link(init_scf, init_scf.parameters.outdir, wann2kc_calc, wann2kc_calc.parameters.outdir)
         self.run_calculator(wann2kc_calc)
 
         # Calculate screening parameters

src/koopmans/workflows/_wannierize.py

@@ -119,18 +119,19 @@ def _run(self):
 
         # Run PW scf and nscf calculations
         # PWscf needs only the valence bands
-        calc_pw = self.new_calculator('pw')
-        calc_pw.parameters.pop('nbnd', None)
-        calc_pw.directory = 'wannier'
-        calc_pw.prefix = 'scf'
+        calc_scf = self.new_calculator('pw')
+        calc_scf.parameters.pop('nbnd', None)
+        calc_scf.directory = 'wannier'
+        calc_scf.prefix = 'scf'
         if self._scf_kgrid:
-            calc_pw.parameters.kpts = self._scf_kgrid
-        self.run_calculator(calc_pw)
+            calc_scf.parameters.kpts = self._scf_kgrid
+        self.run_calculator(calc_scf)
 
-        calc_pw = self.new_calculator('pw', calculation='nscf', nosym=True, noinv=True)
-        calc_pw.directory = 'wannier'
-        calc_pw.prefix = 'nscf'
-        self.run_calculator(calc_pw)
+        calc_nscf = self.new_calculator('pw', calculation='nscf', nosym=True, noinv=True)
+        calc_nscf.directory = 'wannier'
+        calc_nscf.prefix = 'nscf'
+        self.link(calc_scf, calc_scf.parameters.outdir, calc_nscf, calc_nscf.parameters.outdir)
+        self.run_calculator(calc_nscf)
 
         if self.parameters.init_orbitals in ['mlwfs', 'projwfs'] \
                 and self.parameters.init_empty_orbitals in ['mlwfs', 'projwfs']:
@@ -170,8 +171,8 @@ def _run(self):
 
                 # 2) standard pw2wannier90 calculation
                 calc_p2w = self.new_calculator('pw2wannier', directory=w90_dir,
-                                               spin_component=block.spin,
-                                               outdir=calc_pw.parameters.outdir)
+                                               spin_component=block.spin)
+                self.link(calc_nscf, calc_nscf.parameters.outdir, calc_p2w, calc_p2w.parameters.outdir)
                 calc_p2w.prefix = 'pw2wan'
                 self.run_calculator(calc_p2w)
 
@@ -231,17 +232,12 @@ def _run(self):
             calc_pw_bands.parameters.prefix += '_bands'
 
             # Link the save directory so that the bands calculation can use the old density
-            if self.parameters.from_scratch:
-                [src, dest] = [(c.parameters.outdir / c.parameters.prefix).with_suffix('.save')
-                               for c in [calc_pw, calc_pw_bands]]
-
-                if dest.exists():
-                    shutil.rmtree(str(dest))
-                shutil.copytree(src, dest)
+            self.link(calc_nscf, (calc_nscf.parameters.outdir / calc_nscf.parameters.prefix).with_suffix('.save'),
+                      calc_pw_bands, (calc_pw_bands.parameters.outdir / calc_pw_bands.parameters.prefix).with_suffix('.save'))
             self.run_calculator(calc_pw_bands)
 
             # Calculate a projected DOS
-            pseudos = [read_pseudo_file(calc_pw_bands.parameters.pseudo_dir / p) for p in
+            pseudos = [read_pseudo_file(calc_pw_bands.directory / calc_pw_bands.parameters.pseudo_dir / p) for p in
                        self.pseudopotentials.values()]
             if all([p['header']['number_of_wfc'] > 0 for p in pseudos]):
                 calc_dos = self.new_calculator('projwfc', filpdos=self.name)

src/koopmans/workflows/_workflow.py

@@ -612,9 +612,9 @@ def new_calculator(self,
                        calc_type: str,
                        directory: Optional[Path] = None,
                        kpts: Optional[Union[List[int], BandPath]] = None,
-                       **kwargs) -> T:  # type: ignore[type-var, misc]
+                       **kwargs) -> calculators.Calc:  # type: ignore[type-var, misc]
 
-        calc_class: Type[T]
+        calc_class: calculators.CalcType
 
         if calc_type == 'kcp':
             calc_class = calculators.KoopmansCPCalculator
@@ -653,7 +653,7 @@ def new_calculator(self,
             all_kwargs['kpts'] = kpts if kpts is not None else self.kpoints.grid
 
         # Add further information to the calculator as required
-        for kw in ['pseudopotentials', 'pseudo_dir', 'gamma_only', 'kgrid', 'kpath', 'koffset', 'plotting']:
+        for kw in ['pseudopotentials', 'gamma_only', 'kgrid', 'kpath', 'koffset', 'plotting']:
             if kw not in all_kwargs and calculator_parameters.is_valid(kw):
                 val: Any
                 if kw == 'kgrid':
@@ -668,8 +668,6 @@ def new_calculator(self,
                         val = self.kpoints.offset
                 elif kw == 'gamma_only':
                     val = self.kpoints.gamma_only
-                elif kw == 'pseudo_dir':
-                    val = self.parameters.pseudo_directory
                 else:
                     val = getattr(self, kw)
                 all_kwargs[kw] = val
@@ -681,6 +679,12 @@ def new_calculator(self,
         if directory is not None:
             calc.directory = directory
 
+        # Link the pseudopotentials if relevant
+        if calculator_parameters.is_valid('pseudo_dir'):
+            for pseudo in self.pseudopotentials.values():
+                self.link(None, self.parameters.pseudo_directory / pseudo, calc, Path('pseudopotentials') / pseudo)
+            calc.parameters.pseudo_dir = 'pseudopotentials'
+
         return calc
 
     def primitive_to_supercell(self, matrix: Optional[npt.NDArray[np.int_]] = None, **kwargs):
@@ -882,6 +886,21 @@ def load_old_calculator(self, qe_calc: calculators.Calc) -> bool:
 
         return old_calc.is_complete()
 
+    def link(self, src_calc: calculators.Calc | None, src_path: Path, dest_calc: calculators.Calc, dest_path: Path) -> None:
+        """
+        Link a file from one calculator to another
+
+        Paths must be provided relative to the the calculator's directory i.e. calc.directory, unless src_calc is None
+        """
+
+        if src_path.is_absolute() and src_calc is not None:
+            raise ValueError(f'"src_path" in {self.__class__.__name__}.link() must be a relative path if a '
+                             f'"src_calc" is provided')
+        if dest_path.is_absolute():
+            raise ValueError(f'"dest_path" in {self.__class__.__name__}.link() must be a relative path')
+
+        dest_calc.link_file(src_calc, src_path, dest_path)
+
     def print(self, text: str = '', style: str = 'body', **kwargs: Any):
         if style == 'body':
             utils.indented_print(str(text), self.print_indent + 1, **kwargs)