def read_json(filedesc, **kwargs):
"""Reads a JSON-style file with i-pi style comments and creates an Atoms and Cell object.
Args:
filedesc: An open readable file object from a json formatted file with i-PI header comments.
Returns:
An Atoms object with the appropriate atom labels, masses and positions.
A Cell object.
"""
try:
line = json.loads(filedesc.readline())
except ValueError:
raise EOFError("The file descriptor hit EOF.")
atoms = Atoms(line[0])
atoms.q = np.asarray(line[8])
atoms.names = np.asarray(line[9], dtype='|S4')
atoms.m = np.asarray(map(Elements.mass, atoms.names))
a = float(line[1])
b = float(line[2])
c = float(line[3])
alpha = float(line[4]) * np.pi / 180
beta = float(line[5]) * np.pi / 180
gamma = float(line[6]) * np.pi / 180
h = mt.abc2h(a, b, c, alpha, beta, gamma)
cell = Cell(h)
return {"atoms": atoms, "cell": cell}
def create_random_xyz_traj_to_write(request):
natoms, frames, comment, expected_cell, precision = request.param
a, b, c, alpha, beta, gamma = mt.h2abc_deg(expected_cell)
fmt_header = "# CELL(abcABC): %10.5f %10.5f %10.5f %10.5f %10.5f %10.5f %s"
comment = fmt_header % (a, b, c, alpha, beta, gamma, comment)
filedesc, xyz, atom_names = xyz_gen.xyz_traj_filedesc(natoms, frames, comment)
filedesc.seek(0)
masses = [Elements.mass(_am) for _am in atom_names]
cell_list = []
atoms_list = []
for _fr in xrange(frames):
cell = Cell(expected_cell)
atoms = Atoms(natoms)
atoms.q[:] = xyz[_fr * natoms * 3:(_fr + 1) * natoms * 3]
atoms.names = atom_names[_fr * natoms:(_fr + 1) * natoms]
atoms.m[:] = masses[_fr * natoms:(_fr + 1) * natoms]
atoms_list.append(atoms)
cell_list.append(cell)
return (filedesc, atoms_list, cell_list, comment, precision)
def process_units(comment, cell, data, names, masses, natoms, dimension="automatic", units="automatic", cell_units="automatic", mode="xyz"):
"""Convert the data in the file according to the units written in the i-PI format.
Args:
comment:
cell:
data:
names:
masses:
output:
Returns:
"""
dimension, units, cell_units = auto_units(comment, dimension, units, cell_units, mode)
info("Interpreting input with dimension %s, units %s and cell units %s" % (dimension, units, cell_units), verbosity.high)
# Units transformation
cell *= unit_to_internal('length', cell_units, 1) # cell units transformation
data *= unit_to_internal(dimension, units, 1) # units transformation
# Return data as i-PI structures
cell = Cell(cell)
atoms = Atoms(natoms)
atoms.q[:] = data
atoms.names[:] = names
atoms.m[:] = masses
return {
"atoms": atoms,
"cell": cell,
}
def main(prefix, suffix="pos", outprefix="fixcom"):
ipos = []
imode = []
for filename in sorted(glob.glob(prefix + "." + suffix + "*")):
imode.append(filename.split(".")[-1])
ipos.append(open(filename, "r"))
nbeads = len(ipos)
natoms = 0
ifr = 0
lout = []
for b in range(nbeads):
# zero-padded bead number
padb = ("%0" + str(int(1 + np.floor(np.log(nbeads) / np.log(10)))) + "d") % (b)
lout.append(
open(
prefix + "." + outprefix + "." + suffix + "_" + padb + "." + imode[b],
"a",
)
)
while True:
allbeads = []
for i in range(nbeads):
try:
poscell = read_file(imode[i], ipos[i])
cell = poscell["cell"]
pos = poscell["atoms"]
allbeads.append(pos)
if natoms == 0:
natoms = pos.natoms
atoms = Atoms(natoms)
atoms.q += pos.q
atoms.names = pos.names
atoms.m += pos.m
except EOFError: # finished reading files
for ib in range(nbeads):
lout[ib].close()
sys.exit(0)
atoms.q /= nbeads
atoms.m /= nbeads
com = np.zeros(3)
tm = 0
for i in range(atoms.natoms):
com += atoms.m[i] * atoms.q[3 * i : 3 * (i + 1)]
tm += atoms.m[i]
com /= tm
for ib in range(nbeads):
for i in range(allbeads[ib].natoms):
allbeads[ib].q[3 * i : 3 * (i + 1)] -= com
print_file(imode[ib], allbeads[ib], cell, filedesc=lout[ib])
atoms.q[:] = 0.0
atoms.m[:] = 0.0
ifr += 1
def read_pdb(filedesc):
"""Takes a pdb-style file and creates an Atoms and Cell object.
Args:
filedesc: An open readable file object from a pdb formatted file.
Returns:
An Atoms object with the appropriate atom labels, masses and positions,
and a Cell object with the appropriate cell dimensions and an estimate
of a reasonable cell mass.
"""
header = filedesc.readline()
if "TITLE" in header:
header = filedesc.readline() # skip the comment field
if header == "":
raise EOFError("End of file or empty header in PDB file")
a = float(header[6:15])
b = float(header[15:24])
c = float(header[24:33])
alpha = float(header[33:40])
beta = float(header[40:47])
gamma = float(header[47:54])
alpha *= np.pi / 180.0
beta *= np.pi / 180.0
gamma *= np.pi / 180.0
h = mt.abc2h(a, b, c, alpha, beta, gamma)
cell = Cell(h)
natoms = 0
body = filedesc.readline()
qatoms = []
names = []
masses = []
while (body.strip() != "" and body.strip() != "END"):
natoms += 1
name = body[12:16].strip()
names.append(name)
masses.append(Elements.mass(name))
x = float(body[31:39])
y = float(body[39:47])
z = float(body[47:55])
qatoms.append(x)
qatoms.append(y)
qatoms.append(z)
body = filedesc.readline()
atoms = Atoms(natoms)
atoms.q = np.asarray(qatoms)
atoms.names = np.asarray(names, dtype='|S4')
atoms.m = np.asarray(masses)
return atoms, cell
def read_pdb(filedesc):
"""Takes a pdb-style file and creates an Atoms and Cell object.
Args:
filedesc: An open readable file object from a pdb formatted file.
Returns:
An Atoms object with the appropriate atom labels, masses and positions,
and a Cell object with the appropriate cell dimensions and an estimate
of a reasonable cell mass.
"""
header = filedesc.readline()
if "TITLE" in header: header = filedesc.readline() # skip the comment field
if header == "":
raise EOFError("End of file or empty header in PDB file")
a = float(header[6:15])
b = float(header[15:24])
c = float(header[24:33])
alpha = float(header[33:40])
beta = float(header[40:47])
gamma = float(header[47:54])
alpha *= np.pi/180.0
beta *= np.pi/180.0
gamma *= np.pi/180.0
h = mt.abc2h(a, b, c, alpha, beta, gamma)
cell = Cell(h)
natoms = 0
body = filedesc.readline()
qatoms = []
names = []
masses = []
while (body.strip() != "" and body.strip() != "END"):
natoms += 1
name = body[12:16].strip()
names.append(name)
masses.append(Elements.mass(name))
x = float(body[31:39])
y = float(body[39:47])
z = float(body[47:55])
qatoms.append(x)
qatoms.append(y)
qatoms.append(z)
body = filedesc.readline()
atoms = Atoms(natoms)
atoms.q = np.asarray(qatoms)
atoms.names = np.asarray(names,dtype='|S4')
atoms.m = np.asarray(masses)
return atoms, cell
def read_xyz(filedesc):
"""Takes a xyz-style file and creates an Atoms object.
Args:
filedesc: An open readable file object from a xyz formatted file.
Returns:
An Atoms object with the appropriate atom labels, masses and positions.
"""
natoms = filedesc.readline()
if natoms == "":
raise EOFError("The file descriptor hit EOF.")
natoms = int(natoms)
comment = filedesc.readline()
qatoms = []
names = []
masses = []
iat = 0
while (iat < natoms):
body = filedesc.readline()
if body.strip() == "":
break
body = body.split()
name = body[0]
names.append(name)
masses.append(Elements.mass(name))
x = float(body[1])
y = float(body[2])
z = float(body[3])
qatoms.append(x)
qatoms.append(y)
qatoms.append(z)
iat += 1
if natoms != len(names):
raise ValueError(
"The number of atom records does not match the header of the xyz file."
)
atoms = Atoms(natoms)
# for i in range(natoms):
# nat = atoms[i]
# nat.q = qatoms[i]
# nat.name = names[i]
# nat.m = Elements.mass(names[i])
atoms.q = np.asarray(qatoms)
atoms.names = np.asarray(names, dtype='|S4')
atoms.m = np.asarray(masses)
return atoms
def main(filename, natoms):
ipos=open(filename,"r")
imode=filename[-3:]
natoms = int(natoms)
ifr = 0
nn = 2.5
while True:
try:
ret = read_file(imode,ipos,readcell=True)
pos = ret["atoms"]
cell = ret["cell"]
q=depstrip(pos.q).copy()
cell.array_pbc(q)
natin = pos.natoms
q.shape=(natin,3)
s=np.dot(depstrip(cell.ih),q.T).T
# now replicate in scaled coordinates
nrep = int((natoms/natin*nn)**(1./3.))
natrep = natin*(2*nrep+1)**3
ns = np.zeros((natrep,3))
ik = 0
for ix in range(-nrep,nrep+1):
for iy in range(-nrep,nrep+1):
for iz in range(-nrep,nrep+1):
for i in range(natin):
ns[ik] = s[i]+[ix,iy,iz]
ik+=1
ns = np.dot(depstrip(cell.h),ns.T).T
# now removes atoms until we only have natoms
d = np.zeros(natrep)
for i in range(natrep):
d[i] = np.sqrt(np.dot(ns[i],ns[i]))
di = np.argsort(d)
npos = Atoms(natoms)
for i in range(natoms):
npos.q[3*i:3*(i+1)]=ns[di[i]]
except EOFError: # finished reading files
sys.exit(0)
print_file("pdb",npos, cell)
ifr+=1
def main(filename, natoms):
ipos = open(filename, "r")
imode = filename[-3:]
natoms = int(natoms)
ifr = 0
nn = 2.5
while True:
try:
ret = read_file(imode, ipos, readcell=True)
pos = ret["atoms"]
cell = ret["cell"]
q = depstrip(pos.q).copy()
cell.array_pbc(q)
natin = pos.natoms
q.shape = (natin, 3)
s = np.dot(depstrip(cell.ih), q.T).T
# now replicate in scaled coordinates
nrep = int((natoms / natin * nn)**(1. / 3.))
natrep = natin * (2 * nrep + 1)**3
ns = np.zeros((natrep, 3))
ik = 0
for ix in range(-nrep, nrep + 1):
for iy in range(-nrep, nrep + 1):
for iz in range(-nrep, nrep + 1):
for i in range(natin):
ns[ik] = s[i] + [ix, iy, iz]
ik += 1
ns = np.dot(depstrip(cell.h), ns.T).T
# now removes atoms until we only have natoms
d = np.zeros(natrep)
for i in range(natrep):
d[i] = np.sqrt(np.dot(ns[i], ns[i]))
di = np.argsort(d)
npos = Atoms(natoms)
for i in range(natoms):
npos.q[3 * i:3 * (i + 1)] = ns[di[i]]
except EOFError: # finished reading files
sys.exit(0)
print_file("pdb", npos, cell)
ifr += 1
def read_xyz(filedesc):
"""Takes a xyz-style file and creates an Atoms object.
Args:
filedesc: An open readable file object from a xyz formatted file.
Returns:
An Atoms object with the appropriate atom labels, masses and positions.
"""
natoms = filedesc.readline()
if natoms == "":
raise EOFError("The file descriptor hit EOF.")
natoms = int(natoms)
comment = filedesc.readline()
qatoms = []
names = []
masses = []
iat = 0
while iat < natoms:
body = filedesc.readline()
if body.strip() == "":
break
body = body.split()
name = body[0]
names.append(name)
masses.append(Elements.mass(name))
x = float(body[1])
y = float(body[2])
z = float(body[3])
qatoms.append(x)
qatoms.append(y)
qatoms.append(z)
iat += 1
if natoms != len(names):
raise ValueError("The number of atom records does not match the header of the xyz file.")
atoms = Atoms(natoms)
# for i in range(natoms):
# nat = atoms[i]
# nat.q = qatoms[i]
# nat.name = names[i]
# nat.m = Elements.mass(names[i])
atoms.q = np.asarray(qatoms)
atoms.names = np.asarray(names, dtype="|S4")
atoms.m = np.asarray(masses)
return atoms
def process_units(
comment,
cell,
data,
names,
masses,
natoms,
dimension="automatic",
units="automatic",
cell_units="automatic",
mode="xyz",
):
"""Convert the data in the file according to the units written in the i-PI format.
Args:
comment:
cell:
data:
names:
masses:
output:
Returns:
"""
dimension, units, cell_units = auto_units(comment, dimension, units,
cell_units, mode)
info(
" # Interpreting input with dimension %s, units %s and cell units %s" %
(dimension, units, cell_units),
verbosity.high,
)
# Units transformation
cell *= unit_to_internal("length", cell_units,
1) # cell units transformation
data *= unit_to_internal(dimension, units, 1) # units transformation
# Return data as i-PI structures
cell = Cell(cell)
atoms = Atoms(natoms)
atoms.q[:] = data
atoms.names[:] = names
atoms.m[:] = masses
return {
"atoms": atoms,
"cell": cell,
}
def resize(self, natoms, nbeads):
"""Creates all the data arrays needed in the simulation.
Effectively initializes the whole Beads object, according to the
specified number of atoms and beads. Is also used, as the name suggests,
to resize the data to a new number of beads when this is necessary, for
example in initialization from a simulation with a different number of
beads.
Also creates, or recreates, the dependency network, as this requires
the data arrays to be created for it to work.
Args:
natoms: The number of atoms.
nbeads: The number of beads.
"""
self.natoms = natoms
self.nbeads = nbeads
dself = dd(self)
dself.names = depend_array(name="names", value=np.zeros(natoms, np.dtype('|U6')))
# atom masses, and mass-related arrays
dself.m = depend_array(name="m", value=np.zeros(natoms, float)) # this is the prototype mass array (just one independent of bead n)
dself.m3 = depend_array(name="m3", value=np.zeros((nbeads, 3 * natoms), float), # this is m conveniently replicated to be (nb,3*nat)
func=self.mtom3, dependencies=[dself.m])
dself.sm3 = depend_array(name="sm3", value=np.zeros((nbeads, 3 * natoms), float), # this is just the square root of m3
func=self.m3tosm3, dependencies=[dself.m3])
# positions and momenta. bead representation, base storage used everywhere
dself.q = depend_array(name="q", value=np.zeros((nbeads, 3 * natoms), float))
dself.p = depend_array(name="p", value=np.zeros((nbeads, 3 * natoms), float))
# position and momentum of the centroid
dself.qc = depend_array(name="qc", value=np.zeros(3 * natoms, float),
func=self.get_qc, dependencies=[dself.q])
dself.pc = depend_array(name="pc", value=np.zeros(3 * natoms, float),
func=self.get_pc, dependencies=[dself.p])
# path springs potential and force
dself.vpath = depend_value(name="vpath", func=self.get_vpath,
dependencies=[dself.q, dself.m3])
dself.fpath = depend_array(name="fpath", value=np.zeros((nbeads, 3 * natoms), float),
func=self.get_fpath, dependencies=[dself.q])
# create proxies to access the individual beads as Atoms objects
# TODO: ACTUALLY THIS IS ONLY USED HERE METHINK, SO PERHAPS WE COULD REMOVE IT TO DECLUTTER THE CODE.
self._blist = [Atoms(natoms, _prebind=(self.q[i, :], self.p[i, :], self.m, self.names)) for i in range(nbeads)]
# kinetic energies of thhe beads, and total (classical) kinetic stress tensor
dself.kins = depend_array(name="kins", value=np.zeros(nbeads, float),
func=self.kin_gather,
dependencies=[dd(b).kin for b in self._blist])
dself.kin = depend_value(name="kin", func=self.get_kin,
dependencies=[dself.kins])
dself.kstress = depend_array(name="kstress", value=np.zeros((3, 3), float),
func=self.get_kstress,
dependencies=[dd(b).kstress for b in self._blist])
def create_xyz_sample_file(request):
""" Create a fake xyz file and build the atoms and cell object from it.
"""
natoms, frames, comment, expected_cell, units_conv_at, units_conv_cell = request.param
filedesc, xyz, atoms_names = xyz_gen.xyz_traj_filedesc(natoms, frames, comment)
# init_file needs to read from a real file...
tmp_file = tmp.NamedTemporaryFile(mode='wr', prefix='ipi_testing-tmp', delete=False)
tmp_file.seek(0)
tmp_file.write(filedesc.read())
tmp_file.close()
filedesc.close()
masses = np.zeros(natoms * frames)
for _ii, _at in enumerate(atoms_names):
masses[_ii] = Elements.mass(_at)
ratoms = []
for _fr in range(frames):
ratoms.append(Atoms(natoms))
ratoms[-1].q = xyz[_fr * natoms * 3:3 * (_fr + 1) * natoms] * units_conv_at
ratoms[-1].m = masses[_fr * natoms:(_fr + 1) * natoms]
ratoms[-1].names = atoms_names[_fr * natoms:(_fr + 1) * natoms]
cell = Cell(expected_cell * units_conv_cell)
# remove temp file created during this test
def delete_files_after_testing():
if os.path.isfile(tmp_file.name):
os.remove(tmp_file.name)
request.addfinalizer(delete_files_after_testing)
return tmp_file, cell, ratoms
def process_units(comment, cell, qatoms, names, masses, output='objects'):
""" Converts the data in the file according to the units written in the ipi format.
"""
# Extracting trajectory units
family, unit = 'undefined', ''
is_comment_useful = filter(None, [key.search(comment.strip())
for key in traj_re])
if len(is_comment_useful) > 0:
traj = is_comment_useful[0].group()[:-1].split('{')
family, unit = traj_dict[traj[0]]['dimension'], traj[1]
# Extracting cell units
cell_unit = ''
tmp = cell_unit_re.search(comment)
if tmp is not None:
cell_unit = tmp.group(1)
# Units transformation
cell *= unit_to_internal('length', cell_unit, 1) # cell units transformation
qatoms *= unit_to_internal(family, unit, 1) # units transformation
if output == 'objects':
cell = Cell(cell)
atoms = Atoms(len(names))
atoms.q[:] = qatoms
atoms.names[:] = names
atoms.m[:] = masses
return {
"atoms": atoms,
"cell": cell,
}
else:
return {
"data": qatoms,
"masses": masses,
"names": names,
"natoms": len(names),
"cell": cell,
}
def contract_trajectory(fns_in, fn_out_template, n_new, cell_units_in, cell_units_out):
verbosity.level = "low"
n = len(fns_in)
# Generate output file names.
if n_new == 1:
fns_out = [fn_out_template]
else:
fns_out = [fn_out_template.format(i) for i in range(n_new)]
print("Contracting {:d} beads to {:d} beads.".format(n, n_new))
print()
print("input file names:")
for fn in fns_in:
print(fn)
print()
print("output file names:")
for fn in fns_out:
print(fn)
print()
# Open input trajectory iterators.
trjs_in = [iter_file_name_raw(fn) for fn in fns_in]
mode = os.path.splitext(fn)[-1]
# Open output files.
fs_out = [open_backup(fn, "w") for fn in fns_out]
mode_out = os.path.splitext(fn_out_template)[-1]
# prepare ring polymer rescaler
rescale = nm_rescale(n, n_new)
# Loop over all frames.
i_frame = 0
while True:
try:
# Get the frames for all beads.
frames = [trj.next() for trj in trjs_in]
except StopIteration:
# Stop when any of the trajectories runs out of frames.
break
# gets units from first frame
dimension, units, cell_units = auto_units(comment=frames[0]["comment"], cell_units=cell_units_in)
if cell_units_out == "automatic": cell_units_out = cell_units # re-use units unless otherwise specified
# Consistency check.
h = frames[0]["cell"]
natoms = len(frames[0]["data"]) / 3
for i in range(n):
# Check that all the cells are the same.
if (frames[i]["cell"] != h).any():
msg = "Cell for beads {:d} and {:d} differ in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
# Check that the numbers of atoms are the same.
if len(frames[i]["data"]) != 3 * natoms:
msg = "Different numbers of atoms for beads {:d} and {:d} in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
cell = Cell()
cell.h = frames[0]["cell"]
atoms = Atoms(natoms)
atoms.names = frames[0]["names"]
# Compose the ring polymer.
q = np.vstack([frame["data"] for frame in frames]) * unit_to_internal(dimension, units, 1) # units transformation
# Contract the coordinates to `n_new` beads.
q_c = rescale.b1tob2(q)
# Save the output data.
for i, f_out in enumerate(fs_out):
atoms.q = q_c[i, :]
print_file(mode_out, atoms, cell, f_out, dimension=dimension, units=units, cell_units=cell_units_out)
# Count frames and print information on progress.
i_frame += 1
if i_frame % 100 == 0:
print("\rframe {:d}".format(i_frame), end="")
sys.stdout.flush()
for f_out in fs_out:
f_out.close()
print()
print()
print("Processed {:d} frames.".format(i_frame))
def contract_trajectory(fns_in, fn_out_template, n_new, cell_units_in,
cell_units_out):
verbosity.level = "low"
n = len(fns_in)
# Generate output file names.
if n_new == 1:
fns_out = [fn_out_template]
else:
fns_out = [fn_out_template.format(i) for i in range(n_new)]
print("Contracting {:d} beads to {:d} beads.".format(n, n_new))
print()
print("input file names:")
for fn in fns_in:
print(fn)
print()
print("output file names:")
for fn in fns_out:
print(fn)
print()
# Open input trajectory iterators.
trjs_in = [iter_file_name_raw(fn) for fn in fns_in]
mode = os.path.splitext(fn)[-1]
# Open output files.
fs_out = [open_backup(fn, "w") for fn in fns_out]
mode_out = os.path.splitext(fn_out_template)[-1]
# prepare ring polymer rescaler
rescale = nm_rescale(n, n_new)
# Loop over all frames.
i_frame = 0
while True:
try:
# Get the frames for all beads.
frames = [trj.next() for trj in trjs_in]
except StopIteration:
# Stop when any of the trajectories runs out of frames.
break
# gets units from first frame
dimension, units, cell_units = auto_units(comment=frames[0]["comment"],
cell_units=cell_units_in)
if cell_units_out == "automatic":
cell_units_out = cell_units # re-use units unless otherwise specified
# Consistency check.
h = frames[0]["cell"]
natoms = len(frames[0]["data"]) / 3
for i in range(n):
# Check that all the cells are the same.
if (frames[i]["cell"] != h).any():
msg = "Cell for beads {:d} and {:d} differ in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
# Check that the numbers of atoms are the same.
if len(frames[i]["data"]) != 3 * natoms:
msg = "Different numbers of atoms for beads {:d} and {:d} in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
cell = Cell()
cell.h = frames[0]["cell"]
atoms = Atoms(natoms)
atoms.names = frames[0]["names"]
# Compose the ring polymer.
q = np.vstack([frame["data"] for frame in frames]) * unit_to_internal(
dimension, units, 1) # units transformation
# Contract the coordinates to `n_new` beads.
q_c = rescale.b1tob2(q)
# Save the output data.
for i, f_out in enumerate(fs_out):
atoms.q = q_c[i, :]
print_file(mode_out,
atoms,
cell,
f_out,
dimension=dimension,
units=units,
cell_units=cell_units_out)
# Count frames and print information on progress.
i_frame += 1
if i_frame % 100 == 0:
print("\rframe {:d}".format(i_frame), end="")
sys.stdout.flush()
for f_out in fs_out:
f_out.close()
print()
print()
print("Processed {:d} frames.".format(i_frame))