00001
00002
00003
00004
00005
00006
00007
00008
00009
00010
00011
00012
00013
00014
00015
00016
00017
00018
00019
00020 #include "py_commands.h"
00021 #include "AtomSel.h"
00022 #include "VMDApp.h"
00023 #include "MoleculeList.h"
00024 #include "SymbolTable.h"
00025 #include "Measure.h"
00026 #include "SpatialSearch.h"
00027
00028
00029 #if PY_VERSION_HEX < ((2<<24)|(5<<16))
00030 typedef int Py_ssize_t;
00031 #endif
00032
00033 typedef struct {
00034 PyObject_HEAD
00035 AtomSel *atomSel;
00036 VMDApp *app;
00037 } PyAtomSelObject;
00038
00039
00040 static int atomsel_Check(PyObject *obj);
00041
00042
00043 DrawMolecule *get_molecule(PyAtomSelObject *a) {
00044 int molid = a->atomSel->molid();
00045 DrawMolecule *mol = a->app->moleculeList->mol_from_id(molid);
00046 if (!mol) {
00047 PyErr_SetString(PyExc_ValueError, "invalid molid");
00048 }
00049 return mol;
00050 }
00051
00052 static char *atomsel_doc = (char *)
00053 "atomsel( selection, molid = top, frame = now) -> new selection object\n"
00054 ;
00055
00056 static void
00057 atomsel_dealloc( PyAtomSelObject *obj ) {
00058 delete obj->atomSel;
00059 obj->ob_type->tp_free((PyObject *)obj);
00060 }
00061
00062 static PyObject *
00063 atomsel_repr( PyAtomSelObject *obj ) {
00064 AtomSel *sel = obj->atomSel;
00065 char *s = new char[strlen(sel->cmdStr) + 100];
00066 sprintf(s, "atomsel('%s', molid=%d, frame=%d)",
00067 sel->cmdStr, sel->molid(), sel->which_frame);
00068 PyObject *result = PyString_FromString(s);
00069 delete [] s;
00070 return result;
00071 }
00072
00073 static PyObject *
00074 atomsel_str( PyAtomSelObject *obj ) {
00075 return PyString_FromString(obj->atomSel->cmdStr);
00076 }
00077
00078
00079 static PyObject *
00080 atomsel_new( PyTypeObject *type, PyObject *args, PyObject *kwds ) {
00081
00082 const char *sel = "all";
00083 int molid = -1;
00084 int frame = -1;
00085
00086 static char *kwlist[] = { (char *)"selection", (char *)"molid",
00087 (char *)"frame", NULL };
00088
00089 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|sii", kwlist,
00090 &sel, &molid, &frame))
00091 return NULL;
00092
00093 VMDApp *app = get_vmdapp();
00094 if (molid < 0) molid = app->molecule_top();
00095 DrawMolecule *mol = app->moleculeList->mol_from_id(molid);
00096 if (!mol) {
00097 PyErr_SetString(PyExc_ValueError, "invalid molid");
00098 return NULL;
00099 }
00100
00101 AtomSel *atomSel = new AtomSel(app->atomSelParser, mol->id());
00102 atomSel->which_frame = frame;
00103 int parse_result;
00104
00105 Py_BEGIN_ALLOW_THREADS
00106 parse_result = atomSel->change(sel, mol);
00107 Py_END_ALLOW_THREADS
00108
00109 if (parse_result == AtomSel::NO_PARSE) {
00110 PyErr_SetString(PyExc_ValueError, "cannot parse atom selection text");
00111 delete atomSel;
00112 return NULL;
00113 }
00114
00115 PyObject *obj = type->tp_alloc(type, 0);
00116 if (obj == NULL) {
00117 delete atomSel;
00118 return NULL;
00119 }
00120 ((PyAtomSelObject *)obj)->app = app;
00121 ((PyAtomSelObject *)obj)->atomSel = atomSel;
00122 return obj;
00123 }
00124
00125 static char *get_doc = (char *)
00126 "get( attrib ) -> corresponding attrib values for selected atoms."
00127 ;
00128
00129 static PyObject *atomsel_get(PyAtomSelObject *a, PyObject *keyobj) {
00130
00131
00132
00133 VMDApp *app = a->app;
00134 AtomSel *atomSel = a->atomSel;
00135 const int num_atoms = atomSel->num_atoms;
00136 DrawMolecule *mol;
00137 if (!(mol = get_molecule(a))) return NULL;
00138
00139 const char *attr = PyString_AsString(keyobj);
00140 if (!attr) return NULL;
00141
00142
00143
00144
00145 SymbolTable *table = app->atomSelParser;
00146 int attrib_index = table->find_attribute(attr);
00147 if (attrib_index == -1) {
00148 PyErr_SetString(PyExc_ValueError, "unknown atom attribute");
00149 return NULL;
00150 }
00151 SymbolTableElement *elem = table->fctns.data(attrib_index);
00152 if (elem->is_a != SymbolTableElement::KEYWORD &&
00153 elem->is_a != SymbolTableElement::SINGLEWORD) {
00154 PyErr_SetString(PyExc_ValueError, "attribute is not a keyword or singleword");
00155 return NULL;
00156 }
00157
00158
00159
00160
00161
00162 int *flgs = atomSel->on;
00163 atomsel_ctxt context(table, mol, atomSel->which_frame, attr);
00164 PyObject *newlist = PyList_New(atomSel->selected);
00165
00166 if (elem->is_a == SymbolTableElement::SINGLEWORD) {
00167 int *tmp = new int[num_atoms];
00168 memcpy(tmp, atomSel->on, num_atoms*sizeof(int));
00169 elem->keyword_single(&context, num_atoms, tmp);
00170 int j=0;
00171 for (int i=0; i<num_atoms; i++) {
00172 if (flgs[i]) {
00173 PyObject *val = tmp[i] ? Py_True : Py_False;
00174 Py_INCREF(val);
00175 PyList_SET_ITEM(newlist, j++, val);
00176 }
00177 }
00178 delete [] tmp;
00179 } else {
00180 switch(table->fctns.data(attrib_index)->returns_a) {
00181 case (SymbolTableElement::IS_STRING):
00182 {
00183 const char **tmp= new const char *[num_atoms];
00184 elem->keyword_string(&context, num_atoms, tmp, flgs);
00185 int j=0;
00186 for (int i=0; i<num_atoms; i++) {
00187 if (flgs[i]) {
00188 PyList_SET_ITEM(newlist, j++, PyString_FromString(tmp[i]));
00189 }
00190 }
00191 delete [] tmp;
00192 }
00193 break;
00194 case (SymbolTableElement::IS_INT):
00195 {
00196 int *tmp = new int[num_atoms];
00197 elem->keyword_int(&context, num_atoms, tmp, flgs);
00198 int j=0;
00199 for (int i=0; i<num_atoms; i++) {
00200 if (flgs[i]) {
00201 PyList_SET_ITEM(newlist, j++, PyInt_FromLong(tmp[i]));
00202 }
00203 }
00204 delete [] tmp;
00205 }
00206 break;
00207 case (SymbolTableElement::IS_FLOAT):
00208 {
00209 double *tmp = new double[num_atoms];
00210 elem->keyword_double(&context, num_atoms, tmp, flgs);
00211 int j=0;
00212 for (int i=0; i<num_atoms; i++) {
00213 if (flgs[i])
00214 PyList_SET_ITEM(newlist, j++, PyFloat_FromDouble(tmp[i]));
00215 }
00216 delete [] tmp;
00217 }
00218 break;
00219 }
00220 }
00221 return newlist;
00222 }
00223
00224 static char *set_doc = (char *)
00225 "set( attrib, val ) -> set attrib values for selected atoms.\n"
00226 " val must be either a single value, or a sequence of values, one for\n"
00227 " each atom in selection.\n" ;
00228
00229 static PyObject *atomsel_set(PyAtomSelObject *a, PyObject *args) {
00230
00231
00232
00233 VMDApp *app = a->app;
00234 AtomSel *atomSel = a->atomSel;
00235 const int num_atoms = atomSel->num_atoms;
00236 DrawMolecule *mol;
00237 if (!(mol = get_molecule(a))) return NULL;
00238
00239 const char *attr;
00240 PyObject *val;
00241 if (!PyArg_ParseTuple(args, "sO", &attr, &val)) return NULL;
00242
00243
00244
00245
00246 SymbolTable *table = app->atomSelParser;
00247 int attrib_index = table->find_attribute(attr);
00248 if (attrib_index == -1) {
00249 PyErr_SetString(PyExc_ValueError, "unknown atom attribute");
00250 return NULL;
00251 }
00252 SymbolTableElement *elem = table->fctns.data(attrib_index);
00253 if (elem->is_a != SymbolTableElement::KEYWORD &&
00254 elem->is_a != SymbolTableElement::SINGLEWORD) {
00255 PyErr_SetString(PyExc_ValueError, "attribute is not a keyword or singleword");
00256 return NULL;
00257 }
00258 if (!table->is_changeable(attrib_index)) {
00259 PyErr_SetString(PyExc_ValueError, "attribute is not modifiable");
00260 return NULL;
00261 }
00262
00263
00264
00265
00266
00267
00268 const int num_selected = atomSel->selected;
00269 PyObject *fastval = NULL;
00270 int nvals = 1;
00271 if (PySequence_Check(val) && !PyString_Check(val)) {
00272 nvals = PySequence_Size(val);
00273 if (nvals != 1 && nvals != num_selected) {
00274 PyErr_SetString(PyExc_ValueError, "wrong number of items");
00275 return NULL;
00276 }
00277 fastval = PySequence_Fast(val, "Invalid values");
00278 if (!fastval) return NULL;
00279 }
00280
00281 int *flgs = atomSel->on;
00282
00283
00284
00285
00286
00287
00288 atomsel_ctxt context(table, mol, atomSel->which_frame, NULL);
00289 if (elem->returns_a == SymbolTableElement::IS_INT) {
00290 int *list = new int[num_atoms];
00291 if (fastval) {
00292 int j=0;
00293 for (int i=0; i<num_atoms; i++) {
00294 if (!flgs[i]) continue;
00295 int ival = (int)PyInt_AsLong(PySequence_Fast_GET_ITEM(fastval, j++));
00296
00297 list[i] = ival;
00298 }
00299 } else {
00300 int ival = (int)PyInt_AsLong(val);
00301
00302 for (int i=0; i<num_atoms; i++) {
00303 if (flgs[i]) list[i] = ival;
00304 }
00305 }
00306 elem->set_keyword_int(&context, num_atoms, list, flgs);
00307 delete [] list;
00308
00309 } else if (elem->returns_a == SymbolTableElement::IS_FLOAT) {
00310 double *list = new double[num_atoms];
00311 if (fastval) {
00312 int j=0;
00313 for (int i=0; i<num_atoms; i++) {
00314 if (!flgs[i]) continue;
00315 float dval = (float)PyFloat_AsDouble(PySequence_Fast_GET_ITEM(fastval, j++));
00316 list[i] = dval;
00317 }
00318 } else {
00319 float dval = (float)PyFloat_AsDouble(val);
00320 for (int i=0; i<num_atoms; i++) {
00321 if (flgs[i]) list[i] = dval;
00322 }
00323 }
00324 elem->set_keyword_double(&context, num_atoms, list, flgs);
00325 delete [] list;
00326
00327
00328 } else if (elem->returns_a == SymbolTableElement::IS_STRING) {
00329 const char **list = new const char *[num_atoms];
00330 if (fastval) {
00331 int j=0;
00332 for (int i=0; i<num_atoms; i++) {
00333 if (!flgs[i]) continue;
00334 const char *sval = PyString_AsString(PySequence_Fast_GET_ITEM(fastval, j++));
00335 list[i] = sval;
00336 }
00337 } else {
00338 const char *sval = PyString_AsString(val);
00339 for (int i=0; i<num_atoms; i++) {
00340 if (flgs[i]) list[i] = sval;
00341 }
00342 }
00343 elem->set_keyword_string(&context, num_atoms, list, flgs);
00344 delete [] list;
00345 }
00346
00347
00348 if (!strcmp(attr, "name") ||
00349 !strcmp(attr, "type") ||
00350 !strcmp(attr, "resname") ||
00351 !strcmp(attr, "chain") ||
00352 !strcmp(attr, "segid") ||
00353 !strcmp(attr, "segname"))
00354 app->moleculeList->add_color_names(mol->id());
00355
00356 mol->force_recalc(DrawMolItem::SEL_REGEN | DrawMolItem::COL_REGEN);
00357
00358 Py_XDECREF(fastval);
00359 Py_INCREF(Py_None);
00360 return Py_None;
00361 }
00362
00363 static PyObject *getframe(PyAtomSelObject *a, void *) {
00364
00365 AtomSel *atomSel = a->atomSel;
00366 DrawMolecule *mol;
00367 if (!(mol = get_molecule(a))) return NULL;
00368
00369 return PyInt_FromLong(atomSel->which_frame);
00370 }
00371
00372 static int setframe(PyAtomSelObject *a, PyObject *frameobj, void *) {
00373
00374 AtomSel *atomSel = a->atomSel;
00375 DrawMolecule *mol;
00376 if (!(mol = get_molecule(a))) return -1;
00377
00378 int frame = PyInt_AsLong(frameobj);
00379 if (frame < 0 && PyErr_Occurred()) return -1;
00380 if (frame != AtomSel::TS_LAST && frame != AtomSel::TS_NOW &&
00381 (frame < 0 || frame >= mol->numframes())) {
00382 PyErr_SetString(PyExc_ValueError, "Invalid frame");
00383 return -1;
00384 }
00385 atomSel->which_frame = frame;
00386 return 0;
00387 }
00388
00389 static char *frame_doc = (char *)
00390 "frame -- the animation frame referenced by the selection.\n"
00391 " Special values: -1 = current frame, -2 = last frame.\n"
00392 "NOTE: changing the frame does not update the selection;\n"
00393 " use update() to do that.\n";
00394
00395 static char *update_doc = (char *)
00396 "update() -> Recompute atoms in selection; if the selection is\n"
00397 "distance-based (e.g., if it uses 'within'), changes to the frame\n"
00398 "will not be reflected in the selected atoms until this method\n"
00399 "is called.\n" ;
00400
00401 static PyObject *py_update(PyAtomSelObject *a) {
00402
00403 AtomSel *atomSel = a->atomSel;
00404 DrawMolecule *mol;
00405 if (!(mol = get_molecule(a))) return NULL;
00406
00407 Py_BEGIN_ALLOW_THREADS
00408 atomSel->change(NULL, mol);
00409 Py_END_ALLOW_THREADS
00410
00411 Py_INCREF(Py_None);
00412 return Py_None;
00413 }
00414
00415 static char *write_doc = (char *)
00416 "write(filetype, filename) -> None\n"
00417 "Write the atoms in the selection to a file of the given type.";
00418
00419 static PyObject *py_write(PyAtomSelObject *a, PyObject *args, PyObject *kwds) {
00420
00421 AtomSel *atomSel = a->atomSel;
00422 DrawMolecule *mol;
00423 if (!(mol = get_molecule(a))) return NULL;
00424 const char *filetype, *filename;
00425 static char *kwlist[] = { (char *)"filetype", (char *)"filename", NULL };
00426 if (!PyArg_ParseTupleAndKeywords(args, kwds, "ss", kwlist,
00427 &filetype, &filename))
00428 return NULL;
00429
00430 int frame=-1;
00431 switch (atomSel -> which_frame) {
00432 case AtomSel::TS_NOW: frame = mol->frame(); break;
00433 case AtomSel::TS_LAST: frame = mol->numframes()-1; break;
00434 default: frame = atomSel->which_frame; break;
00435 }
00436 if (frame < 0) frame = 0;
00437 else if (frame >= mol->numframes()) frame = mol->numframes()-1;
00438
00439 FileSpec spec;
00440 spec.first = frame;
00441 spec.last = frame;
00442 spec.stride = 1;
00443 spec.waitfor = -1;
00444 spec.selection = atomSel->on;
00445 int rc;
00446 VMDApp *app = get_vmdapp();
00447 Py_BEGIN_ALLOW_THREADS
00448 rc = app->molecule_savetrajectory(mol->id(), filename, filetype, &spec);
00449 Py_END_ALLOW_THREADS
00450
00451 if (rc < 0) {
00452 PyErr_SetString(PyExc_ValueError, "Unable to write selection to file.");
00453 return NULL;
00454 }
00455 Py_INCREF(Py_None);
00456 return Py_None;
00457 }
00458
00459 static PyObject *getbonds(PyAtomSelObject *a, void *) {
00460
00461 AtomSel *atomSel = a->atomSel;
00462 DrawMolecule *mol;
00463 if (!(mol = get_molecule(a))) return NULL;
00464
00465 PyObject *newlist = PyList_New(atomSel->selected);
00466
00467 int k=0;
00468 for (int i=0; i< atomSel->num_atoms; i++) {
00469 if (!atomSel->on[i]) continue;
00470 const MolAtom *atom = mol->atom(i);
00471 PyObject *bondlist = PyList_New(atom->bonds);
00472 for (int j=0; j<atom->bonds; j++) {
00473 PyList_SET_ITEM(bondlist, j, PyInt_FromLong(atom->bondTo[j]));
00474 }
00475 PyList_SET_ITEM(newlist, k++, bondlist);
00476 }
00477 return newlist;
00478 }
00479
00480 static int setbonds(PyAtomSelObject *a, PyObject *obj, void *) {
00481
00482 AtomSel *atomSel = a->atomSel;
00483 DrawMolecule *mol;
00484 if (!(mol = get_molecule(a))) return -1;
00485
00486 if (!PySequence_Check(obj) || PySequence_Size(obj) != atomSel->selected) {
00487 PyErr_SetString(PyExc_ValueError,
00488 (char *)"setbonds: atomlist and bondlist must have the same size");
00489 return -1;
00490 }
00491 PyObject *fastbonds = PySequence_Fast(obj, "Argument to setbonds must be sequence");
00492 if (!fastbonds) return -1;
00493
00494 int ibond = 0;
00495 mol->force_recalc(DrawMolItem::MOL_REGEN);
00496 for (int i=0; i<atomSel->num_atoms; i++) {
00497 if (!atomSel->on[i]) continue;
00498 MolAtom *atom = mol->atom(i);
00499
00500 PyObject *atomids = PySequence_Fast_GET_ITEM(fastbonds, ibond++);
00501 if (!PyList_Check(atomids)) continue;
00502 int numbonds = PyList_Size(atomids);
00503 int k=0;
00504 for (int j=0; j<numbonds; j++) {
00505 int bond = PyInt_AsLong(PyList_GET_ITEM(atomids, j));
00506 if (bond >= 0 && bond < mol->nAtoms) {
00507 atom->bondTo[k++] = bond;
00508 }
00509 }
00510 atom->bonds = k;
00511 }
00512 Py_DECREF(fastbonds);
00513 return 0;
00514 }
00515
00516 static char *bonds_doc = (char *)
00517 "bonds - for each atom in selection, a list of the indices\n"
00518 " of the bonded atoms.\n";
00519
00520 static PyObject *getmolid(PyAtomSelObject *a, void *) {
00521 AtomSel *atomSel = a->atomSel;
00522 return PyInt_FromLong(atomSel->molid());
00523 }
00524
00525 static char *molid_doc = (char *)
00526 "molid - The id of the molecule this selection is associated with.\n";
00527
00528
00529 static PyGetSetDef atomsel_getset[] = {
00530 { (char *)"frame", (getter)getframe, (setter)setframe, frame_doc, NULL },
00531 { (char *)"bonds", (getter)getbonds, (setter)setbonds, bonds_doc, NULL },
00532 { (char *)"molid", (getter)getmolid, (setter)NULL, molid_doc, NULL },
00533 { NULL },
00534 };
00535
00536
00537
00538 static PyObject *macro(PyObject *self, PyObject *args, PyObject *keywds) {
00539 char *name = NULL, *selection = NULL;
00540 static char *kwlist[] = {
00541 (char *)"name", (char *)"selection", NULL
00542 };
00543 if (!PyArg_ParseTupleAndKeywords(args, keywds, (char *)"|ss:atomsel.macro", kwlist, &name, &selection))
00544 return NULL;
00545
00546 if (selection && !name) {
00547 PyErr_SetString(PyExc_ValueError, (char *)"Must specify name for macro");
00548 return NULL;
00549 }
00550 SymbolTable *table = get_vmdapp()->atomSelParser;
00551 if (!name && !selection) {
00552
00553 PyObject *macrolist = PyList_New(0);
00554 for (int i=0; i<table->num_custom_singleword(); i++) {
00555 const char *s = table->custom_singleword_name(i);
00556 if (s && strlen(s))
00557 PyList_Append(macrolist, PyString_FromString(s));
00558 }
00559 return macrolist;
00560 }
00561 if (name && !selection) {
00562
00563 const char *s = table->get_custom_singleword(name);
00564 if (!s) {
00565 PyErr_SetString(PyExc_ValueError, (char *)"No macro for given name");
00566 return NULL;
00567 }
00568 return PyString_FromString(s);
00569 }
00570
00571 if (!table->add_custom_singleword(name, selection)) {
00572 PyErr_SetString(PyExc_ValueError, (char *)"Unable to create new macro");
00573 return NULL;
00574 }
00575 Py_INCREF(Py_None);
00576 return Py_None;
00577 }
00578
00579
00580 static PyObject *delmacro(PyObject *self, PyObject *args) {
00581 char *name;
00582 if (!PyArg_ParseTuple(args, (char *)"s:atomsel.delmacro", &name))
00583 return NULL;
00584 if (!get_vmdapp()->atomSelParser->remove_custom_singleword(name)) {
00585 PyErr_SetString(PyExc_ValueError, (char *)"Unable to remove macro");
00586 return NULL;
00587 }
00588 Py_INCREF(Py_None);
00589 return Py_None;
00590 }
00591
00592
00593 static int py_get_vector(PyObject *matobj, int n, float *vec) {
00594
00595 if (!PySequence_Check(matobj) || PySequence_Size(matobj) != n) {
00596 PyErr_SetString(PyExc_ValueError, "vector has incorrect size");
00597 return 0;
00598 }
00599 PyObject *fastval = PySequence_Fast(matobj, "Invalid sequence");
00600 if (!fastval) return 0;
00601
00602 for (int i=0; i<n; i++) {
00603 vec[i] = PyFloat_AsDouble(PySequence_Fast_GET_ITEM(fastval, i));
00604 if (PyErr_Occurred()) {
00605 Py_DECREF(fastval);
00606 return 0;
00607 }
00608 }
00609 Py_DECREF(fastval);
00610 return 1;
00611 }
00612
00613 static PyObject *inverse(PyObject *self, PyObject *matobj) {
00614 Matrix4 mat;
00615 if (!py_get_vector(matobj, 16, mat.mat)) return NULL;
00616 if (mat.inverse()) {
00617 PyErr_SetString(PyExc_ValueError, "Matrix is singular.");
00618 return NULL;
00619 }
00620 PyObject *result = PyTuple_New(16);
00621 for (int i=0; i<16; i++) {
00622 PyTuple_SET_ITEM(result, i, PyFloat_FromDouble(mat.mat[i]));
00623 }
00624 return result;
00625 }
00626
00627 #define SYMBOL_TABLE_FUNC(funcname, elemtype) \
00628 static PyObject *funcname(PyObject *self) { \
00629 VMDApp *app = get_vmdapp(); \
00630 PyObject *result = PyList_New(0); \
00631 SymbolTable *table = app->atomSelParser; \
00632 int i, n = table->fctns.num(); \
00633 for (i=0; i<n; i++) \
00634 if (table->fctns.data(i)->is_a == elemtype) \
00635 PyList_Append(result, PyString_FromString(table->fctns.name(i))); \
00636 return result; \
00637 }
00638
00639 SYMBOL_TABLE_FUNC(keywords, SymbolTableElement::KEYWORD)
00640 SYMBOL_TABLE_FUNC(booleans, SymbolTableElement::SINGLEWORD)
00641 SYMBOL_TABLE_FUNC(functions, SymbolTableElement::FUNCTION)
00642 SYMBOL_TABLE_FUNC(stringfunctions, SymbolTableElement::STRINGFCTN)
00643
00644 #undef SYMBOL_TABLE_FUNC
00645
00646
00647
00648 static float *parse_weight(AtomSel *sel, PyObject *wtobj) {
00649 float *weight = new float[sel->selected];
00650 if (!wtobj || wtobj == Py_None) {
00651 for (int i=0; i<sel->selected; i++) weight[i] = 1.0f;
00652 return weight;
00653 }
00654
00655 PyObject *seq = PySequence_Fast(wtobj, (char *)"weight must be a sequence.");
00656 if (!seq) return NULL;
00657 if (PySequence_Size(seq) != sel->selected) {
00658 Py_DECREF(seq);
00659 PyErr_SetString(PyExc_ValueError, "weight must be same size as selection.");
00660 delete [] weight;
00661 return NULL;
00662 }
00663 for (int i=0; i<sel->selected; i++) {
00664 double tmp = PyFloat_AsDouble(PySequence_Fast_GET_ITEM(seq, i));
00665 if (PyErr_Occurred()) {
00666 PyErr_SetString(PyExc_ValueError, "non-floating point value found in weight.");
00667 Py_DECREF(seq);
00668 delete [] weight;
00669 return NULL;
00670 }
00671 weight[i] = (float)tmp;
00672 }
00673 Py_DECREF(seq);
00674 return weight;
00675 }
00676
00677 static char *minmax_doc = (char *)
00678 "minmax() -> (min, max)\n"
00679 "Return minimum and maximum coordinates for selected atoms.";
00680
00681 static PyObject *minmax(PyAtomSelObject *a, PyObject *withradii) {
00682
00683 AtomSel *atomSel = a->atomSel;
00684 DrawMolecule *mol;
00685 if (!(mol = get_molecule(a))) return NULL;
00686 const float *radii = NULL;
00687 if (withradii && PyObject_IsTrue(withradii)) {
00688 radii = mol->extraflt.data("radius");
00689 }
00690
00691 const float *pos = atomSel->coordinates(a->app->moleculeList);
00692 float min[3], max[3];
00693 int rc = measure_minmax(atomSel->num_atoms, atomSel->on, pos, radii, min, max);
00694 if (rc < 0) {
00695 PyErr_SetString(PyExc_ValueError, measure_error(rc));
00696 return NULL;
00697 }
00698 return Py_BuildValue("[f,f,f],[f,f,f]",
00699 min[0], min[1], min[2], max[0], max[1], max[2]);
00700 }
00701
00702 static char *center_doc = (char *)
00703 "center(weight=None) -> (x, y, z)\n"
00704 "Return a tuple corresponding to the center of the selection,\n"
00705 "optionally weighted by weight.";
00706
00707 static PyObject *center(PyAtomSelObject *a, PyObject *args, PyObject *kwds) {
00708
00709 AtomSel *atomSel = a->atomSel;
00710 DrawMolecule *mol;
00711 if (!(mol = get_molecule(a))) return NULL;
00712 PyObject *weightobj = NULL;
00713 static char *kwlist[] = { "weight", NULL };
00714
00715 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", kwlist, &weightobj))
00716 return NULL;
00717 float *weight = parse_weight(atomSel, weightobj);
00718 if (!weight) return NULL;
00719
00720 float cen[3];
00721
00722 int ret_val = measure_center(atomSel,
00723 atomSel->coordinates(a->app->moleculeList),
00724 weight, cen);
00725 delete [] weight;
00726 if (ret_val < 0) {
00727 PyErr_SetString(PyExc_ValueError, measure_error(ret_val));
00728 return NULL;
00729 }
00730 return Py_BuildValue("(f,f,f)", cen[0], cen[1], cen[2]);
00731 }
00732
00733 static float *parse_two_selections_return_weight(PyAtomSelObject *a,
00734 PyObject *args, AtomSel **othersel) {
00735
00736 PyObject *other;
00737 PyObject *weightobj = NULL;
00738
00739 AtomSel *atomSel = a->atomSel;
00740 DrawMolecule *mol;
00741 if (!(mol = get_molecule(a))) return NULL;
00742
00743 if (!PyArg_ParseTuple(args, "O|O", &other, &weightobj))
00744 return NULL;
00745 float *weight = parse_weight(atomSel, weightobj);
00746 if (!weight) return NULL;
00747 if (!atomsel_Check(other)) return NULL;
00748 AtomSel *sel2 = ((PyAtomSelObject *)other)->atomSel;
00749 if (atomSel->selected != sel2->selected) {
00750 PyErr_SetString(PyExc_ValueError, "Selections must have same number of atoms");
00751 return NULL;
00752 }
00753 *othersel = sel2;
00754 return weight;
00755 }
00756
00757 static char *rmsd_doc = (char *)
00758 "rmsd(sel, weight=None) -> rms distance between selections.\n"
00759 " Selections must have the same number of atoms.\n"
00760 " Weight must be None or list of same size as selections.";
00761
00762 static PyObject *py_rmsd(PyAtomSelObject *a, PyObject *args) {
00763
00764 AtomSel *sel2;
00765 float *weight = parse_two_selections_return_weight(a, args, &sel2);
00766 if (!weight) return NULL;
00767
00768 float rmsd;
00769 int rc = measure_rmsd(a->atomSel, sel2, a->atomSel->selected,
00770 a->atomSel->coordinates(a->app->moleculeList),
00771 sel2->coordinates(a->app->moleculeList),
00772 weight, &rmsd);
00773 delete [] weight;
00774 if (rc < 0) {
00775 PyErr_SetString(PyExc_ValueError, measure_error(rc));
00776 return NULL;
00777 }
00778 return PyFloat_FromDouble(rmsd);
00779 }
00780
00781 static char *fit_doc = (char *)
00782 "fit(sel, weight=None) -> transformation matrix\n"
00783 " Compute and return the transformation matrix for the RMS alignment\n"
00784 " of the selection to sel. The format of the matrix is a 16-element\n"
00785 " tuple suitable for passing to the move() method.\n"
00786 " Weight must be None or list of same size as selections.";
00787
00788 static PyObject *py_fit(PyAtomSelObject *a, PyObject *args) {
00789
00790 AtomSel *sel2;
00791 float *weight = parse_two_selections_return_weight(a, args, &sel2);
00792 if (!weight) return NULL;
00793
00794 Matrix4 mat;
00795 int rc = measure_fit(a->atomSel, sel2,
00796 a->atomSel->coordinates(a->app->moleculeList),
00797 sel2->coordinates(a->app->moleculeList),
00798 weight, NULL, &mat);
00799 delete [] weight;
00800 if (rc < 0) {
00801 PyErr_SetString(PyExc_ValueError, measure_error(rc));
00802 return NULL;
00803 }
00804 PyObject *result = PyTuple_New(16);
00805 for (int i=0; i<16; i++) {
00806 PyTuple_SET_ITEM(result, i, PyFloat_FromDouble(mat.mat[i]));
00807 }
00808 return result;
00809 }
00810
00811 static char *moveby_doc = (char *)
00812 "moveby( vec ) -> shift the selection by the three-element vector vec.";
00813
00814 static PyObject *py_moveby(PyAtomSelObject *a, PyObject *vecobj) {
00815
00816 AtomSel *atomSel = a->atomSel;
00817 DrawMolecule *mol;
00818 if (!(mol = get_molecule(a))) return NULL;
00819
00820 float offset[3];
00821 if (!py_get_vector(vecobj, 3, offset)) return NULL;
00822 float *pos = atomSel->coordinates(a->app->moleculeList);
00823 if (!pos) {
00824 PyErr_SetString(PyExc_ValueError, "No coordinates");
00825 return NULL;
00826 }
00827 for (int i=0; i<atomSel->num_atoms; i++) {
00828 if (atomSel->on[i]) {
00829 vec_add(pos, pos, offset);
00830 }
00831 pos += 3;
00832 }
00833 mol->force_recalc(DrawMolItem::MOL_REGEN);
00834 Py_INCREF(Py_None);
00835 return Py_None;
00836 }
00837
00838 static char *move_doc = (char *)
00839 "move( matrix ) -> apply coordinate transformation to selection.\n"
00840 " matrix should be of the form returned by fit().";
00841
00842 static PyObject *py_move(PyAtomSelObject *a, PyObject *matobj) {
00843
00844 AtomSel *atomSel = a->atomSel;
00845 DrawMolecule *mol;
00846 if (!(mol = get_molecule(a))) return NULL;
00847
00848 Matrix4 mat;
00849 if (!py_get_vector(matobj, 16, mat.mat)) return NULL;
00850
00851 int err;
00852 if ((err = measure_move(
00853 atomSel,
00854 atomSel->coordinates(a->app->moleculeList),
00855 mat)) != MEASURE_NOERR) {
00856 PyErr_SetString(PyExc_ValueError, measure_error(err));
00857 return NULL;
00858 }
00859 mol->force_recalc(DrawMolItem::MOL_REGEN);
00860 Py_INCREF(Py_None);
00861 return Py_None;
00862 }
00863
00864 static char *contacts_doc = (char *)
00865 "contacts(sel, cutoff) -> contact pairs\n"
00866 "Return two lists, whose corresponding elements contain atom indices\n"
00867 "in selection that are within cutoff of sel, but not directly bonded.\n";
00868
00869
00870 static PyObject *contacts(PyAtomSelObject *a, PyObject *args) {
00871 AtomSel *sel1 = a->atomSel;
00872 DrawMolecule *mol;
00873 if (!(mol = get_molecule(a))) return NULL;
00874
00875 PyObject *obj2;
00876 float cutoff;
00877 if (!PyArg_ParseTuple(args, "Of", &obj2, &cutoff))
00878 return NULL;
00879 if (!atomsel_Check(obj2)) return NULL;
00880 if (!(get_molecule((PyAtomSelObject *)obj2))) return NULL;
00881 AtomSel *sel2 = ((PyAtomSelObject *)obj2)->atomSel;
00882
00883 const float *ts1 = sel1->coordinates(a->app->moleculeList);
00884 const float *ts2 = sel2->coordinates(a->app->moleculeList);
00885 if (!ts1 || !ts2) {
00886 PyErr_SetString(PyExc_ValueError, "No coordinates in selection");
00887 return NULL;
00888 }
00889
00890 GridSearchPair *pairlist = vmd_gridsearch3(
00891 ts1, sel1->num_atoms, sel1->on,
00892 ts2, sel2->num_atoms, sel2->on,
00893 cutoff, -1, (sel1->num_atoms + sel2->num_atoms) * 27);
00894
00895 GridSearchPair *p, *tmp;
00896 PyObject *list1 = PyList_New(0);
00897 PyObject *list2 = PyList_New(0);
00898 PyObject *tmp1;
00899 PyObject *tmp2;
00900 for (p=pairlist; p != NULL; p=tmp) {
00901
00902 MolAtom *a1 = mol->atom(p->ind1);
00903 if (sel1->molid() != sel2->molid() || !a1->bonded(p->ind2)) {
00904
00905
00906
00907
00908 tmp1 = PyInt_FromLong(p->ind1);
00909 tmp2 = PyInt_FromLong(p->ind2);
00910 PyList_Append(list1, tmp1);
00911 PyList_Append(list2, tmp2);
00912 Py_DECREF(tmp1);
00913 Py_DECREF(tmp2);
00914 }
00915 tmp = p->next;
00916 free(p);
00917 }
00918 PyObject *result = PyList_New(2);
00919 PyList_SET_ITEM(result, 0, list1);
00920 PyList_SET_ITEM(result, 1, list2);
00921 return result;
00922 }
00923
00924 static char *sasa_doc = (char *)
00925 "sasa(srad, sel, ... ) -> solvent accessible surface area.\n"
00926 "srad gives solvent radius.\n"
00927 "Optional keyword arguments:\n"
00928 " samples=500 -- specifies number of sample points used per atom.\n"
00929 " points=None -- If points is a list, coordinates of surface points\n"
00930 " will be appended to the list.\n"
00931 " restrict=None -- Pass an atom selection as argument to restrict\n"
00932 " to find contributions coming from just atoms in restrict.\n";
00933
00934 static PyObject *sasa(PyAtomSelObject *a, PyObject *args, PyObject *keywds) {
00935 float srad = 0;
00936 int samples = -1;
00937 const int *sampleptr = NULL;
00938 PyObject *pointsobj = NULL;
00939 PyObject *restrictobj = NULL;
00940
00941 AtomSel *sel = a->atomSel;
00942 DrawMolecule *mol;
00943 if (!(mol = get_molecule(a))) return NULL;
00944
00945 static char *kwlist[] = {
00946 "srad", "samples", "points", "restrict", NULL
00947 };
00948 if (!PyArg_ParseTupleAndKeywords(args, keywds,
00949 "f|iOO:atomsel.sasa", kwlist,
00950 &srad, &samples, &pointsobj, &restrictobj))
00951 return NULL;
00952
00953
00954 if (srad < 0) {
00955 PyErr_SetString(PyExc_ValueError, "atomselect.sasa: srad must be non-negative.");
00956 return NULL;
00957 }
00958
00959
00960 const float *radii = mol->extraflt.data("radius");
00961 const float *coords = sel->coordinates(a->app->moleculeList);
00962
00963
00964 if (samples > 1) sampleptr = &samples;
00965
00966
00967 AtomSel *restrictsel = NULL;
00968 if (restrictobj) {
00969 if (!atomsel_Check(restrictobj)) return NULL;
00970 if (!get_molecule((PyAtomSelObject *)restrictobj)) return NULL;
00971 restrictsel = ((PyAtomSelObject *)restrictobj)->atomSel;
00972 }
00973
00974
00975 ResizeArray<float> sasapts;
00976 ResizeArray<float> *sasaptsptr = pointsobj ? &sasapts : NULL;
00977
00978
00979 float sasa = 0;
00980 int rc = measure_sasa(sel, coords, radii, srad, &sasa,
00981 sasaptsptr, restrictsel, sampleptr);
00982 if (rc) {
00983 PyErr_SetString(PyExc_ValueError, measure_error(rc));
00984 return NULL;
00985 }
00986
00987
00988 if (pointsobj) {
00989 for (int i=0; i<sasapts.num(); i++) {
00990 PyList_Append(pointsobj, PyFloat_FromDouble(sasapts[i]));
00991 }
00992 }
00993
00994
00995 return PyFloat_FromDouble(sasa);
00996 }
00997
00998
00999
01000
01001
01002 static Py_ssize_t
01003 atomselection_length( PyObject *a ) {
01004 return ((PyAtomSelObject *)a)->atomSel->selected;
01005 }
01006
01007
01008 static PyObject *
01009 atomselection_subscript( PyAtomSelObject * a, PyObject * keyobj ) {
01010 long ind = PyInt_AsLong(keyobj);
01011 if (ind < 0 && PyErr_Occurred()) return NULL;
01012 PyObject *result = Py_False;
01013 if (ind >= 0 && ind < a->atomSel->num_atoms &&
01014 a->atomSel->on[ind]) {
01015 result = Py_True;
01016 }
01017 Py_INCREF(result);
01018 return result;
01019 }
01020
01021 static PyMappingMethods atomsel_mapping = {
01022 atomselection_length,
01023 (binaryfunc)atomselection_subscript,
01024 0
01025 };
01026
01027
01028 static PyMethodDef atomselection_methods[] = {
01029 { "get", (PyCFunction)atomsel_get, METH_O, get_doc },
01030 { "set", (PyCFunction)atomsel_set, METH_VARARGS, set_doc },
01031 { "update", (PyCFunction)py_update, METH_NOARGS, update_doc },
01032 { "write", (PyCFunction)py_write, METH_VARARGS|METH_KEYWORDS, write_doc },
01033 { "minmax", (PyCFunction)minmax, METH_VARARGS, minmax_doc },
01034 { "center", (PyCFunction)center, METH_VARARGS|METH_KEYWORDS, center_doc },
01035 { "rmsd", (PyCFunction)py_rmsd, METH_VARARGS, rmsd_doc, },
01036 { "fit", (PyCFunction)py_fit, METH_VARARGS, fit_doc, },
01037 { "move", (PyCFunction)py_move, METH_O, move_doc },
01038 { "moveby", (PyCFunction)py_moveby, METH_O, moveby_doc },
01039 { "contacts", (PyCFunction)contacts, METH_VARARGS, contacts_doc },
01040 { "sasa", (PyCFunction)sasa, METH_VARARGS | METH_KEYWORDS, sasa_doc },
01041 { NULL, NULL }
01042 };
01043
01044
01045
01046 namespace {
01047 typedef struct {
01048 PyObject_HEAD
01049 int index;
01050 PyAtomSelObject * a;
01051 } iterobject;
01052
01053 PyObject *atomsel_iter(PyObject *);
01054
01055 PyObject *iter_next(iterobject *it) {
01056 for ( ; it->index < it->a->atomSel->num_atoms; ++it->index) {
01057 if (it->a->atomSel->on[it->index])
01058 return PyInt_FromLong(it->index++);
01059 }
01060 return NULL;
01061 }
01062
01063 void iter_dealloc(iterobject *it) {
01064 Py_XDECREF(it->a);
01065 }
01066 PyObject *iter_len(iterobject *it) {
01067 return PyInt_FromLong(it->a->atomSel->selected);
01068 }
01069
01070 PyMethodDef iter_methods[] = {
01071 {"__length_hint__", (PyCFunction)iter_len, METH_NOARGS },
01072 {NULL, NULL}
01073 };
01074
01075 PyTypeObject itertype = {
01076 PyObject_HEAD_INIT(&PyType_Type)
01077 0,
01078 "atomsel.iterator",
01079 sizeof(iterobject),
01080 0,
01081
01082 (destructor)iter_dealloc,
01083 0,
01084 0,
01085 0,
01086 0,
01087 0,
01088 0,
01089 0,
01090 0,
01091 0,
01092 0,
01093 0,
01094 PyObject_GenericGetAttr,
01095 0,
01096 0,
01097 Py_TPFLAGS_DEFAULT,
01098 0,
01099 0,
01100 0,
01101 0,
01102 0,
01103 PyObject_SelfIter,
01104 (iternextfunc)iter_next,
01105 iter_methods,
01106 0,
01107 };
01108
01109 PyObject *atomsel_iter(PyObject *self) {
01110 iterobject * iter = PyObject_New(iterobject, &itertype);
01111 if (!iter) return NULL;
01112 Py_INCREF( iter->a = (PyAtomSelObject *)self );
01113 iter->index = 0;
01114 return (PyObject *)iter;
01115 }
01116 }
01117
01118
01119
01120 PyTypeObject atomsel_type = {
01121 PyObject_HEAD_INIT(0)
01122 0,
01123 "atomsel.atomsel",
01124 sizeof(PyAtomSelObject),
01125 0,
01126 (destructor)atomsel_dealloc,
01127 0,
01128 0,
01129 0,
01130 0,
01131 (reprfunc)atomsel_repr,
01132 0,
01133 0,
01134 &atomsel_mapping,
01135 0,
01136 0,
01137 (reprfunc)atomsel_str,
01138 PyObject_GenericGetAttr,
01139 0,
01140 0,
01141 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE|Py_TPFLAGS_HAVE_CLASS,
01142 atomsel_doc,
01143 0,
01144 0,
01145 0,
01146 0,
01147 atomsel_iter,
01148 0,
01149 atomselection_methods,
01150 0,
01151 atomsel_getset,
01152 0,
01153 0,
01154 0,
01155 0,
01156 0,
01157 0,
01158 PyType_GenericAlloc,
01159 atomsel_new,
01160 _PyObject_Del,
01161 };
01162
01163 static int atomsel_Check(PyObject *obj) {
01164 if (PyObject_TypeCheck(obj, &atomsel_type)) return 1;
01165 PyErr_SetString(PyExc_TypeError, "expected atomsel");
01166 return 0;
01167 }
01168
01169 AtomSel * atomsel_AsAtomSel( PyObject *obj) {
01170 if (!atomsel_Check(obj)) return NULL;
01171 return ((PyAtomSelObject *)obj)->atomSel;
01172 }
01173
01174
01175 static PyMethodDef atomsel_methods[] = {
01176 {"macro", (PyCFunction)macro, METH_VARARGS | METH_KEYWORDS,
01177 "macro(name=None, selection=None) -- create and query selection macros.\n"
01178 "If both name and selection are None, return list of macro names.\n"
01179 "If selection is None, return definition for name.\n"
01180 "If both name and selection are given, define new macro.\n" },
01181 {"delmacro", (vmdPyMethod)delmacro, METH_VARARGS,
01182 "delmacro(name) -> Delete atom selection macro with given name." },
01183 {"inverse", (PyCFunction)inverse, METH_O,
01184 "inverse(matrix) -> Inverse of matrix returned by atomsel.fit(...)"},
01185 {"keywords", (PyCFunction)keywords, METH_NOARGS,
01186 "keywords() -> List of available atom selection keywords."},
01187 {"booleans", (PyCFunction)booleans, METH_NOARGS,
01188 "booleans() -> List of available atom selection boolean tokens."},
01189 {"functions", (PyCFunction)functions, METH_NOARGS,
01190 "functions() -> List of available atom selection functions."},
01191 {"stringfunctions", (PyCFunction)stringfunctions, METH_NOARGS,
01192 "stringfunctions() -> List of available atom selection string functions."},
01193 { NULL, NULL }
01194 };
01195
01196 static char *module_doc = (char *)
01197 "Methods for creating, updating, querying, and modifying\n"
01198 "selections of atoms.\n"
01199 "\n"
01200 "Example of usage:\n"
01201 ">>> from atomsel import *\n"
01202 ">>> s1 = atomsel('residue 1 to 10 and backbone')\n"
01203 ">>> s1.get('resid')\n"
01204 " <snip> \n"
01205 ">>> s1.set('beta', 5') # set B value to 5 for atoms in s1\n"
01206 ">>> # Mass-weighted RMS alignment:\n"
01207 ">>> mass = s1.get('mass')\n"
01208 ">>> s2 = atomsel('residue 21 to 30 and backbone')\n"
01209 ">>> mat = s1.fit(s2, mass)\n"
01210 ">>> s1.move(mat)\n"
01211 ">>> print s1.rmsd(s2)\n" ;
01212
01213 void
01214 initatomsel(void) {
01215 PyObject *m;
01216
01217 atomsel_type.ob_type = &PyType_Type;
01218 m = Py_InitModule3( "atomsel", atomsel_methods, module_doc );
01219
01220 Py_INCREF((PyObject *)&atomsel_type);
01221 if (PyModule_AddObject(m, "atomsel",
01222 (PyObject *)&atomsel_type) !=0)
01223 return;
01224 if (PyType_Ready(&atomsel_type) < 0)
01225 return;
01226
01227 return;
01228 }
01229