<commit_before>"""Classes for use with Quantum ESPRESSO

Representation of a spectrum
"""

import re
import os
import copy  as cp
import numpy as np

import atk.atomistic.fundamental as fu
import atk.atomistic.constants as atc

class Dispersion(object):
    """Holds k-points belonging to one spin"""
    def __init__(self, kpoints=None, kvectors=None):
        if kpoints is None:
            self.kpoints = []
            self.kvectors = []
        else:
            self.kpoints = kpoints
            self.kvectors = kvectors

    def addkpoint(self, kpoint, kvector):
        self.kpoints.append(kpoint)
        self.kvectors.append(kvector)

    def energylevels(self):
        s = fu.EnergyLevels()
        for l in self.kpoints:
            s += l
        return s

    def mergekpoints(self):
        self.kpoints = energylevels(self.kpoints)

    def __iadd__(self, s):
        """Merging two spins"""
        if len(self.kpoints) != len(s.kpoints):
            print "Unable to merge due to different number of kpoints"
        for i in range(len(self.kpoints)):
            self.kpoints[i] += s.kpoints[i]
        return self

    def shift(self, e):
        for l in self.kpoints:
            l.shift(e)


class Spectrum(object):
    """A collection of dispersions, grouped by spin"""

    def __init__(self, dispersions=None, spins=None):
        """Set up spectrum from list of dispersions."""
        self.dispersions = dispersions
        self.spins = spins

    @classmethod
    def from_save(cls, fname):
        """Creates Spectrum from QE save directory"""
        tmp = Spectrum()
        tmp.read_from_save(fname)
        return tmp

    @classmethod
    def from_output(cls, fname):
        """Creates Spectrum from QE output"""
        tmp = Spectrum()
        tmp.read_from_output(fname)
        return tmp

    @property
    def energies(self):
        """Returns list of all energy levels of all spins."""
        list = [d.energies for d in self.dispersions]
        return np.concatenate(list)

    @property
    def occupations(self):
        """Returns list of level occupations of all spins."""
        os = np.array([])
        for d in self.dispersions:
            os = np.concatenate( (os, d.occupations))
        return os

    @property
    def fermi(self):
        """Returns Fermi energy."""
        fermis = [el.fermi for d in self.dispersions]

        fermi = np.unique(fermis)

        if len( np.unique(fermis) ) != 1:
            print "There are Fermi energies {}".format(fermis)
            print "Using the mean {}".format(np.mean(fermis))

        return np.mean(fermis)

    def copy(self, spectrum):
        """Performs deep copy of spectrum."""
        self.dispersions = [ d.copy() for d in spectrum.dispersions ]
        self.spins = cp.copy(spectrum.spins)

    def shift(self, de):
        for d in self.dispersions:
            d.shift(de)

    def __str__(self):
        text  = "Spectrum containing {} spins\n".format(len(self.dispersions))
        for i in range(len(self.dispersions)):
            d = self.dispersions[i]
            s = self.spins[i]
            text += 'spin {} : {}\n'.format(s+1, d.__str__())
        return text

    def __getitem__(self, index):
        return self.dispersions[index]

    def read_from_save(self, prefix):
        """Reads Spectrum from QE save directory"""
        dispersion = Dispersion()

        savedir = prefix + '.save'
        if not os.path.exists(savedir):
            print "Error: directory {s} not found.".format(s=savedir)
            return

        os.chdir(savedir)
        dataxml = open('data-file.xml', 'r').read()

        # should be able to match scientific and non-scientific notation
        floatregex = '-?\d+\.\d+(?:[Ee][+\-]?\d+)?'

        # get fermi energy
        fermiregex = '<FERMI_ENERGY.*?>\s*({f})'.format(f=floatregex)
        fermi = float( re.search(fermiregex, dataxml, re.DOTALL).group(1) )
        fermi *= atc.Ha / atc.eV

        #get lattice parameter
        alatregex = '<LATTICE_PARAMETER.*?>\s*({f})'.format(f=floatregex)
        alat = float( re.search(alatregex, dataxml, re.DOTALL).group(1) )
        alat *= atc.a0  


        # get k-points
        kptregex = '<K-POINT\.(\d+)\s+XYZ=\"({f}) ({f}) ({f})\"\s+WEIGHT=\"({f})\"/>'\
                .format(f=floatregex)
        kptdatas = re.findall(kptregex, dataxml)

        for kpt in kptdatas:
            kindex = int(kpt[0])
            kvec   = np.array([ kpt[1], kpt[2], kpt[3] ], dtype=float)
            #kvec  *= 2*np.pi / alat

            kdir = 'K{k:05d}'.format(k=kindex)
            if not os.path.exists(kdir):
                print "Error: directory {s} not found.".format(s=kdir)
                return

            # Read energy levels
            os.chdir(kdir)

            eigenvalxml = open('eigenval.xml', 'r').read()

            eigregex  = '<EIGENVALUES.*?>(.*?)<'
            eigstring = re.search(eigregex, eigenvalxml, re.DOTALL).group(1)

            levelregex  = '\-?\d+.*'
            levelregex  = floatregex
            levels = np.array(re.findall(levelregex, eigstring), dtype = float)
            levels *= atc.Ha / atc.eV

            dispersion.addkpoint( 
                    fu.EnergyLevels(energies=levels, fermi=fermi), 
                    kvec)

            os.chdir('..')

        os.chdir('..')

        self.dispersions = [dispersion]

    def read_from_output(self, prefix):
        return 0;

<commit_msg>Added capability to import spin-polarized QE calculations<commit_after>"""Classes for use with Quantum ESPRESSO

Representation of a spectrum
"""

import re
import os
import copy  as cp
import numpy as np

import atk.atomistic.fundamental as fu
import atk.atomistic.constants as atc

class Dispersion(object):
    """Holds k-points belonging to one spin"""
    def __init__(self, kpoints=None, kvectors=None):
        if kpoints is None:
            self.kpoints = []
            self.kvectors = []
        else:
            self.kpoints = kpoints
            self.kvectors = kvectors

    def addkpoint(self, kpoint, kvector):
        self.kpoints.append(kpoint)
        self.kvectors.append(kvector)

    def energylevels(self):
        s = fu.EnergyLevels()
        for l in self.kpoints:
            s += l
        return s

    def mergekpoints(self):
        self.kpoints = energylevels(self.kpoints)

    def __iadd__(self, s):
        """Merging two spins"""
        if len(self.kpoints) != len(s.kpoints):
            print "Unable to merge due to different number of kpoints"
        for i in range(len(self.kpoints)):
            self.kpoints[i] += s.kpoints[i]
        return self

    def shift(self, e):
        for l in self.kpoints:
            l.shift(e)


class Spectrum(object):
    """A collection of dispersions, grouped by spin"""

    def __init__(self, dispersions=None, spins=None):
        """Set up spectrum from list of dispersions."""
        self.dispersions = dispersions
        self.spins = spins

    @classmethod
    def from_save(cls, fname):
        """Creates Spectrum from QE save directory"""
        tmp = Spectrum()
        tmp.read_from_save(fname)
        return tmp

    @classmethod
    def from_output(cls, fname):
        """Creates Spectrum from QE output"""
        tmp = Spectrum()
        tmp.read_from_output(fname)
        return tmp

    @property
    def energies(self):
        """Returns list of all energy levels of all spins."""
        list = [d.energies for d in self.dispersions]
        return np.concatenate(list)

    @property
    def occupations(self):
        """Returns list of level occupations of all spins."""
        os = np.array([])
        for d in self.dispersions:
            os = np.concatenate( (os, d.occupations))
        return os

    @property
    def fermi(self):
        """Returns Fermi energy."""
        fermis = [el.fermi for d in self.dispersions]

        fermi = np.unique(fermis)

        if len( np.unique(fermis) ) != 1:
            print "There are Fermi energies {}".format(fermis)
            print "Using the mean {}".format(np.mean(fermis))

        return np.mean(fermis)

    def copy(self, spectrum):
        """Performs deep copy of spectrum."""
        self.dispersions = [ d.copy() for d in spectrum.dispersions ]
        self.spins = cp.copy(spectrum.spins)

    def shift(self, de):
        for d in self.dispersions:
            d.shift(de)

    def __str__(self):
        text  = "Spectrum containing {} spins\n".format(len(self.dispersions))
        for i in range(len(self.dispersions)):
            d = self.dispersions[i]
            s = self.spins[i]
            text += 'spin {} : {}\n'.format(s+1, d.__str__())
        return text

    def __getitem__(self, index):
        return self.dispersions[index]

    def read_from_save(self, prefix):
        """Reads Spectrum from QE save directory"""

        savedir = prefix + '.save'
        if not os.path.exists(savedir):
            print "Error: directory {s} not found.".format(s=savedir)
            return

        os.chdir(savedir)
        dataxml = open('data-file.xml', 'r').read()

        nspinregex = '<NUMBER_OF_SPIN_COMPONENTS.*?>\s*(\d+)'
        nspin = int( re.search(nspinregex, dataxml, re.DOTALL).group(1) )

        # should be able to match scientific and non-scientific notation
        floatregex = '-?\d+\.\d+(?:[Ee][+\-]?\d+)?'

        # get fermi energy
        fermiregex = '<FERMI_ENERGY.*?>\s*({f})'.format(f=floatregex)
        fermi = float( re.search(fermiregex, dataxml, re.DOTALL).group(1) )
        fermi *= atc.Ha / atc.eV

        #get lattice parameter
        alatregex = '<LATTICE_PARAMETER.*?>\s*({f})'.format(f=floatregex)
        alat = float( re.search(alatregex, dataxml, re.DOTALL).group(1) )
        alat *= atc.a0  


        # get k-points
        kptregex = '<K-POINT\.(\d+)\s+XYZ=\"({f}) ({f}) ({f})\"\s+WEIGHT=\"({f})\"/>'\
                .format(f=floatregex)
        kptdatas = re.findall(kptregex, dataxml)

        self.dispersions = []
        for spin in range(1, nspin+1):
            dispersion = Dispersion()

            for kpt in kptdatas:
                kindex = int(kpt[0])
                kvec   = np.array([ kpt[1], kpt[2], kpt[3] ], dtype=float)
                #kvec  *= 2*np.pi / alat

                kdir = 'K{k:05d}'.format(k=kindex)
                if not os.path.exists(kdir):
                    print "Error: directory {s} not found.".format(s=kdir)
                    return

                # Read energy levels
                os.chdir(kdir)

                # get correct file name
                if nspin == 1:
                    eigf = 'eigenval.xml'
                elif nspin == 2:
                    eigf = 'eigenval{}.xml'.format(spin)
                else:
                    print "Error: Can only handle nspin=1, 2"

                if not os.path.exists(eigf):
                    print "Error: Cannot find file {}".format(eigf)

                eigenvalxml = open(eigf, 'r').read()

                eigregex  = '<EIGENVALUES.*?>(.*?)<'
                eigstring = re.search(eigregex, eigenvalxml, re.DOTALL).group(1)

                levelregex  = '\-?\d+.*'
                levelregex  = floatregex
                levels = np.array(re.findall(levelregex, eigstring), dtype = float)
                levels *= atc.Ha / atc.eV

                dispersion.addkpoint( 
                        fu.EnergyLevels(energies=levels, fermi=fermi), 
                        kvec)

                os.chdir('..')

            self.dispersions.append(dispersion)    

        os.chdir('..')

    def read_from_output(self, prefix):
        return 0;


{'fraction_non_alphanumeric': 0.12916488222698072, 'fraction_numerical': 0.0029978586723768737, 'mean_word_length': 1.9153558052434456, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 18}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5783653', 'n_tokens_mistral': 3822, 'n_tokens_neox': 3621, 'n_words': 900}