WIP: Tutorial 3 running (results not yet checked)

src/koopmans/workflows/_koopmans_dfpt.py

@@ -15,12 +15,20 @@
 from koopmans.calculators import (KoopmansHamCalculator, PWCalculator,
                                   Wann2KCCalculator, Wannier90Calculator)
 from koopmans.files import FilePointer
+from koopmans.outputs import OutputModel
 
 from ._workflow import Workflow
 
 
+class KoopmansDFPTOutputs(OutputModel):
+    pass
+
+
 class KoopmansDFPTWorkflow(Workflow):
 
+    output_model = KoopmansDFPTOutputs  # type: ignore
+    outputs: KoopmansDFPTOutputs
+
     def __init__(self, scf_kgrid=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -139,6 +147,7 @@ def _run(self):
             coarse_wf.run(subdirectory='coarse_grid')
             self.print('Regular grid calculations', style='heading')
 
+        wannier_files_to_link = {}
         if all(self.atoms.pbc):
             # Run PW and Wannierization
             for key in self.calculator_parameters.keys():
@@ -152,6 +161,20 @@ def _run(self):
             init_pw = [c for c in self.calculations if isinstance(
                 c, PWCalculator) and c.parameters.calculation == 'nscf'][-1]
 
+            # Populate a list of files to link to subsequent calculations
+            for f in [wf_workflow.outputs.u_matrices_files["occ"],
+                      wf_workflow.outputs.hr_files["occ"],
+                      wf_workflow.outputs.u_dis_file,
+                      wf_workflow.outputs.centers_files["occ"]]:
+                assert f is not None
+                wannier_files_to_link[f.name] = f
+
+            for f in [wf_workflow.outputs.u_matrices_files["emp"],
+                      wf_workflow.outputs.hr_files["emp"],
+                      wf_workflow.outputs.centers_files["emp"]]:
+                assert f is not None
+                wannier_files_to_link[f.parent.prefix + '_emp' + f.name[len(f.parent.prefix):]] = f
+
         else:
             # Run PW
             self.print('Initialization of density and variational orbitals', style='heading')
@@ -172,12 +195,15 @@ def _run(self):
         # Convert from wannier to KC
         wann2kc_calc = self.new_calculator('wann2kc')
         self.link(init_pw, init_pw.parameters.outdir, wann2kc_calc, wann2kc_calc.parameters.outdir)
+        for dst, f in wannier_files_to_link.items():
+            self.link(f.parent, f.name, wann2kc_calc, dst, symlink=True)
         self.run_calculator(wann2kc_calc)
 
         # Calculate screening parameters
         if self.parameters.calculate_alpha:
             if self.parameters.dfpt_coarse_grid is None:
-                screen_wf = ComputeScreeningViaDFPTWorkflow.fromparent(self)
+                screen_wf = ComputeScreeningViaDFPTWorkflow.fromparent(
+                    self, wannier_files_to_link=wannier_files_to_link)
                 screen_wf.run(subdirectory='screening')
             else:
                 self.bands.alphas = coarse_wf.bands.alphas
@@ -197,6 +223,8 @@ def _run(self):
                       kc_ham_calc, kc_ham_calc.parameters.outdir / (kc_ham_calc.parameters.prefix + '.xml'), symlink=True)
             self.link(wann2kc_calc, wann2kc_calc.parameters.outdir / 'kcw', kc_ham_calc,
                       kc_ham_calc.parameters.outdir / 'kcw', recursive_symlink=True)
+            for dst, f in wannier_files_to_link.items():
+                self.link(f.parent, f.name, kc_ham_calc, dst, symlink=True)
             self.run_calculator(kc_ham_calc)
 
             # Postprocessing
@@ -233,12 +261,24 @@ def new_calculator(self, calc_presets, **kwargs):
         return internal_new_calculator(self, calc_presets, **kwargs)
 
 
+class ComputeScreeningViaDFPTOutputs(OutputModel):
+    pass
+
+
 class ComputeScreeningViaDFPTWorkflow(Workflow):
+    output_model = ComputeScreeningViaDFPTOutputs
+    outputs: ComputeScreeningViaDFPTOutputs
+
+    def __init__(self, *args, wannier_files_to_link: Dict[str, FilePointer], **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._wannier_files_to_link = wannier_files_to_link
+
     def _run(self):
         # Group the bands by spread
         self.bands.assign_groups(sort_by='spread', allow_reassignment=True)
 
-        if len(self.bands.to_solve) == len(self.bands) and False:
+        if len(self.bands.to_solve) == len(self.bands):
             # If there is no orbital grouping, do all orbitals in one calculation
 
             # 1) Create the calculator
@@ -283,6 +323,9 @@ def new_calculator(self, calc_type: str, *args, **kwargs):
         self.link(wann2kc_calc, wann2kc_calc.parameters.outdir / 'kcw', calc,
                   calc.parameters.outdir / 'kcw', recursive_symlink=True)
 
+        for dst, f in self._wannier_files_to_link.items():
+            self.link(f.parent, f.name, calc, dst, symlink=True)
+
         return calc
 
 
@@ -330,11 +373,6 @@ def internal_new_calculator(workflow, calc_presets, **kwargs):
     calc.parameters.have_empty = have_empty
     calc.parameters.has_disentangle = has_disentangle
 
-    # Provide the rotation matrices and the wannier centers if required
-    if all(workflow.atoms.pbc):
-        for process in workflow.processes:
-            workflow.link(process, process.outputs.dst_file, calc, process.outputs.dst_file, symlink=True)
-
     # Apply any additional calculator keywords passed as kwargs
     for k, v in kwargs.items():
         setattr(calc.parameters, k, v)

src/koopmans/workflows/_koopmans_dscf.py

@@ -1149,14 +1149,18 @@ def _run(self) -> None:
             wannier_workflow.run()
 
             # Store the Hamitonian files
+            hr_file_keys: List[Tuple[str, str | None]]
             if self.parameters.spin_polarized:
-                wannier_hamiltonian_files = {('occ', "up"): wannier_workflow.outputs.hr_files['occ_up'],
-                                             ('emp', "up"): wannier_workflow.outputs.hr_files['emp_up'],
-                                             ('occ', "down"): wannier_workflow.outputs.hr_files['occ_down'],
-                                             ('emp', "down"): wannier_workflow.outputs.hr_files['emp_down']}
+                hr_file_keys = [('occ', 'up'), ('emp', 'up'), ('occ', 'down'), ('emp', 'down')]
             else:
-                wannier_hamiltonian_files = {('occ', None): wannier_workflow.outputs.hr_files['occ'],
-                                             ('emp', None): wannier_workflow.outputs.hr_files['emp']}
+                hr_file_keys = [('occ', None), ('emp', None)]
+
+            wannier_hamiltonian_files = {}
+            for key in hr_file_keys:
+                key_str = '_'.join([k for k in key if k is not None])
+                hr_file = wannier_workflow.outputs.hr_files[key_str]
+                assert hr_file is not None
+                wannier_hamiltonian_files[key] = hr_file
 
             # Convert the files over from w90 format to kcp format
             nscf_calc = wannier_workflow.steps[1]

src/koopmans/workflows/_wannierize.py

@@ -57,6 +57,7 @@ class Config:
 class WannierizeWorkflow(Workflow):
 
     output_model = WannierizeOutput  # type: ignore
+    outputs: WannierizeOutput
 
     def __init__(self, *args, force_nspin2=False, scf_kgrid=None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -188,32 +189,34 @@ def _run(self):
                         # Fetching the list of calculations for this block
                         src_calcs: List[calculators.Wannier90Calculator] = [
                             b.w90_calc for b in block if b.w90_calc is not None]
-                        emp_label = '_emp' if filling_label == 'emp' else ''
+                        emp_label = '_emp' if label == 'emp' else ''
+                        prefix = src_calcs[-1].prefix
 
                         # Merging the wannier_hr (Hamiltonian) files
                         merge_hr_proc = MergeProcess(merge_function=merge_wannier_hr_file_contents,
-                                                     src_files=[(calc, Path(prefix + '_hr.dat')) for calc in src_calcs],
+                                                     src_files=[(calc, Path(calc.prefix + '_hr.dat'))
+                                                                for calc in src_calcs],
                                                      dst_file=prefix + f'{emp_label}_hr.dat')
-                        merge_hr_proc.name = f'merge_{filling_label}_wannier_hamiltonian'
+                        merge_hr_proc.name = f'merge_{label}_wannier_hamiltonian'
                         self.run_process(merge_hr_proc)
                         hr_files[label] = FilePointer(merge_hr_proc, merge_hr_proc.outputs.dst_file)
 
                         if self.parameters.method == 'dfpt' and self.parent is not None:
                             # Merging the U (rotation matrix) files
                             merge_u_proc = MergeProcess(merge_function=merge_wannier_u_file_contents,
-                                                        src_files=[(calc, Path(prefix + '_u.mat'))
+                                                        src_files=[(calc, Path(calc.prefix + '_u.mat'))
                                                                    for calc in src_calcs],
                                                         dst_file=prefix + f'{emp_label}_u.mat')
-                            merge_u_proc.name = f'merge_{filling_label}_wannier_u'
+                            merge_u_proc.name = f'merge_{label}_wannier_u'
                             self.run_process(merge_u_proc)
                             u_matrices_files[label] = FilePointer(merge_u_proc, merge_u_proc.outputs.dst_file)
 
                             # Merging the wannier centers files
                             merge_centers_proc = MergeProcess(merge_function=partial(merge_wannier_centers_file_contents, atoms=self.atoms),
-                                                              src_files=[(calc, Path(prefix + '_centres.xyz'))
+                                                              src_files=[(calc, Path(calc.prefix + '_centres.xyz'))
                                                                          for calc in src_calcs],
                                                               dst_file=prefix + f'{emp_label}_centres.xyz')
-                            merge_centers_proc.name = f'merge_{filling_label}_wannier_centers'
+                            merge_centers_proc.name = f'merge_{label}_wannier_centers'
                             self.run_process(merge_centers_proc)
                             centers_files[label] = FilePointer(merge_centers_proc, merge_centers_proc.outputs.dst_file)
 

src/koopmans/workflows/_workflow.py

@@ -1043,12 +1043,7 @@ def _parent_context(self, subdirectory: Optional[str] = None,
 
         # Link the bands
         if hasattr(self.parent, 'bands'):
-            self.bands = copy.deepcopy(self.parent.bands)
-            # Only include the most recent screening parameter and wipe the error history
-            for b in self.bands:
-                if len(b.alpha_history) > 0:
-                    b.alpha_history = [b.alpha]
-                b.error_history = []
+            self.bands = self.parent.bands
 
         subdirectory = self.name.replace(' ', '-').lower() if subdirectory is None else subdirectory
 
@@ -1077,20 +1072,9 @@ def _parent_context(self, subdirectory: Optional[str] = None,
             if from_scratch is None:
                 self.parent.parameters.from_scratch = self.parameters.from_scratch
 
-            # Copy back over the bands
-            if hasattr(self, 'bands'):
-                if hasattr(self.parent, 'bands'):
-                    # Add the alpha and error history if the length of the bands match
-                    if len(self.parent.bands) == len(self.bands):
-                        for b, b_sub in zip(self.parent.bands, self.bands):
-                            b.alpha_history += b_sub.alpha_history[1:]
-                            b.error_history += b_sub.error_history[1:]
-                            b.self_hartree = b_sub.self_hartree
-                            b.spread = b_sub.spread
-                            b.center = b_sub.center
-                else:
-                    # Copy the entire bands object
-                    self.parent.bands = self.bands
+            if hasattr(self, 'bands') and not hasattr(self.parent, 'bands'):
+                # Copy the entire bands object
+                self.parent.bands = self.bands
 
         self.parent.steps.append(self.steps)