Performing DFPT on a coarser grid (#194)

This PR implements the ability to calculate the screening parameters via DFPT on a coarser grid than that which is used to construct the Hamiltonian. This is controlled by the keyword `dfpt_coarse_grid` in the `workflow` block, which is a direct analogue of the `grid` keyword in the `kpoints` block.

It also introduces the ability to group variational orbitals based on their spread (provided that they are Wannier functions). This is required in order to perform orbital grouping for DFPT workflows.

Finally, this PR makes `koopmans` suppress traceback by default. This can be re-enabled by running `koopmans` with the flag `--traceback` or `-t`.

README.rst

@@ -133,7 +133,7 @@ Written and maintained by Edward Linscott, Riccardo De Gennaro, and Nicola Colon
 
 For help and feedback email [email protected]
 
-.. |GH Actions| image:: https://img.shields.io/github/workflow/status/epfl-theos/koopmans/Run%20tests/master?label=master&logo=github
+.. |GH Actions| image:: https://img.shields.io/github/actions/workflow/status/epfl-theos/koopmans/tests.yml?master&label=master&logo=github
    :target: https://github.com/epfl-theos/koopmans/actions?query=branch%3Amaster
 .. |Coverage Status| image:: https://img.shields.io/codecov/c/gh/epfl-theos/koopmans/master?logo=codecov
    :target: https://codecov.io/gh/epfl-theos/koopmans

README_pypi.rst

@@ -27,7 +27,7 @@ Written and maintained by Edward Linscott, Riccardo De Gennaro, and Nicola Colon
 
 For help and feedback email [email protected]
 
-.. |GH Actions| image:: https://img.shields.io/github/workflow/status/epfl-theos/koopmans/Run%20tests/master?label=master&logo=github
+.. |GH Actions| image:: https://img.shields.io/github/actions/workflow/status/epfl-theos/koopmans/tests.yml?master&label=master&logo=github
    :target: https://github.com/epfl-theos/koopmans/actions?query=branch%3Amaster
 .. |Coverage Status| image:: https://img.shields.io/codecov/c/gh/epfl-theos/koopmans/master?logo=codecov
    :target: https://codecov.io/gh/epfl-theos/koopmans

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "koopmans"
-version = "1.0.0-beta.6"
+version = "1.0.0-rc.1"
 description = "Koopmans spectral functional calculations with python and Quantum ESPRESSO"
 readme = "README_pypi.rst"
 requires-python = ">=3.7"

src/koopmans/bands.py

@@ -1,5 +1,5 @@
 import itertools
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -75,7 +75,7 @@ def error(self, value: Optional[float]):
 
 class Bands(object):
     def __init__(self, n_bands: Union[int, List[int]], n_spin: int = 1, spin_polarized: bool = False,
-                 self_hartree_tol=None, **kwargs):
+                 tolerances: Dict[str, float] = {}, **kwargs):
         if isinstance(n_bands, int):
             self.n_bands = [n_bands for _ in range(n_spin)]
         else:
@@ -95,7 +95,7 @@ def __init__(self, n_bands: Union[int, List[int]], n_spin: int = 1, spin_polariz
             self._bands = [Band(i_band + 1, i_spin, group=i_band) for i_spin in range(n_spin)
                            for i_band in range(self.n_bands[i_spin])]
 
-        self.self_hartree_tol = self_hartree_tol
+        self.tolerances = tolerances
         for k, v in kwargs.items():
             assert hasattr(self, k)
             if v:
@@ -187,15 +187,18 @@ def groups(self, value: List[List[int]]):
         for b, v in zip(self, [v for subarray in value for v in subarray]):
             b.group = v
 
-    def assign_groups(self, sh_tol: Optional[float] = None, allow_reassignment: bool = False):
+    def assign_groups(self, sort_by: str = 'self_hartree', tol: Optional[float] = None, allow_reassignment: bool = False):
         # Basic clustering algorithm for assigning groups
 
-        if self.self_hartree_tol is None:
+        if self.tolerances == {}:
             # Do not perform clustering
             return
 
+        if sort_by not in self.tolerances:
+            return ValueError(f'Cannot sort bands according to {sort_by}; valid choices are' + '/'.join(self.tolerances.keys()))
+
         # By default use the settings provided when Bands() was initialized
-        sh_tol = sh_tol if sh_tol is not None else self.self_hartree_tol
+        tol = tol if tol is not None else self.tolerances[sort_by]
 
         # Separate the orbitals into different subsets, where we don't want any grouping of orbitals belonging to
         # different subsets
@@ -209,10 +212,8 @@ def assign_groups(self, sh_tol: Optional[float] = None, allow_reassignment: bool
 
         def points_are_close(p0: Band, p1: Band, factor: Union[int, float] = 1) -> bool:
             # Determine if two bands are "close"
-            for obs, tol in (('self_hartree', sh_tol),):
-                if tol is not None and abs(getattr(p0, obs) - getattr(p1, obs)) > tol * factor:
-                    return False
-            return True
+            assert tol is not None
+            return abs(getattr(p0, sort_by) - getattr(p1, sort_by)) < tol * factor
 
         group = 0
         for unassigned in unassigned_sets:
@@ -224,8 +225,8 @@ def points_are_close(p0: Band, p1: Band, factor: Union[int, float] = 1) -> bool:
                 neighborhood = [b for b in unassigned if points_are_close(guess, b)]
 
                 # Find the center of that neighborhood
-                av_sh = np.mean([b.self_hartree for b in neighborhood])
-                center = Band(self_hartree=av_sh)
+                av = np.mean([getattr(b, sort_by) for b in neighborhood])
+                center = Band(**{sort_by: av})
 
                 # Find a revised neighborhood close to the center (using a factor of 0.5 because we want points that
                 # are on opposite sides of the neighborhood to be within "tol" of each other which means they can be
@@ -236,11 +237,11 @@ def points_are_close(p0: Band, p1: Band, factor: Union[int, float] = 1) -> bool:
                 wider_neighborhood = [b for b in unassigned if points_are_close(center, b)]
 
                 if neighborhood != wider_neighborhood:
-                    if self.self_hartree_tol and sh_tol < 0.01 * self.self_hartree_tol:
+                    if self.tolerances[sort_by] and tol < 0.01 * self.tolerances[sort_by]:
                         # We have recursed too deeply, abort
                         raise Exception('Clustering algorithm failed')
                     else:
-                        self.assign_groups(sh_tol=0.9 * sh_tol if sh_tol else None,
+                        self.assign_groups(sort_by, tol=0.9 * tol if tol else None,
                                            allow_reassignment=allow_reassignment)
                         return
 
@@ -266,10 +267,10 @@ def points_are_close(p0: Band, p1: Band, factor: Union[int, float] = 1) -> bool:
                 [match] = [b_op for b_op in self.get(spin=0) if b_op.index == b.index]
                 b.group = match.group
 
-        if sh_tol != self.self_hartree_tol:
-            warn(f'It was not possible to group orbitals with the self-Hartree tolerance of {self.self_hartree_tol:.2e} eV. '
-                 f'A grouping was found for a self-Hartree tolerance of {sh_tol:.2e} eV.\n'
-                 f'Try a larger self-Hartree tolerance to group more orbitals together')
+        if tol != self.tolerances[sort_by]:
+            warn(f'It was not possible to group orbitals with the {sort_by} tolerance of {self.tolerances[sort_by]:.2e} eV. '
+                 f'A grouping was found for a tolerance of {tol:.2e} eV.\n'
+                 f'Try a larger tolerance to group more orbitals together')
 
         return
 

src/koopmans/cli/main.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import sys
 import argparse
 import textwrap
 
@@ -16,12 +17,18 @@ def main():
         epilog='See https://koopmans-functionals.org for more details')
     parser.add_argument('json', metavar='system.json', type=str,
                         help='a single JSON file containing the workflow and code settings')
+    parser.add_argument('-t', '--traceback', action='store_true', help='enable traceback')
 
     # Parse arguments
     args = parser.parse_args()
 
     # Reading in JSON file
     workflow = read(args.json)
 
+    # Set traceback behaviour
+    if not args.traceback:
+        sys.tracebacklimit = 0
+
+
     # Run workflow
     workflow.run()

src/koopmans/settings/_workflow.py

@@ -104,6 +104,10 @@ def __init__(self, **kwargs) -> None:
                     'together only if their self-Hartree energy is within this '
                     'threshold',
                     float, None, None),
+            Setting('orbital_groups_spread_tol',
+                    'when calculating alpha parameters, the code will group orbitals '
+                    'together only if their spread is within this threshold',
+                    float, None, None),
             Setting('convergence_observable',
                     'System observable of interest which we converge',
                     str, 'total energy', None),
@@ -114,6 +118,9 @@ def __init__(self, **kwargs) -> None:
                     'The observable of interest will be converged with respect to this/these '
                     'simulation parameter(s)',
                     (list, str), ['ecutwfc'], None),
+            Setting('dfpt_coarse_grid',
+                    'The coarse k-point grid on which to perform the DFPT calculations',
+                    list, None, None),
             Setting('eps_cavity',
                     'a list of epsilon_infinity values for the cavity in dscf calculations',
                     list, [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20], None)]

src/koopmans/workflows/_koopmans_dfpt.py

@@ -8,6 +8,7 @@
 
 import shutil
 from pathlib import Path
+from typing import Dict
 
 from koopmans import utils, pseudopotentials
 from koopmans.bands import Bands
@@ -39,6 +40,9 @@ def __init__(self, *args, **kwargs):
                 raise ValueError(
                     'Calculating screening parameters with DFPT for a non-periodic system is only possible '
                     'with Kohn-Sham orbitals as the variational orbitals')
+        if self.parameters.dfpt_coarse_grid is not None and not self.parameters.calculate_alpha:
+            raise ValueError('Using a DFPT coarse grid to calculate the screening parameters is not possible '
+                             'when calculate_alpha = False')
         for params in self.calculator_parameters.values():
             if all(self.atoms.pbc):
                 # Gygi-Baldereschi
@@ -93,14 +97,21 @@ def __init__(self, *args, **kwargs):
         nemp = ntot - nocc
         if self.parameters.orbital_groups is None:
             self.parameters.orbital_groups = [list(range(nocc + nemp))]
+        tols: Dict[str, float] = {}
+        for key in ['self_hartree', 'spread']:
+            val = self.parameters.get(f'orbital_groups_{key}_tol', None)
+            if val is not None:
+                tols[key] = val
         self.bands = Bands(n_bands=nocc + nemp, filling=[[True] * nocc + [False] * nemp],
-                           groups=self.parameters.orbital_groups)
+                           groups=self.parameters.orbital_groups, tolerances=tols)
 
         # Populating kpoints if absent
         if not all(self.atoms.pbc):
             for key in ['pw', 'kc_screen']:
                 self.calculator_parameters[key].kpts = [1, 1, 1]
 
+        self._perform_ham_calc: bool = True
+
     def _run(self):
         '''
         This function runs the workflow from start to finish
@@ -116,6 +127,15 @@ def _run(self):
                 if output_directory.exists():
                     shutil.rmtree(output_directory)
 
+        if self.parameters.dfpt_coarse_grid is not None:
+            self.print('Coarse grid calculations', style='heading')
+            coarse_wf = self.__class__.fromparent(self)
+            coarse_wf.parameters.dfpt_coarse_grid = None
+            coarse_wf.kpoints.grid = self.parameters.dfpt_coarse_grid
+            coarse_wf._perform_ham_calc = False
+            coarse_wf.run(subdirectory='coarse_grid')
+            self.print('Regular grid calculations', style='heading')
+
         if all(self.atoms.pbc):
             # Run PW and Wannierization
             for key in self.calculator_parameters.keys():
@@ -156,33 +176,39 @@ def _run(self):
 
         # Calculate screening parameters
         if self.parameters.calculate_alpha:
-            self.print('Calculation of screening parameters', style='heading')
-            all_groups = [g for gspin in self.parameters.orbital_groups for g in gspin]
-            if len(set(all_groups)) == len(all_groups):
-                # If there is no orbital grouping, do all orbitals in one calculation
-                # 1) Create the calculator
-                kc_screen_calc = self.new_calculator('kc_screen')
-
-                # 2) Run the calculator
-                self.run_calculator(kc_screen_calc)
-
-                # 3) Store the computed screening parameters
-                self.bands.alphas = kc_screen_calc.results['alphas']
-            else:
-                # If there is orbital grouping, do the orbitals one-by-one
-                for band in self.bands.to_solve:
-                    # 1) Create the calculator (in a subdirectory)
-                    kc_screen_calc = self.new_calculator('kc_screen', i_orb=band.index)
-                    kc_screen_calc.directory /= f'band_{band.index}'
+            if self.parameters.dfpt_coarse_grid is None:
+                self.print('Calculation of screening parameters', style='heading')
+
+                # Group the bands by spread
+                self.bands.assign_groups(sort_by='spread', allow_reassignment=True)
+
+                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
+                    kc_screen_calc = self.new_calculator('kc_screen')
 
                     # 2) Run the calculator
                     self.run_calculator(kc_screen_calc)
 
-                    # 3) Store the computed screening parameter (accounting for band groupings)
-                    for b in self.bands:
-                        if b.group == band.group:
-                            alpha = kc_screen_calc.results['alphas'][band.spin]
-                            b.alpha = alpha[band.spin]
+                    # 3) Store the computed screening parameters
+                    self.bands.alphas = kc_screen_calc.results['alphas']
+                else:
+                    # If there is orbital grouping, do the orbitals one-by-one
+                    for band in self.bands.to_solve:
+                        # 1) Create the calculator (in a subdirectory)
+                        kc_screen_calc = self.new_calculator('kc_screen', i_orb=band.index)
+                        kc_screen_calc.directory /= f'band_{band.index}'
+
+                        # 2) Run the calculator
+                        self.run_calculator(kc_screen_calc)
+
+                        # 3) Store the computed screening parameter (accounting for band groupings)
+                        for b in self.bands:
+                            if b.group == band.group:
+                                alpha = kc_screen_calc.results['alphas'][band.spin]
+                                b.alpha = alpha[band.spin]
+            else:
+                self.bands.alphas = coarse_wf.bands.alphas
         else:
             # Load the alphas
             if self.parameters.alpha_from_file:
@@ -191,31 +217,33 @@ def _run(self):
                 self.bands.alphas = self.parameters.alpha_guess
 
         # Calculate the Hamiltonian
-        self.print('Construction of the Hamiltonian', style='heading')
-        kc_ham_calc = self.new_calculator('kc_ham', kpts=self.kpoints.path)
-
-        if self.parameters.calculate_alpha and kc_ham_calc.parameters.lrpa != kc_screen_calc.parameters.lrpa:
-            raise ValueError('Do not set "lrpa" to different values in the "screen" and "ham" blocks')
-        self.run_calculator(kc_ham_calc)
-
-        # Postprocessing
-        if all(self.atoms.pbc) and self.projections and self.kpoints.path is not None \
-                and self.calculator_parameters['ui'].do_smooth_interpolation:
-            from koopmans.workflows import UnfoldAndInterpolateWorkflow
-            self.print(f'\nPostprocessing', style='heading')
-            ui_workflow = UnfoldAndInterpolateWorkflow.fromparent(self)
-            ui_workflow.run(subdirectory='postproc')
-
-        # Plotting
-        self.plot_bandstructure()
+        if self._perform_ham_calc:
+            self.print('Construction of the Hamiltonian', style='heading')
+            kc_ham_calc = self.new_calculator('kc_ham', kpts=self.kpoints.path)
+
+            if self.parameters.calculate_alpha and self.parameters.dfpt_coarse_grid is None:
+                if kc_ham_calc.parameters.lrpa != kc_screen_calc.parameters.lrpa:
+                    raise ValueError('Do not set "lrpa" to different values in the "screen" and "ham" blocks')
+            self.run_calculator(kc_ham_calc)
+
+            # Postprocessing
+            if all(self.atoms.pbc) and self.projections and self.kpoints.path is not None \
+                    and self.calculator_parameters['ui'].do_smooth_interpolation:
+                from koopmans.workflows import UnfoldAndInterpolateWorkflow
+                self.print(f'\nPostprocessing', style='heading')
+                ui_workflow = UnfoldAndInterpolateWorkflow.fromparent(self)
+                ui_workflow.run(subdirectory='postproc')
+
+            # Plotting
+            self.plot_bandstructure()
 
     def plot_bandstructure(self):
         if not all(self.atoms.pbc):
             return
 
         # Identify the relevant calculators
-        [wann2kc_calc] = [c for c in self.calculations if isinstance(c, Wann2KCCalculator)]
-        [kc_ham_calc] = [c for c in self.calculations if isinstance(c, KoopmansHamCalculator)]
+        wann2kc_calc = [c for c in self.calculations if isinstance(c, Wann2KCCalculator)][-1]
+        kc_ham_calc = [c for c in self.calculations if isinstance(c, KoopmansHamCalculator)][-1]
 
         # Plot the bandstructure if the band path has been specified
         bs = kc_ham_calc.results['band structure']