Extend VASP parsing of charges and magnetism

electronicparsers/vasp/parser.py

@@ -66,6 +66,7 @@
 )
 from runschema.calculation import (
     Calculation,
+    Charges,
     Energy,
     EnergyEntry,
     Forces,
@@ -399,6 +400,8 @@ def str_to_eigenvalues(val_in):
                 val.extend(['nan' if '*' in v else v for v in line.split()])
             return np.array(val, np.float64)
 
+        atomic_quantities = r'# of ion\s+s\s+p\s+d\s+tot\s*-+((?:\s*-?[\d\.]+)+)+'
+
         scf_iteration = [
             Quantity(
                 'energy_total',
@@ -509,6 +512,18 @@ def str_to_eigenvalues(val_in):
                 r'LOOP\+\: +cpu time +([\d\.]+)\: +real time +([\d\.]+)',
                 dtype=np.dtype(np.float64),
             ),
+            Quantity(
+                'atomic_charges',
+                r'total charge[?\s\S]+' + atomic_quantities,
+                str_operation=lambda x: x.strip().split(),
+                convert=True,
+            ),
+            Quantity(
+                'atomic_magnetisms',
+                r'magnetization \(x\)[?\s\S]+' + atomic_quantities,
+                str_operation=lambda x: x.strip().split(),
+                convert=True,
+            ),
         ]
 
         self._quantities = [
@@ -1716,6 +1731,7 @@ def parse_method(self):
                 break
 
         # Atom Parameters
+
         atomtypes = self.parser.atom_info.get('atomtypes', {})
         element = atomtypes.get('element', [])
         atom_counts = {e: 0 for e in element}
@@ -1866,7 +1882,7 @@ def parse_method(self):
             # correct based on core-holes
             # since ZVAL information is centrally reported, it's all or nothing
             try:
-                species_electrons = param.n_electrons
+                species_electrons = param.n_valence_electrons
             except AttributeError:
                 break
             try:
@@ -2231,6 +2247,19 @@ def parse_dos(n_calc):
             # dos
             parse_dos(n)
 
+            # atomic charges and magentizations
+            charge_groups = np.array(
+                self.parser.calculations[n].results['atomic_charges']
+            ).reshape(-1, 5)
+            total_charges = charge_groups[:, -1]
+            sec_scc.charges.append(
+                Charges(
+                    # analysis_method =
+                    n_atoms=len(total_charges),
+                    value=total_charges,
+                )
+            )
+
             # convergence
             converged = self.parser.is_converged(n)
             if converged is not None: