Updating convergence machinery (#205)

This PR updates the machinery behind convergence calculations to make them much more flexible and much less dependent on what has been implemented inside of `koopmans`.

Users can now perform arbitrary convergence tests by defining their own `ConvergenceVariable`s and/or defining an `observable` function that tells `koopmans` how to extract the observable from the subworkflow. The custom `ConvergenceWorkflow` can easily be generated using the new `ConvergenceWorkflowFactory`.

For example, see the new `custom_convergence.py` script in Tutorial 4.

CITATION.cff

@@ -70,4 +70,4 @@ preferred-citation:
 repository-code: https://github.com/epfl-theos/koopmans/
 title: koopmans
 type: software
-version: 1.0.0
+version: 1.0.1

docs/_static/keywords/convergence_keywords.html

@@ -0,0 +1,26 @@
+<table id="convTable" style="width:100%; text-align:left">
+   <tr>
+      <th>Keyword</th>
+      <th>Description</th>
+      <th>Type</th>
+      <th>Default</th>
+   </tr>
+   <tr>
+      <td><code>observable</code></td>
+      <td>System observable of interest which we converge</td>
+      <td><code>str</code></td>
+      <td><code>total energy</code></td>
+   </tr>
+   <tr>
+      <td><code>threshold</code></td>
+      <td>threshold for the convergence of the observable of interest</td>
+      <td><code>str</code>/<code>float</code></td>
+      <td></td>
+   </tr>
+   <tr>
+      <td><code>variables</code></td>
+      <td>The observable of interest will be converged with respect to this (these) simulation variable(s)</td>
+      <td><code>list</code>/<code>str</code></td>
+      <td><code>['ecutwfc']</code></td>
+   </tr>
+</table>

docs/_static/keywords/generate_keywords.py

@@ -1,4 +1,5 @@
-from koopmans.settings import (MLSettingsDict, PlotSettingsDict,
+from koopmans.settings import (ConvergenceSettingsDict, MLSettingsDict,
+                               PlotSettingsDict,
                                UnfoldAndInterpolateSettingsDict,
                                WorkflowSettingsDict)
 
@@ -53,3 +54,6 @@ def print_settings_as_html(settings, table_id, fd):
 
     with open('ml_keywords.html', 'w') as fd:
         print_settings_as_html(MLSettingsDict(), 'mlTable', fd)
+
+    with open('convergence_keywords.html', 'w') as fd:
+        print_settings_as_html(ConvergenceSettingsDict(), 'convTable', fd)

docs/_static/keywords/workflow_keywords.html

@@ -38,7 +38,7 @@
    <tr>
       <td><code>pseudo_directory</code></td>
       <td>the folder containing the pseudopotentials to use (mutually exclusive with "pseudo_library")</td>
-      <td><code>str</code>/<code>Path</code></td>
+      <td><code>Path</code></td>
       <td></td>
    </tr>
    <tr>
@@ -131,6 +131,12 @@
       <td><code>float</code>/<code>list</code></td>
       <td><code>0.6</code></td>
    </tr>
+   <tr>
+      <td><code>alpha_mixing</code></td>
+      <td>mixing parameter for updating alpha</td>
+      <td><code>float</code></td>
+      <td><code>1.0</code></td>
+   </tr>
    <tr>
       <td><code>alpha_from_file</code></td>
       <td>if True, uses the file_alpharef.txt from the base directory as a starting guess</td>
@@ -162,22 +168,22 @@
       <td></td>
    </tr>
    <tr>
-      <td><code>convergence_observable</code></td>
-      <td>System observable of interest which we converge</td>
-      <td><code>str</code></td>
-      <td><code>total energy</code></td>
+      <td><code>orbital_groups_spread_tol</code></td>
+      <td>when calculating alpha parameters, the code will group orbitals together only if their spread is within this threshold</td>
+      <td><code>float</code></td>
+      <td></td>
    </tr>
    <tr>
-      <td><code>convergence_threshold</code></td>
-      <td>Convergence threshold for the system observable of interest</td>
-      <td><code>str</code>/<code>float</code></td>
-      <td></td>
+      <td><code>converge</code></td>
+      <td>If True, repeat the workflow increasing the convergence_parameters until the convergence_observable converges within the convergence_threshold</td>
+      <td><code>bool</code></td>
+      <td><code>False</code></td>
    </tr>
    <tr>
-      <td><code>convergence_parameters</code></td>
-      <td>The observable of interest will be converged with respect to this/these simulation parameter(s)</td>
-      <td><code>list</code>/<code>str</code></td>
-      <td><code>['ecutwfc']</code></td>
+      <td><code>dfpt_coarse_grid</code></td>
+      <td>The coarse k-point grid on which to perform the DFPT calculations</td>
+      <td><code>list</code></td>
+      <td></td>
    </tr>
    <tr>
       <td><code>eps_cavity</code></td>

docs/input_file.rst

@@ -163,6 +163,15 @@ This subblock controls the unfolding and interpolation procedure for generating
 
   List of valid keywords <input_file/ui_keywords>
 
+The convergence block
+^^^^^^^^^^^^^^^^^^^^^
+This block can be used to customize a convergence calculation.
+
+.. toctree::
+
+  List of valid keywords <input_file/convergence_keywords>
+
+See :ref:`here for a more detailed tutorial on performing convergence calculations <tutorials/tutorial_4:Tutorial 4: running convergence tests>`.
 
 The plotting block
 ^^^^^^^^^^^^^^^^^^

docs/input_file/convergence_keywords.rst

@@ -0,0 +1,15 @@
+Valid keywords
+==============
+
+.. raw:: html
+
+   <center>
+   <input type="text" id="convInput" onkeyup="lookupTable(convInput, convTable)" placeholder="Search for keywords...", style="width:50%">
+   </center>
+   <br>
+
+.. raw:: html
+   :file: ../_static/keywords/convergence_keywords.html
+
+.. raw:: html
+   :file: ../_static/keywords/lookup_table.html

docs/tutorials/tutorial_4.rst

@@ -1,30 +1,40 @@
 Tutorial 4: running convergence tests
 =====================================
-While ``koopmans`` is a package primarily oriented towards performing Koopmans spectral functional calculations, it does have a couple of other useful functionalities. In this tutorial, we will make use of its ``convergence`` task to determine how large a cell size and energy cutoff is required to converge the PBE energy of the highest occupied molecular orbital (HOMO) of a water molecule.
+While ``koopmans`` is a package primarily oriented towards performing Koopmans functional calculations, it does have a couple of other useful functionalities. Among these functionalities is the ability to perform arbitrary covnergence tests.
+
+
+A simple convergence test
+-------------------------
+
+In this tutorial, we will make use of its ``convergence`` task to determine how large a cell size and energy cutoff is required to converge the PBE energy of the highest occupied molecular orbital (HOMO) of a water molecule. This workflow was chosen for its simplicity; it is possible to run convergence tests on any workflow implemented in the ``koopmans`` package.
 
 The input file
---------------
-In order to run this calculation, our ``workflow`` block needs a few particular keywords:
+''''''''''''''
+In order to run this calculation, in the ``workflow`` block we need to set the ``converge`` parameter to true:
 
 .. literalinclude:: ../../tutorials/tutorial_4/h2o_conv.json
-  :lines: 1-11
-  :linenos:
-  :emphasize-lines: 4, 6-10
+  :lines: 1-9
+  :emphasize-lines: 6
 
-The important lines are highlighted. ``"task": "convergence"`` means that we will be performing a convergence test. The other three highlighted keywords specifying that we are going to converge the HOMO energy to within 0.01 eV, with respect to *both* the energy cutoff ``ecutwfc`` and the cell size. The full input file can be found :download:`here <../../tutorials/tutorial_4/h2o_conv.json>`.
+and then provide the convergence information in the ``convergence`` block:
+
+.. literalinclude:: ../../tutorials/tutorial_4/h2o_conv.json
+  :lines: 10-17
+
+These settings state that we are going to converge the HOMO energy to within 0.01 eV, with respect to *both* the energy cutoff ``ecutwfc`` and the cell size. The full input file can be found :download:`here <../../tutorials/tutorial_4/h2o_conv.json>`.
 
 The output file
----------------
+'''''''''''''''
 When you run the calculation, you should see something like this after the header:
 
-.. literalinclude:: ../../tutorials/tutorial_4/h2o_conv.out
-  :lines: 14-32
+.. literalinclude:: ../../tutorials/tutorial_4/h2o_conv.log
+  :lines: 15-34
   :language: text
 
-Here, the code is attempting to use progressively larger energy cutoffs and cell sizes. It will ultimately arrive at a converged solution, with a ``ecutwfc`` of 50.0 Ha and a cell 1.3 times larger than that provided in the ``.json`` input file.
+Here, the code is attempting to use progressively larger energy cutoffs and cell sizes. It will ultimately arrive at a converged solution, with a ``ecutwfc`` of 50.0 Ha and a cell slightly larger than that provided in the ``.json`` input file.
 
 Plotting
---------
+''''''''
 
 The individual ``Quantum ESPRESSO`` calculations reside in nested subdirectories. If you plot the HOMO energies from each of these, you should get something like this:
 
@@ -35,4 +45,30 @@ The individual ``Quantum ESPRESSO`` calculations reside in nested subdirectories
 
   Plot of the HOMO energy of water with respect to the energy cutoff and cell size (generated using :download:`this script <../../tutorials/tutorial_4/plot.py>`)
 
-We can see that indeed the calculation with ``ecutwfc = 50.0`` and ``cell_size = 1.3`` is the one where the energy of the HOMO goes within (and stays within) 0.01 eV of the most accurate calculation.
+We can see that indeed the calculation with ``ecutwfc = 50.0`` and ``celldm(1) = 13.3`` is the one where the energy of the HOMO goes within (and stays within) 0.01 eV of the most accurate calculation.
+
+A custom convergence test
+-------------------------
+
+In the previous example, we performed a convergence test with respect to ``ecutwfc`` and ``celldm1``. A full list of supported convergence variables can be found :ref:`here <input_file/convergence_keywords:Valid keywords>`. You will see that only a couple of variables are implemented by default. But don't let that limit you! It is possible to perform a convergence test on arbitrary keywords using factories.
+
+First, try taking the input file from the first part of the tutorial, and switch the pseudopotential library to ``pseudo_dojo_standard``. What do you notice?
+
+Hopefully, the first thing you will see is that there are now some warnings about the small box parameters ``nrb`` like this,
+
+.. code-block:: text
+
+  UserWarning: Small box parameters "nrb" not provided in input: these will be automatically set to safe default values. These values can probably be decreased, but this would require convergence tests.
+  Estimated real mesh dimension (nr1, nr2, nr3) = ...
+  Small box mesh dimension (nr1b, nr2b, nr3b) = ...
+
+These parameters are associated with the way ``Quantum ESPRESSO`` handles non-local core corrections in pseudopotentials, and corrections are present in the new set of pseudopotentials but absent in the SG15 pseudopotentials.
+
+So, let's perform a convergence test! The ``nrb`` parameters are not directly implemented as a convergence variable in ``koopmans``, but we can use a factory to perform a convergence test on them, by making use of the ``ConvergenceVariable`` and ``ConvergenceWorkflowfactory`` classes.
+
+.. literalinclude:: ../../tutorials/tutorial_4/custom_convergence.py
+
+Running this script will perform a convergence test with respect to ``nrb 1-3``.
+
+.. warning:: 
+  This tutorial performs convergence tests in a slightly incorrect way, To see this, add the keyword ``length = 10`` to the ``ConvergenceVarialbe`` in the above script. How is the behaviour different? Which behaviour is correct? Why?

pyproject.toml

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

src/koopmans/cell.py

@@ -1,5 +1,5 @@
 import copy
-from typing import Dict, Union
+from typing import Dict, TypedDict, Union
 
 import numpy as np
 from ase.cell import Cell
@@ -8,6 +8,11 @@
 from ase.units import Bohr
 
 
+class CellParams(TypedDict):
+    ibrav: int
+    celldms: Dict[int, float]
+
+
 def parameters_to_cell(ibrav: int, celldms: Dict[int, float]) -> Cell:
     '''
     Converts QE cell parameters to the corresponding ASE `Cell` object.
@@ -129,7 +134,7 @@ def parameters_to_cell(ibrav: int, celldms: Dict[int, float]) -> Cell:
     return Cell(new_array)
 
 
-def cell_to_parameters(cell: Cell) -> Dict[str, Union[int, Dict[int, float]]]:
+def cell_to_parameters(cell: Cell) -> CellParams:
     '''
     Identifies a cell in the language of Quantum ESPRESSO
 
@@ -140,8 +145,8 @@ def cell_to_parameters(cell: Cell) -> Dict[str, Union[int, Dict[int, float]]]:
 
     Returns
     -------
-    Dict
-        a dictionary containing the ibrav and celldms
+    CellParams
+        a typed dictionary containing the ibrav and celldms
 
     '''
 
@@ -254,7 +259,7 @@ def cell_to_parameters(cell: Cell) -> Dict[str, Union[int, Dict[int, float]]]:
     # Convert to Bohr radii
     celldms[1] /= Bohr
 
-    return {'ibrav': ibrav, 'celldms': celldms}
+    return CellParams(ibrav=ibrav, celldms=celldms)
 
 
 def cell_follows_qe_conventions(cell: Cell) -> bool:

src/koopmans/io/_json.py

@@ -34,8 +34,6 @@ def read_json(fd: TextIO, override: Dict[str, Any] = {}):
     workflow_cls: Type[workflows.Workflow]
     if parameters.task == 'singlepoint':
         workflow_cls = workflows.SinglepointWorkflow
-    elif parameters.task == 'convergence':
-        workflow_cls = workflows.ConvergenceWorkflow
     elif parameters.task == 'wannierize':
         workflow_cls = workflows.WannierizeWorkflow
     elif parameters.task == 'environ_dscf':

src/koopmans/io/_kwf.py

@@ -6,6 +6,7 @@
 
 """
 
+import functools
 import inspect
 import json
 import os
@@ -16,6 +17,7 @@
 from ase.io import jsonio as ase_json
 
 import koopmans.workflows as workflows
+from koopmans import utils
 
 
 class KoopmansEncoder(ase_json.MyEncoder):
@@ -26,6 +28,14 @@ def default(self, obj) -> dict:
             return {'__path__': os.path.relpath(obj, '.')}
         elif inspect.isclass(obj):
             return {'__class__': {'__name__': obj.__name__, '__module__': obj.__module__}}
+        elif isinstance(obj, functools.partial):
+            return {'__partial__': {'func': obj.func, 'args': obj.args, 'keywords': obj.keywords}}
+        elif inspect.isfunction(obj):
+            module = import_module(obj.__module__)
+            if hasattr(module, obj.__name__):
+                return {'__function__': {'__name__': obj.__name__, '__module__': obj.__module__}}
+            else:
+                return {'__function__': {}}
         elif hasattr(obj, 'todict'):
             d = obj.todict()
             if '__koopmans_name__' in d:
@@ -49,6 +59,17 @@ def object_hook(dct):
         subdct = dct['__class__']
         module = import_module(subdct['__module__'])
         return getattr(module, subdct['__name__'])
+    elif '__partial__' in dct:
+        subdct = dct['__partial__']
+        return functools.partial(subdct['func'], **subdct['keywords'])
+    elif '__function__' in dct:
+        subdct = dct['__function__']
+        if '__module__' in subdct:
+            module = import_module(subdct['__module__'])
+            if hasattr(module, subdct['__name__']):
+                return getattr(module, subdct['__name__'])
+        utils.warn('Function was not able to be serialized')
+        return f'<placeholder for function {subdct["__name__"]} lost during serialization>'
     else:
         # Patching bug in ASE where allocating an np.empty(dtype=str) will assume a particular length for each
         # string. dtype=object allows for individual strings to be different lengths

src/koopmans/settings/__init__.py

@@ -6,6 +6,7 @@
 
 '''
 
+from ._convergence import ConvergenceSettingsDict
 from ._koopmans_cp import KoopmansCPSettingsDict
 from ._koopmans_ham import KoopmansHamSettingsDict
 from ._koopmans_screen import KoopmansScreenSettingsDict

src/koopmans/settings/_convergence.py

@@ -0,0 +1,36 @@
+import os
+from pathlib import Path
+from typing import Any
+
+from koopmans import pseudopotentials
+
+from ._utils import Setting, SettingsDictWithChecks
+
+
+class ConvergenceSettingsDict(SettingsDictWithChecks):
+
+    def __init__(self, **kwargs) -> None:
+        settings = [
+            Setting('observable',
+                    'System observable of interest which we converge',
+                    str, 'total energy', None),
+            Setting('threshold',
+                    'threshold for the convergence of the observable of interest',
+                    (str, float), None, None),
+            Setting('variables',
+                    'The observable of interest will be converged with respect to this (these) '
+                    'simulation variable(s)',
+                    (list, str), ['ecutwfc'], None)]
+
+        super().__init__(settings=settings, physicals=['threshold'], **kwargs)
+
+    @property
+    def _other_valid_keywords(self):
+        return []
+
+    def __setitem__(self, key: str, value: Any):
+        # Convert parameters to a list
+        if key == 'parameters' and isinstance(value, str):
+            value = [value]
+
+        return super().__setitem__(key, value)