From 0cb32cd0f3e76159d6058ac5b3cee5b714a21930 Mon Sep 17 00:00:00 2001
From: Aliaksandr Yakutovich <yakutovicha@gmail.com>
Date: Tue, 29 Aug 2023 21:47:59 +0200
Subject: [PATCH] Allow `UPF` pseudos input. (#195)

- Add `aiida-pseudo` and `upf-to-json` dependencies.
- Update pre-commit settings.
---
 .flake8                                 |  6 +++---
 .pre-commit-config.yaml                 | 12 ++++++------
 aiida_cp2k/calculations/__init__.py     | 17 +++++++++++++++++
 examples/gaussian_datatypes/gdt_data.py |  4 +++-
 pyproject.toml                          |  2 ++
 5 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/.flake8 b/.flake8
index c04df776..50b1c1d5 100644
--- a/.flake8
+++ b/.flake8
@@ -3,8 +3,8 @@
 # rather than using the flake8 default of 79:
 max-line-length = 88
 extend-ignore =
-    E501  # Line length handled by black.
-    W503  # Line break before binary operator, preferred formatting for black.
-    E203  # Whitespace before ':', preferred formatting for black.
+    E501
+    W503
+    E203
 exclude =
     docs/source/conf.py
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d9d66788..d08d6934 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,7 +1,7 @@
 ---
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.2.0
+    rev: v4.4.0
     hooks:
 
         - id: check-json
@@ -16,32 +16,32 @@ repos:
           exclude: *exclude_pre_commit_hooks
 
   - repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 23.7.0
     hooks:
         - id: black
           language_version: python3 # Should be a command that runs python3.6+
           exclude: aiida_cp2k/workchains/base.py
 
   - repo: https://github.com/PyCQA/flake8
-    rev: 4.0.1
+    rev: 6.1.0
     hooks:
         - id: flake8
           args: [--count, --show-source, --statistics]
 
   - repo: https://github.com/pycqa/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
         - id: isort
           args: [--profile, black, --filter-files]
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.32.1
+    rev: v3.10.1
     hooks:
         - id: pyupgrade
           args: [--py37-plus]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.961
+    rev: v1.5.1
     hooks:
         - id: mypy
           additional_dependencies:
diff --git a/aiida_cp2k/calculations/__init__.py b/aiida_cp2k/calculations/__init__.py
index b79fb7ad..9eadca6d 100644
--- a/aiida_cp2k/calculations/__init__.py
+++ b/aiida_cp2k/calculations/__init__.py
@@ -6,12 +6,14 @@
 ###############################################################################
 """AiiDA-CP2K input plugin."""
 
+import json
 from operator import add
 
 from aiida.common import CalcInfo, CodeInfo, InputValidationError
 from aiida.engine import CalcJob
 from aiida.orm import Dict, RemoteData, SinglefileData
 from aiida.plugins import DataFactory
+from upf_to_json import upf_to_json
 
 from ..utils import Cp2kInput
 from ..utils.datatype_helpers import (
@@ -26,6 +28,7 @@
 BandsData = DataFactory("array.bands")  # pylint: disable=invalid-name
 StructureData = DataFactory("structure")  # pylint: disable=invalid-name
 KpointsData = DataFactory("array.kpoints")  # pylint: disable=invalid-name
+UpfData = DataFactory("pseudo.upf")  # pylint: disable=invalid-name
 
 
 class Cp2kCalculation(CalcJob):
@@ -105,6 +108,14 @@ def define(cls, spec):
             ),
         )
 
+        spec.input_namespace(
+            "pseudos_upf",
+            valid_type=UpfData,
+            dynamic=True,
+            required=True,
+            help="A mapping of `UpfData` nodes onto the kind name to which they should apply.",
+        )
+
         # Specify default parser.
         spec.input(
             "metadata.options.parser_name",
@@ -262,6 +273,12 @@ def prepare_for_submission(self, folder):
             )
             write_pseudos(inp, self.inputs.pseudos, folder)
 
+        if "pseudos_upf" in self.inputs:
+            for atom_kind, pseudo in self.inputs.pseudos_upf.items():
+                pseudo_dict = upf_to_json(pseudo.get_content(), atom_kind)
+                with folder.open(atom_kind + ".json", "w") as fobj:
+                    fobj.write(json.dumps(pseudo_dict, indent=2))
+
         # Kpoints.
         if "kpoints" in self.inputs:
             try:
diff --git a/examples/gaussian_datatypes/gdt_data.py b/examples/gaussian_datatypes/gdt_data.py
index 74f0466e..3ae4619a 100644
--- a/examples/gaussian_datatypes/gdt_data.py
+++ b/examples/gaussian_datatypes/gdt_data.py
@@ -74,7 +74,9 @@ def load_data(prefix="MY-"):
             p.element: p
             for p in Pseudo.from_cp2k(fhandle_pseudo, duplicate_handling="error")
         }
-    except UniquenessError:  # if the user already ran the script, fetch the data from the db instead
+    except (
+        UniquenessError
+    ):  # if the user already ran the script, fetch the data from the db instead
         bsets = {
             "H": BasisSet.get("H", f"{prefix}AUTO-DZVP-MOLOPT-GTH"),
             "O": BasisSet.get("O", f"{prefix}AUTO-DZVP-MOLOPT-SR-GTH"),
diff --git a/pyproject.toml b/pyproject.toml
index 52ffb55c..7841d3e2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,6 +25,8 @@ dependencies = [
     "ase",
     "ruamel.yaml>=0.16.5",
     "cp2k-output-tools",
+    "aiida-pseudo~=0.6",
+    "upf-to-json>=0.9",
 ]
 
 [[project.authors]]