diff --git a/pymolfile/molfile/libpymolfile.i b/pymolfile/molfile/libpymolfile.i index e1abfa430208aeb212d77e4422439569e6d3cc25..e3438c550f6f63e28878605ad8310f244902312f 100644 --- a/pymolfile/molfile/libpymolfile.i +++ b/pymolfile/molfile/libpymolfile.i @@ -20,7 +20,7 @@ %{ /* Python SWIG interface to libpymolfile - Copyright (c) 2017 Berk Onat + Copyright (c) 2018 Berk Onat Published with UIUC LICENSE */ #define SWIG_FILE_WITH_INIT @@ -79,7 +79,7 @@ extern PyObject * molfile_plugin_info(PyObject* molcapsule, int plugin_no); %inline %{ PyObject * my_open_file_read(PyObject* molcapsule, char* fname, char* ftype, int natoms) { if (PyType_Ready(&MolObjectType) < 0) - return NULL; + Py_RETURN_NONE; PyTypeObject *type = &MolObjectType; MolObject *plugin_c; molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(molcapsule); @@ -97,6 +97,82 @@ PyObject * my_open_file_read(PyObject* molcapsule, char* fname, char* ftype, int } %} +%feature("autodoc", "0") my_open_file_write; +%rename (open_file_write) my_open_file_write; +%exception my_open_file_write { + $action + if (PyErr_Occurred()) SWIG_fail; +} +%inline %{ +PyObject * my_open_file_write(PyObject* molcapsule, char* fname, char* ftype, int natoms) { + if (PyType_Ready(&MolObjectType) < 0) + Py_RETURN_NONE; + if(natoms < 1) + Py_RETURN_NONE; + PyTypeObject *type = &MolObjectType; + MolObject *plugin_c; + molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(molcapsule); + plugin_c = (MolObject *)type->tp_alloc(type, 0); + plugin_c->plugin = molcapsule; + void *file_handle = plugin->open_file_write(fname, ftype, natoms); + plugin_c->file_handle = PyMolfileCapsule_FromVoidPtr(file_handle, del_molfile_file_handle); + if (!plugin_c->file_handle) { + Py_RETURN_NONE; + } else { + plugin_c->natoms = natoms; + return (PyObject *)plugin_c; + } + } +%} + +/* +%feature("autodoc", "0") my_set_file_write_stdout; +%rename (set_file_write_stdout) my_set_file_write_stdout; +%exception my_set_file_write_stdout { + $action + if (PyErr_Occurred()) SWIG_fail; +} +%inline %{ +PyObject * my_set_file_write_stdout(PyObject* molpack) { + MolObject* plugin_handle = (MolObject*) molpack; + PyObject* filecapsule = plugin_handle->file_handle; + void *file_handle = (void*) PyMolfileCapsule_AsVoidPtr(filecapsule); + if(file_handle->fd){ + fclose(file_handle->fd); + file_handle->fd = stdout; + } + else if(file_handle->fp){ + fclose(file_handle->fp); + file_handle->fp = stdout; + } + else if(file_handle->file){ + fclose(file_handle->file); + file_handle->file = stdout; + } + else if(file_handle->mf){ + fclose(file_handle->mf); + file_handle->mf = stdout; + } + else if(file_handle->writer){ + if(file_handle->writer->fd){ + fclose(file_handle->writer->fd); + file_handle->writer->fd = stdout; + } + } + else (file_handle){ + fclose(file_handle); + file_handle = stdout; + } + plugin_handle->file_handle = PyMolfileCapsule_FromVoidPtr(file_handle, del_molfile_file_handle); + if (!plugin_c->file_handle) { + Py_RETURN_NONE; + } else { + return (PyObject *)plugin_handle; + } + } +%} +*/ + %feature("autodoc", "0") my_close_file_read; %rename (close_file_read) my_close_file_read; %exception my_close_file_read { @@ -115,18 +191,49 @@ PyObject * my_close_file_read(PyObject* molpack) { } %} +%feature("autodoc", "0") my_close_file_write; +%rename (close_file_write) my_close_file_write; +%exception my_close_file_write { + $action + if (PyErr_Occurred()) SWIG_fail; +} +%inline %{ +PyObject * my_close_file_write(PyObject* molpack) { + MolObject* plugin_handle = (MolObject*) molpack; + PyObject* plugincapsule = plugin_handle->plugin; + PyObject* filecapsule = plugin_handle->file_handle; + molfile_plugin_t* plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugincapsule); + void *file_handle = (void*) PyMolfileCapsule_AsVoidPtr(filecapsule); + plugin->close_file_write(file_handle); + Py_RETURN_TRUE; + } +%} + + %feature("autodoc", "0") read_fill_structure; extern PyObject * read_fill_structure(PyObject* molpack, PyObject* prototype); +%feature("autodoc", "0") write_fill_structure; +extern PyObject * write_fill_structure(PyObject* molpack, PyObject* molarray); + %feature("autodoc", "0") read_fill_bonds; extern PyObject * read_fill_bonds(PyObject* molpack); +%feature("autodoc", "0") write_fill_bonds; +extern PyObject * write_fill_bonds(PyObject* molpack, PyObject* moldict); + %feature("autodoc", "0") read_fill_angles; extern PyObject * read_fill_angles(PyObject* molpack); +%feature("autodoc", "0") write_fill_angles; +extern PyObject * write_fill_angles(PyObject* molpack, PyObject* moldict); + %feature("autodoc", "0") read_fill_next_timestep; extern PyObject * read_fill_next_timestep(PyObject* molpack); +%feature("autodoc", "0") write_fill_timestep; +extern PyObject * write_fill_timestep(PyObject* molpack, PyObject* moldict); + %feature("autodoc", "0") are_plugins_same; PyObject* are_plugins_same(PyObject* molpack_a, PyObject* molpack_b); diff --git a/pymolfile/molfile/pymolfile.cxx b/pymolfile/molfile/pymolfile.cxx index e3b3e07d71d717195b2543ede60500e163e1cc52..5fd7a032edc94d7375a0ae4d5242695f5321543b 100644 --- a/pymolfile/molfile/pymolfile.cxx +++ b/pymolfile/molfile/pymolfile.cxx @@ -209,6 +209,62 @@ PyObject * molfile_plugin_info(PyObject* molcapsule, int plugin_no) { return tuple; } +PyObject* write_fill_structure(PyObject* molpack, PyObject* molarray) +{ + //Py_Initialize(); + initNumpyArray(); + int options = 0; + options = MOLFILE_INSERTION | MOLFILE_OCCUPANCY | MOLFILE_BFACTOR | + MOLFILE_ALTLOC | MOLFILE_ATOMICNUMBER | MOLFILE_BONDSSPECIAL | + MOLFILE_MASS | MOLFILE_CHARGE; + molfile_plugin_t* plugin; + void* file_handle; + molfile_atom_t* data; + int numatoms, status; + // Access plugin_handle values + MolObject* plugin_handle = (MolObject*) molpack; + if (plugin_handle->plugin) { + plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugin_handle->plugin); + } else { + PyErr_Format(PyExc_IOError, "molfile plugin is not active."); + Py_RETURN_NONE; + } + if (plugin_handle->file_handle) { + file_handle = (void*) PyMolfileCapsule_AsVoidPtr(plugin_handle->file_handle); + } else { + PyErr_Format(PyExc_IOError, "no file handle in molfile plugin handle."); + Py_RETURN_NONE; + } + numatoms = (int) PyArray_DIM((PyArrayObject*)molarray, 1); + if (numatoms<0){ + if (plugin_handle->natoms) { + numatoms = plugin_handle->natoms; + if (numatoms<0){ + PyErr_Format(PyExc_IOError, "no assigned number of atoms in molfile plugin handle."); + Py_RETURN_NONE; + } + } else { + PyErr_Format(PyExc_AttributeError, "plugin does not have number of atoms information."); + Py_RETURN_NONE; + } + } + // Aquire memory pointer of molfile_atom_t struct from numpy array + data = (molfile_atom_t*) PyArray_DATA((PyArrayObject*)molarray); + // Write array values in molfile_atom_t + if (plugin->write_structure) { + status = plugin->write_structure(file_handle, options, data); + // Check if the status is ok + if (status!=0){ + PyErr_Format(PyExc_IOError, "Error in write_structure function of plugin."); + Py_RETURN_FALSE; + } + Py_RETURN_TRUE; + } else { + PyErr_Format(PyExc_AttributeError, "molfile plugin does not have write_structure function."); + Py_RETURN_FALSE; + } +} + PyObject* read_fill_structure(PyObject* molpack, PyObject* prototype) { //Py_Initialize(); @@ -224,14 +280,12 @@ PyObject* read_fill_structure(PyObject* molpack, PyObject* prototype) MolObject* plugin_handle = (MolObject*) molpack; if (plugin_handle->plugin) { plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugin_handle->plugin); - //plugin = plugin_handle->plugin; } else { PyErr_Format(PyExc_IOError, "molfile plugin is not active."); return NULL; } if (plugin_handle->file_handle) { file_handle = (void*) PyMolfileCapsule_AsVoidPtr(plugin_handle->file_handle); - //file_handle = plugin_handle->file_handle; } else { PyErr_Format(PyExc_IOError, "no file handle in molfile plugin handle."); return NULL; @@ -368,6 +422,81 @@ PyObject* read_fill_bonds(PyObject* molpack) } } +PyObject* write_fill_bonds(PyObject* molpack, PyObject* moldict) +{ + if(!PyDict_Check(moldict)) { + PyErr_Format(PyExc_IOError, "argument 2 is not a Python dictionary."); + Py_RETURN_FALSE; + } + initNumpyArray(); + molfile_plugin_t* plugin; + void* file_handle; + molfile_atom_t* data; + int numatoms, status; + PyObject *ret = NULL; + // Access plugin_handle values + MolObject* plugin_handle = (MolObject*) molpack; + if (plugin_handle->plugin) { + plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugin_handle->plugin); + } else { + PyErr_Format(PyExc_IOError, "molfile plugin is not set."); + return NULL; + } + if (plugin_handle->file_handle) { + file_handle = (void*) PyMolfileCapsule_AsVoidPtr(plugin_handle->file_handle); + } else { + PyErr_Format(PyExc_IOError, "no file handle in molfile plugin handle."); + return NULL; + } + numatoms = plugin_handle->natoms; + if (plugin->write_bonds) { + int nbonds, *from, *to, nbondtypes; + int *bondtype = NULL; + float *bondorder = NULL; + char **bondtypename = NULL; + // Lets see whether dictionary includes a numpy array for coords. + PyObject *from_arr = PyDict_GetItemString(moldict, "from"); + PyObject *to_arr = PyDict_GetItemString(moldict, "to"); + PyObject *bondorder_arr = PyDict_GetItemString(moldict, "bondorder"); + PyObject *bondtype_arr = PyDict_GetItemString(moldict, "bondtype"); + PyObject *bondtypename_arr = PyDict_GetItemString(moldict, "bondtypename"); + if(from_arr && to_arr) { + nbonds = (int) PyArray_DIMS((PyArrayObject*)from_arr)[0]; + from = (int*) PyArray_DATA((PyArrayObject*)from_arr); + to = (int*) PyArray_DATA((PyArrayObject*)to_arr); + } + if(bondorder_arr) { + bondorder = (float*) PyArray_DATA((PyArrayObject*)bondorder_arr); + } + nbondtypes = 0; + if(bondtype_arr) { + nbondtypes = (int) PyArray_DIMS((PyArrayObject*)bondtype_arr)[0]; + bondtype = (int*) PyArray_DATA((PyArrayObject*)bondtype_arr); + } + if(bondtypename_arr) { + nbondtypes = (int) PyArray_DIMS((PyArrayObject*)bondtypename_arr)[0]; + bondtypename = (char**) PyArray_DATA((PyArrayObject*)bondtypename_arr); + } + if (nbonds>0) { + if ((status = plugin->write_bonds(file_handle, nbonds, from, to, + bondorder, bondtype, nbondtypes, bondtypename))) { + PyErr_Format(PyExc_IOError, "Error accessing write_bonds function of plugin."); + Py_RETURN_NONE; + } + if (status!=0){ + PyErr_Format(PyExc_IOError, "Error in write_bonds function of plugin."); + Py_RETURN_NONE; + } + Py_RETURN_TRUE; + } else { + Py_RETURN_FALSE; + } + } else { + PyErr_Format(PyExc_AttributeError, "molfile plugin does not have write_bonds function."); + Py_RETURN_NONE; + } +} + PyObject* read_fill_angles(PyObject* molpack) { initNumpyArray(); @@ -567,6 +696,391 @@ PyObject* read_fill_angles(PyObject* molpack) } +PyObject* write_fill_angles(PyObject* molpack, PyObject* moldict) +{ + if(!PyDict_Check(moldict)) { + PyErr_Format(PyExc_IOError, "argument 2 is not a Python dictionary."); + Py_RETURN_FALSE; + } + initNumpyArray(); + int options = 0; + molfile_plugin_t* plugin; + void* file_handle; + molfile_atom_t* data; + int numatoms, status; + // Access plugin_handle values + MolObject* plugin_handle = (MolObject*) molpack; + if (plugin_handle->plugin) { + plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugin_handle->plugin); + } else { + PyErr_Format(PyExc_IOError, "molfile plugin is not active."); + return NULL; + } + if (plugin_handle->file_handle) { + file_handle = (void*) PyMolfileCapsule_AsVoidPtr(plugin_handle->file_handle); + } else { + PyErr_Format(PyExc_IOError, "no file handle in molfile plugin handle."); + return NULL; + } + numatoms = plugin_handle->natoms; + // Check if there is write_angles support in this plugin + if (plugin->write_angles) { + // Angles + int numangles; + int *angles = NULL; + int *angletypes = NULL; + int numangletypes; + const char **angletypenames = NULL; + // Dihedrals + int numdihedrals; + int *dihedrals = NULL; + int *dihedraltypes = NULL; + int numdihedraltypes; + const char **dihedraltypenames = NULL; + // Impropers + int numimpropers; + int *impropers = NULL; + int *impropertypes = NULL; + int numimpropertypes; + const char **impropertypenames = NULL; + // Cterms + int ndimcterms = 2; + int numcterms, ctermcols, ctermrows; + int *cterms = NULL; + // Initilize zeros to number of angles, dihedrals, so on ... + numangles = 0; + numangletypes = 0; + numdihedrals = 0; + numdihedraltypes = 0; + numimpropers = 0; + numimpropertypes = 0; + numcterms = 0; + // Lets see whether dictionary includes a numpy arrays. + PyObject *angles_arr = PyDict_GetItemString(moldict, "angles"); + PyObject *angletypes_arr = PyDict_GetItemString(moldict, "angletypes"); + PyObject *angletypenames_arr = PyDict_GetItemString(moldict, "angletypenames"); + PyObject *dihedrals_arr = PyDict_GetItemString(moldict, "dihedrals"); + PyObject *dihedraltypes_arr = PyDict_GetItemString(moldict, "dihedraltypes"); + PyObject *dihedraltypenames_arr = PyDict_GetItemString(moldict, "dihedraltypenames"); + PyObject *impropers_arr = PyDict_GetItemString(moldict, "impropers"); + PyObject *impropertypes_arr = PyDict_GetItemString(moldict, "impropertypes"); + PyObject *impropertypenames_arr = PyDict_GetItemString(moldict, "impropertypenames"); + PyObject *cterms_arr = PyDict_GetItemString(moldict, "cterms"); + // Even if there is no info for angles/dihedrals/impropers, this function will let the + // the arrays to be NULL on plugin level. + // We will do the checking one-by-one for all available numpy arrays + if(angles_arr) { + numangles = (int) PyArray_DIMS((PyArrayObject*)angles_arr)[0]; + angles = (int*) PyArray_DATA((PyArrayObject*)angles_arr); + } + if(angletypes_arr) { + numangletypes = (int) PyArray_DIMS((PyArrayObject*)angletypes_arr)[0]; + angletypes = (int*) PyArray_DATA((PyArrayObject*)angletypes_arr); + } + if(angletypenames_arr) { + numangletypes = (int) PyArray_DIMS((PyArrayObject*)angletypenames_arr)[0]; + angletypenames = (const char**) PyArray_DATA((PyArrayObject*)angletypenames_arr); + } + if(dihedrals_arr) { + numdihedrals = (int) PyArray_DIMS((PyArrayObject*)dihedrals_arr)[0]; + dihedrals = (int*) PyArray_DATA((PyArrayObject*)dihedrals_arr); + } + if(dihedraltypes_arr) { + numdihedraltypes = (int) PyArray_DIMS((PyArrayObject*)dihedraltypes_arr)[0]; + dihedraltypes = (int*) PyArray_DATA((PyArrayObject*)dihedraltypes_arr); + } + if(dihedraltypenames_arr) { + numdihedraltypes = (int) PyArray_DIMS((PyArrayObject*)dihedraltypenames_arr)[0]; + dihedraltypenames = (const char**) PyArray_DATA((PyArrayObject*)dihedraltypenames_arr); + } + if(impropers_arr) { + numimpropers = (int) PyArray_DIMS((PyArrayObject*)impropers_arr)[0]; + impropers = (int*) PyArray_DATA((PyArrayObject*)impropers_arr); + } + if(impropertypes_arr) { + numimpropertypes = (int) PyArray_DIMS((PyArrayObject*)impropertypes_arr)[0]; + impropertypes = (int*) PyArray_DATA((PyArrayObject*)impropertypes_arr); + } + if(impropertypenames_arr) { + numimpropertypes = (int) PyArray_DIMS((PyArrayObject*)impropertypenames_arr)[0]; + impropertypenames = (const char**) PyArray_DATA((PyArrayObject*)impropertypenames_arr); + } + // Cterms + if(cterms_arr) { + ndimcterms = (int) PyArray_NDIM((PyArrayObject*)cterms_arr); + numcterms = (int) PyArray_SIZE((PyArrayObject*)cterms_arr); + if (ndimcterms>1) { + ctermrows = (int) PyArray_DIMS((PyArrayObject*)cterms_arr)[0]; + ctermcols = (int) PyArray_DIMS((PyArrayObject*)cterms_arr)[1]; + } else { + ctermrows = 0; + ctermcols = 0; + } + cterms = (int*) PyArray_DATA((PyArrayObject*)cterms_arr); + } + // Calling write_angles to write the information + if ((status = plugin->write_angles(file_handle, numangles, angles, angletypes, + numangletypes, angletypenames, numdihedrals, + dihedrals, dihedraltypes, numdihedraltypes, + dihedraltypenames, numimpropers, impropers, + impropertypes, numimpropertypes, impropertypenames, + numcterms, cterms, ctermcols, ctermrows))) { + PyErr_Format(PyExc_IOError, "Error accessing read_angles function of plugin."); + Py_RETURN_NONE; + } + Py_RETURN_TRUE; + } else { + PyErr_Format(PyExc_AttributeError, "molfile plugin does not have read_angles function."); + Py_RETURN_NONE; + } +} + + +PyObject* write_fill_timestep(PyObject* molpack, PyObject* moldict) +{ + if(!PyDict_Check(moldict)) { + PyErr_Format(PyExc_IOError, "argument 2 is not a Python dictionary."); + Py_RETURN_FALSE; + } + initNumpyArray(); + molfile_plugin_t* plugin; + void* file_handle; + int numatoms, status; + int nd; + int i, d; + // Access plugin_handle values + MolObject* plugin_handle = (MolObject*) molpack; + if (plugin_handle->plugin) { + plugin = (molfile_plugin_t*) PyMolfileCapsule_AsVoidPtr(plugin_handle->plugin); + } else { + PyErr_Format(PyExc_IOError, "molfile plugin is not active."); + Py_RETURN_NONE; + } + if (plugin_handle->file_handle) { + file_handle = (void*) PyMolfileCapsule_AsVoidPtr(plugin_handle->file_handle); + } else { + PyErr_Format(PyExc_IOError, "no file handle in molfile plugin handle."); + Py_RETURN_NONE; + } + if (plugin_handle->natoms) { + numatoms = plugin_handle->natoms; + } else { + PyErr_Format(PyExc_IOError, "no assigned number of atoms in molfile plugin handle."); + Py_RETURN_NONE; + } + if (plugin->write_timestep) { + PyObject *coords_arr = NULL; + npy_intp n, m, i, j; + int ndim = 1; + // Lets see whether dictionary includes a numoy array for coords. + // We will use the first dimension as the loop over write_timestep. + coords_arr = PyDict_GetItemString(moldict, "coords"); + if(coords_arr) { + ndim = (int) PyArray_NDIM((PyArrayObject*)coords_arr); + } else { + // It seams someone forgot to put coords in dictionary. + // Nothing to write to output file + // Return False in this case + Py_RETURN_FALSE; + } + int numsteps = 0; + int numatoms = 0; + // Check if dimension is correct for numpy array + // If the array dimension is not correct return False. + if(ndim>2){ + n = PyArray_DIMS((PyArrayObject*)coords_arr)[0]; + m = PyArray_DIMS((PyArrayObject*)coords_arr)[1]; + numsteps = (int)n; + numatoms = (int)m; + } + else if(ndim>1){ + n = 1; + m = PyArray_DIMS((PyArrayObject*)coords_arr)[0]; + numsteps = (int)n; + numatoms = (int)m; + } else { + Py_RETURN_FALSE; + } + //if(numsteps>0){ + // molfile_timestep_t timestep; + //} else { + // Py_RETURN_FALSE; + //} + // It seams we have coordinates in a numpy array and + // if we have at least one snapshot of coordinates, we can write it. + // Set if the velocities can be written with this plugin + int has_velocities = 0; + unsigned int total_steps = 1; + unsigned int bytes_per_step = 0; + double a_sca, b_sca, c_sca, alpha_sca, beta_sca, gamma_sca, time_sca; + //molfile_timestep_metadata_t timestep_metadata; + PyObject *bytes_per_step_arr = PyDict_GetItemString(moldict, "bytes_per_step"); + if(bytes_per_step_arr) { + bytes_per_step = (unsigned int)PyLong_AsLong(bytes_per_step_arr); + //timestep_metadata.avg_bytes_per_timestep = bytes_per_step; + } + PyObject *total_steps_arr = PyDict_GetItemString(moldict, "total_steps"); + if(total_steps_arr){ + total_steps = (unsigned int)PyLong_AsLong(total_steps_arr); + //timestep_metadata.count = total_steps; + } + PyObject *has_velocities_arr = PyDict_GetItemString(moldict, "has_velocities"); + if(has_velocities_arr){ + has_velocities = (int)PyLong_AsLong(has_velocities_arr); + //timestep_metadata.has_velocities = has_velocities; + } + PyObject *velocities_arr = NULL; + if(has_velocities > 0) { + velocities_arr = PyDict_GetItemString(moldict, "velocities"); + } + // All support arrays' sizes should match the size of coords array if supplied. + PyObject *a_arr = PyDict_GetItemString(moldict, "A"); + if(a_arr) + if(PyArray_Check(a_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)a_arr)[0]) + Py_RETURN_FALSE; + } else { + a_sca = PyFloat_AsDouble(a_arr); + } + PyObject *b_arr = PyDict_GetItemString(moldict, "B"); + if(b_arr) + if(PyArray_Check(b_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)b_arr)[0]) + Py_RETURN_FALSE; + } else { + b_sca = PyFloat_AsDouble(b_arr); + } + PyObject *c_arr = PyDict_GetItemString(moldict, "C"); + if(c_arr) + if(PyArray_Check(c_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)c_arr)[0]) + Py_RETURN_FALSE; + } else { + c_sca = PyFloat_AsDouble(c_arr); + } + PyObject *alpha_arr = PyDict_GetItemString(moldict, "alpha"); + if(alpha_arr) + if(PyArray_Check(alpha_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)alpha_arr)[0]) + Py_RETURN_FALSE; + } else { + alpha_sca = PyFloat_AsDouble(alpha_arr); + } + PyObject *beta_arr = PyDict_GetItemString(moldict, "beta"); + if(beta_arr) + if(PyArray_Check(beta_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)beta_arr)[0]) + Py_RETURN_FALSE; + } else { + beta_sca = PyFloat_AsDouble(beta_arr); + } + PyObject *gamma_arr = PyDict_GetItemString(moldict, "gamma"); + if(gamma_arr) + if(PyArray_Check(gamma_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)gamma_arr)[0]) + Py_RETURN_FALSE; + } else { + gamma_sca = PyFloat_AsDouble(gamma_arr); + } + PyObject *pt_arr = PyDict_GetItemString(moldict, "physical_time"); + if(pt_arr) + if(PyArray_Check(pt_arr)){ + if(n != PyArray_DIMS((PyArrayObject*)pt_arr)[0]) + Py_RETURN_FALSE; + } else { + time_sca = PyFloat_AsDouble(pt_arr); + } + molfile_timestep_t *timestep; + // Good old for loop over first dimension of numpy array will do the writing out. + for (i = 0; i < n; i++) { + //if (plugin->write_timestep_metadata) plugin->write_timestep_metadata(file_handle, ×tep_metadata); + if(a_arr){ + if(PyArray_Check(a_arr)){ + timestep->A = *(float*)(PyArray_BYTES((PyArrayObject*)a_arr) + i*PyArray_STRIDES((PyArrayObject*)a_arr)[0]); + } else { + timestep->A = (float) a_sca; + } + } else { + timestep->A = 0.0; + } + if(b_arr){ + if(PyArray_Check(b_arr)){ + timestep->B = *(float*)(PyArray_BYTES((PyArrayObject*)b_arr) + i*PyArray_STRIDES((PyArrayObject*)b_arr)[0]); + } else { + timestep->B = (float) b_sca; + } + } else { + timestep->B = (float) 0.0; + } + if(c_arr){ + if(PyArray_Check(c_arr)){ + timestep->C = *(float*)(PyArray_BYTES((PyArrayObject*)c_arr) + i*PyArray_STRIDES((PyArrayObject*)c_arr)[0]); + } else { + timestep->C = (float) c_sca; + } + } else { + timestep->C = (float) 0.0; + } + if(alpha_arr){ + if(PyArray_Check(alpha_arr)){ + timestep->alpha = *(float*)(PyArray_BYTES((PyArrayObject*)alpha_arr) + i*PyArray_STRIDES((PyArrayObject*)alpha_arr)[0]); + } else { + timestep->alpha = (float) alpha_sca; + } + } else { + timestep->alpha = NULL; + } + if(beta_arr){ + if(PyArray_Check(beta_arr)){ + timestep->beta = *(float*)(PyArray_BYTES((PyArrayObject*)beta_arr) + i*PyArray_STRIDES((PyArrayObject*)beta_arr)[0]); + } else { + timestep->beta = (float) beta_sca; + } + } else { + timestep->beta = NULL; + } + if(gamma_arr){ + if(PyArray_Check(gamma_arr)){ + timestep->gamma = *(float*)(PyArray_BYTES((PyArrayObject*)gamma_arr) + i*PyArray_STRIDES((PyArrayObject*)gamma_arr)[0]); + } else { + timestep->gamma = (float) gamma_sca; + } + } else { + timestep->gamma = NULL; + } + if(pt_arr){ + if(PyArray_Check(pt_arr)){ + timestep->physical_time = *(float*)(PyArray_BYTES((PyArrayObject*)pt_arr) + i*PyArray_STRIDES((PyArrayObject*)pt_arr)[0]); + } else { + timestep->physical_time = (float) time_sca; + } + } else { + timestep->physical_time = i; + } + if(has_velocities > 0) { + if(velocities_arr){ + if(PyArray_Check(velocities_arr)){ + timestep->velocities = (float*)(PyArray_BYTES((PyArrayObject*)velocities_arr) + i*PyArray_STRIDES((PyArrayObject*)velocities_arr)[0]); + } + } + } + timestep->coords = (float*)(PyArray_BYTES((PyArrayObject*)coords_arr) + i*PyArray_STRIDES((PyArrayObject*)coords_arr)[0]); + status = plugin->write_timestep(file_handle, timestep); + if (status == MOLFILE_EOF) { + Py_RETURN_FALSE; + } + else if (status != MOLFILE_SUCCESS) { + PyErr_Format(PyExc_AttributeError, "Failed in calling write_timestep function of plugin."); + Py_RETURN_FALSE; + } + } + Py_RETURN_TRUE; + } else { + PyErr_Format(PyExc_AttributeError, "molfile plugin does not have write_timestep function."); + Py_RETURN_NONE; + } +} + PyObject* read_fill_next_timestep(PyObject* molpack) { initNumpyArray(); diff --git a/pymolfile/molfile/pymolfile.h b/pymolfile/molfile/pymolfile.h index b0c31fb3c2096ce487f0bf1da58b7927f8cd6c27..fe803879326cfd036bdc358053ac7851520a46e2 100644 --- a/pymolfile/molfile/pymolfile.h +++ b/pymolfile/molfile/pymolfile.h @@ -263,9 +263,13 @@ static PyTypeObject MolObjectType = { PyObject * molfile_plugin_info(PyObject* molcapsule, int plugin_no); PyObject * read_fill_structure(PyObject* molpack, PyObject* prototype); +PyObject * write_fill_structure(PyObject* molpack, PyObject* molarray); PyObject * read_fill_bonds(PyObject* molpack); +PyObject * write_fill_bonds(PyObject* molpack, PyObject* moldict); PyObject * read_fill_angles(PyObject* molpack); +PyObject * write_fill_angles(PyObject* molpack, PyObject* moldict); PyObject * read_fill_next_timestep(PyObject* molpack); +PyObject * write_fill_timestep(PyObject* molpack, PyObject* moldict); PyObject * are_plugins_same(PyObject* molpack_a, PyObject* molpack_b); PyObject * are_filehandles_same(PyObject* molpack_a, PyObject* molpack_b); PyObject * get_plugin(PyObject* molcapsule, int plug_no); diff --git a/pymolfile/pymolfile.py b/pymolfile/pymolfile.py index 1d22139172465e5621d621d1d09d9866bcee982a..fa25e44fa4bfdc75850b996ea1655e62c531ffc5 100644 --- a/pymolfile/pymolfile.py +++ b/pymolfile/pymolfile.py @@ -7,6 +7,9 @@ import re import sys import numpy as np from contextlib import contextmanager +import platform +import ctypes +import tempfile import warnings if sys.version_info > (3,): @@ -19,6 +22,18 @@ except ImportError: from .plugin_list import plugins, byte_str_decode, MOLFILE_PLUGINS, C_MOLFILE_PLUGINS +if platform.system() == "Windows": + libc = ctypes.cdll.msvcrt +else: + libc = ctypes.CDLL(None) + +if platform.system() == "Windows": + c_stdout = ctypes.c_void_p.in_dll(libc, 'stdout') +elif "Darwin" in platform.system(): + c_stdout = ctypes.c_void_p.in_dll(libc, '__stdoutp') +else: + c_stdout = ctypes.c_void_p.in_dll(libc, 'stdout') + def list_plugins(): global MOLFILE_PLUGINS return MOLFILE_PLUGINS[:] @@ -49,6 +64,8 @@ def stdout_redirected(to=os.devnull): ####assert libc.fileno(ctypes.c_void_p.in_dll(libc, "stdout")) == fd == 1 def _redirect_stdout(to): + # Flush the C-level buffer stdout + libc.fflush(c_stdout) sys.stdout.close() # + implicit flush() os.dup2(to.fileno(), fd) # fd writes to 'to' file sys.stdout = os.fdopen(fd, 'w') # Python writes to fd @@ -63,6 +80,42 @@ def stdout_redirected(to=os.devnull): # buffering and flags such as # CLOEXEC may be different +# The function stdout_redirect_stream is taken from +#https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ + +@contextmanager +def stdout_redirect_stream(stream): + # The original fd stdout points to. Usually 1 on POSIX systems. + original_stdout_fd = sys.stdout.fileno() + + def _redirect_stdout(to_fd): + """Redirect stdout to the given file descriptor.""" + # Flush the C-level buffer stdout + libc.fflush(c_stdout) + # Flush and close sys.stdout - also closes the file descriptor (fd) + sys.stdout.close() + # Make original_stdout_fd point to the same file as to_fd + os.dup2(to_fd, original_stdout_fd) + # Create a new sys.stdout that points to the redirected fd + sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, 'wb')) + + # Save a copy of the original stdout fd in saved_stdout_fd + saved_stdout_fd = os.dup(original_stdout_fd) + try: + # Create a temporary file and redirect stdout to it + tfile = tempfile.TemporaryFile(mode='w+b') + _redirect_stdout(tfile.fileno()) + # Yield to caller, then redirect stdout back to the saved fd + yield + _redirect_stdout(saved_stdout_fd) + # Copy contents of temporary file to the given stream + tfile.flush() + tfile.seek(0, io.SEEK_SET) + stream.write(tfile.read()) + finally: + tfile.close() + os.close(saved_stdout_fd) + def get_dir_base_extension(file_name): """ Splits directory, file base and file extensions @@ -346,6 +399,64 @@ class Trajectory(object): finally: del iter_obj +def molfile_prototype(file_format=None): + """ Generates a prototype numpy array that has the same struct + definition as in molfile_atom_t arrays. + + Returns: numpy.array + """ + if file_format is None: + file_format = '' + #if 0 && vmdplugin_ABIVERSION > 17 + # /* The new PDB file formats allows for much larger structures, */ + # /* which can therefore require longer chain ID strings. The */ + # /* new PDBx/mmCIF file formats do not have length limits on */ + # /* fields, so PDB chains could be arbitrarily long strings */ + # /* in such files. At present, we know we need at least 3-char */ + # /* chains for existing PDBx/mmCIF files. */ + # char chain[4]; /**< required chain name, or "" */ + #else + # char chain[2]; /**< required chain name, or "" */ + #endif + # + # Change 'chain', S2 to S4 + # + #if('pdb' in file_format or + # 'psf' in file_format): + chain_size = 'S4' + #else: + # chain_size = 'S2' + if('dtr' in file_format.lower() or + 'stk' in file_format.lower() or + 'atr' in file_format.lower()): + prototype = np.array([ + ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0), + ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0) + ], + dtype=[ + ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'), + ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size), + ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'), + ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'), + ('radius', 'f4'), ('atomicnumber', 'i4'), ('ctnumber', 'i4') + ] + ) + else: + prototype = np.array([ + ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6), + ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6) + ], + dtype=[ + ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'), + ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size), + ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'), + ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'), + ('radius', 'f4'), ('atomicnumber', 'i4') + ] + ) + + return prototype + def read_topology(file_name, file_format, plugin, silent): """ Reads structure, bonds, angles, dihedrals, impropers and additional informations through molfile_plugin if the @@ -373,53 +484,7 @@ def read_topology(file_name, file_format, plugin, silent): plugin is not None): natoms=0 topo = Topology(plugin, file_name, file_format, natoms, silent) - #if 0 && vmdplugin_ABIVERSION > 17 - # /* The new PDB file formats allows for much larger structures, */ - # /* which can therefore require longer chain ID strings. The */ - # /* new PDBx/mmCIF file formats do not have length limits on */ - # /* fields, so PDB chains could be arbitrarily long strings */ - # /* in such files. At present, we know we need at least 3-char */ - # /* chains for existing PDBx/mmCIF files. */ - # char chain[4]; /**< required chain name, or "" */ - #else - # char chain[2]; /**< required chain name, or "" */ - #endif - # - # Change 'chain', S2 to S4 - # - #if('pdb' in file_format or - # 'psf' in file_format): - chain_size = 'S4' - #else: - # chain_size = 'S2' - if('dtr' in file_format or - 'stk' in file_format or - 'atr' in file_format): - prototype = np.array([ - ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0), - ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6,0) - ], - dtype=[ - ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'), - ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size), - ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'), - ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'), - ('radius', 'f4'), ('atomicnumber', 'i4'), ('ctnumber', 'i4') - ] - ) - else: - prototype = np.array([ - ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6), - ('','','',0,'','','','',1.0,1.0,1.0,1.0,1.0,6) - ], - dtype=[ - ('name', 'S16'), ('type', 'S16'), ('resname', 'S8'), - ('resid', 'i4'), ('segid', 'S8'), ('chain', chain_size), - ('altloc', 'S2'), ('insertion', 'S2'), ('occupancy', 'f4'), - ('bfactor', 'f4'), ('mass', 'f4'), ('charge', 'f4'), - ('radius', 'f4'), ('atomicnumber', 'i4') - ] - ) + prototype = molfile_prototype(file_format) if topo.read_structure(prototype) is not None: topo.structure = decode_array(topo._structure) diff --git a/setup.py b/setup.py index 010b9645f3c383d0ea08724117f9be1e53ed9100..3f0a8832ed5d9c55b9b4447d15f2430d389c8ad7 100644 --- a/setup.py +++ b/setup.py @@ -98,7 +98,7 @@ def check_tcl_version(): val = None return val -VERSION = "0.0.4" +VERSION = "0.0.5" CLASSIFIERS = [ "Development Status :: 1 - Alpha", "Intended Audience :: Science/Research",