<commit_before>import sys
import numpy as np
from phonopy.structure.atoms import Atoms
from phonopy.interface import vasp
from phonopy.hphonopy.file_IO import write_FORCE_SETS_vasp, read_force_constant_vasprun_xml

###########
#
# File I/O
#
###########

#
# Output
#
def write_supercells_with_displacements(supercell,
                                        double_displacements,
                                        amplitude=None):
    if amplitude==None:
        distance = 0.01
    else:
        distance = amplitude
    
    # YAML
    file = open('disp.yaml', 'w')
    file.write("natom: %d\n" %  supercell.get_number_of_atoms())
    file.write("num_first_displacements: %d\n" %  len(double_displacements))
    num_second = 0
    for d1 in double_displacements:
        for d2 in d1['second_atoms']:
            num_second += len(d2['directions'])
    file.write("num_second_displacements: %d\n" %  num_second)
    file.write("first_atoms:\n")
    lattice = supercell.get_cell()
    lattice_inv = np.linalg.inv(lattice)
    count1 = 1
    count2 = len(double_displacements) + 1
    for disp1 in double_displacements:
        disp_cartesian3 = np.dot(disp1['direction'], lattice)
        disp_cartesian3 = disp_cartesian3 / np.linalg.norm(disp_cartesian3) * distance
        positions = supercell.get_positions()
        positions[disp1['number']] += disp_cartesian3
        atoms = Atoms(numbers = supercell.get_atomic_numbers(),
                      masses = supercell.get_masses(),
                      positions = positions,
                      cell = lattice,
                      pbc = True)
        vasp.write_vasp('POSCAR-%04d' % count1, atoms, direct=True)

        # YAML
        file.write("- number: %5d\n" % (disp1['number'] + 1))
        file.write("  displacement:\n")
        file.write("    [ %20.16f,%20.16f,%20.16f ] # %d \n" %
                   (disp_cartesian3[0], disp_cartesian3[1], disp_cartesian3[2], count1))
        file.write("  second_atoms:\n")
        count1 += 1

        for disp2 in disp1['second_atoms']:
            # YAML
            file.write("  - number: %5d\n" % (disp2['number'] + 1))
            file.write("    displacements:\n")
            
            for direction in disp2['directions']:
                disp_cartesian2 = np.dot(direction, lattice)
                disp_cartesian2 = disp_cartesian2 / np.linalg.norm(disp_cartesian2) * distance
                positions = supercell.get_positions()
                positions[disp1['number']] += disp_cartesian3
                positions[disp2['number']] += disp_cartesian2
                atoms = Atoms(numbers = supercell.get_atomic_numbers(),
                               masses = supercell.get_masses(),
                               positions = positions,
                               cell = lattice,
                               pbc = True)
                vasp.write_vasp('POSCAR-%04d' % count2, atoms, direct=True)

                # YAML
                file.write("    - [ %20.16f,%20.16f,%20.16f ] # %d \n" %
                           (disp_cartesian2[0], disp_cartesian2[1], disp_cartesian2[2], count2))
                count2 += 1

    file.write("lattice:\n")
    for axis in supercell.get_cell():
        file.write("- [ %20.15f,%20.15f,%20.15f ]\n" % tuple(axis))
    symbols = supercell.get_chemical_symbols()
    positions = supercell.get_scaled_positions()
    file.write("atoms:\n")
    for i, (s, v) in enumerate(zip(symbols, positions)):
        file.write("- symbol: %-2s # %d\n" % (s, i+1))
        file.write("  position: [ %18.14f,%18.14f,%18.14f ]\n" % \
                       (v[0], v[1], v[2]))

    file.close()

def write_FORCES_THIRD(vaspruns,
                       displacements,
                       forces_third='FORCES_THIRD',
                       forces_second='FORCES_SECOND',
                       force_sets='FORCE_SETS'):
    natom = displacements['natom']
    num_disp1 = len(displacements['first_atoms'])
    disp_datasets = []
    set_of_forces = get_forces_from_vasprun_xmls(vaspruns, natom)
    filethird = open(forces_third, 'w')
    filesecond = open(forces_second, 'w')

    for i, disp1 in enumerate(displacements['first_atoms']):
        filesecond.write("# File: %-5d\n" % (i + 1))
        filesecond.write("# %-5d " % disp1['number'] + 1)
        filesecond.write("%20.16f %20.16f %20.16f\n" %
                         tuple(disp1['displacement']))
        for f in set_of_forces[i]:
            filesecond.write("%15.10f %15.10f %15.10f\n" % (tuple(f)))
        
        disp_datasets.append([disp1['number'], disp1['displacement']])

    count = num_disp1
    for i, disp1 in enumerate(displacements['first_atoms']):
        for disp2 in disp1['second_atoms']:
            for d in disp2['displacements']:
                filethird.write("# File: %-5d\n" % (count + 1))
                filethird.write("# %-5d " % disp1['number'] + 1)
                filethird.write("%20.16f %20.16f %20.16f\n" %
                                 tuple(disp1['displacement']))
                filethird.write("# %-5d " % disp2['number'] + 1)
                filethird.write("%20.16f %20.16f %20.16f\n" % tuple(d))

                for f in set_of_forces[count]:
                    filethird.write("%15.10f %15.10f %15.10f\n" % (tuple(f)))
                count += 1

    write_FORCE_SETS_vasp(vaspruns[:num_disp1],
                          disp_datasets,
                          natom,
                          force_sets)

def write_DELTA_FC2_SETS(vaspruns,
                         displacements,
                         fc2_file='fc2.dat',
                         dfc2_file='DELTA_FC2_SETS'):
    fc2_set = get_force_constants_from_vasprun_xmls(vaspruns)
    w = open(dfc2_file, 'w')
    perfect_fc2 = fc2_set.pop(0)
    write_fc2_dat(perfect_fc2, fc2_file)
    for i, (dfc2, first_disp) in enumerate(
        zip(fc2_set, displacements['first_atoms'])):
        w.write("# File: %d\n" % (i + 1))
        w.write("# %-5d " % (first_disp['number'] + 1))
        w.write("%20.16f %20.16f %20.16f\n" %
                tuple(first_disp['displacement']))<commit_msg>Update anharmonic code<commit_after>import sys
import numpy as np
from phonopy.structure.atoms import Atoms
from phonopy.interface import vasp
from phonopy.hphonopy.file_IO import write_FORCE_SETS_vasp, read_force_constant_vasprun_xml

###########
#
# File I/O
#
###########

#
# Output
#
def write_supercells_with_displacements(supercell,
                                        double_displacements,
                                        amplitude=None):
    if amplitude==None:
        distance = 0.01
    else:
        distance = amplitude
    
    # YAML
    file = open('disp.yaml', 'w')
    file.write("natom: %d\n" %  supercell.get_number_of_atoms())
    file.write("num_first_displacements: %d\n" %  len(double_displacements))
    num_second = 0
    for d1 in double_displacements:
        for d2 in d1['second_atoms']:
            num_second += len(d2['directions'])
    file.write("num_second_displacements: %d\n" %  num_second)
    file.write("first_atoms:\n")
    lattice = supercell.get_cell()
    lattice_inv = np.linalg.inv(lattice)
    count1 = 1
    count2 = len(double_displacements) + 1
    for disp1 in double_displacements:
        disp_cartesian3 = np.dot(disp1['direction'], lattice)
        disp_cartesian3 = disp_cartesian3 / np.linalg.norm(disp_cartesian3) * distance
        positions = supercell.get_positions()
        positions[disp1['number']] += disp_cartesian3
        atoms = Atoms(numbers = supercell.get_atomic_numbers(),
                      masses = supercell.get_masses(),
                      positions = positions,
                      cell = lattice,
                      pbc = True)
        vasp.write_vasp('POSCAR-%04d' % count1, atoms, direct=True)

        # YAML
        file.write("- number: %5d\n" % (disp1['number'] + 1))
        file.write("  displacement:\n")
        file.write("    [ %20.16f,%20.16f,%20.16f ] # %d \n" %
                   (disp_cartesian3[0], disp_cartesian3[1], disp_cartesian3[2], count1))
        file.write("  second_atoms:\n")
        count1 += 1

        for disp2 in disp1['second_atoms']:
            # YAML
            file.write("  - number: %5d\n" % (disp2['number'] + 1))
            file.write("    displacements:\n")
            
            for direction in disp2['directions']:
                disp_cartesian2 = np.dot(direction, lattice)
                disp_cartesian2 = disp_cartesian2 / np.linalg.norm(disp_cartesian2) * distance
                positions = supercell.get_positions()
                positions[disp1['number']] += disp_cartesian3
                positions[disp2['number']] += disp_cartesian2
                atoms = Atoms(numbers = supercell.get_atomic_numbers(),
                               masses = supercell.get_masses(),
                               positions = positions,
                               cell = lattice,
                               pbc = True)
                vasp.write_vasp('POSCAR-%04d' % count2, atoms, direct=True)

                # YAML
                file.write("    - [ %20.16f,%20.16f,%20.16f ] # %d \n" %
                           (disp_cartesian2[0], disp_cartesian2[1], disp_cartesian2[2], count2))
                count2 += 1

    file.write("lattice:\n")
    for axis in supercell.get_cell():
        file.write("- [ %20.15f,%20.15f,%20.15f ]\n" % tuple(axis))
    symbols = supercell.get_chemical_symbols()
    positions = supercell.get_scaled_positions()
    file.write("atoms:\n")
    for i, (s, v) in enumerate(zip(symbols, positions)):
        file.write("- symbol: %-2s # %d\n" % (s, i+1))
        file.write("  position: [ %18.14f,%18.14f,%18.14f ]\n" % \
                       (v[0], v[1], v[2]))

    file.close()

def write_FORCES_THIRD(vaspruns,
                       displacements,
                       forces_third='FORCES_THIRD',
                       forces_second='FORCES_SECOND',
                       force_sets='FORCE_SETS'):
    natom = displacements['natom']
    num_disp1 = len(displacements['first_atoms'])
    disp_datasets = []
    set_of_forces = get_forces_from_vasprun_xmls(vaspruns, natom)
    filethird = open(forces_third, 'w')
    filesecond = open(forces_second, 'w')

    for i, disp1 in enumerate(displacements['first_atoms']):
        filesecond.write("# File: %-5d\n" % (i + 1))
        filesecond.write("# %-5d " % (disp1['number'] + 1))
        filesecond.write("%20.16f %20.16f %20.16f\n" %
                         tuple(disp1['displacement']))
        for f in set_of_forces[i]:
            filesecond.write("%15.10f %15.10f %15.10f\n" % (tuple(f)))
        
        disp_datasets.append([disp1['number'], disp1['displacement']])

    count = num_disp1
    for i, disp1 in enumerate(displacements['first_atoms']):
        for disp2 in disp1['second_atoms']:
            for d in disp2['displacements']:
                filethird.write("# File: %-5d\n" % (count + 1))
                filethird.write("# %-5d " % (disp1['number'] + 1))
                filethird.write("%20.16f %20.16f %20.16f\n" %
                                 tuple(disp1['displacement']))
                filethird.write("# %-5d " % (disp2['number'] + 1))
                filethird.write("%20.16f %20.16f %20.16f\n" % tuple(d))

                for f in set_of_forces[count]:
                    filethird.write("%15.10f %15.10f %15.10f\n" % (tuple(f)))
                count += 1

    write_FORCE_SETS_vasp(vaspruns[:num_disp1],
                          disp_datasets,
                          natom,
                          force_sets)

def write_DELTA_FC2_SETS(vaspruns,
                         displacements,
                         fc2_file='fc2.dat',
                         dfc2_file='DELTA_FC2_SETS'):
    fc2_set = get_force_constants_from_vasprun_xmls(vaspruns)
    w = open(dfc2_file, 'w')
    perfect_fc2 = fc2_set.pop(0)
    write_fc2_dat(perfect_fc2, fc2_file)
    for i, (dfc2, first_disp) in enumerate(
        zip(fc2_set, displacements['first_atoms'])):
        w.write("# File: %d\n" % (i + 1))
        w.write("# %-5d " % (first_disp['number'] + 1))
        w.write("%20.16f %20.16f %20.16f\n" %
                tuple(first_disp['displacement']))
{'fraction_non_alphanumeric': 0.13689265076910556, 'fraction_numerical': 0.037437942540896885, 'mean_word_length': 1.9531362653208364, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2932058', 'n_tokens_mistral': 4322, 'n_tokens_neox': 4004, 'n_words': 683}