<commit_before>import os
import numpy as np
import h5py
from phonopy.interface.vasp import get_forces_from_vasprun_xmls

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

def write_disp_fc4_yaml(dataset, supercell, filename='disp_fc4.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)

    num_second = 0
    for d1 in dataset['first_atoms']:
        num_second += len(d1['second_atoms'])
    w.write("num_second_displacements: %d\n" %  num_second)

    num_third = 0
    for d1 in dataset['first_atoms']:
        for d2 in d1['second_atoms']:
            num_third += len(d2['third_atoms'])
    w.write("num_third_displacements: %d\n" %  num_third)

    w.write("first_atoms:\n")
    count1 = 1
    count2 = num_first + 1
    count3 = num_first + num_second + 1
    for disp1 in dataset['first_atoms']:
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], count1))
        w.write("  second_atoms:\n")
        count1 += 1
        for disp2 in disp1['second_atoms']:
            w.write("  - number: %5d\n" % (disp2['number'] + 1))
            w.write("    displacement:\n")
            disp_cart2 = disp2['displacement']
            w.write("      [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                    (disp_cart2[0], disp_cart2[1], disp_cart2[2], count2))
            w.write("    third_atoms:\n")
            count2 += 1
            atom3 = -1
            for disp3 in disp2['third_atoms']:
                if atom3 != disp3['number']:
                    atom3 = disp3['number']
                    w.write("    - number: %5d\n" % (atom3 + 1))
                    w.write("      displacements:\n")
                disp_cart3 = disp3['displacement']
                w.write("      - [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                        (disp_cart3[0], disp_cart3[1], disp_cart3[2], count3))
                count3 += 1

    write_cell_yaml(w, supercell)

    w.close()

    return num_first + num_second + num_third

def write_disp_fc3_yaml(dataset, supercell, filename='disp_fc3.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)
    if 'cutoff_distance' in dataset:
        w.write("cutoff_distance: %f\n" %  dataset['cutoff_distance'])

    num_second = 0
    num_disp_files = 0
    for d1 in dataset['first_atoms']:
        num_disp_files += 1
        num_second += len(d1['second_atoms'])
        for d2 in d1['second_atoms']:
            if 'included' in d2:
                if d2['included']:
                    num_disp_files += 1
            else:
                num_disp_files += 1

    w.write("num_second_displacements: %d\n" %  num_second)
    w.write("num_displacements_created: %d\n" %  num_disp_files)

    w.write("first_atoms:\n")
    count1 = 1
    count2 = num_first + 1
    for disp1 in dataset['first_atoms']:
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], count1))
        w.write("  second_atoms:\n")
        count1 += 1

        included = None
        distance = 0.0
        atom2 = -1
        for disp2 in disp1['second_atoms']:
            if atom2 != disp2['number']:
                atom2 = disp2['number']
                if 'included' in disp2:
                    included = disp2['included']
                pair_distance = disp2['pair_distance']
                w.write("  - number: %5d\n" % (atom2 + 1))
                w.write("    distance: %f\n" % pair_distance)
                if included is not None:
                    if included:
                        w.write("    included: %s\n" % "true")
                    else:
                        w.write("    included: %s\n" % "false")
                w.write("    displacements:\n")

            disp_cart2 = disp2['displacement']
            w.write("    - [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                    (disp_cart2[0], disp_cart2[1], disp_cart2[2], count2))
            count2 += 1

    write_cell_yaml(w, supercell)

    w.close()

    return num_first + num_second, num_disp_files

def write_disp_fc2_yaml(dataset, supercell, filename='disp_fc2.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)
    w.write("first_atoms:\n")
    for i, disp1 in enumerate(dataset['first_atoms']):
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], i + 1))

    write_cell_yaml(w, supercell)

    w.close()

    return num_first

def write_FORCES_FC4_vasp(vaspruns,
                          disp_dataset,
                          filename='FORCES_FC4'):
    natom = disp_dataset['natom']
    forces = get_forces_from_vasprun_xmls(vaspruns, natom)
    w = open(filename, 'w')
    write_FORCES_FC4(disp_dataset, forces, fp=w)
    w.close()

def write_FORCES_FC3_vasp(vaspruns,
                          disp_dataset,
                          filename='FORCES_FC3'):
    natom = disp_dataset['natom']
    forces = get_forces_from_vasprun_xmls(vaspruns, natom)
    w = open(filename, 'w')
    write_FORCES_FC3(disp_dataset, forces, fp=w)
    w.close()

def write_FORCES_FC2_vasp(vaspruns,
                          disp_dataset,
                          filename='FORCES_FC2'):
    natom = disp_dataset['natom']
    forces_fc2 = get_forces_from_vasprun_xmls(vaspruns, natom)
    w = open(filename, 'w')
    write_FORCES_FC2(disp_dataset, forces_fc2=forces_fc2, fp=w)
    w.close()

def write_FORCES_FC2(disp_dataset,
                     forces_fc2=None,
                     fp=None,
                     filename="FORCES_FC2"):
    if fp is None:
        w = open(filename, 'w')
    else:
        w = fp

    for i, disp1 in enumerate(disp_dataset['first_atoms']):
        w.write("# File: %-5d\n" % (i + 1))
        w.write("# %-5d " % (disp1['number'] + 1))
        w.write("%20.16f %20.16f %20.16f\n" % tuple(disp1['displacement']))
        if forces_fc2 is None:
            force_set = disp1['forces']
        else:
            force_set = forces_fc2[i]
        for forces in force_set:<commit_msg>Remove useless method to write FORCE_FCX<commit_after>import os
import numpy as np
import h5py

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

def write_disp_fc4_yaml(dataset, supercell, filename='disp_fc4.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)

    num_second = 0
    for d1 in dataset['first_atoms']:
        num_second += len(d1['second_atoms'])
    w.write("num_second_displacements: %d\n" %  num_second)

    num_third = 0
    for d1 in dataset['first_atoms']:
        for d2 in d1['second_atoms']:
            num_third += len(d2['third_atoms'])
    w.write("num_third_displacements: %d\n" %  num_third)

    w.write("first_atoms:\n")
    count1 = 1
    count2 = num_first + 1
    count3 = num_first + num_second + 1
    for disp1 in dataset['first_atoms']:
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], count1))
        w.write("  second_atoms:\n")
        count1 += 1
        for disp2 in disp1['second_atoms']:
            w.write("  - number: %5d\n" % (disp2['number'] + 1))
            w.write("    displacement:\n")
            disp_cart2 = disp2['displacement']
            w.write("      [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                    (disp_cart2[0], disp_cart2[1], disp_cart2[2], count2))
            w.write("    third_atoms:\n")
            count2 += 1
            atom3 = -1
            for disp3 in disp2['third_atoms']:
                if atom3 != disp3['number']:
                    atom3 = disp3['number']
                    w.write("    - number: %5d\n" % (atom3 + 1))
                    w.write("      displacements:\n")
                disp_cart3 = disp3['displacement']
                w.write("      - [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                        (disp_cart3[0], disp_cart3[1], disp_cart3[2], count3))
                count3 += 1

    write_cell_yaml(w, supercell)

    w.close()

    return num_first + num_second + num_third

def write_disp_fc3_yaml(dataset, supercell, filename='disp_fc3.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)
    if 'cutoff_distance' in dataset:
        w.write("cutoff_distance: %f\n" %  dataset['cutoff_distance'])

    num_second = 0
    num_disp_files = 0
    for d1 in dataset['first_atoms']:
        num_disp_files += 1
        num_second += len(d1['second_atoms'])
        for d2 in d1['second_atoms']:
            if 'included' in d2:
                if d2['included']:
                    num_disp_files += 1
            else:
                num_disp_files += 1

    w.write("num_second_displacements: %d\n" %  num_second)
    w.write("num_displacements_created: %d\n" %  num_disp_files)

    w.write("first_atoms:\n")
    count1 = 1
    count2 = num_first + 1
    for disp1 in dataset['first_atoms']:
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], count1))
        w.write("  second_atoms:\n")
        count1 += 1

        included = None
        distance = 0.0
        atom2 = -1
        for disp2 in disp1['second_atoms']:
            if atom2 != disp2['number']:
                atom2 = disp2['number']
                if 'included' in disp2:
                    included = disp2['included']
                pair_distance = disp2['pair_distance']
                w.write("  - number: %5d\n" % (atom2 + 1))
                w.write("    distance: %f\n" % pair_distance)
                if included is not None:
                    if included:
                        w.write("    included: %s\n" % "true")
                    else:
                        w.write("    included: %s\n" % "false")
                w.write("    displacements:\n")

            disp_cart2 = disp2['displacement']
            w.write("    - [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                    (disp_cart2[0], disp_cart2[1], disp_cart2[2], count2))
            count2 += 1

    write_cell_yaml(w, supercell)

    w.close()

    return num_first + num_second, num_disp_files

def write_disp_fc2_yaml(dataset, supercell, filename='disp_fc2.yaml'):
    w = open(filename, 'w')
    w.write("natom: %d\n" %  dataset['natom'])

    num_first = len(dataset['first_atoms'])
    w.write("num_first_displacements: %d\n" %  num_first)
    w.write("first_atoms:\n")
    for i, disp1 in enumerate(dataset['first_atoms']):
        disp_cart1 = disp1['displacement']
        w.write("- number: %5d\n" % (disp1['number'] + 1))
        w.write("  displacement:\n")
        w.write("    [%20.16f,%20.16f,%20.16f ] # %05d\n" %
                (disp_cart1[0], disp_cart1[1], disp_cart1[2], i + 1))

    write_cell_yaml(w, supercell)

    w.close()

    return num_first

def write_FORCES_FC2(disp_dataset,
                     forces_fc2=None,
                     fp=None,
                     filename="FORCES_FC2"):
    if fp is None:
        w = open(filename, 'w')
    else:
        w = fp

    for i, disp1 in enumerate(disp_dataset['first_atoms']):
        w.write("# File: %-5d\n" % (i + 1))
        w.write("# %-5d " % (disp1['number'] + 1))
        w.write("%20.16f %20.16f %20.16f\n" % tuple(disp1['displacement']))
        if forces_fc2 is None:
            force_set = disp1['forces']
        else:
            force_set = forces_fc2[i]
        for forces in force_set:
{'fraction_non_alphanumeric': 0.1521434368774441, 'fraction_numerical': 0.042057109127130524, 'mean_word_length': 2.068598596332352, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 4}, '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': '1734562', 'n_tokens_mistral': 5151, 'n_tokens_neox': 4815, 'n_words': 896}