<commit_before>import numpy as np
from phonopy.structure.symmetry import Symmetry
from phonopy.structure.cells import get_supercell, get_primitive
from phonopy.structure.atoms import Atoms
from phonopy.units import VaspToTHz
from anharmonic.phonon3.imag_self_energy import get_imag_self_energy, \
     write_imag_self_energy, get_linewidth, write_linewidth
from anharmonic.phonon3.frequency_shift import FrequencyShift
from anharmonic.phonon3.interaction import Interaction
from anharmonic.phonon3.conductivity_RTA import get_thermal_conductivity
from anharmonic.phonon3.joint_dos import JointDos
from anharmonic.phonon3.gruneisen import Gruneisen
from anharmonic.phonon3.displacement_fc3 import get_third_order_displacements, \
     direction_to_displacement
from anharmonic.file_IO import write_frequency_shift, write_joint_dos
from anharmonic.other.isotope import Isotope
from phonopy.harmonic.force_constants import get_fc2, set_permutation_symmetry, \
     set_translational_invariance
from anharmonic.phonon3.fc3 import get_fc3, set_permutation_symmetry_fc3, \
     set_translational_invariance_fc3, cutoff_fc3_by_zero

class Phono3py:
    def __init__(self,
                 unitcell,
                 supercell_matrix,
                 primitive_matrix=None,
                 phonon_supercell_matrix=None,
                 mesh=None,
                 band_indices=None,
                 sigmas=[],
                 cutoff_frequency=1e-4,
                 frequency_factor_to_THz=VaspToTHz,
                 is_symmetry=True,
                 is_nosym=False,
                 symmetrize_fc3_q=False,
                 symprec=1e-5,
                 log_level=0,
                 lapack_zheev_uplo='L'):
        self._symprec = symprec
        self._sigmas = sigmas
        self._frequency_factor_to_THz = frequency_factor_to_THz
        self._is_symmetry = is_symmetry
        self._is_nosym = is_nosym
        self._lapack_zheev_uplo =  lapack_zheev_uplo
        self._symmetrize_fc3_q = symmetrize_fc3_q
        self._cutoff_frequency = cutoff_frequency
        self._log_level = log_level

        # Create supercell and primitive cell
        self._unitcell = unitcell
        self._supercell_matrix = supercell_matrix
        self._primitive_matrix = primitive_matrix
        self._phonon_supercell_matrix = phonon_supercell_matrix # optional
        self._supercell = None
        self._primitive = None
        self._phonon_supercell = None
        self._phonon_primitive = None
        self._build_supercell()
        self._build_primitive_cell()
        self._build_phonon_supercell()
        self._build_phonon_primitive_cell()

        # Set supercell, primitive, and phonon supercell symmetries
        self._symmetry = None
        self._primitive_symmetry = None
        self._phonon_supercell_symmetry = None
        self._search_symmetry()
        self._search_primitive_symmetry()
        self._search_phonon_supercell_symmetry()

        # Displacements and supercells
        self._supercells_with_displacements = None
        self._displacement_dataset = None
        
        # Thermal conductivity
        self._thermal_conductivity = None # conductivity_RTA object

        # Imaginary part of self energy at frequency points
        self._imag_self_energy = None

        # Linewidth (Imaginary part of self energy x 2) at temperatures
        self._linewidth = None

        self._grid_points = None
        self._frequency_points = None
        self._temperatures = None

        # Other variables
        self._fc2 = None
        self._fc3 = None
        
        # Setup interaction
        self._interaction = None
        self._mesh = None
        self._band_indices = None
        self._band_indices_flatten = None
        if mesh is not None:
            self._mesh = np.array(mesh, dtype='intc')
        if band_indices is None:
            num_band = self._primitive.get_number_of_atoms() * 3
            self._band_indices = [np.arange(num_band)]
        else:
            self._band_indices = band_indices
        self._band_indices_flatten = np.hstack(self._band_indices).astype('intc')

    def set_phph_interaction(self,
                             nac_params=None,
                             nac_q_direction=None,
                             frequency_scale_factor=None):
        self._interaction = Interaction(
            self._supercell,
            self._primitive,
            self._mesh,
            self._primitive_symmetry,
            fc3=self._fc3,
            band_indices=self._band_indices_flatten,
            frequency_factor_to_THz=self._frequency_factor_to_THz,
            cutoff_frequency=self._cutoff_frequency,
            is_nosym=self._is_nosym,
            symmetrize_fc3_q=self._symmetrize_fc3_q,
            lapack_zheev_uplo=self._lapack_zheev_uplo)
        self._interaction.set_dynamical_matrix(
            self._fc2,
            self._phonon_supercell,
            self._phonon_primitive,
            nac_params=nac_params,
            frequency_scale_factor=frequency_scale_factor)
        self._interaction.set_nac_q_direction(nac_q_direction=nac_q_direction)

    def generate_displacements(self,
                               distance=0.03,
                               cutoff_pair_distance=None,
                               is_plusminus='auto',
                               is_diagonal=True):
        direction_dataset = get_third_order_displacements(
            self._supercell,
            self._symmetry,
            is_plusminus=is_plusminus,
            is_diagonal=is_diagonal)
        self._displacement_dataset = direction_to_displacement(
            direction_dataset,
            distance,
            self._supercell,
            cutoff_distance=cutoff_pair_distance)

    def produce_fc2(self,
                    forces_fc2,
                    displacement_dataset=None,
                    is_permutation_symmetry=False,
                    is_translational_symmetry=False):
        if displacement_dataset is None:
            disp_dataset = self._displacement_dataset
        else:
            disp_dataset = displacement_dataset
            
        if self._log_level:
            print "Solving fc2"
        for forces, disp1 in zip(forces_fc2, disp_dataset['first_atoms']):
            disp1['forces'] = forces
        self._fc2 = get_fc2(self._phonon_supercell,
                            self._phonon_supercell_symmetry,
                            disp_dataset)
        if is_permutation_symmetry:
            set_permutation_symmetry(self._fc2)
        if is_translational_symmetry:
            set_translational_invariance(self._fc2)

    def produce_fc3(self,
                    forces_fc3,
                    displacement_dataset=None,
                    cutoff_distance=None, # set fc3 zero
                    is_translational_symmetry=False,
                    is_permutation_symmetry=False):
        if displacement_dataset is None:
            disp_dataset = self._displacement_dataset
        else:
            disp_dataset = displacement_dataset
        
        if self._log_level:
            print "Solving fc2"
        for forces, disp1 in zip(forces_fc3, disp_dataset['first_atoms']):
            disp1['forces'] = forces
        fc2 = get_fc2(self._supercell, self._symmetry, disp_dataset)
        if is_permutation_symmetry:
            set_permutation_symmetry(fc2)
        if is_translational_symmetry:
            set_translational_invariance(fc2)
        
        if self._log_level:
            print "Solving fc3:"
        count = len(disp_dataset['first_atoms'])
        for disp1 in disp_dataset['first_atoms']:
            for disp2 in disp1['second_atoms']:
                disp2['delta_forces'] = forces_fc3[count] - disp1['forces']
                count += 1
        self._fc3 = get_fc3(
            self._supercell,
            disp_dataset,
            fc2,
            self._symmetry,
            is_translational_symmetry=is_translational_symmetry,
            is_permutation_symmetry=is_permutation_symmetry,
            verbose=self._log_level)<commit_msg>Minor change<commit_after>import numpy as np
from phonopy.structure.symmetry import Symmetry
from phonopy.structure.cells import get_supercell, get_primitive
from phonopy.structure.atoms import Atoms
from phonopy.units import VaspToTHz
from anharmonic.phonon3.imag_self_energy import get_imag_self_energy, \
     write_imag_self_energy, get_linewidth, write_linewidth
from anharmonic.phonon3.frequency_shift import FrequencyShift
from anharmonic.phonon3.interaction import Interaction
from anharmonic.phonon3.conductivity_RTA import get_thermal_conductivity
from anharmonic.phonon3.joint_dos import JointDos
from anharmonic.phonon3.gruneisen import Gruneisen
from anharmonic.phonon3.displacement_fc3 import get_third_order_displacements, \
     direction_to_displacement
from anharmonic.file_IO import write_frequency_shift, write_joint_dos
from anharmonic.other.isotope import Isotope
from phonopy.harmonic.force_constants import get_fc2, set_permutation_symmetry, \
     set_translational_invariance
from anharmonic.phonon3.fc3 import get_fc3, set_permutation_symmetry_fc3, \
     set_translational_invariance_fc3, cutoff_fc3_by_zero

class Phono3py:
    def __init__(self,
                 unitcell,
                 supercell_matrix,
                 primitive_matrix=None,
                 phonon_supercell_matrix=None,
                 mesh=None,
                 band_indices=None,
                 sigmas=[],
                 cutoff_frequency=1e-4,
                 frequency_factor_to_THz=VaspToTHz,
                 is_symmetry=True,
                 is_nosym=False,
                 symmetrize_fc3_q=False,
                 symprec=1e-5,
                 log_level=0,
                 lapack_zheev_uplo='L'):
        self._symprec = symprec
        self._sigmas = sigmas
        self._frequency_factor_to_THz = frequency_factor_to_THz
        self._is_symmetry = is_symmetry
        self._is_nosym = is_nosym
        self._lapack_zheev_uplo =  lapack_zheev_uplo
        self._symmetrize_fc3_q = symmetrize_fc3_q
        self._cutoff_frequency = cutoff_frequency
        self._log_level = log_level

        # Create supercell and primitive cell
        self._unitcell = unitcell
        self._supercell_matrix = supercell_matrix
        self._primitive_matrix = primitive_matrix
        self._phonon_supercell_matrix = phonon_supercell_matrix # optional
        self._supercell = None
        self._primitive = None
        self._phonon_supercell = None
        self._phonon_primitive = None
        self._build_supercell()
        self._build_primitive_cell()
        self._build_phonon_supercell()
        self._build_phonon_primitive_cell()

        # Set supercell, primitive, and phonon supercell symmetries
        self._symmetry = None
        self._primitive_symmetry = None
        self._phonon_supercell_symmetry = None
        self._search_symmetry()
        self._search_primitive_symmetry()
        self._search_phonon_supercell_symmetry()

        # Displacements and supercells
        self._supercells_with_displacements = None
        self._displacement_dataset = None
        
        # Thermal conductivity
        self._thermal_conductivity = None # conductivity_RTA object

        # Imaginary part of self energy at frequency points
        self._imag_self_energy = None

        # Linewidth (Imaginary part of self energy x 2) at temperatures
        self._linewidth = None

        self._grid_points = None
        self._frequency_points = None
        self._temperatures = None

        # Other variables
        self._fc2 = None
        self._fc3 = None
        
        # Setup interaction
        self._interaction = None
        self._mesh = None
        self._band_indices = None
        self._band_indices_flatten = None
        if mesh is not None:
            self._mesh = np.array(mesh, dtype='intc')
        if band_indices is None:
            num_band = self._primitive.get_number_of_atoms() * 3
            self._band_indices = [np.arange(num_band)]
        else:
            self._band_indices = band_indices
        self._band_indices_flatten = np.hstack(self._band_indices).astype('intc')

    def set_phph_interaction(self,
                             nac_params=None,
                             nac_q_direction=None,
                             frequency_scale_factor=None):
        self._interaction = Interaction(
            self._supercell,
            self._primitive,
            self._mesh,
            self._primitive_symmetry,
            fc3=self._fc3,
            band_indices=self._band_indices_flatten,
            frequency_factor_to_THz=self._frequency_factor_to_THz,
            cutoff_frequency=self._cutoff_frequency,
            is_nosym=self._is_nosym,
            symmetrize_fc3_q=self._symmetrize_fc3_q,
            lapack_zheev_uplo=self._lapack_zheev_uplo)
        self._interaction.set_dynamical_matrix(
            self._fc2,
            self._phonon_supercell,
            self._phonon_primitive,
            nac_params=nac_params,
            frequency_scale_factor=frequency_scale_factor)
        self._interaction.set_nac_q_direction(nac_q_direction=nac_q_direction)

    def generate_displacements(self,
                               distance=0.03,
                               cutoff_pair_distance=None,
                               is_plusminus='auto',
                               is_diagonal=True):
        direction_dataset = get_third_order_displacements(
            self._supercell,
            self._symmetry,
            is_plusminus=is_plusminus,
            is_diagonal=is_diagonal)
        self._displacement_dataset = direction_to_displacement(
            direction_dataset,
            distance,
            self._supercell,
            cutoff_distance=cutoff_pair_distance)

    def produce_fc2(self,
                    forces_fc2,
                    displacement_dataset=None,
                    is_permutation_symmetry=False,
                    is_translational_symmetry=False):
        if displacement_dataset is None:
            disp_dataset = self._displacement_dataset
        else:
            disp_dataset = displacement_dataset
            
        if self._log_level:
            print "Solving fc2"
        for forces, disp1 in zip(forces_fc2, disp_dataset['first_atoms']):
            disp1['forces'] = forces
        self._fc2 = get_fc2(self._phonon_supercell,
                            self._phonon_supercell_symmetry,
                            disp_dataset)
        if is_permutation_symmetry:
            set_permutation_symmetry(self._fc2)
        if is_translational_symmetry:
            set_translational_invariance(self._fc2)

    def produce_fc3(self,
                    forces_fc3,
                    displacement_dataset=None,
                    cutoff_distance=None, # set fc3 zero
                    is_translational_symmetry=False,
                    is_permutation_symmetry=False):
        if displacement_dataset is None:
            disp_dataset = self._displacement_dataset
        else:
            disp_dataset = displacement_dataset
        
        if self._log_level:
            print "Solving fc2,",
        for forces, disp1 in zip(forces_fc3, disp_dataset['first_atoms']):
            disp1['forces'] = forces
        fc2 = get_fc2(self._supercell, self._symmetry, disp_dataset)
        if is_permutation_symmetry:
            set_permutation_symmetry(fc2)
        if is_translational_symmetry:
            set_translational_invariance(fc2)
        
        if self._log_level:
            print "fc3:"
        count = len(disp_dataset['first_atoms'])
        for disp1 in disp_dataset['first_atoms']:
            for disp2 in disp1['second_atoms']:
                disp2['delta_forces'] = forces_fc3[count] - disp1['forces']
                count += 1
        self._fc3 = get_fc3(
            self._supercell,
            disp_dataset,
            fc2,
            self._symmetry,
            is_translational_symmetry=is_translational_symmetry,
            is_permutation_symmetry=is_permutation_symmetry,
            verbose=self._log_level)
{'fraction_non_alphanumeric': 0.05767453324296013, 'fraction_numerical': 0.00838006038572925, 'mean_word_length': 2.1038439472174413, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '4414947', 'n_tokens_mistral': 4767, 'n_tokens_neox': 4573, 'n_words': 769}