auto-calculate ADDED_MOS based on basissets

aiida_cp2k/calculations/__init__.py

@@ -22,6 +22,7 @@
     validate_pseudos_namespace,
     write_basissets,
     write_pseudos,
+    estimate_added_mos,
 )
 from ..utils import Cp2kInput
 
@@ -139,7 +140,7 @@ def prepare_for_submission(self, folder):
         :return: `aiida.common.datastructures.CalcInfo` instance
         """
 
-        # pylint: disable=too-many-statements,too-many-branches
+        # pylint: disable=too-many-statements,too-many-branches,too-many-locals
 
         # Create cp2k input file.
         inp = Cp2kInput(self.inputs.parameters.get_dict())
@@ -167,6 +168,21 @@ def prepare_for_submission(self, folder):
                                self.inputs.structure if 'structure' in self.inputs else None)
             write_basissets(inp, self.inputs.basissets, folder)
 
+        # if we have both basissets and structure we can start helping the user :)
+        if 'basissets' in self.inputs and 'structure' in self.inputs:
+            try:
+                scf_section = inp.get_section_dict('FORCE_EVAL/DFT/SCF')
+
+                if 'SMEAR' in scf_section and 'ADDED_MOS' not in scf_section:
+                    # now is our time to shine!
+                    added_mos = estimate_added_mos(self.inputs.basissets, self.inputs.structure)
+                    inp.add_keyword('FORCE_EVAL/DFT/SCF/ADDED_MOS', added_mos)
+                    self.logger.info(f'The FORCE_EVAL/DFT/SCF/ADDED_MOS was added with an automatically estimated value'
+                                     f' of {added_mos}')
+
+            except (KeyError, TypeError):  # no SCF, no smearing, or multiple FORCE_EVAL, nothing to do (yet)
+                pass
+
         if 'pseudos' in self.inputs:
             validate_pseudos(inp, self.inputs.pseudos, self.inputs.structure if 'structure' in self.inputs else None)
             write_pseudos(inp, self.inputs.pseudos, folder)

aiida_cp2k/utils/datatype_helpers.py

@@ -160,6 +160,31 @@ def validate_basissets(inp, basissets, structure):
             kind_sec["ELEMENT"] = bset.element
 
 
+def estimate_added_mos(basissets, structure, fraction=0.3):
+    """Calculate an estimate for ADDED_MOS based on used basis sets"""
+
+    symbols = [structure.get_kind(s.kind_name).get_symbols_string() for s in structure.sites]
+    n_mos = 0
+
+    # We are currently overcounting in the following cases:
+    #   * if we get a mix of ORB basissets for the same chemical symbol but different sites
+    #   * if we get multiple basissets for one element (merged within CP2K)
+
+    for label, bset in _unpack(basissets):
+        try:
+            _, bstype = label.split("_", maxsplit=1)
+        except ValueError:
+            bstype = "ORB"
+
+        if bstype != "ORB":  # ignore non-ORB basissets
+            continue
+
+        n_mos += symbols.count(bset.element) * bset.n_orbital_functions
+
+    # at least one additional MO per site, otherwise a fraction of the total number of orbital functions
+    return max(len(symbols), int(fraction * n_mos))
+
+
 def write_basissets(inp, basissets, folder):
     """Writes the unified BASIS_SETS file with the used basissets"""
     _write_gdt(inp, basissets, folder, "BASIS_SET_FILE_NAME", "BASIS_SETS")

aiida_cp2k/utils/input_generator.py

@@ -56,6 +56,37 @@ def add_keyword(self, kwpath, value, override=True, conflicting_keys=None):
 
         Cp2kInput._add_keyword(kwpath, value, self._params, ovrd=override, cfct=conflicting_keys)
 
+    def get_section_dict(self, kwpath):
+        return self._get_section_or_kw(kwpath, True)
+
+    def get_keyword_value(self, kwpath):
+        return self._get_section_or_kw(kwpath, False)
+
+    def _get_section_or_kw(self, kwpath, section_requested):
+        """Retrieve either a section or a keyword given a path"""
+
+        if isinstance(kwpath, str):  # turn to list of sections
+            kwpath = kwpath.split("/")
+        kwpath = [k.upper() for k in kwpath]  # accept any case, but internally we use uppercase
+        orig_kwpath = kwpath
+
+        current = self._params
+
+        try:
+            while kwpath:
+                current = current[kwpath.pop(0)]
+            return current
+        except KeyError:
+            raise KeyError("Section '{}' not found in parameters".format("/".join(orig_kwpath)))
+
+        if isinstance(current, Mapping):
+            if not section_requested:
+                raise TypeError("Section '{}' requested, but keyword found".format("/".join(orig_kwpath)))
+        elif section_requested:
+            raise TypeError("Section '{}' requested, but keyword found".format("/".join(orig_kwpath)))
+
+        return deepcopy(current)
+
     def render(self):
         output = [self.DISCLAIMER]
         self._render_section(output, deepcopy(self._params))

test/test_gaussian_datatypes.py

@@ -778,3 +778,83 @@ def test_without_kinds(cp2k_code, cp2k_basissets, cp2k_pseudos, clear_database):
 
     _, calc_node = run_get_node(CalculationFactory("cp2k"), **inputs)
     assert calc_node.exit_status == 0
+
+
+def test_added_mos(cp2k_code, cp2k_basissets, cp2k_pseudos, clear_database):  # pylint: disable=unused-argument
+    """Testing CP2K with the Basis Set stored in gaussian.basisset and a smearing section"""
+
+    # structure
+    atoms = ase.build.molecule("H2O")
+    atoms.center(vacuum=2.0)
+    structure = StructureData(ase=atoms)
+
+    # parameters
+    parameters = Dict(
+        dict={
+            'GLOBAL': {
+                'RUN_TYPE': 'ENERGY',
+            },
+            'FORCE_EVAL': {
+                'METHOD': 'Quickstep',
+                'DFT': {
+                    "XC": {
+                        "XC_FUNCTIONAL": {
+                            "_": "PBE",
+                        },
+                    },
+                    "MGRID": {
+                        "CUTOFF": 1000.0,
+                        "REL_CUTOFF": 100.0,
+                    },
+                    "QS": {
+                        "METHOD": "GPW",
+                        "EXTRAPOLATION": "USE_GUESS",
+                    },
+                    "SCF": {
+                        "EPS_SCF": 1e-08,
+                        "MAX_SCF": 200,
+                        "MIXING": {
+                            "METHOD": "BROYDEN_MIXING",
+                            "ALPHA": 0.4,
+                        },
+                        "SMEAR": {
+                            "METHOD": "FERMI_DIRAC",
+                            "ELECTRONIC_TEMPERATURE": 300.0,
+                        },
+                    },
+                    "KPOINTS": {
+                        "SCHEME": "MONKHORST-PACK 4 4 2",
+                        "FULL_GRID": False,
+                        "SYMMETRY": False,
+                        "PARALLEL_GROUP_SIZE": -1,
+                    },
+                },
+            },
+        })
+
+    options = {
+        "resources": {
+            "num_machines": 1,
+            "num_mpiprocs_per_machine": 1
+        },
+        "max_wallclock_seconds": 1 * 3 * 60,
+    }
+
+    inputs = {
+        "structure": structure,
+        "parameters": parameters,
+        "code": cp2k_code,
+        "metadata": {
+            "options": options,
+        },
+        "basissets": cp2k_basissets,
+        "pseudos": cp2k_pseudos,
+    }
+
+    _, calc_node = run_get_node(CalculationFactory("cp2k"), **inputs)
+
+    # check that the ADDED_MOS keyword was added within the calculation
+    with calc_node.open("aiida.inp") as fhandle:
+        assert any("ADDED_MOS" in line for line in fhandle), "ADDED_MOS not found in the generated CP2K input file"
+
+    assert calc_node.exit_status == 0