<commit_before>'''
    Script to create .xyz file.
'''
import sys
import re
import commands
import argparse

from vaspy import atomco, matstudio

if "__main__" == __name__:

    # Set argument parser.
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--step",
                        help="step number on which the data is extracted.")
    args = parser.parse_args()

    if args.step:
        # extract certain step data to .xyz file
        step = args.step
        step_regex = re.compile(r'^STEP\s+=\s+' + step + r'$')
        with open('OUT.ANI', 'r') as f:
            natom = int(f.readline().strip())
            # match step
            content = ''
            count = 0
            line = f.readline()
            while line:
                if step_regex.match(line):
                    # get data content
                    for i in xrange(natom):
                        content += f.readline()
                    break
                line = f.readline()
            if not content:
                print 'Step: %s is out of range.' % step
                sys.exit(1)
        # write to .xyz file
        with open('ts.xyz', 'w') as f:
            head = '%12d\nSTEP = %8s\n' % (natom, step)
            content = head + content
            f.write(content)

        # coordinate transformation
        xyz = atomco.XyzFile('ts.xyz')
        poscar = atomco.PosCar()
        direct_coordinates = xyz.coordinate_transform(poscar.bases)
        suffix = '-' + step + '.xsd'
    else:  # the last step data
        contcar = atomco.ContCar()
        direct_coordinates = contcar.data
        suffix = '-y.xsd'

# create .xsd file
status, output = commands.getstatusoutput('ls *.xsd | head -1')
if not output.endswith('.xsd'):
    print "No .xsd file in current directory."
    sys.exit(1)
xsd = matstudio.XsdFile(filename=output)
xsd.data = direct_coordinates

jobname = output.split('.')[0]
xsd.tofile(filename=jobname+suffix)<commit_msg>Add energy and force update.<commit_after>'''
    Script to create .xyz file.
'''
import argparse
import commands
import logging
import re
import sys

from vaspy import atomco, matstudio
from vaspy.iter import OutCar, OsziCar

if "__main__" == __name__:

    # Set argument parser.
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--step",
                        help="step number on which the data is extracted.")
    args = parser.parse_args()

    if args.step:
        # extract certain step data to .xyz file
        step = args.step
        step_regex = re.compile(r'^STEP\s+=\s+' + step + r'$')
        with open('OUT.ANI', 'r') as f:
            natom = int(f.readline().strip())
            # match step
            content = ''
            count = 0
            line = f.readline()
            while line:
                if step_regex.match(line):
                    # get data content
                    for i in xrange(natom):
                        content += f.readline()
                    break
                line = f.readline()
            if not content:
                print 'Step: %s is out of range.' % step
                sys.exit(1)
        # write to .xyz file
        with open('ts.xyz', 'w') as f:
            head = '%12d\nSTEP = %8s\n' % (natom, step)
            content = head + content
            f.write(content)

        # coordinate transformation
        xyz = atomco.XyzFile('ts.xyz')
        poscar = atomco.PosCar()
        direct_coordinates = xyz.coordinate_transform(poscar.bases)
        suffix = '-' + step + '.xsd'
    else:  # the last step data
        contcar = atomco.ContCar()
        direct_coordinates = contcar.data
        suffix = '-y.xsd'

# create .xsd file
status, output = commands.getstatusoutput('ls *.xsd | head -1')
if not output.endswith('.xsd'):
    print "No .xsd file in current directory."
    sys.exit(1)
xsd = matstudio.XsdFile(filename=output)
xsd.data = direct_coordinates

# Get energy and force.
oszicar = OsziCar()
outcar = OutCar()
xsd.force = outcar.total_forces[-1]
logging.info("Total Force --> {}".format(xsd.force))
xsd.energy = oszicar.E0[-1]
logging.info("Total Energy --> {}".format(xsd.energy))

jobname = output.split('.')[0]
xsd.tofile(filename=jobname+suffix)
{'fraction_non_alphanumeric': 0.12300094073377234, 'fraction_numerical': 0.004468485418626529, 'mean_word_length': 1.8833898305084746, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '315425', 'n_tokens_mistral': 1327, 'n_tokens_neox': 1249, 'n_words': 354}