from ase.io import write
from ase.io import read
from ase import Atom, Atoms
from math import floor
import numpy as np
import sys
__author__ = "Alexander Gabourie"
__email__ = "gabourie@stanford.edu"
#########################################
# Helper Functions
#########################################
def __get_atom_line(atom, velocity, layer, groups, type_dict, info):
'''
Constructs an atom's line in an xyz.in file.
Args:
atom (ase.Atom):
Atom object to write to file.
velocity (bool):
If velocities need to be added.
layer (bool):
If the layer number needs to be added.
groups (bool):
If the groups need to be added.
type_dict (dict):
Dictionary to convert symbol to type number.
info (dict):
Dictionary that stores all velocity, layer, and groups data.
Returns:
out (str)
The line to be printed to file.
'''
optional = ''
if info:
try:
option = info[atom.index]
if velocity:
optional += ' ' + ' '.join([str(val) for val in option['velocity']])
if layer:
optional += ' ' + str(option['layer'])
if groups:
optional += ' ' + ' '.join([str(val) for val in option['groups']])
except KeyError:
pass
required = ' '.join([str(type_dict[atom.symbol])] + \
[str(val) for val in list(atom.position)] + \
[str(atom.mass)])
return required + optional
def __set_atoms(atoms, types):
"""
Sets the atom symbols for atoms loaded from GPUMD where in.xyz does not
contain that information
Args:
atoms (ase.Atoms):
Atoms object to change symbols in
types (list(str)):
List of strings to assign to atomic symbols
"""
for atom in atoms:
atom.symbol = types[atom.number]
def __atom_type_sortkey(atom, order=None):
"""
Used as a key for sorting atom type for GPUMD in.xyz files
Args:
atom (ase.Atom):
Atom object
order (list(str)):
A list of atomic symbol strings in the desired order.
"""
if order:
for i, sym in enumerate(order):
if sym == atom.symbol:
return i
else:
ValueError('type sortkey error: Missing order.')
def __atom_group_sortkey(atom, info=None, group_index=None, order=None):
"""
Used as a key for sorting atom groups for GPUMD in.xyz files
Args:
atom (ase.Atom):
Atom object
info (dict):
Info dictionary for Atoms object that 'atom' belongs to. Stores velocity,
groups, & layer information
group_index (int):
Index of the grouping list that is part of the 'groups' key for the atom.index
element from the info dictionary.
order (list(int)):
A list of ints in desired order for groups at group_index
"""
if not (info and not group_index is None):
ValueError('group sortkey error: Missing either info or group_index.')
if order:
for i, group in enumerate(order):
if group == info[atom.index]['groups'][group_index]:
return i
else:
return info[atom.index]['groups'][group_index]
return sys.maxsize
def __atom_layer_sortkey(atom, info=None, order=None):
"""
Used as a key for sorting atom layers for GPUMD in.xyz files
Args:
atom (ase.Atom):
Atom object
info (dict):
Info dictionary for Atoms object that 'atom' belongs to. Stores velocity,
groups, & layer information
order (list(int)):
A list of ints in desired order for layers
"""
if not info:
ValueError('layer sortkey error: Missing info.')
if order:
for i, layer in enumerate(order):
if layer == info[atom.index]['layer']:
return i
else:
return info[atom.index]['layer']
return sys.maxsize
#########################################
# Read Related
#########################################
[docs]def load_xyz(filename='xyz.in', atom_types=None):
"""
Reads and returns the structure input file from GPUMD.
Args:
filename (str):
Name of structure file
atom_types (list(str)):
List of atom types (elements).
Returns:
atoms (ase.Atoms):
ASE atoms object with x,y,z, mass, group, type, cell, and PBCs
from input file. group is stored in tag, atom type may not
correspond to correct atomic symbol
M (int):
Max number of neighbor atoms
cutoff (float):
Initial cutoff for neighbor list build
"""
# read file
with open(filename) as f:
xyz_lines = f.readlines()
# get global structure params
l1 = tuple(xyz_lines[0].split()) # first line
N, M, use_triclinic, has_velocity, \
has_layer, num_of_groups = [int(val) for val in l1[:2]+l1[3:]]
cutoff = float(l1[2])
l2 = tuple(xyz_lines[1].split()) # second line
if use_triclinic:
pbc, cell = [int(val) for val in l2[:3]], [float(val) for val in l2[3:]]
else:
pbc, L = [int(val) for val in l2[:3]], [float(val) for val in l2[3:]]
# get atomic params
info = dict()
atoms = Atoms()
atoms.set_pbc((pbc[0], pbc[1], pbc[2]))
if use_triclinic:
atoms.set_cell(np.array(cell).reshape((3,3)))
else:
atoms.set_cell([(L[0], 0, 0), (0, L[1], 0), (0, 0, L[2])])
for index, line in enumerate(xyz_lines[2:]):
data = dict()
lc = tuple(line.split()) # current line
type_, mass = int(lc[0]), float(lc[4])
position = [float(val) for val in lc[1:4]]
atom = Atom(type_, position, mass=mass)
lc = lc[5:] # reduce list length for easier indexing
if has_velocity:
velocity = [float(val) for val in lc[:3]]
lc = lc[3:]
data['velocity'] = velocity
if has_layer:
layer = int(lc[0])
lc = lc[1:]
data['layer'] = layer
if num_of_groups:
groups = [int(group) for group in lc]
data['groups'] = groups
atoms.append(atom)
info[index] = data
atoms.info = info
if atom_types:
__set_atoms(atoms, atom_types)
return atoms, M, cutoff
[docs]def import_trajectory(filename='movie.xyz', in_file=None, atom_types=None):
"""
Reads the trajectory from GPUMD run and creates a list of ASE atoms.
Args:
filename (str):
Name of the file that holds the GPUMD trajectory.
in_file (str):
Name of the original structure input file. Not required, but
can help load extra information not included in trajectory output.
atom_types (list(str)):
List of atom types (elements).
Returns:
traj (list(ase.Atoms)):
A list of ASE atoms objects.
"""
# get extra information about system if wanted
if in_file:
atoms, _, _ = load_xyz(in_file,atom_types)
pbc = atoms.get_pbc()
else:
pbc = None
with open(filename, 'r') as f:
xyz_line = f.readlines()
num_atoms = int(xyz_line[0])
block_size = num_atoms + 2
num_blocks = len(xyz_line)/block_size
traj = list()
for block in range(num_blocks):
types = []
positions = []
## TODO Loop may be inefficient in accessing xyz_line
for entry in xyz_line[block_size*block+2:block_size*(block+1)]:
# type_ can be an atom number or index to atom_types
type_, x, y, z = entry.split()[:4]
positions.append([float(x), float(y), float(z)])
if atom_types:
types.append(atom_types[int(type_)])
else:
types.append(int(type_))
if atom_types:
traj.append(Atoms(symbols=types, positions=positions, pbc=pbc))
else:
traj.append(Atoms(numbers=types, positions=positions, pbc=pbc))
return traj
#########################################
# Write Related
#########################################
[docs]def convert_gpumd_atoms(in_file='xyz.in', out_filename='in.xyz',
format='xyz', atom_types=None):
"""
Converts the GPUMD input structure file to any compatible ASE
output structure file.
Warning: Info dictionary may not be preserved.
Args:
in_file (str):
GPUMD position file to get structure from
out_filename (str):
Name of output file after conversion
format (str):
ASE supported output format
atom_types (list(str)):
List of atom types (elements).
"""
atoms, M, cutoff = load_xyz(in_file, atom_types)
write(out_filename, atoms, format)
return
[docs]def convert_gpumd_trajectory(traj_file='xyz.out', out_filename='out.xyz',
in_file='xyz.in', format='xyz'):
"""
Converts GPUMD trajectory to any compatible ASE output. Default: xyz
Args:
traj_file (str):
Trajetory from GPUMD
out_filename (str):
File in which final trajectory should be saved
in_file (str):
Original stucture input file to GPUMD. Needed to get atom
numbers/types
format (str):
ASE supported format
"""
traj = load_traj(traj_file, in_file)
write(out_filename, traj, format)
return
[docs]def lammps_atoms_to_gpumd(filename, M, cutoff, style='atomic',
gpumd_file='xyz.in'):
"""
Converts a lammps data file to GPUMD compatible position file
Args:
filename (str):
LAMMPS data file name
M (int):
Maximum number of neighbors for one atom
cutoff (float):
Initial cutoff distance for building the neighbor list
style (str):
Atom style used in LAMMPS data file
gpumd_file (str):
File to save the structure data to
"""
# Load atoms
atoms = read(filename, format='lammps-data', style=style)
ase_atoms_to_gpumd(atoms, M, cutoff, gpumd_file=gpumd_file)
return
[docs]def ase_atoms_to_gpumd(atoms, M, cutoff, gpumd_file='xyz.in', sort_key=None,
order=None, group_index=None):
"""
Converts ASE atoms to GPUMD compatible position file
Args:
atoms (ase.Atoms):
Atoms to write to gpumd file
M (int):
Maximum number of neighbors for one atom
cutoff (float):
Initial cutoff distance for building the neighbor list
gpumd_file (str):
File to save the structure data to
sort_key (str):
How to sort atoms ('group', 'type', 'layer'). Default is None.
order (list(type)):
List to sort by. Provide str for 'type', and int for 'group'
and 'layer'
group_index (int):
Selects the group to sort in the output.
"""
info = atoms.info # info dictionary that stores velocities, groups, layers
# sort atoms by desired property
if sort_key == 'type':
atoms_list = sorted(atoms, key=lambda x: __atom_type_sortkey(x, order))
elif sort_key == 'group':
atoms_list = sorted(atoms, key=lambda x: __atom_group_sortkey(x, info, group_index, order))
elif sort_key == 'layer':
atoms_list = sorted(atoms, key=lambda x: __atom_layer_sortkey(x, info, order))
else:
atoms_list = atoms
# set order of types
if sort_key=='type' and atom_order:
types = atom_order
else:
types = list(set(atoms.get_chemical_symbols()))
type_dict = dict()
for i, type_ in enumerate(types):
type_dict[type_] = i
# assume info[0] has same keys and number of groups as all other indices
num_groups = 0
if info and (0 in info.keys()):
infokeys = info[0].keys()
velocity = 'velocity' in infokeys
layer = 'layer' in infokeys
if 'groups' in infokeys:
groups = True
num_groups = str(len(info[0]['groups']))
else:
groups = False
else:
velocity = 0
layer = 0
groups = 0
# prepare cell to write
N = len(atoms)
pbc = [str(1) if val else str(0) for val in atoms.get_pbc()]
lx, ly, lz, a1, a2, a3 = tuple(atoms.get_cell_lengths_and_angles())
summary = ' '.join([str(N), str(M), str(cutoff), '@', \
'1' if velocity else '0', \
'1' if layer else '0', \
num_groups if groups else '0','\n'])
# if orthorhombic
if (a1 == a2 == a3):
summary = summary.replace('@', '0')
lx, ly, lz = str(lx), str(ly), str(lz)
summary += ' '.join(pbc + [lx, ly, lz] + ['\n'])
else: # if triclinic
summary = summary.replace('@', '1')
cell_str_vec = [str(val) for val in atoms.get_cell().flatten()]
summary += ' '.join(pbc + cell_str_vec + ['\n'])
# write structure
with open(gpumd_file, 'w') as f:
f.writelines(summary)
for atom in atoms_list[:-1]:
line = __get_atom_line(atom, velocity, layer, groups, type_dict, info)
f.writelines(line + '\n')
# Last line
atom = atoms_list[-1]
line = __get_atom_line(atom, velocity, layer, groups, type_dict, info)
f.writelines(line)
return