Commit fe013941 authored by Berk Onat's avatar Berk Onat

Adding WriteTrajMol class for python interface.

parent 2c31d23f
......@@ -309,6 +309,123 @@ class Topology(object):
#if self.pluginhandle is not None:
# libpymolfile.close_file_read(self.pluginhandle)
class WriteMolfile(object):
def __init__(self, molfile_handle):
self.handle = molfile_handle
def __enter__(self):
if self.handle.foutOpen is False:
global MOLFILE_PLUGINS
global C_MOLFILE_PLUGINS
#numlist = libpymolfile.molfile_init()
self.handle.wplugin = libpymolfile.get_plugin(C_MOLFILE_PLUGINS,
self.handle.wfplugin[0])
self.handle.wpluginhandle = None
if self.handle.topology:
self.handle.natoms = self.handle.topology.natoms
#if self.handle.trajectory:
# rtn = self.handle.trajectory.iread()
# if self.handle.trajectory.atoms is not None and rtn is not None:
# self.handle.natoms = len(self.handle.trajectory.atoms["coords"])
if self.handle.kwords["silent"]:
with stdout_redirected():
self.handle.wpluginhandle = libpymolfile.open_file_write(
self.handle.wplugin, self.handle.wfname,
self.handle.wtype, self.handle.natoms)
self.handle.foutOpen = True
else:
self.handle.wpluginhandle = libpymolfile.open_file_write(
self.handle.wplugin, self.handle.wfname,
self.handle.wtype, self.handle.natoms)
self.handle.foutOpen = True
if self.handle.wpluginhandle is not None:
return self
else:
return None
def __exit__(self, type, value, traceback):
if self.handle.foutEOF:
if self.handle.foutOpen:
if self.handle.wpluginhandle is not None:
rtn = libpymolfile.close_file_write(self.handle.wpluginhandle)
if rtn:
self.handle.wpluginhandle = None
def write_topology(self):
status = False
bstatus = False
astatus = False
if self.handle.topology.bonds is not None:
if self.handle.wfplugin[1][7]>0:
bstatus = libpymolfile.write_fill_bonds(self.handle.wpluginhandle, self.handle.topology.bonds)
if self.handle.topology.angles is not None:
if self.handle.wfplugin[1][8]>0:
astatus = libpymolfile.write_fill_angles(self.handle.wpluginhandle, self.handle.topology.angles)
if self.handle.wfplugin[1][6]>0 and self.handle.topology.structure is not None:
prototype = molfile_prototype(self.handle.wtype)
newstructure = [
(str(item[0]),
str(item[1]),
str(item[2]),
int(item[3]),
str(item[4]),
str(item[5]),
str(item[6]),
str(item[7]),
float(item[8]),
float(item[9]),
float(item[10]),
float(item[11]),
float(item[12]),
int(item[13])
) if len(item)<15 else (
str(item[0]),
str(item[1]),
str(item[2]),
int(item[3]),
str(item[4]),
str(item[5]),
str(item[6]),
str(item[7]),
float(item[8]),
float(item[9]),
float(item[10]),
float(item[11]),
float(item[12]),
int(item[13]),
int(item[14])
) for item in self.handle.topology.structure
]
thestructure = np.array(newstructure, dtype=prototype.dtype)
status = libpymolfile.write_fill_structure(self.handle.wpluginhandle, thestructure)
else:
status = True
return status
def write_trajectory(self):
if self.handle.wfplugin[1][9]>0 and self.handle.trajectory:
rtnOut = None
done = False
print("Writing Trajectory")
if self.handle.trajectory.atoms is not None:
rtnOut = libpymolfile.write_fill_timestep(self.handle.wpluginhandle, self.handle.trajectory.atoms)
while done is False:
rtn = self.handle.trajectory.iread()
if self.handle.trajectory.atoms is not None and rtn is not None:
print("Writing Trajectory")
rtnOut = libpymolfile.write_fill_timestep(self.handle.wpluginhandle, self.handle.trajectory.atoms)
else:
print("Done")
done = True
self.handle.foutEOF = True
return True
else:
self.handle.foutEOF = True
return True
class OpenMolTraj(object):
def __init__(self, traj_handle):
......@@ -496,8 +613,12 @@ def read_topology(file_name, file_format, plugin, silent):
if(topo.structure is not None or
topo.bonds is not None or
topo.angles is not None):
if topo.pluginhandle is not None:
libpymolfile.close_file_read(topo.pluginhandle)
return topo
else:
if topo.pluginhandle is not None:
libpymolfile.close_file_read(topo.pluginhandle)
del topo
return None
......@@ -651,11 +772,39 @@ class OpenMolfile(object):
warnings.warn("Pymolfile can not find a plugin to open the '" + self.kwords["file_format"] +
"' file format of the file " + self.kwords["file_name"])
def save(self, wfname, wtype=None):
self.wfname = wfname
if wtype is None:
file_dir, file_base, file_ext = get_dir_base_extension(wfname)
if file_ext:
self.wtype = file_ext
else:
self.wtype = file_base
else:
self.wtype = wtype
self.wfplugin = get_plugin_with_ext(wtype)
if self.wfplugin:
done = False
structOK = False
with WriteMolfile(self) as f:
if f:
if self.wfplugin[1][6]>0:
structOK = f.write_topology()
if self.wfplugin[1][9]>0:
done = f.write_trajectory()
def initialize_settings(self):
self.foutEOF = False
self.foutOpen = False
self.trajectory = None
self.topology = None
self.fplugin = None
self.tplugin = None
self.wfname = None
self.wtype = None
self.wplugin = None
self.wfplugin = None
self.natoms = None
self.smolplugin = None
self.cmolplugin = None
self.kwords = {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment