Skip to content
Snippets Groups Projects
Commit fe013941 authored by Berk Onat's avatar Berk Onat
Browse files

Adding WriteTrajMol class for python interface.

parent 2c31d23f
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 149 additions and 0 deletions
pymolfile/pymolfile.py
+ 149
− 0
View file @ fe013941
...@@ -309,6 +309,123 @@ class Topology(object): ...@@ -309,6 +309,123 @@ class Topology(object):
#if self.pluginhandle is not None: #if self.pluginhandle is not None:
# libpymolfile.close_file_read(self.pluginhandle) # 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): class OpenMolTraj(object):
def __init__(self, traj_handle): def __init__(self, traj_handle):
...@@ -496,8 +613,12 @@ def read_topology(file_name, file_format, plugin, silent): ...@@ -496,8 +613,12 @@ def read_topology(file_name, file_format, plugin, silent):
if(topo.structure is not None or if(topo.structure is not None or
topo.bonds is not None or topo.bonds is not None or
topo.angles is not None): topo.angles is not None):
if topo.pluginhandle is not None:
libpymolfile.close_file_read(topo.pluginhandle)
return topo return topo
else: else:
if topo.pluginhandle is not None:
libpymolfile.close_file_read(topo.pluginhandle)
del topo del topo
return None return None
...@@ -651,11 +772,39 @@ class OpenMolfile(object): ...@@ -651,11 +772,39 @@ class OpenMolfile(object):
warnings.warn("Pymolfile can not find a plugin to open the '" + self.kwords["file_format"] + warnings.warn("Pymolfile can not find a plugin to open the '" + self.kwords["file_format"] +
"' file format of the file " + self.kwords["file_name"]) "' 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): def initialize_settings(self):
self.foutEOF = False
self.foutOpen = False
self.trajectory = None self.trajectory = None
self.topology = None self.topology = None
self.fplugin = None self.fplugin = None
self.tplugin = None self.tplugin = None
self.wfname = None
self.wtype = None
self.wplugin = None
self.wfplugin = None
self.natoms = None
self.smolplugin = None self.smolplugin = None
self.cmolplugin = None self.cmolplugin = None
self.kwords = { self.kwords = {
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment