<commit_before># Copyright (C) 2018 <NAME>, <NAME>
# All rights reserved.
#
# This file is part of yambopy
#
#
import os
import re
import shutil
from math import sqrt
from .pseudo import get_pseudo_path
from .tools import fortran_bool
from .lattice import red_car, car_red

class PwIn(object):
    """
    Class to generate an manipulate Quantum Espresso input files
    Can be initialized either reading from a file or starting from a new file.
    This class is not meant to be comprehensive but a lightweight version 
    capable of basic input/ouput of PW files.
    For a comprehensive class use the tools provoded by the AiiDa package: http://www.aiida.net/

    Examples of use:

    To read a local file with name "mos2.in"

        .. code-block :: python

            qe = PwIn.from_file('mos2.scf')
            print qe

    To start a file from scratch

        .. code-block :: python

            qe = PwIn()
            qe.atoms = [['N',[ 0.0, 0.0,0.0]],
                        ['B',[1./3,2./3,0.0]]]
            qe.atypes = {'B': [10.811, "B.pbe-mt_fhi.UPF"],
                         'N': [14.0067,"N.pbe-mt_fhi.UPF"]}

            qe.control['prefix'] = "'%s'"%prefix
            qe.control['verbosity'] = "'high'"
            qe.control['wf_collect'] = '.true.'
            qe.control['pseudo_dir'] = "'../pseudos/'"
            qe.system['celldm(1)'] = 4.7
            qe.system['celldm(3)'] = layer_separation/qe.system['celldm(1)']
            qe.system['ecutwfc'] = 60
            qe.system['occupations'] = "'fixed'"
            qe.system['nat'] = 2
            qe.system['ntyp'] = 2
            qe.system['ibrav'] = 4
            qe.kpoints = [6, 6, 1]
            qe.electrons['conv_thr'] = 1e-8

            print qe

    Special care should be taken with string variables e.g. "'high'"

    """
    _pw = 'pw.x'

    def __init__(self):
        """ TODO: specify the required parameters """
        #kpoints
        self.ktype = "automatic"
        self.kpoints = [1,1,1]
        self.shiftk = [0,0,0]
        self.klist = []
        #dictionaries
        self.control = dict(prefix="'pw'",wf_collect='.true.')
        self.system = dict()
        self.electrons = dict(conv_thr=1e-8)
        self.ions = dict()
        self.cell = dict()
        self.atypes = dict()
        self.cell_units = 'bohr'
        self.atomic_pos_type = 'crystal'

    @classmethod
    def from_file(cls,filename):
        """ Initialize the QE structure from a file """
        new = cls()
 
        with open(filename,"r") as f:
            new.file_lines = f.readlines() #set file lines
            new.store(new.control,"control")     #read &control
            new.store(new.system,"system")      #read &system
            new.store(new.electrons,"electrons")   #read &electrons
            new.store(new.ions,"ions")        #read &ions
            new.store(new.cell,"cell")        #read &ions
            #read ATOMIC_SPECIES
            new.read_atomicspecies()
            #read ATOMIC_POSITIONS
            new.read_atoms()
            #read K_POINTS
            new.read_kpoints()
            #read CELL_PARAMETERS
            new.read_cell_parameters()

        return new

    @classmethod
    def from_structure_dict(cls,structure,kpoints=None,ecut=None,pseudo_dir='.'):
        pwi = cls()
        pwi.set_structure(structure)
        if kpoints: pwi.set_kpoints(kpoints)
        if ecut: pwi.set_ecut(ecut)
        if pseudo_dir: pwi.pseudo_dir = pseudo_dir
        return pwi
      
    @property
    def natoms(self): 
        return len(self.atoms)

    @property
    def pseudo_dir(self):
        if 'pseudo_dir' not in self.control: return None
        return self.control['pseudo_dir'].replace("'",'')
    
    @pseudo_dir.setter
    def pseudo_dir(self,value):
        self.control['pseudo_dir'] = "'%s'"%value.replace("'",'')

    @property
    def prefix(self): 
        return self.control['prefix'].replace("'",'')

    @prefix.setter
    def prefix(self,value):
        self.control['prefix'] = "'%s'"%value.replace("'",'')

    def set_ecut(self,ecut):
        self.system['ecutwfc'] = ecut

    def set_structure(self,structure):
        """
        Set the structure from a structure dictionary
        
        Example:
            .. code-block :: python
            
            structure = dict(ibrav=4,celldm1=4.7,celldm3=12)
            atypes = dict(Si=[28.086,"Si.pbe-mt_fhi.UPF"])
            atoms = [['N',[ 0.0, 0.0,0.5]],
                     ['B',[1./3,2./3,0.5]]]
            BN = dict(structure=structure,atypes=atypes,atoms=atoms)

            pwi = PwIn()
            pwi.set_structure(BN)
        """
        if 'lattice' in structure: self.set_lattice(**structure['lattice'])
        if 'atypes'  in structure: self.set_atypes(structure['atypes'])
        if 'atoms'   in structure: self.set_atoms(structure['atoms'])

    def get_structure(self):
        """ Return an instance of a structure dictionary
        """
        lattice = self.get_lattice()
        return dict(lattice=lattice,atypes=self.atypes,atoms=self.atoms)

    def change_cell_parameters(self):
        """
        Convert the atomic postions to cartesian, change the lattice and convert
        the atomic positions to reduced
        """
        raise NotImplementedError('TODO')

    def set_lattice(self,ibrav=None,celldm1=None,celldm2=None,celldm3=None,
                      celldm4=None,celldm5=None,celldm6=None,cell_parameters=None):
        """Set the structure using the typical QE input variables"""
        if ibrav == 0 and cell_parameters is None:
            raise ValueError('ibrav = 0 implies that the cell_parameters variable is set')
        if cell_parameters: self.cell_parameters = cell_parameters
        if ibrav is not None: self.ibrav = ibrav
        if celldm1 is not None: self.system['celldm(1)'] = celldm1
        if celldm2 is not None: self.system['celldm(2)'] = celldm2
        if celldm3 is not None: self.system['celldm(3)'] = celldm3
        if celldm4 is not None: self.system['celldm(4)'] = celldm4
        if celldm5 is not None: self.system['celldm(5)'] = celldm5
        if celldm6 is not None: self.system['celldm(6)'] = celldm6

    def get_lattice(self):
        lattice_dict = {}
        if 'ibrav' in self.system: lattice_dict['ibrav'] = self.ibrav 
        if 'celldm(1)' in self.system: lattice_dict['celldm1'] = self.system['celldm(1)']
        if 'celldm(2)' in self.system: lattice_dict['celldm2'] = self.system['celldm(2)']
        if 'celldm(3)' in self.system: lattice_dict['celldm3'] = self.system['celldm(3)']
        if 'celldm(4)' in self.system: lattice_dict['celldm4'] = self.system['celldm(4)']
        if 'celldm(5)' in self.system: lattice_dict['celldm5'] = self.system['celldm(5)']
        if 'celldm(6)' in self.system: lattice_dict['celldm6'] = self.system['celldm(6)']
        return lattice_dict 

    def set_atoms(self,atoms,coordtype="reduced"):
        """
        Set the atoms.
        Internally the atoms are always stored in reduced coordinates.
        The positions are written/read in the format from atomic_pos_type

        Arguments:
            coordtype: ("reduced"/"crystal") or cartesian

        Example:
            .. code-block :: python
   
                pwi = PwIn()
                pwi.set_atoms( [['Si',[0.125,0.125,0.125]],
                                ['Si',[-.125,-.125,-.125]]])
        """
        self.system['nat'] = len(atoms)
        #TODO: add consistency check 

        #convert the positions if needed
        if coordtype in ["reduced","crystal"]:
            self._atoms = atoms
        else:
            red_atoms = []
            for atype,apos in atoms:<commit_msg>Fixed get_lattice in PwIn<commit_after># Copyright (C) 2018 <NAME>, <NAME>
# All rights reserved.
#
# This file is part of yambopy
#
#
import os
import re
import shutil
from math import sqrt
from .pseudo import get_pseudo_path
from .tools import fortran_bool
from .lattice import red_car, car_red

class PwIn(object):
    """
    Class to generate an manipulate Quantum Espresso input files
    Can be initialized either reading from a file or starting from a new file.
    This class is not meant to be comprehensive but a lightweight version 
    capable of basic input/ouput of PW files.
    For a comprehensive class use the tools provoded by the AiiDa package: http://www.aiida.net/

    Examples of use:

    To read a local file with name "mos2.in"

        .. code-block :: python

            qe = PwIn.from_file('mos2.scf')
            print qe

    To start a file from scratch

        .. code-block :: python

            qe = PwIn()
            qe.atoms = [['N',[ 0.0, 0.0,0.0]],
                        ['B',[1./3,2./3,0.0]]]
            qe.atypes = {'B': [10.811, "B.pbe-mt_fhi.UPF"],
                         'N': [14.0067,"N.pbe-mt_fhi.UPF"]}

            qe.control['prefix'] = "'%s'"%prefix
            qe.control['verbosity'] = "'high'"
            qe.control['wf_collect'] = '.true.'
            qe.control['pseudo_dir'] = "'../pseudos/'"
            qe.system['celldm(1)'] = 4.7
            qe.system['celldm(3)'] = layer_separation/qe.system['celldm(1)']
            qe.system['ecutwfc'] = 60
            qe.system['occupations'] = "'fixed'"
            qe.system['nat'] = 2
            qe.system['ntyp'] = 2
            qe.system['ibrav'] = 4
            qe.kpoints = [6, 6, 1]
            qe.electrons['conv_thr'] = 1e-8

            print qe

    Special care should be taken with string variables e.g. "'high'"

    """
    _pw = 'pw.x'

    def __init__(self):
        """ TODO: specify the required parameters """
        #kpoints
        self.ktype = "automatic"
        self.kpoints = [1,1,1]
        self.shiftk = [0,0,0]
        self.klist = []
        #dictionaries
        self.control = dict(prefix="'pw'",wf_collect='.true.')
        self.system = dict()
        self.electrons = dict(conv_thr=1e-8)
        self.ions = dict()
        self.cell = dict()
        self.atypes = dict()
        self.cell_units = 'bohr'
        self.atomic_pos_type = 'crystal'

    @classmethod
    def from_file(cls,filename):
        """ Initialize the QE structure from a file """
        new = cls()
 
        with open(filename,"r") as f:
            new.file_lines = f.readlines() #set file lines
            new.store(new.control,"control")     #read &control
            new.store(new.system,"system")      #read &system
            new.store(new.electrons,"electrons")   #read &electrons
            new.store(new.ions,"ions")        #read &ions
            new.store(new.cell,"cell")        #read &ions
            #read ATOMIC_SPECIES
            new.read_atomicspecies()
            #read ATOMIC_POSITIONS
            new.read_atoms()
            #read K_POINTS
            new.read_kpoints()
            #read CELL_PARAMETERS
            new.read_cell_parameters()

        return new

    @classmethod
    def from_structure_dict(cls,structure,kpoints=None,ecut=None,pseudo_dir='.'):
        pwi = cls()
        pwi.set_structure(structure)
        if kpoints: pwi.set_kpoints(kpoints)
        if ecut: pwi.set_ecut(ecut)
        if pseudo_dir: pwi.pseudo_dir = pseudo_dir
        return pwi
      
    @property
    def natoms(self): 
        return len(self.atoms)

    @property
    def pseudo_dir(self):
        if 'pseudo_dir' not in self.control: return None
        return self.control['pseudo_dir'].replace("'",'')
    
    @pseudo_dir.setter
    def pseudo_dir(self,value):
        self.control['pseudo_dir'] = "'%s'"%value.replace("'",'')

    @property
    def prefix(self): 
        return self.control['prefix'].replace("'",'')

    @prefix.setter
    def prefix(self,value):
        self.control['prefix'] = "'%s'"%value.replace("'",'')

    def set_ecut(self,ecut):
        self.system['ecutwfc'] = ecut

    def set_structure(self,structure):
        """
        Set the structure from a structure dictionary
        
        Example:
            .. code-block :: python
            
            structure = dict(ibrav=4,celldm1=4.7,celldm3=12)
            atypes = dict(Si=[28.086,"Si.pbe-mt_fhi.UPF"])
            atoms = [['N',[ 0.0, 0.0,0.5]],
                     ['B',[1./3,2./3,0.5]]]
            BN = dict(structure=structure,atypes=atypes,atoms=atoms)

            pwi = PwIn()
            pwi.set_structure(BN)
        """
        if 'lattice' in structure: self.set_lattice(**structure['lattice'])
        if 'atypes'  in structure: self.set_atypes(structure['atypes'])
        if 'atoms'   in structure: self.set_atoms(structure['atoms'])

    def get_structure(self):
        """ Return an instance of a structure dictionary
        """
        lattice = self.get_lattice()
        return dict(lattice=lattice,atypes=self.atypes,atoms=self.atoms)

    def change_cell_parameters(self):
        """
        Convert the atomic postions to cartesian, change the lattice and convert
        the atomic positions to reduced
        """
        raise NotImplementedError('TODO')

    def set_lattice(self,ibrav=None,celldm1=None,celldm2=None,celldm3=None,
                      celldm4=None,celldm5=None,celldm6=None,cell_parameters=None):
        """Set the structure using the typical QE input variables"""
        if ibrav == 0 and cell_parameters is None:
            raise ValueError('ibrav = 0 implies that the cell_parameters variable is set')
        if cell_parameters: self.cell_parameters = cell_parameters
        if ibrav is not None: self.ibrav = ibrav
        if celldm1 is not None: self.system['celldm(1)'] = celldm1
        if celldm2 is not None: self.system['celldm(2)'] = celldm2
        if celldm3 is not None: self.system['celldm(3)'] = celldm3
        if celldm4 is not None: self.system['celldm(4)'] = celldm4
        if celldm5 is not None: self.system['celldm(5)'] = celldm5
        if celldm6 is not None: self.system['celldm(6)'] = celldm6

    def get_lattice(self):
        lattice_dict = {}
        if 'ibrav' in self.system: lattice_dict['ibrav'] = self.ibrav 
        if 'celldm(1)' in self.system: lattice_dict['celldm1'] = self.system['celldm(1)']
        if 'celldm(2)' in self.system: lattice_dict['celldm2'] = self.system['celldm(2)']
        if 'celldm(3)' in self.system: lattice_dict['celldm3'] = self.system['celldm(3)']
        if 'celldm(4)' in self.system: lattice_dict['celldm4'] = self.system['celldm(4)']
        if 'celldm(5)' in self.system: lattice_dict['celldm5'] = self.system['celldm(5)']
        if 'celldm(6)' in self.system: lattice_dict['celldm6'] = self.system['celldm(6)']
        if self.ibrav == 0: lattice_dict['cell_parameters'] = self.cell_parameters
        return lattice_dict 

    def set_atoms(self,atoms,coordtype="reduced"):
        """
        Set the atoms.
        Internally the atoms are always stored in reduced coordinates.
        The positions are written/read in the format from atomic_pos_type

        Arguments:
            coordtype: ("reduced"/"crystal") or cartesian

        Example:
            .. code-block :: python
   
                pwi = PwIn()
                pwi.set_atoms( [['Si',[0.125,0.125,0.125]],
                                ['Si',[-.125,-.125,-.125]]])
        """
        self.system['nat'] = len(atoms)
        #TODO: add consistency check 

        #convert the positions if needed
        if coordtype in ["reduced","crystal"]:
            self._atoms = atoms
        else:
            red_atoms = []
            for atype,apos in atoms:
{'fraction_non_alphanumeric': 0.13209614021522006, 'fraction_numerical': 0.018235711063857206, 'mean_word_length': 2.1654089333061393, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, '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': '3766978', 'n_tokens_mistral': 5190, 'n_tokens_neox': 5042, 'n_words': 1153}