Enable spin orbit calculation (#35)

* Enable spin-orbit coupling calculation in `Wannier90BandsWorkChain`.

* Adjust the selection of bands and semicores in `get_wannier_number_of_bands`
  and `get_semicore_list` respectively, when `spin_type` is set to `SpinType.SPIN_ORBIT`
* Add support for `PseudoDojo/0.4/PBE/FR/standard/upf`, and `PSlibrary` (should be installed by the user)
* Add tests for `spin_type=SpinType.SPIN_ORBIT`

---------

Co-authored-by: Yuhao Jiang <[email protected]>
Co-authored-by: npaulish <[email protected]>

src/aiida_wannier90_workflows/utils/pseudo/__init__.py

@@ -84,6 +84,10 @@ def get_pseudo_orbitals(pseudos: ty.Mapping[str, PseudoPotentialData]) -> dict:
     pseudo_data.append(
         load_pseudo_metadata("semicore/PseudoDojo_0.4_LDA_SR_stringent_upf.json")
     )
+    pseudo_data.append(
+        load_pseudo_metadata("semicore/PseudoDojo_0.4_PBE_FR_standard_upf.json")
+    )
+    pseudo_data.append(load_pseudo_metadata("semicore/pslibrary_paw_relpbe_1.0.0.json"))
 
     pseudo_orbitals = {}
     for element in pseudos:
@@ -99,11 +103,14 @@ def get_pseudo_orbitals(pseudos: ty.Mapping[str, PseudoPotentialData]) -> dict:
     return pseudo_orbitals
 
 
-def get_semicore_list(structure: orm.StructureData, pseudo_orbitals: dict) -> list:
+def get_semicore_list(
+    structure: orm.StructureData, pseudo_orbitals: dict, spin_orbit_coupling: bool
+) -> list:
     """Get semicore states (a subset of pseudo wavefunctions) in the pseudopotential.
 
     :param structure: [description]
     :param pseudo_orbitals: [description]
+    :param spin_orbit_coupling: [description]
     :return: [description]
     """
     from copy import deepcopy
@@ -122,6 +129,8 @@ def get_semicore_list(structure: orm.StructureData, pseudo_orbitals: dict) -> li
     #     "semicores": ["5S", "5P"]
     # }
     label2num = {"S": 1, "P": 3, "D": 5, "F": 7}
+    # for spin-orbit-coupling, every orbit contains 2 electrons
+    nspin = 2 if spin_orbit_coupling else 1
 
     semicore_list = []  # index should start from 1
     num_pswfcs = 0
@@ -132,7 +141,7 @@ def get_semicore_list(structure: orm.StructureData, pseudo_orbitals: dict) -> li
         site_semicores = deepcopy(pseudo_orbitals[site.kind_name]["semicores"])
 
         for orb in site_pswfcs:
-            num_orbs = label2num[orb[-1]]
+            num_orbs = label2num[orb[-1]] * nspin
             if orb in site_semicores:
                 site_semicores.remove(orb)
                 semicore_list.extend(
@@ -182,7 +191,7 @@ def get_wannier_number_of_bands(
 
     num_electrons = get_number_of_electrons(structure, pseudos)
     num_projections = get_number_of_projections(structure, pseudos, spin_orbit_coupling)
-    nspin = 2 if spin_polarized else 1
+    nspin = 2 if (spin_polarized or spin_orbit_coupling) else 1
     # TODO check nospin, spin, soc  # pylint: disable=fixme
     if only_valence:
         num_bands = int(0.5 * num_electrons * nspin)

src/aiida_wannier90_workflows/utils/pseudo/data/__init__.py

@@ -2,6 +2,7 @@
 import json
 import os
 import typing as ty
+import xml.sax
 
 __all__ = ("load_pseudo_metadata",)
 
@@ -29,17 +30,110 @@ def md5(filename):
     return hash_md5.hexdigest()
 
 
+class PSHandler(xml.sax.ContentHandler):
+    """Read xml files in pslibrary using xml.sax.
+
+    This script can generate cutoff energy/rho, pswfcs and semicores for
+    FULLY RELATIVISTIC pseudopotentials (rel-*).
+    Use SSSP or Pesudo-dojo for scalar/non relativistic pseudos.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        # define the p-block
+        self.pblock = list(range(31, 37))
+        self.pblock.extend(list(range(49, 55)))
+        self.pblock.extend(list(range(81, 87)))
+        self.pblock.extend(list(range(113, 119)))
+        # in p-block ns, np are planewaves and n-1d are semicores
+        # otherwise ns, np and n-1d are pws, n-1s n-1p are semicores
+        self.readWFC = False
+        self.pswfcs = []
+        self.pswfcs_shell = []
+        self.semicores = []
+        self.znum = 0
+
+    def startElement(self, name, attrs):
+        """Instructions for the beginning of different sections.
+
+        If in PP_MESH section, read element's index in periodic table;
+        If in PP_SPIN_ORB, set read wavefunctions flag to True;
+        If in PP_RELWFC, read the orbital labels (5S, 5D, etc.)
+        and calculate the "weight" by taking the principal quantum number
+        and adding 1 for D orbitals, S and P orbitals for the elements in P block,
+        and 2 for F orbitals. If not in P block, S and P weight is 0.
+        This will later be used to identify semicores.
+        """
+
+        if name == "PP_MESH":
+            try:
+                self.znum = int(float(attrs["zmesh"]))
+            except ValueError:
+                print(f"z = {attrs['zmesh']} is not acceptable")
+
+        if name == "PP_SPIN_ORB":
+            # <PP_SPIN_ORB>
+            #   <PP_RELWFC.1 index="1" els="5S" nn="1" lchi="0" jchi="0.500000000000000" oc="2.00000000000000"/>
+            #   <PP_RELWFC.2 index="2" els="6S" nn="2" lchi="0" jchi="0.500000000000000" oc="1.00000000000000"/>
+            #   ...
+            #   <PP_RELBETA.1 index="1" lll="0" jjj="0.500000000000000"/>
+            #   <PP_RELBETA.2 index="2" lll="0" jjj="0.500000000000000"/>
+            #   <PP_RELBETA.3 index="3" lll="0" jjj="0.500000000000000"/>
+            #   ...
+            # </PP_SPIN_ORB>
+
+            self.readWFC = True
+        if "PP_RELWFC" in name and self.readWFC:
+            orb = attrs["els"]
+            if not orb in self.pswfcs:
+                self.pswfcs.append(orb)
+                nn = int(orb[0])
+                orbtype = orb[-1]
+                if orbtype in ("S", "P"):
+                    ll = 1 if self.znum in self.pblock else 0
+                if orbtype == "D":
+                    ll = 1
+                if orbtype == "F":
+                    ll = 2
+                self.pswfcs_shell.append(nn + ll)
+
+    def endElement(self, name):
+        """Select semicores.
+
+        When the PP_SPIN_ORB ends, all pswfcs are recorded in list,
+        then we can determine semicores using the rules introduced in init
+        """
+
+        if name == "PP_SPIN_ORB":
+            self.readWFC = False
+            self.znum = 0
+            maxshell = max(self.pswfcs_shell)
+            for iorb in enumerate(self.pswfcs_shell):
+                if iorb[1] < maxshell:
+                    self.semicores.append(self.pswfcs[iorb[0]])
+
+
 def get_metadata(filename):
     """Return metadata."""
+
+    # this part reads the upf file twice, but it do not take too much time
     result = {"filename": filename, "md5": md5(filename), "pseudopotential": "100PAW"}
     with open(filename, encoding="utf-8") as handle:
         for line in handle:
             if "Suggested minimum cutoff for wavefunctions" in line:
                 wave = float(line.strip().split()[-2])
             if "Suggested minimum cutoff for charge density" in line:
                 charge = float(line.strip().split()[-2])
-        result["cutoff"] = wave
-        result["dual"] = charge / wave
+    # use xml.sax to parse upf file
+    parser = xml.sax.make_parser()
+    Handler = PSHandler()
+    parser.setContentHandler(Handler)
+    parser.parse(filename)
+    # cutoffs in unit: Ry
+    result["cutoff_wfc"] = wave
+    result["cutoff_rho"] = charge
+    result["pswfcs"] = Handler.pswfcs
+    result["semicores"] = Handler.semicores
     return result
 
 
@@ -49,10 +143,11 @@ def generate_pslibrary_metadata(dirname=None):
     :param dirname: folder to be scanned, if None download from QE website
     :type dirname: str
     """
+    import shutil
     import urllib.request
 
     output_filename = "pslibrary_paw_relpbe_1.0.0.json"
-    qe_site = "https://www.quantum-espresso.org/upf_files/"
+    qe_site = "https://pseudopotentials.quantum-espresso.org/upf_files/"
 
     # these are the suggested PP from https://dalcorso.github.io/pslibrary/PP_list.html (2020.07.21)
     suggested = r"""H:  H.$fct-*_psl.1.0.0
@@ -152,6 +247,15 @@ def generate_pslibrary_metadata(dirname=None):
     # use PAW, PBE, SOC
     star = "kjpaw"
     fct = "rel-pbe"
+    # For conventionally installed pslibrary, the dirname folder will contain
+    # more pseudos than what we expect. It is essential to create a sub-folder
+    # to store the pseudos we need if we want to install the pseudos to aiida-pseudo.
+    # If dirname == None, the pslibrary_install folder will be placed in workdir.
+    if dirname is not None:
+        dir_pseudos_install = os.path.join(dirname, "pslibrary_install")
+    else:
+        dir_pseudos_install = "pslibrary_install"
+    os.makedirs(dir_pseudos_install, exist_ok=True)
     result = {}
     suggested = suggested.replace("*", star).replace("$fct", fct)
     suggested = suggested.split("\n")
@@ -164,19 +268,49 @@ def generate_pslibrary_metadata(dirname=None):
             filename = "Co.rel-pbe-spn-kjpaw_psl.0.3.1.UPF"
         if filename == "Au.rel-pbe-n-kjpaw_psl.1.0.1.UPF":
             filename = "Au.rel-pbe-n-kjpaw_psl.1.0.0.UPF"
+        # Cs.rel-pbe-spnl-kjpaw_psl.1.0.0.UPF gives 'z_valence=-5.00000000000000'
+        # it is not a legal value for aiida-pseudo install
+        if filename == "Cs.rel-pbe-spnl-kjpaw_psl.1.0.0.UPF":
+            filename = "Cs.rel-pbe-spn-kjpaw_psl.1.0.0.UPF"
+        # these file do not exist in qe_site
+        if filename == "Ar.rel-pbe-nl-kjpaw_psl.1.0.0.UPF":
+            filename = "Ar.rel-pbe-n-kjpaw_psl.1.0.0.UPF"
+        if filename == "Fe.rel-pbe-n-kjpaw_psl.1.0.0.UPF":
+            filename = "Fe.rel-pbe-spn-kjpaw_psl.1.0.0.UPF"
+        if filename == "Ni.rel-pbe-n-kjpaw_psl.1.0.0.UPF":
+            filename = "Ni.rel-pbe-spn-kjpaw_psl.1.0.0.UPF"
+        if filename == "Zn.rel-pbe-dn-kjpaw_psl.1.0.0.UPF":
+            filename = "Zn.rel-pbe-dnl-kjpaw_psl.1.0.0.UPF"
         print(filename)
+        # copy the pseudopotentials to another folder, which only contain the pp we need
+        dst_filename = os.path.join(dir_pseudos_install, filename)
         if dirname is not None:
-            filename = os.path.join(dirname, filename)
+            shutil.copy2(os.path.join(dirname, filename), dst_filename)
         else:
             if not os.path.exists(filename):
                 # download from QE website
                 url = qe_site + filename
-                urllib.request.urlretrieve(url, filename)
-        result[element] = get_metadata(filename)
-
+                urllib.request.urlretrieve(url, dst_filename)
+        result[element] = get_metadata(dst_filename)
+        result[element]["filename"] = filename
+    # the output file (as well as the pseudo_install if dirname==None) will be placed in the workdir
     with open(output_filename, "w", encoding="utf-8") as handle:
         json.dump(result, handle, indent=2)
 
+    print(
+        "Use following commands to install pslibrary as a cutoffs family in aiida-pseudo"
+    )
+    print(
+        f"\taiida-pseudo install family '{dir_pseudos_install}' <LABEL> -F pseudo.family.cutoffs -P pseudo.upf"
+    )
+    print(
+        f"\taiida-pseudo family cutoffs set -s <STRINGENCY(standard)> -u Ry <FAMILY> '{output_filename}'"
+    )
+    print(
+        "Please move the json file into "
+        + "'.../aiida-wannier90-workflows/src/aiida_wannier90_workflows/utils/pseudo/data/semicore/'"
+    )
+
 
 def generate_dojo_metadata():
     """Generate metadata for pseduo-dojo SOC pseudos.
@@ -230,5 +364,5 @@ def _print_exclude_semicore():
 
 
 # if __name__ == '__main__':
-#     # generate_pslibrary_metadata()
-#     generate_dojo_metadata()
+#     generate_pslibrary_metadata()
+#     # generate_dojo_metadata()