allow symbol choosing for GTH PP parser

fishjojo · Feb 4, 2024 · 543bca3 · 543bca3
1 parent 6f16c30
commit 543bca3
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/pyscf/gto/basis/parse_cp2k_pp.py b/pyscf/gto/basis/parse_cp2k_pp.py
@@ -21,6 +21,7 @@
 '''
 
 import sys
+import re
 from pyscf.lib.exceptions import BasisNotFoundError
 import numpy as np
 
@@ -45,8 +46,14 @@ def parse(string, symb=None):
     ...      0.28637912    0
     ... """)}
     '''
-    pseudotxt = [x.strip() for x in string.splitlines()
-                 if x.strip() and 'END' not in x and '#PSEUDOPOTENTIAL' not in x]
+    if symb is not None:
+        raw_data = list(filter(None, re.split('#PSEUDOPOTENTIAL', string)))
+        pseudotxt = _search_gthpp_block(raw_data, symb)
+        if not pseudotxt:
+            raise BasisNotFoundError(f'Pseudopotential not found for {symb}.')
+    else:
+        pseudotxt = [x.strip() for x in string.splitlines()
+                     if x.strip() and 'END' not in x and '#PSEUDOPOTENTIAL' not in x]
     return _parse(pseudotxt)
 
 def load(pseudofile, symb, suffix=None):
@@ -95,7 +102,13 @@ def search_seg(pseudofile, symb, suffix=None):
     fin = open(pseudofile, 'r')
     fdata = fin.read().split('#PSEUDOPOTENTIAL')
     fin.close()
-    for dat in fdata[1:]:
+    dat = _search_gthpp_block(fdata[1:], symb, suffix)
+    if not dat:
+        raise BasisNotFoundError(f'Pseudopotential for {symb} in {pseudofile}')
+    return dat
+
+def _search_gthpp_block(raw_data, symb, suffix=None):
+    for dat in raw_data:
         dat0 = dat.split(None, 1)
         if dat0 and dat0[0] == symb:
             dat = [x.strip() for x in dat.splitlines()
@@ -107,7 +120,7 @@ def search_seg(pseudofile, symb, suffix=None):
             else:
                 if any(suffix == x.split('-')[-1] for x in dat[0].split()):
                     return dat
-    raise BasisNotFoundError(f'Pseudopotential for {symb} in {pseudofile}')
+    return None
 
 if __name__ == '__main__':
     args = sys.argv[1:]

diff --git a/pyscf/gto/test/test_basis_parser.py b/pyscf/gto/test/test_basis_parser.py
@@ -21,7 +21,7 @@
 from pyscf import lib
 from pyscf.gto.basis import parse_molpro
 from pyscf.gto.basis import parse_gaussian
-from pyscf.gto.basis import parse_cp2k
+from pyscf.gto.basis import parse_cp2k, parse_cp2k_pp
 from pyscf.lib.exceptions import BasisNotFoundError
 
 class KnownValues(unittest.TestCase):
@@ -505,6 +505,29 @@ def test_parse_gth_basis(self):
         ref = gto.basis.load('gth-dzv', 'C')
         self.assertEqual(ref, basis1)
 
+    def test_parse_gth_pp(self):
+        pp_str = '''
+            #PSEUDOPOTENTIAL
+            B GTH-PADE-q3 GTH-LDA-q3 GTH-PADE GTH-LDA
+                2    1
+                 0.43392956    2    -5.57864173     0.80425145
+                2
+                 0.37384326    1     6.23392822
+                 0.36039317    0
+            #PSEUDOPOTENTIAL
+            C GTH-PADE-q4 GTH-LDA-q4 GTH-PADE GTH-LDA
+                2    2
+                 0.34883045    2    -8.51377110     1.22843203
+                2
+                 0.30455321    1     9.52284179
+                 0.23267730    0'''
+        pp1 = parse_cp2k_pp.parse(pp_str, 'B')
+        ref = gto.basis.load_pseudo('gth-pade', 'B')
+        self.assertEqual(ref, pp1)
+        pp1 = parse_cp2k_pp.parse(pp_str, 'C')
+        ref = gto.basis.load_pseudo('gth-pade', 'C')
+        self.assertEqual(ref, pp1)
+
 if __name__ == "__main__":
     print("test basis module")
     unittest.main()