text: # coding: utf-8
  # Copyright (c) Pymatgen Development Team.
  # Distributed under the terms of the MIT License
  
  """
  Module for reading Lobster output files. For more information
  on LOBSTER see www.cohp.de.
  """
  
  import collections
  import fnmatch
  import os
  import re
  import warnings
  from collections import defaultdict
  from typing import Any, Dict, List, Optional
  
  import numpy as np
  from monty.io import zopen
  from pymatgen.core.structure import Structure
  from pymatgen.electronic_structure.bandstructure import LobsterBandStructureSymmLine
  from pymatgen.electronic_structure.core import Orbital, Spin
  from pymatgen.electronic_structure.dos import Dos, LobsterCompleteDos
  from pymatgen.io.vasp.inputs import Kpoints
  from pymatgen.io.vasp.outputs import Vasprun, VolumetricData
  
  __author__ = "Janine George, Marco Esters"
  __copyright__ = "Copyright 2017, The Materials Project"
  __version__ = "0.2"
  __maintainer__ = "Janine George, Marco Esters "
  __email__ = "janine.george@uclouvain.be, esters@uoregon.edu"
  __date__ = "Dec 13, 2017"
  
  MODULE_DIR = os.path.dirname(os.path.abspath(__file__))
  
  
  class Cohpcar:
      """
      Class to read COHPCAR/COOPCAR files generated by LOBSTER.
  
      .. attribute: cohp_data
  
           Dict that contains the COHP data of the form:
             {bond: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                     "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                     "length": bond length,
                     "sites": sites corresponding to the bond}
           Also contains an entry for the average, which does not have
           a "length" key.
  
      .. attribute: efermi
  
           The Fermi energy in eV.
  
      .. attribute: energies
  
           Sequence of energies in eV. Note that LOBSTER shifts the energies
           so that the Fermi energy is at zero.
  
      .. attribute: is_spin_polarized
  
           Boolean to indicate if the calculation is spin polarized.
  
      .. attribute: orb_res_cohp
  
          orb_cohp[label] = {bond_data["orb_label"]: {"COHP": {Spin.up: cohps, Spin.down:cohps},
                                                       "ICOHP": {Spin.up: icohps, Spin.down: icohps},
                                                       "orbitals": orbitals,
                                                       "length": bond lengths,
                                                       "sites": sites corresponding to the bond}}
  
      """
  
      def __init__(self, are_coops: bool = False, filename: str = None):
          """
          Args:
              are_coops: Determines if the file is a list of COHPs or COOPs.
                Default is False for COHPs.
              filename: Name of the COHPCAR file. If it is None, the default
                file name will be chosen, depending on the value of are_coops.
          """
          self.are_coops = are_coops
          if filename is None:
              filename = "COOPCAR.lobster" if are_coops else "COHPCAR.lobster"
  
          with zopen(filename, "rt") as f:
              contents = f.read().split("\n")
  
          # The parameters line is the second line in a COHPCAR file. It
          # contains all parameters that are needed to map the file.
          parameters = contents[1].split()
          # Subtract 1 to skip the average
          num_bonds = int(parameters[0]) - 1
          self.efermi = float(parameters[-1])
          if int(parameters[1]) == 2:
              spins = [Spin.up, Spin.down]
              self.is_spin_polarized = True
          else:
              spins = [Spin.up]
              self.is_spin_polarized = False
  
          # The COHP data start in row num_bonds + 3
          data = np.array([np.array(row.split(), dtype=float) for row in contents[num_bonds + 3 :]]).transpose()
          data = np.array([np.array(row.split(), dtype=float) for row in contents[num_bonds + 3 :]]).transpose()
          self.energies = data[0]
          cohp_data = {
              "average": {
                  "COHP": {spin: data[1 + 2 * s * (num_bonds + 1)] for s, spin in enumerate(spins)},
                  "ICOHP": {spin: data[2 + 2 * s * (num_bonds + 1)] for s, spin in enumerate(spins)},
              }
          }  # type: Dict[Any, Any]
  
          orb_cohp = {}  # type: Dict[str, Any]
          # present for Lobster versions older than Lobster 2.2.0
          veryold = False
          # the labeling had to be changed: there are more than one COHP for each atom combination
          # this is done to make the labeling consistent with ICOHPLIST.lobster
          bondnumber = 0
          for bond in range(num_bonds):
              bond_data = self._get_bond_data(contents[3 + bond])
  
              label = str(bondnumber)
  
              orbs = bond_data["orbitals"]
              cohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 3] for s, spin in enumerate(spins)}
  
              icohp = {spin: data[2 * (bond + s * (num_bonds + 1)) + 4] for s, spin in enumerate(spins)}
              if orbs is None:
                  bondnumber = bondnumber + 1
                  label = str(bondnumber)
                  cohp_data[label] = {
                      "COHP": cohp,
                      "ICOHP": icohp,
                      "length": bond_data["length"],
                      "sites": bond_data["sites"],
                  }
  
              elif label in orb_cohp:
                  orb_cohp[label].update(
                      {
                          bond_data["orb_label"]: {
                              "COHP": cohp,
                              "ICOHP": icohp,
                              "orbitals": orbs,
                              "length": bond_data["length"],
                              "sites": bond_data["sites"],
                          }
                      }
                  )
              else:
                  # present for Lobster versions older than Lobster 2.2.0
                  if bondnumber == 0:
                      veryold = True
                  if veryold:
                      bondnumber += 1
                      label = str(bondnumber)
  
                  orb_cohp[label] = {
                      bond_data["orb_label"]: {
                          "COHP": cohp,
                          "ICOHP": icohp,
                          "orbitals": orbs,
                          "length": bond_data["length"],
                          "sites": bond_data["sites"],
                      }
                  }
  
          # present for lobster older than 2.2.0
          if veryold:
              for bond_str in orb_cohp:
                  cohp_data[bond_str] = {
                      "COHP": None,
                      "ICOHP": None,
                      "length": bond_data["length"],
                      "sites": bond_data["sites"],
                  }
  
          self.orb_res_cohp = orb_cohp if orb_cohp else None
          self.cohp_data = cohp_data
  
      @staticmethod
      def _get_bond_data(line: str) -> dict:
          """
          Subroutine to extract bond label, site indices, and length from
          a LOBSTER header line. The site indices are zero-based, so they
          can be easily used with a Structure object.
  
          Example header line: No.4:Fe1->Fe9(2.4524893531900283)
          Example header line for orbtial-resolved COHP:
              No.1:Fe1[3p_x]->Fe2[3d_x^2-y^2](2.456180552772262)
  
          Args:
              line: line in the COHPCAR header describing the bond.
  
          Returns:
              Dict with the bond label, the bond length, a tuple of the site
              indices, a tuple containing the orbitals (if orbital-resolved),
              and a label for the orbitals (if orbital-resolved).
          """
  
          orb_labs = [
              "s",
              "p_y",
              "p_z",
              "p_x",
              "d_xy",
              "d_yz",
              "d_z^2",
              "d_xz",
              "d_x^2-y^2",
              "f_y(3x^2-y^2)",
              "f_xyz",
              "f_yz^2",
              "f_z^3",
              "f_xz^2",
              "f_z(x^2-y^2)",
              "f_x(x^2-3y^2)",
          ]
  
          line_new = line.rsplit("(", 1)
          # bondnumber = line[0].replace("->", ":").replace(".", ":").split(':')[1]
          length = float(line_new[-1][:-1])
  
          sites = line_new[0].replace("->", ":").split(":")[1:3]
          site_indices = tuple(int(re.split(r"\D+", site)[1]) - 1 for site in sites)
  
          # species = tuple(re.split(r"\d+", site)[0] for site in sites)
          if "[" in sites[0]:
              orbs = [re.findall(r"\[(.*)\]", site)[0] for site in sites]
              orbitals = [tuple((int(orb[0]), Orbital(orb_labs.index(orb[1:])))) for orb in orbs]  # type: Any
              orb_label = "%d%s-%d%s" % (
                  orbitals[0][0],
                  orbitals[0][1].name,
                  orbitals[1][0],
                  orbitals[1][1].name,
              )  # type: Any
  
          else:
              orbitals = None
              orb_label = None
  
          # a label based on the species alone is not feasible, there can be more than one bond for each atom combination
          # label = "%s" % (bondnumber)
  
          bond_data = {
              "length": length,
              "sites": site_indices,
              "orbitals": orbitals,
              "orb_label": orb_label,
          }
          return bond_data
  
  
  class Icohplist:
      """
      Class to read ICOHPLIST/ICOOPLIST files generated by LOBSTER.
  
      .. attribute: are_coops
           Boolean to indicate if the populations are COOPs or COHPs.
  
      .. attribute: is_spin_polarized
           Boolean to indicate if the calculation is spin polarized.
  
      .. attribute: Icohplist
          Dict containing the listfile data of the form:
             {bond: "length": bond length,
                    "number_of_bonds": number of bonds
                    "icohp": {Spin.up: ICOHP(Ef) spin up, Spin.down: ...}}
  
      .. attribute: IcohpCollection
          IcohpCollection Object
  
      """
  
      def __init__(self, are_coops: bool = False, filename: str = None):
          """
          Args:
              are_coops: Determines if the file is a list of ICOHPs or ICOOPs.
                Defaults to False for ICOHPs.
              filename: Name of the ICOHPLIST file. If it is None, the default
                file name will be chosen, depending on the value of are_coops.
          """
          self.are_coops = are_coops
          if filename is None:
              filename = "ICOOPLIST.lobster" if are_coops else "ICOHPLIST.lobster"
  
          # LOBSTER list files have an extra trailing blank line
          # and we don't need the header.
          with zopen(filename, "rt") as f:
              data = f.read().split("\n")[1:-1]
          if len(data) == 0:
              raise IOError("ICOHPLIST file contains no data.")
  
          # Which Lobster version?
          if len(data[0].split()) == 8:
              version = "3.1.1"
          elif len(data[0].split()) == 6:
              version = "2.2.1"
              warnings.warn("Please consider using the new Lobster version. See www.cohp.de.")
          else:
              raise ValueError
  
          # If the calculation is spin polarized, the line in the middle
          # of the file will be another header line.
          if "distance" in data[len(data) // 2]:
              num_bonds = len(data) // 2
              if num_bonds == 0:
                  raise IOError("ICOHPLIST file contains no data.")
              self.is_spin_polarized = True
          else:
              num_bonds = len(data)
              self.is_spin_polarized = False
  
          list_labels = []
          list_atom1 = []
          list_atom2 = []
          list_length = []
          list_translation = []
          list_num = []
          list_icohp = []
          for bond in range(num_bonds):
              line = data[bond].split()
              icohp = {}
              if version == "2.2.1":
                  label = "%s" % (line[0])
                  atom1 = str(line[1])
                  atom2 = str(line[2])
                  length = float(line[3])
                  icohp[Spin.up] = float(line[4])
                  num = int(line[5])
                  translation = [0, 0, 0]
                  if self.is_spin_polarized:
                      icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[4])
  
              elif version == "3.1.1":
                  label = "%s" % (line[0])
                  atom1 = str(line[1])
                  atom2 = str(line[2])
                  length = float(line[3])
                  translation = [int(line[4]), int(line[5]), int(line[6])]
                  icohp[Spin.up] = float(line[7])
                  num = int(1)
  
                  if self.is_spin_polarized:
                      icohp[Spin.down] = float(data[bond + num_bonds + 1].split()[7])
  
              list_labels.append(label)
              list_atom1.append(atom1)
              list_atom2.append(atom2)
              list_length.append(length)
              list_translation.append(translation)
              list_num.append(num)
              list_icohp.append(icohp)
  
          # to avoid circular dependencies
          from pymatgen.electronic_structure.cohp import IcohpCollection
  
          self._icohpcollection = IcohpCollection(
              are_coops=are_coops,
              list_labels=list_labels,
              list_atom1=list_atom1,
              list_atom2=list_atom2,
              list_length=list_length,
              list_translation=list_translation,
              list_num=list_num,
              list_icohp=list_icohp,
              is_spin_polarized=self.is_spin_polarized,
          )
  
      @property
      def icohplist(self) -> Dict[Any, Dict[str, Any]]:
          """
          Returns: icohplist compatible with older version of this class
          """
          icohplist_new = {}
          for key, value in self._icohpcollection._icohplist.items():
              icohplist_new[key] = {
                  "length": value._length,
                  "number_of_bonds": value._num,
                  "icohp": value._icohp,
                  "translation": value._translation,
              }
          return icohplist_new
  
      @property
      def icohpcollection(self):
          """
          Returns: IcohpCollection object
          """
          return self._icohpcollection
  
  
  class Doscar:
      """
      Class to deal with Lobster's projected DOS and local projected DOS.
      The beforehand quantum-chemical calculation was performed with VASP
  
      .. attribute:: completedos
  
          LobsterCompleteDos Object
  
      .. attribute:: pdos
          List of Dict including numpy arrays with pdos. Access as pdos[atomindex]['orbitalstring']['Spin.up/Spin.down']
  
      .. attribute:: tdos
          Dos Object of the total density of states
  
      .. attribute:: energies
          numpy array of the energies at which the DOS was calculated (in eV, relative to Efermi)
  
      .. attribute:: tdensities
          tdensities[Spin.up]: numpy array of the total density of states for the Spin.up contribution at each of the
              energies
          tdensities[Spin.down]: numpy array of the total density of states for the Spin.down contribution at each of the
              energies
  
          if is_spin_polarized=False:
          tdensities[Spin.up]: numpy array of the total density of states
  
  
      .. attribute:: itdensities:
          itdensities[Spin.up]: numpy array of the total density of states for the Spin.up contribution at each of the
              energies
          itdensities[Spin.down]: numpy array of the total density of states for the Spin.down contribution at each of the
              energies
  
          if is_spin_polarized=False:
          itdensities[Spin.up]: numpy array of the total density of states
  
  
      .. attribute:: is_spin_polarized
          Boolean. Tells if the system is spin polarized
  
  
      """
  
      def __init__(
          self,
          doscar: str = "DOSCAR.lobster",
          structure_file: str = "POSCAR",
          dftprogram: str = "Vasp",
      ):
          """
          Args:
              doscar: DOSCAR filename, typically "DOSCAR.lobster"
              structure_file: for vasp, this is typically "POSCAR"
              dftprogram: so far only "vasp" is implemented
          """
          self._doscar = doscar
          if dftprogram == "Vasp":
              self._final_structure = Structure.from_file(structure_file)
  
          self._parse_doscar()
  
      def _parse_doscar(self):
          doscar = self._doscar
  
          tdensities = {}
          itdensities = {}
          f = open(doscar)
          natoms = int(f.readline().split()[0])
          efermi = float([f.readline() for nn in range(4)][3].split()[17])
          dos = []
          orbitals = []
          for atom in range(natoms + 1):
              line = f.readline()
              ndos = int(line.split()[2])
              orbitals.append(line.split(";")[-1].split())
              line = f.readline().split()
              cdos = np.zeros((ndos, len(line)))
              cdos[0] = np.array(line)
              for nd in range(1, ndos):
                  line = f.readline().split()
                  cdos[nd] = np.array(line)
              dos.append(cdos)
          f.close()
          doshere = np.array(dos[0])
          if len(doshere[0, :]) == 5:
              self._is_spin_polarized = True
          elif len(doshere[0, :]) == 3:
              self._is_spin_polarized = False
          else:
              raise ValueError("There is something wrong with the DOSCAR. Can't extract spin polarization.")
          energies = doshere[:, 0]
          if not self._is_spin_polarized:
              tdensities[Spin.up] = doshere[:, 1]
              itdensities[Spin.up] = doshere[:, 2]
              pdoss = []
              spin = Spin.up
              for atom in range(natoms):
                  pdos = defaultdict(dict)
                  data = dos[atom + 1]
                  _, ncol = data.shape
                  orbnumber = 0
                  for j in range(1, ncol):
                      orb = orbitals[atom + 1][orbnumber]
                      pdos[orb][spin] = data[:, j]
                      orbnumber = orbnumber + 1
                  pdoss.append(pdos)
          else:
              tdensities[Spin.up] = doshere[:, 1]
              tdensities[Spin.down] = doshere[:, 2]
              itdensities[Spin.up] = doshere[:, 3]
              itdensities[Spin.down] = doshere[:, 4]
              pdoss = []
              for atom in range(natoms):
                  pdos = defaultdict(dict)
                  data = dos[atom + 1]
                  _, ncol = data.shape
                  orbnumber = 0
                  for j in range(1, ncol):
                      if j % 2 == 0:
                          spin = Spin.down
                      else:
                          spin = Spin.up
                      orb = orbitals[atom + 1][orbnumber]
                      pdos[orb][spin] = data[:, j]
                      if j % 2 == 0:
                          orbnumber = orbnumber + 1
                  pdoss.append(pdos)
  
          self._efermi = efermi
          self._pdos = pdoss
          self._tdos = Dos(efermi, energies, tdensities)
          self._energies = energies
          self._tdensities = tdensities
          self._itdensities = itdensities
          final_struct = self._final_structure
  
          pdossneu = {final_struct[i]: pdos for i, pdos in enumerate(self._pdos)}
  
          self._completedos = LobsterCompleteDos(final_struct, self._tdos, pdossneu)
  
      @property
      def completedos(self) -> LobsterCompleteDos:
          """
          :return: CompleteDos
          """
          return self._completedos
  
      @property
      def pdos(self) -> list:
          """
          :return: Projected DOS
          """
          return self._pdos
  
      @property
      def tdos(self) -> Dos:
          """
          :return: Total DOS
          """
          return self._tdos
  
      @property
      def energies(self) -> np.ndarray:
          """
          :return: Energies
          """
          return self._energies
  
      @property
      def tdensities(self) -> np.ndarray:
          """
          :return: total densities as a np.ndarray
          """
          return self._tdensities
  
      @property
      def itdensities(self) -> np.ndarray:
          """
          :return: integrated total densities as a np.ndarray
          """
          return self._itdensities
  
      @property
      def is_spin_polarized(self) -> bool:
          """
          :return: Whether run is spin polarized.
          """
          return self._is_spin_polarized
  
  
  class Charge:
      """
      Class to read CHARGE files generated by LOBSTER
  
      .. attribute: atomlist
          List of atoms in CHARGE.lobster
      .. attribute: types
          List of types of atoms in CHARGE.lobster
      .. attribute: Mulliken
          List of Mulliken charges of atoms in CHARGE.lobster
      .. attribute: Loewdin
          List of Loewdin charges of atoms in CHARGE.Loewdin
      .. attribute: num_atoms
          Number of atoms in CHARGE.lobster
  
      """
  
      def __init__(self, filename: str = "CHARGE.lobster"):
          """
          Args:
              filename: filename for the CHARGE file, typically "CHARGE.lobster"
          """
          with zopen(filename, "rt") as f:
              data = f.read().split("\n")[3:-3]
          if len(data) == 0:
              raise IOError("CHARGES file contains no data.")
  
          self.num_atoms = len(data)
          self.atomlist = []  # type: List[str]
          self.types = []  # type: List[str]
          self.Mulliken = []  # type: List[float]
          self.Loewdin = []  # type: List[float]
          for atom in range(0, self.num_atoms):
              line = data[atom].split()
              self.atomlist.append(line[1] + line[0])
              self.types.append(line[1])
              self.Mulliken.append(float(line[2]))
              self.Loewdin.append(float(line[3]))
  
      def get_structure_with_charges(self, structure_filename):
          """
          get a Structure with Mulliken and Loewdin charges as site properties
          Args:
              structure_filename: filename of POSCAR
          Returns:
              Structure Object with Mulliken and Loewdin charges as site properties
          """
  
          struct = Structure.from_file(structure_filename)
          Mulliken = self.Mulliken
          Loewdin = self.Loewdin
          site_properties = {"Mulliken Charges": Mulliken, "Loewdin Charges": Loewdin}
          new_struct = struct.copy(site_properties=site_properties)
          return new_struct
  
  
  class Lobsterout:
      """
      Class to read in the lobsterout and evaluate the spilling, save the basis, save warnings, save infos
  
        .. attribute: basis_functions
          list of basis functions that were used in lobster run as strings
  
        .. attribute: basis_type
           list of basis type that were used in lobster run as strings
  
        .. attribute: chargespilling
           list of charge spilling (first entry: result for spin 1, second entry: result for spin 2 or not present)
  
        .. attribute: dftprogram
           string representing the dft program used for the calculation of the wave function
  
        .. attribute: elements
           list of strings of elements that were present in lobster calculation
  
        .. attribute: has_CHARGE
           Boolean, indicates that CHARGE.lobster is present
  
        .. attribute: has_COHPCAR
          Boolean, indicates that COHPCAR.lobster and ICOHPLIST.lobster are present
  
        .. attribute: has_COOPCAR
          Boolean, indicates that COOPCAR.lobster and ICOOPLIST.lobster are present
  
        .. attribute: has_DOSCAR
          Boolean, indicates that DOSCAR.lobster is present
  
        .. attribute: has_Projection
          Boolean, indcates that projectionData.lobster is present
  
        .. attribute: has_bandoverlaps
          Boolean, indcates that bandOverlaps.lobster is present
  
        .. attribute: has_density_of_energies
          Boolean, indicates that DensityOfEnergy.lobster is present
  
        .. attribute: has_fatbands
          Boolean, indicates that fatband calculation was performed
  
        .. attribute: has_grosspopulation
          Boolean, indicates that GROSSPOP.lobster is present
  
        .. attribute: info_lines
          string with additional infos on the run
  
        .. attribute: info_orthonormalization
          string with infos on orthonormalization
  
        .. attribute: is_restart_from_projection
          Boolean that indicates that calculation was restartet from existing projection file
  
        .. attribute: lobster_version
          string that indicates Lobster version
  
        .. attribute: number_of_spins
          Integer indicating the number of spins
  
        .. attribute: number_of_threads
         integer that indicates how many threads were used
  
        .. attribute: timing
          dict with infos on timing
  
        .. attribute: totalspilling
          list of values indicating the total spilling for spin channel 1 (and spin channel 2)
  
        .. attribute: warninglines
          string with all warnings
  
  
      """
  
      def __init__(self, filename="lobsterout"):
          """
          Args:
              filename: filename of lobsterout
          """
          warnings.warn("Make sure the lobsterout is read in correctly. This is a brand new class.")
          # read in file
          with zopen(filename, "rt") as f:
              data = f.read().split("\n")  # [3:-3]
          if len(data) == 0:
              raise IOError("lobsterout does not contain any data")
  
          # check if Lobster starts from a projection
          self.is_restart_from_projection = "loading projection from projectionData.lobster..." in data
  
          self.lobster_version = self._get_lobster_version(data=data)
  
          self.number_of_threads = int(self._get_threads(data=data))
          self.dftprogram = self._get_dft_program(data=data)
  
          self.number_of_spins = self._get_number_of_spins(data=data)
          chargespilling, totalspilling = self._get_spillings(data=data, number_of_spins=self.number_of_spins)
          self.chargespilling = chargespilling
          self.totalspilling = totalspilling
  
          (
              elements,
              basistype,
              basisfunctions,
          ) = self._get_elements_basistype_basisfunctions(data=data)
          self.elements = elements
          self.basis_type = basistype
          self.basis_functions = basisfunctions
  
          wall_time, user_time, sys_time = self._get_timing(data=data)
          timing = {}
          timing["walltime"] = wall_time
          timing["usertime"] = user_time
          timing["sys_time"] = sys_time
          self.timing = timing
  
          warninglines = self._get_all_warning_lines(data=data)
          self.warninglines = warninglines
  
          orthowarning = self._get_warning_orthonormalization(data=data)
          self.info_orthonormalization = orthowarning
  
          infos = self._get_all_info_lines(data=data)
          self.info_lines = infos
  
          self.has_DOSCAR = "writing DOSCAR.lobster..." in data and "SKIPPING writing DOSCAR.lobster..." not in data
          self.has_COHPCAR = (
              "writing COOPCAR.lobster and ICOOPLIST.lobster..." in data
              and "SKIPPING writing COOPCAR.lobster and ICOOPLIST.lobster..." not in data
          )
          self.has_COOPCAR = (
              "writing COHPCAR.lobster and ICOHPLIST.lobster..." in data
              and "SKIPPING writing COHPCAR.lobster and ICOHPLIST.lobster..." not in data
          )
          self.has_CHARGE = "SKIPPING writing CHARGE.lobster..." not in data
          self.has_Projection = "saving projection to projectionData.lobster..." in data
          self.has_bandoverlaps = "WARNING: I dumped the band overlap matrices to the file bandOverlaps.lobster." in data
          self.has_fatbands = self._has_fatband(data=data)
          self.has_grosspopulation = "writing CHARGE.lobster and GROSSPOP.lobster..." in data
          self.has_density_of_energies = "writing DensityOfEnergy.lobster..." in data
  
      def get_doc(self):
          """
  
          Returns: LobsterDict with all the information stored in lobsterout
          """
          LobsterDict = {}
          # check if Lobster starts from a projection
          LobsterDict["restart_from_projection"] = self.is_restart_from_projection
          LobsterDict["lobster_version"] = self.lobster_version
          LobsterDict["threads"] = self.number_of_threads
          LobsterDict["Dftprogram"] = self.dftprogram
  
          LobsterDict["chargespilling"] = self.chargespilling
          LobsterDict["totalspilling"] = self.totalspilling
  
          LobsterDict["elements"] = self.elements
          LobsterDict["basistype"] = self.basis_type
          LobsterDict["basisfunctions"] = self.basis_functions
  
          LobsterDict["timing"] = self.timing
  
          LobsterDict["warnings"] = self.warninglines
  
          LobsterDict["orthonormalization"] = self.info_orthonormalization
  
          LobsterDict["infos"] = self.info_lines
  
          LobsterDict["hasDOSCAR"] = self.has_DOSCAR
          LobsterDict["hasCOHPCAR"] = self.has_COHPCAR
          LobsterDict["hasCOOPCAR"] = self.has_COOPCAR
          LobsterDict["hasCHARGE"] = self.has_CHARGE
          LobsterDict["hasProjection"] = self.has_Projection
          LobsterDict["hasbandoverlaps"] = self.has_bandoverlaps
          LobsterDict["hasfatband"] = self.has_fatbands
          LobsterDict["hasGrossPopuliation"] = self.has_grosspopulation
          LobsterDict["hasDensityOfEnergies"] = self.has_density_of_energies
  
          return LobsterDict
  
      @staticmethod
      def _get_lobster_version(data):
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 1:
                  if splitrow[0] == "LOBSTER":
                      return splitrow[1]
          raise RuntimeError("Version not found.")
  
      @staticmethod
      def _has_fatband(data):
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 1:
                  if splitrow[1] == "FatBand":
                      return True
          return False
  
      @staticmethod
      def _get_dft_program(data):
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 4:
                  if splitrow[3] == "program...":
                      return splitrow[4]
          return None
  
      @staticmethod
      def _get_number_of_spins(data):
          if "spillings for spin channel 2" in data:
              return 2
          return 1
  
      @staticmethod
      def _get_threads(data):
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 11:
                  if (splitrow[11]) == "threads" or (splitrow[11] == "thread"):
                      return splitrow[10]
          raise ValueError("Threads not found.")
  
      @staticmethod
      def _get_spillings(data, number_of_spins):
          charge_spilling = []
          total_spilling = []
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 2:
                  if splitrow[2] == "spilling:":
                      if splitrow[1] == "charge":
                          charge_spilling.append(np.float_(splitrow[3].replace("%", "")) / 100.0)
                      if splitrow[1] == "total":
                          total_spilling.append(np.float_(splitrow[3].replace("%", "")) / 100.0)
  
              if len(charge_spilling) == number_of_spins and len(total_spilling) == number_of_spins:
                  break
  
          return charge_spilling, total_spilling
  
      @staticmethod
      def _get_elements_basistype_basisfunctions(data):
          begin = False
          end = False
          elements = []
          basistype = []
          basisfunctions = []
          for row in data:
  
              if begin and not end:
                  splitrow = row.split()
                  if splitrow[0] not in [
                      "INFO:",
                      "WARNING:",
                      "setting",
                      "calculating",
                      "post-processing",
                      "saving",
                      "spillings",
                      "writing",
                  ]:
  
                      elements.append(splitrow[0])
                      basistype.append(splitrow[1].replace("(", "").replace(")", ""))
                      # last sign is a ''
                      basisfunctions.append(splitrow[2:])
                  else:
                      end = True
              if "setting up local basis functions..." in row:
                  begin = True
          return elements, basistype, basisfunctions
  
      @staticmethod
      def _get_timing(data):
          # will give back wall, user and sys time
          begin = False
          # end=False
          # time=[]
  
          for row in data:
              splitrow = row.split()
              if "finished" in splitrow:
                  begin = True
              if begin:
                  if "wall" in splitrow:
                      wall_time = splitrow[2:10]
                  if "user" in splitrow:
                      user_time = splitrow[0:8]
                  if "sys" in splitrow:
                      sys_time = splitrow[0:8]
  
          wall_time_dict = {
              "h": wall_time[0],
              "min": wall_time[2],
              "s": wall_time[4],
              "ms": wall_time[6],
          }
          user_time_dict = {
              "h": user_time[0],
              "min": user_time[2],
              "s": user_time[4],
              "ms": user_time[6],
          }
          sys_time_dict = {
              "h": sys_time[0],
              "min": sys_time[2],
              "s": sys_time[4],
              "ms": sys_time[6],
          }
  
          return wall_time_dict, user_time_dict, sys_time_dict
  
      @staticmethod
      def _get_warning_orthonormalization(data):
          orthowarning = []
          for row in data:
              splitrow = row.split()
              if "orthonormalized" in splitrow:
                  orthowarning.append(" ".join(splitrow[1:]))
          return orthowarning
  
      @staticmethod
      def _get_all_warning_lines(data):
          ws = []
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 0:
                  if splitrow[0] == "WARNING:":
                      ws.append(" ".join(splitrow[1:]))
          return ws
  
      @staticmethod
      def _get_all_info_lines(data):
          infos = []
          for row in data:
              splitrow = row.split()
              if len(splitrow) > 0:
                  if splitrow[0] == "INFO:":
                      infos.append(" ".join(splitrow[1:]))
          return infos
  
  
  class Fatband:
      """
      Reads in FATBAND_x_y.lobster files
  
      .. attribute: efermi
  
          efermi that was read in from vasprun.xml
  
      .. attribute: eigenvals
          {Spin.up:[][],Spin.down:[][]}, the first index of the array
              [][] refers to the band and the second to the index of the
              kpoint. The kpoints are ordered according to the order of the
              kpoints array. If the band structure is not spin polarized, we
              only store one data set under Spin.up.
  
      .. attribute: is_spinpolarized
  
          Boolean that tells you whether this was a spin-polarized calculation
  
      .. attribute: kpoints_array
  
          list of kpoint as numpy arrays, in frac_coords of the given lattice by default
  
      .. attribute: label_dict
  
           (dict) of {} this link a kpoint (in frac coords or cartesian coordinates depending on the coords).
  
      .. attribute: lattice
  
          lattice object of reciprocal lattice as read in from vasprun.xml
  
      .. attribute: nbands
  
          number of bands used in the calculation
  
      .. attribute: p_eigenvals
  
          dict of orbital projections as {spin: array of dict}.
          The indices of the array are [band_index, kpoint_index].
          The dict is then built the following way:
          {"string of element": "string of orbital as read in from FATBAND file"}
          If the band structure is not spin polarized, we only store one data set under Spin.up.
  
      .. attribute: structure
  
          structure read in from vasprun.xml
      """
  
      def __init__(self, filenames=".", vasprun="vasprun.xml", Kpointsfile="KPOINTS"):
          """
          Args:
              filenames (list or string): can be a list of file names or a path to a folder folder from which all
                  "FATBAND_*" files will be read
              vasprun: corresponding vasprun file
              Kpointsfile: KPOINTS file for bandstructure calculation, typically "KPOINTS"
          """
          warnings.warn("Make sure all relevant FATBAND files were generated and read in!")
          warnings.warn("Use Lobster 3.2.0 or newer for fatband calculations!")
  
          VASPRUN = Vasprun(
              filename=vasprun,
              ionic_step_skip=None,
              ionic_step_offset=0,
              parse_dos=True,
              parse_eigen=False,
              parse_projected_eigen=False,
              parse_potcar_file=False,
              occu_tol=1e-8,
              exception_on_bad_xml=True,
          )
          self.structure = VASPRUN.final_structure
          self.lattice = self.structure.lattice.reciprocal_lattice
          self.efermi = VASPRUN.efermi
          kpoints_object = Kpoints.from_file(Kpointsfile)
  
          atomtype = []
          atomnames = []
          orbital_names = []
  
          if not isinstance(filenames, list) or filenames is None:
              filenames_new = []
              if filenames is None:
                  filenames = "."
              for file in os.listdir(filenames):
                  if fnmatch.fnmatch(file, "FATBAND_*.lobster"):
                      filenames_new.append(os.path.join(filenames, file))
              filenames = filenames_new
          if len(filenames) == 0:
              raise ValueError("No FATBAND files in folder or given")
          for ifilename, filename in enumerate(filenames):
              with zopen(filename, "rt") as f:
                  contents = f.read().split("\n")
  
              atomnames.append(os.path.split(filename)[1].split("_")[1].capitalize())
              parameters = contents[0].split()
              atomtype.append(re.split(r"[0-9]+", parameters[3])[0].capitalize())
              orbital_names.append(parameters[4])
  
          # get atomtype orbital dict
          atom_orbital_dict = {}
          for iatom, atom in enumerate(atomnames):
              if atom not in atom_orbital_dict:
                  atom_orbital_dict[atom] = []
              atom_orbital_dict[atom].append(orbital_names[iatom])
          # test if there are the same orbitals twice or if two different formats were used or if all necessary orbitals
          # are there
          for key, items in atom_orbital_dict.items():
              if len(set(items)) != len(items):
                  raise ValueError("The are two FATBAND files for the same atom and orbital. The program will stop.")
              split = []
              for item in items:
                  split.append(item.split("_")[0])
              for orb, number in collections.Counter(split).items():
                  if number not in (1, 3, 5, 7):
                      raise ValueError(
                          "Make sure all relevant orbitals were generated and that no duplicates (2p and 2p_x) are "
                          "present"
                      )
  
          kpoints_array = []
          for ifilename, filename in enumerate(filenames):
              with zopen(filename, "rt") as f:
                  contents = f.read().split("\n")
  
              if ifilename == 0:
                  self.nbands = int(parameters[6])
                  self.number_kpts = kpoints_object.num_kpts - int(contents[1].split()[2]) + 1
  
              if len(contents[1:]) == self.nbands + 2:
                  self.is_spinpolarized = False
              elif len(contents[1:]) == self.nbands * 2 + 2:
                  self.is_spinpolarized = True
              else:
                  linenumbers = []
                  for iline, line in enumerate(contents[1 : self.nbands * 2 + 4]):
                      if line.split()[0] == "#":
                          linenumbers.append(iline)
  
                  if ifilename == 0:
                      if len(linenumbers) == 2:
                          self.is_spinpolarized = True
                      else:
                          self.is_spinpolarized = False
  
              if ifilename == 0:
                  eigenvals = {}
                  eigenvals[Spin.up] = [
                      [collections.defaultdict(float) for i in range(self.number_kpts)] for j in range(self.nbands)
                  ]
                  if self.is_spinpolarized:
                      eigenvals[Spin.down] = [
                          [collections.defaultdict(float) for i in range(self.number_kpts)] for j in range(self.nbands)
                      ]
  
                  p_eigenvals = {}
                  p_eigenvals[Spin.up] = [
                      [
                          {
                              str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                              for e in atomnames
                          }
                          for i in range(self.number_kpts)
                      ]
                      for j in range(self.nbands)
                  ]
  
                  if self.is_spinpolarized:
                      p_eigenvals[Spin.down] = [
                          [
                              {
                                  str(e): {str(orb): collections.defaultdict(float) for orb in atom_orbital_dict[e]}
                                  for e in atomnames
                              }
                              for i in range(self.number_kpts)
                          ]
                          for j in range(self.nbands)
                      ]
  
              ikpoint = -1
              for iline, line in enumerate(contents[1:-1]):
                  if line.split()[0] == "#":
                      KPOINT = np.array(
                          [
                              float(line.split()[4]),
                              float(line.split()[5]),
                              float(line.split()[6]),
                          ]
                      )
                      if ifilename == 0:
                          kpoints_array.append(KPOINT)
  
                      linenumber = 0
                      iband = 0
                      ikpoint += 1
                  if linenumber == self.nbands:
                      iband = 0
                  if line.split()[0] != "#":
  
                      if linenumber < self.nbands:
                          if ifilename == 0:
                              eigenvals[Spin.up][iband][ikpoint] = float(line.split()[1]) + self.efermi
  
                          p_eigenvals[Spin.up][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
                              line.split()[2]
                          )
                      if linenumber >= self.nbands and self.is_spinpolarized:
                          if ifilename == 0:
                              eigenvals[Spin.down][iband][ikpoint] = float(line.split()[1]) + self.efermi
                          p_eigenvals[Spin.down][iband][ikpoint][atomnames[ifilename]][orbital_names[ifilename]] = float(
                              line.split()[2]
                          )
  
                      linenumber += 1
                      iband += 1
  
          self.kpoints_array = kpoints_array
          self.eigenvals = eigenvals
          self.p_eigenvals = p_eigenvals
  
          label_dict = {}
          for ilabel, label in enumerate(kpoints_object.labels[-self.number_kpts :], start=0):
  
              if label is not None:
                  label_dict[label] = kpoints_array[ilabel]
  
          self.label_dict = label_dict
  
      def get_bandstructure(self):
          """
          returns a LobsterBandStructureSymmLine object which can be plotted with a normal BSPlotter
          """
  
          return LobsterBandStructureSymmLine(
              kpoints=self.kpoints_array,
              eigenvals=self.eigenvals,
              lattice=self.lattice,
              efermi=self.efermi,
              labels_dict=self.label_dict,
              structure=self.structure,
              projections=self.p_eigenvals,
          )
  
  
  class Bandoverlaps:
      """
      Class to read in bandOverlaps.lobster files. These files are not created during every Lobster run.
      .. attribute: bandoverlapsdict is a dict of the following form:
          {spin:{"kpoint as string": {"maxDeviation": float that describes the max deviation, "matrix": 2D
          array of the size number of bands times number of bands including the overlap matrices with } }}
  
      .. attribute: maxDeviation is a list of floats describing the maximal Deviation for each problematic kpoint
  
      """
  
      def __init__(self, filename: str = "bandOverlaps.lobster"):
          """
          Args:
              filename: filename of the "bandOverlaps.lobster" file
          """
          with zopen(filename, "rt") as f:
              contents = f.read().split("\n")
  
          self._read(contents)
  
      def _read(self, contents: list):
          """
          will read in all contents of the file
          Args:
           contents: list of strings
          """
          self.bandoverlapsdict = {}  # type: Dict
          self.max_deviation = []  # type: List
          # This has to be done like this because there can be different numbers of problematic k-points per spin
          for line in contents:
              if "Overlap Matrix (abs) of the orthonormalized projected bands for spin 0" in line:
                  spin = Spin.up
              elif "Overlap Matrix (abs) of the orthonormalized projected bands for spin 1" in line:
                  spin = Spin.down
              elif "k-point" in line:
                  kpoint = line.split(" ")
                  kpoint_array = []
                  for kpointel in kpoint:
                      if kpointel not in ["at", "k-point", ""]:
                          kpoint_array.append(str(kpointel))
  
              elif "maxDeviation" in line:
                  if spin not in self.bandoverlapsdict:
                      self.bandoverlapsdict[spin] = {}
                  if not " ".join(kpoint_array) in self.bandoverlapsdict[spin]:
                      self.bandoverlapsdict[spin][" ".join(kpoint_array)] = {}
                  maxdev = line.split(" ")[2]
                  self.bandoverlapsdict[spin][" ".join(kpoint_array)]["maxDeviation"] = float(maxdev)
                  self.max_deviation.append(float(maxdev))
                  self.bandoverlapsdict[spin][" ".join(kpoint_array)]["matrix"] = []
  
              else:
                  overlaps = []
                  for el in line.split(" "):
                      if el not in [""]:
                          overlaps.append(float(el))
                  self.bandoverlapsdict[spin][" ".join(kpoint_array)]["matrix"].append(overlaps)
  
      def has_good_quality_maxDeviation(self, limit_maxDeviation: float = 0.1) -> bool:
          """
          will check if the maxDeviation from the ideal bandoverlap is smaller or equal to limit_maxDeviation
          Args:
           limit_maxDeviation: limit of the maxDeviation
          Returns:
               Boolean that will give you information about the quality of the projection
          """
  
          for deviation in self.max_deviation:
              if deviation > limit_maxDeviation:
                  return False
          return True
  
      def has_good_quality_check_occupied_bands(
          self,
          number_occ_bands_spin_up: int,
          number_occ_bands_spin_down: Optional[int] = None,
          spin_polarized: bool = False,
          limit_deviation: float = 0.1,
      ) -> bool:
          """
          will check if the deviation from the ideal bandoverlap of all occupied bands is smaller or equal to
          limit_deviation
  
          Args:
          number_occ_bands_spin_up (int): number of occupied bands of spin up
          number_occ_bands_spin_down (int): number of occupied bands of spin down
          spin_polarized (bool):  If True, then it was a spin polarized calculation
          limit_deviation (float): limit of the maxDeviation
          Returns:
               Boolean that will give you information about the quality of the projection
          """
  
          for matrix in self.bandoverlapsdict[Spin.up].values():
              for iband1, band1 in enumerate(matrix["matrix"]):
                  for iband2, band2 in enumerate(band1):
                      if iband1 < number_occ_bands_spin_up and iband2 < number_occ_bands_spin_up:
                          if iband1 == iband2:
                              if abs(band2 - 1.0) > limit_deviation:
                                  return False
                          else:
                              if band2 > limit_deviation:
                                  return False
  
          if spin_polarized:
              for matrix in self.bandoverlapsdict[Spin.down].values():
                  for iband1, band1 in enumerate(matrix["matrix"]):
                      for iband2, band2 in enumerate(band1):
                          if number_occ_bands_spin_down is not None:
                              if iband1 < number_occ_bands_spin_down and iband2 < number_occ_bands_spin_down:
                                  if iband1 == iband2:
                                      if abs(band2 - 1.0) > limit_deviation:
                                          return False
                                  else:
                                      if band2 > limit_deviation:
                                          return False
                          else:
                              ValueError("number_occ_bands_spin_down has to be specified")
          return True
  
  
  class Grosspop:
      """
      Class to read in GROSSPOP.lobster files.
  
      .. attribute: list_dict_grosspop
         which is a list of dicts including all information about the grosspopulations, one sample dict looks like this:
          {'element': 'O', 'Mulliken GP': {'2s': '1.80', '2p_y': '1.83', '2p_z': '1.79', '2p_x': '1.75', 'total': '7.18'},
           'Loewdin GP': {'2s': '1.60', '2p_y': '1.82', '2p_z': '1.77', '2p_x': '1.73', 'total': '6.92'}}
          The 0. entry of the list refers to the first atom in GROSSPOP.lobster and so on.
      """
  
      def __init__(self, filename: str = "GROSSPOP.lobster"):
          """
          Args:
              filename: filename of the "GROSSPOP.lobster" file
          """
          # opens file
          with zopen(filename, "rt") as f:
              contents = f.read().split("\n")
  
          self.list_dict_grosspop = []  # type: List[Any]
          # transfers content of file to list of dict
          for line in contents[3:]:
              cleanline = [i for i in line.split(" ") if not i == ""]
              if len(cleanline) == 5:
                  smalldict = {}
                  smalldict["element"] = cleanline[1]
                  smalldict["Mulliken GP"] = {}
                  smalldict["Loewdin GP"] = {}
                  smalldict["Mulliken GP"][cleanline[2]] = float(cleanline[3])
                  smalldict["Loewdin GP"][cleanline[2]] = float(cleanline[4])
              elif len(cleanline) > 0:
                  smalldict["Mulliken GP"][cleanline[0]] = float(cleanline[1])
                  smalldict["Loewdin GP"][cleanline[0]] = float(cleanline[2])
                  if "total" in cleanline[0]:
                      self.list_dict_grosspop.append(smalldict)
  
      def get_structure_with_total_grosspop(self, structure_filename: str) -> Structure:
          """
          get a Structure with Mulliken and Loewdin total grosspopulations as site properties
          Args:
              structure_filename (str): filename of POSCAR
          Returns:
              Structure Object with Mulliken and Loewdin total grosspopulations as site properties
          """
  
          struct = Structure.from_file(structure_filename)
          site_properties = {}  # type: Dict[str, Any]
          mullikengp = []
          loewdingp = []
          for grosspop in self.list_dict_grosspop:
              mullikengp.append(grosspop["Mulliken GP"]["total"])
              loewdingp.append(grosspop["Loewdin GP"]["total"])
  
          site_properties = {
              "Total Mulliken GP": mullikengp,
              "Total Loewdin GP": loewdingp,
          }
          new_struct = struct.copy(site_properties=site_properties)
          return new_struct
  
  
  class Wavefunction:
      """
      Class to read in wave function files from Lobster and transfer them into an object of the type VolumetricData
  
      .. attribute: grid
  
          grid for the wave function [Nx+1,Ny+1,Nz+1]
  
      .. attribute: points
  
          list of points
  
      .. attribute: real
  
          list of real part of wave function
  
      .. attribute: imaginary
  
          list of imaginary part of wave function
  
      .. attribute: distance
  
          list of distance to first point in wave function file
  
  
      """
  
      def __init__(self, filename, structure):
          """
  
          Args:
              filename: filename of wavecar file from Lobster
              structure: Structure object (e.g., created by Structure.from_file(""))
          """
  
          self.filename = filename
          self.structure = structure
  
          (
              self.grid,
              self.points,
              self.real,
              self.imaginary,
              self.distance,
          ) = Wavefunction._parse_file(filename)
  
      @staticmethod
      def _parse_file(filename):
          with zopen(filename, "rt") as f:
              contents = f.read().split("\n")
          points = []
          distance = []
          real = []
          imaginary = []
          splitline = contents[0].split()
          grid = [int(splitline[7]), int(splitline[8]), int(splitline[9])]
          for line in contents[1:]:
              splitline = line.split()
              if len(splitline) >= 6:
                  points.append([float(splitline[0]), float(splitline[1]), float(splitline[2])])
                  distance.append(float(splitline[3]))
                  real.append(float(splitline[4]))
                  imaginary.append(float(splitline[5]))
  
          if not len(real) == grid[0] * grid[1] * grid[2]:
              raise ValueError("Something went wrong while reading the file")
          if not len(imaginary) == grid[0] * grid[1] * grid[2]:
              raise ValueError("Something went wrong while reading the file")
          return grid, points, real, imaginary, distance
  
      def set_volumetric_data(self, grid, structure):
          """
          Will create the VolumetricData Objects
  
          Args:
              grid: grid on which wavefunction was calculated, e.g. [1,2,2]
              structure: Structure object
  
          Returns:
  
          """
          Nx = grid[0] - 1
          Ny = grid[1] - 1
          Nz = grid[2] - 1
          a = structure.lattice.matrix[0]
          b = structure.lattice.matrix[1]
          c = structure.lattice.matrix[2]
          new_x = []
          new_y = []
          new_z = []
          new_real = []
          new_imaginary = []
          new_density = []
  
          runner = 0
          for x in range(0, Nx + 1):
              for y in range(0, Ny + 1):
                  for z in range(0, Nz + 1):
                      x_here = x / float(Nx) * a[0] + y / float(Ny) * b[0] + z / float(Nz) * c[0]
                      y_here = x / float(Nx) * a[1] + y / float(Ny) * b[1] + z / float(Nz) * c[1]
                      z_here = x / float(Nx) * a[2] + y / float(Ny) * b[2] + z / float(Nz) * c[2]
  
                      if x != Nx:
                          if y != Ny:
                              if z != Nz:
                                  if not np.isclose(self.points[runner][0], x_here, 1e-3):
                                      if not np.isclose(self.points[runner][1], y_here, 1e-3):
                                          if not np.isclose(self.points[runner][2], z_here, 1e-3):
                                              raise ValueError(
                                                  "The provided wavefunction from Lobster does not contain all relevant"
                                                  " points. "
                                                  "Please use a line similar to: printLCAORealSpaceWavefunction kpoint 1 "
                                                  "coordinates 0.0 0.0 0.0 coordinates 1.0 1.0 1.0 box bandlist 1 "
                                              )
  
                                  new_x.append(x_here)
                                  new_y.append(y_here)
                                  new_z.append(z_here)
  
                                  new_real.append(self.real[runner])
                                  new_imaginary.append(self.imaginary[runner])
                                  new_density.append(self.real[runner] ** 2 + self.imaginary[runner] ** 2)
  
                      runner += 1
  
          self.final_real = np.reshape(new_real, [Nx, Ny, Nz])
          self.final_imaginary = np.reshape(new_imaginary, [Nx, Ny, Nz])
          self.final_density = np.reshape(new_density, [Nx, Ny, Nz])
  
          self.volumetricdata_real = VolumetricData(structure, {"total": self.final_real})
          self.volumetricdata_imaginary = VolumetricData(structure, {"total": self.final_imaginary})
          self.volumetricdata_density = VolumetricData(structure, {"total": self.final_density})
  
      def get_volumetricdata_real(self):
          """
          will return a VolumetricData object including the real part of the wave function
  
          Returns: VolumetricData object
          """
  
          if not hasattr(self, "volumetricdata_real"):
              self.set_volumetric_data(self.grid, self.structure)
          return self.volumetricdata_real
  
      def get_volumetricdata_imaginary(self):
          """
          will return a VolumetricData object including the imaginary part of the wave function
  
          Returns: VolumetricData object
          """
          if not hasattr(self, "volumetricdata_imaginary"):
              self.set_volumetric_data(self.grid, self.structure)
          return self.volumetricdata_imaginary
  
      def get_volumetricdata_density(self):
          """
          will return a VolumetricData object including the imaginary part of the wave function
  
          Returns: VolumetricData object
  
          """
          if not hasattr(self, "volumetricdata_density"):
              self.set_volumetric_data(self.grid, self.structure)
          return self.volumetricdata_density
  
      def write_file(self, filename="WAVECAR.vasp", part="real"):
          """
          will save the wavefunction in a file format that can be read by VESTA
          This will only work if the wavefunction from lobster was constructed with:
          "printLCAORealSpaceWavefunction kpoint 1 coordinates 0.0 0.0 0.0 coordinates 1.0 1.0 1.0 box bandlist 1 2 3 4
          5 6 "
          or similar (the whole unit cell has to be covered!)
  
          Args:
              filename: Filename for the output, e.g., WAVECAR.vasp
              part: which part of the wavefunction will be saved ("real" or "imaginary")
  
          Returns:
  
          """
          if not (
              hasattr(self, "volumetricdata_real")
              and hasattr(self, "volumetricdata_imaginary")
              and hasattr(self, "volumetricdata_density")
          ):
              self.set_volumetric_data(self.grid, self.structure)
          if part == "real":
              self.volumetricdata_real.write_file(filename)
          elif part == "imaginary":
              self.volumetricdata_imaginary.write_file(filename)
          elif part == "density":
              self.volumetricdata_density.write_file(filename)
          else:
              raise ValueError('part can be only "real" or "imaginary" or "density"')
  