From 92ee82b932e8d198b17e3555780d8af009794343 Mon Sep 17 00:00:00 2001
From: Jamal Mustafa <jimustafa@gmail.com>
Date: Tue, 23 Aug 2022 23:17:00 -0700
Subject: [PATCH] parse nnkpts block in NNKP file

---
 src/wannier90io/_nnkp.py   | 12 ++++++++++++
 src/wannier90io/_schema.py |  1 +
 2 files changed, 13 insertions(+)

diff --git a/src/wannier90io/_nnkp.py b/src/wannier90io/_nnkp.py
index 9ab93c0..da91e0b 100644
--- a/src/wannier90io/_nnkp.py
+++ b/src/wannier90io/_nnkp.py
@@ -26,6 +26,11 @@
         r'(?P<projections>.+)',
         re.IGNORECASE | re.DOTALL
     ),
+    'nnkpts': re.compile(
+        r'\d+\s+'
+        r'(?P<nnkpts>.+)',
+        re.IGNORECASE | re.DOTALL
+    ),
     'exclude_bands': re.compile(
         r'\d+\s*'
         r'(?P<exclude_bands>.*)',
@@ -115,6 +120,12 @@ def parse_spinor_projections(string: str) -> dict:
     ]
 
 
+def parse_nnkpts(string: str) -> list:
+    match = patterns['nnkpts'].search(string)
+
+    return [[int(x) for x in line.split()] for line in match.group('nnkpts').splitlines()]
+
+
 def parse_exclude_bands(string: str) -> dict:
     match = patterns['exclude_bands'].search(string)
 
@@ -149,6 +160,7 @@ def parse_nnkp_raw(string: str) -> dict:
         'direct_lattice': parse_direct_lattice(blocks['real_lattice']),
         'reciprocal_lattice': parse_reciprocal_lattice(blocks['recip_lattice']),
         'kpoints': parse_kpoints(blocks['kpoints']),
+        'nnkpts': parse_nnkpts(blocks['nnkpts']),
         'exclude_bands': parse_exclude_bands(blocks['exclude_bands']),
     }
     if 'projections' in blocks:
diff --git a/src/wannier90io/_schema.py b/src/wannier90io/_schema.py
index 1ea7e12..bb34f10 100644
--- a/src/wannier90io/_schema.py
+++ b/src/wannier90io/_schema.py
@@ -61,6 +61,7 @@ class NNKP(pydantic.BaseModel):
     comments: typing.List[str]
     parameters: dict
     blocks: dict
+    nnkpts: typing.List[typing.List[int]]
     direct_lattice: DirectLattice
     reciprocal_lattice: ReciprocalLattice
     kpoints: Kpoints