00001 /*************************************************************************** 00002 *cr 00003 *cr (C) Copyright 1995-2019 The Board of Trustees of the 00004 *cr University of Illinois 00005 *cr All Rights Reserved 00006 *cr 00007 ***************************************************************************/ 00008 00009 /*************************************************************************** 00010 * RCS INFORMATION: 00011 * 00012 * $RCSfile: py_numeric.C,v $ 00013 * $Author: johns $ $Locker: $ $State: Exp $ 00014 * $Revision: 1.28 $ $Date: 2020年07月23日 03:27:52 $ 00015 * 00016 *************************************************************************** 00017 * DESCRIPTION: 00018 * Python interface to the Python Numeric module. 00019 ***************************************************************************/ 00020 00021 #include "py_commands.h" 00022 #ifdef VMDNUMPY 00023 00024 // Promise we're not using any deprecated stuff to squash a compiler warning 00025 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION 00026 #include "numpy/ndarrayobject.h" 00027 00028 #include "AtomSel.h" 00029 #include "VMDApp.h" 00030 #include "MoleculeList.h" 00031 00032 /* 00033 * Note: changes to the timestep method now return an N x 3 array 00034 * rather than flat 3N arrays. This may break older plugins like IED 00035 * if they are not kept up-to-date with VMD. 00036 */ 00037 static const char timestep_doc[] = 00038 "Return a zero-copy reference to atomic coordinates\n\n" 00039 "Args:\n" 00040 " molid (int): Molecule ID. Defaults to -1 (top molecule)\n" 00041 " frame (int): Frame to select. Defaults to -1 (current frame)\n" 00042 "Returns:\n" 00043 " (N x 3 float32 ndarray): Reference to coordinates, where N is the\n" 00044 " number of atoms in the molecule"; 00045 static PyObject *py_timestep(PyObject *self, PyObject *args, PyObject *kwargs) { 00046 const char *kwlist[] = {"molid", "frame", NULL}; 00047 int molid = -1, frame = -1; 00048 00049 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|ii:vmdnumpy.timestep", 00050 (char**) kwlist, &molid, &frame)) 00051 return NULL; 00052 00053 VMDApp *app; 00054 if (!(app = get_vmdapp())) 00055 return NULL; 00056 00057 if (molid == -1) 00058 molid = app->molecule_top(); 00059 if (!valid_molid(molid, app)) 00060 return NULL; 00061 00062 Timestep *ts = parse_timestep(app, molid, frame); 00063 if (!ts) 00064 return NULL; 00065 00066 npy_intp dims[2] = {0, 3}; 00067 dims[0] = ts->num; 00068 PyArrayObject *result; 00069 result = (PyArrayObject *)PyArray_SimpleNewFromData(2, dims, NPY_FLOAT, 00070 (char *)ts->pos); 00071 return PyArray_Return(result); 00072 } 00073 00074 00075 static const char velocities_doc[] = 00076 "Return a zero-copy reference to atomic velocites, or None if no velocity\n" 00077 "information is present\n\n" 00078 "Args:\n" 00079 " molid (int): Molecule ID. Defaults to -1 (top molecule)\n" 00080 " frame (int): Frame to select. Defaults to -1 (current frame)\n" 00081 "Returns:\n" 00082 " (N x 3 float32 ndarray): Reference to velocities, where N is the\n" 00083 " number of atoms in the molecule"; 00084 static PyObject *py_velocities(PyObject *self, PyObject *args, 00085 PyObject *kwargs) { 00086 const char *kwlist[] = {"molid", "frame", NULL}; 00087 int molid = -1, frame = -1; 00088 00089 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|ii:vmdnumpy.velocities", 00090 (char**) kwlist, &molid, &frame)) 00091 return NULL; 00092 00093 VMDApp *app; 00094 if (!(app = get_vmdapp())) 00095 return NULL; 00096 00097 if (molid == -1) 00098 molid = app->molecule_top(); 00099 if (!valid_molid(molid, app)) 00100 return NULL; 00101 00102 Timestep *ts = parse_timestep(app, molid, frame); 00103 if (!ts) 00104 return NULL; 00105 00106 // If no velocities, return None 00107 if (!ts->vel) { 00108 Py_INCREF(Py_None); 00109 return Py_None; 00110 } 00111 00112 npy_intp dims[2] = {0, 3}; 00113 dims[0] = ts->num; 00114 PyArrayObject *result; 00115 result = (PyArrayObject *)PyArray_SimpleNewFromData(2, dims, NPY_FLOAT, 00116 (char *)ts->vel); 00117 return PyArray_Return(result); 00118 } 00119 00120 00121 static const char atomselect_doc[] = 00122 "Return an array of ints representing flags for on/off atoms in an atom\n" 00123 "selection\n\n" 00124 "Args:\n" 00125 " selection (str): Atom selection string\n" 00126 " molid (int): Molecule ID. Defaults to -1 (top molecule)\n" 00127 " frame (int): Frame to select. Defaults to -1 (current frame)\n" 00128 "Returns:\n" 00129 " (N, int ndarray): Flags for atoms, where N is the number of atoms\n" 00130 " in the molecule. The value for an atom will be 1 if it is in the\n" 00131 " selection, 0 otherwise"; 00132 static PyObject *py_atomselect(PyObject *self, PyObject *args, 00133 PyObject *kwargs) { 00134 const char *kwlist[] = {"selection", "molid", "frame", NULL}; 00135 int molid = -1, frame = -1; 00136 char *sel = NULL; 00137 00138 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|ii:vmdnumpy.atomselect", 00139 (char**) kwlist, &sel, &molid, &frame)) 00140 return NULL; 00141 00142 VMDApp *app; 00143 if (!(app = get_vmdapp())) 00144 return NULL; 00145 00146 if (molid == -1) 00147 molid = app->molecule_top(); 00148 if (!valid_molid(molid, app)) 00149 return NULL; 00150 00151 // Create a new atom selection 00152 DrawMolecule *mol = app->moleculeList->mol_from_id(molid); 00153 AtomSel *atomSel = new AtomSel(app, app->atomSelParser, mol->id()); 00154 atomSel->which_frame = frame; 00155 00156 if (atomSel->change(sel, mol) == AtomSel::NO_PARSE) { 00157 PyErr_Format(PyExc_ValueError, "can't parse atom selection text '%s'", sel); 00158 delete atomSel; 00159 return NULL; 00160 } 00161 00162 // As the atom selection object is about to be deleted, we need to copy 00163 // the atom masks to a new numpy array so they remain valid. 00164 npy_intp dims[1] = {0}; 00165 dims[0] = mol->nAtoms; 00166 PyArrayObject *result; 00167 result = (PyArrayObject *) PyArray_SimpleNew(1, dims, NPY_INT); 00168 memcpy(PyArray_DATA(result), atomSel->on, dims[0]*sizeof(NPY_INT)); 00169 delete atomSel; 00170 00171 return PyArray_Return(result); 00172 } 00173 00174 00175 static PyMethodDef methods[] = { 00176 {"timestep", (PyCFunction)py_timestep, METH_VARARGS | METH_KEYWORDS, timestep_doc}, 00177 {"positions", (PyCFunction)py_timestep, METH_VARARGS | METH_KEYWORDS, timestep_doc}, 00178 {"velocities", (PyCFunction)py_velocities, METH_VARARGS | METH_KEYWORDS, velocities_doc}, 00179 {"atomselect", (PyCFunction)py_atomselect, METH_VARARGS | METH_KEYWORDS, atomselect_doc}, 00180 {NULL, NULL} 00181 }; 00182 00183 00184 static const char vnumpy_moddoc[] = 00185 "Methods for interacting with atomic positions or velocities as numpy arrays. " 00186 "This can offer significant performance gains over using atomsel types"; 00187 00188 00189 #if PY_MAJOR_VERSION >= 3 00190 static struct PyModuleDef vmdnumpydef = { 00191 PyModuleDef_HEAD_INIT, 00192 "vmdnumpy", 00193 vnumpy_moddoc, 00194 -1, 00195 methods, 00196 }; 00197 #endif 00198 00199 00200 PyObject* initvmdnumpy() { 00201 #if PY_MAJOR_VERSION >= 3 00202 PyObject *module = PyModule_Create(&vmdnumpydef); 00203 #else 00204 PyObject *module = Py_InitModule3("vmdnumpy", methods, vnumpy_moddoc); 00205 #endif 00206 _import_array(); // Don't use import_array macro as it expands to return void 00207 00208 if (PyErr_Occurred()) 00209 return NULL; 00210 00211 return module; 00212 } 00213 00214 #endif // -DVMDNUMPY