Main Page   Class Hierarchy   Alphabetical List   Compound List   File List   Compound Members   File Members   Related Pages  

py_atomselection.C

Go to the documentation of this file.
00001 /***************************************************************************
00002  *cr
00003  *cr            (C) Copyright 1995-2008 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_atomselection.C,v $
00013  *      $Author: johns $        $Locker:  $             $State: Exp $
00014  *      $Revision: 1.46 $       $Date: 2008/03/27 19:36:51 $
00015  *
00016  ***************************************************************************
00017  * DESCRIPTION:
00018  *  Python atom selection interface.
00019  ***************************************************************************/
00020 
00021 #include "py_commands.h"
00022 #include "AtomSel.h"
00023 #include "VMDApp.h"
00024 #include "MoleculeList.h"
00025 #include "SymbolTable.h"
00026 #include "Measure.h"
00027 #include "SpatialSearch.h"
00028 
00029 static char create_doc[] = "create(molid, frame, selection) -> tuple\nFind atoms in atom selection.";
00030 static PyObject *create(PyObject *self, PyObject *args) {
00031 
00032   int molid = 0;
00033   int frame = 0;
00034   char *sel = 0;
00035 
00036   if (!PyArg_ParseTuple(args, (char *)"iis", &molid, &frame, &sel))
00037     return NULL;
00038 
00039   VMDApp *app = get_vmdapp();
00040 
00041   DrawMolecule *mol = app->moleculeList->mol_from_id(molid);
00042   if (!mol) {
00043     PyErr_SetString(PyExc_ValueError, "invalid molid");
00044     return NULL;
00045   }
00046 
00047   AtomSel *atomSel = new AtomSel(app->atomSelParser, mol->id());
00048   atomSel->which_frame = frame;
00049   if (atomSel->change(sel, mol) == AtomSel::NO_PARSE) {
00050     PyErr_SetString(PyExc_ValueError, "cannot parse atom selection text");
00051     delete atomSel;
00052     return NULL;
00053   }
00054   
00055   // construct a Python tuple to return
00056   PyObject *newlist = PyTuple_New(atomSel->selected);
00057   int j=0;
00058   for (int i=0; i<atomSel->num_atoms; i++) {
00059     if (atomSel->on[i])
00060       PyTuple_SET_ITEM(newlist, j++, PyInt_FromLong(i));
00061   }
00062   delete atomSel;
00063   return newlist;
00064 }
00065 
00066 static char get_doc[] = "get(molid, frame, tuple, attribute) -> value list\nGet selected atom values for the given attribute.";
00067 static PyObject *get(PyObject *self, PyObject *args) {
00068   int i, molid, frame;
00069   PyObject *selected;
00070   int num_selected;
00071   char *attr = 0;
00072 
00073   //
00074   // get molid, list, and attribute
00075   //
00076   if (!PyArg_ParseTuple(args, (char *)"iiO!s", 
00077                         &molid, &frame, &PyTuple_Type, &selected, &attr))
00078     return NULL;  // bad args
00079 
00080   //
00081   // check molecule
00082   //
00083   VMDApp *app = get_vmdapp();
00084   Molecule *mol = app->moleculeList->mol_from_id(molid);
00085   if (!mol) {
00086     PyErr_SetString(PyExc_ValueError, "molecule no longer exists");
00087     return NULL;
00088   }
00089   const int num_atoms = mol->nAtoms;
00090  
00091   // 
00092   // Check for a valid attribute
00093   //
00094   SymbolTable *table = app->atomSelParser;
00095   int attrib_index = table->find_attribute(attr);
00096   if (attrib_index == -1) {
00097     PyErr_SetString(PyExc_ValueError, "unknown atom attribute");
00098     return NULL;
00099   }
00100   SymbolTableElement *elem = table->fctns.data(attrib_index);
00101   if (elem->is_a != SymbolTableElement::KEYWORD &&
00102       elem->is_a != SymbolTableElement::SINGLEWORD) {
00103     PyErr_SetString(PyExc_ValueError, "attribute is not a keyword or singleword");
00104     return NULL;
00105   }
00106  
00107   // 
00108   // fetch the data
00109   //
00110 
00111   atomsel_ctxt context(table, mol, frame, attr);
00112   num_selected = PyTuple_Size(selected);
00113   PyObject *newlist = PyList_New(num_selected); 
00114 
00115   // XXX should check that selected contains valid indices
00116   int *flgs = new int[num_atoms];
00117   memset(flgs,0,num_atoms*sizeof(int));
00118   for (i=0; i<num_selected; i++) 
00119     flgs[PyInt_AsLong(PyTuple_GET_ITEM(selected,i))] = 1;
00120 
00121   if (elem->is_a == SymbolTableElement::SINGLEWORD) {
00122     int *tmp = new int[num_atoms];
00123     memcpy(tmp, flgs, num_atoms*sizeof(int));
00124     elem->keyword_single(&context, num_atoms, tmp);
00125     int j=0;
00126     for (i=0; i<num_atoms; i++) {
00127       if (flgs[i]) {
00128         if (tmp[i]) {
00129           PyList_SET_ITEM(newlist, j++, PyInt_FromLong(1));       
00130         } else {
00131           PyList_SET_ITEM(newlist, j++, PyInt_FromLong(0));       
00132         }
00133       }
00134     }
00135     delete [] tmp;
00136   } else {
00137     switch(table->fctns.data(attrib_index)->returns_a) {
00138       case (SymbolTableElement::IS_STRING):
00139       {
00140         const char **tmp= new const char *[num_atoms];
00141         elem->keyword_string(&context, num_atoms, tmp, flgs);
00142         int j=0;
00143         for (int i=0; i<num_atoms; i++) {
00144           if (flgs[i]) {
00145             PyList_SET_ITEM(newlist, j++, PyString_FromString(tmp[i]));
00146           }
00147         }
00148         delete [] tmp;
00149       }
00150       break;
00151       case (SymbolTableElement::IS_INT):
00152       {
00153         int *tmp = new int[num_atoms];
00154         elem->keyword_int(&context, num_atoms, tmp, flgs);
00155         int j=0;
00156         for (int i=0; i<num_atoms; i++) {
00157           if (flgs[i]) {
00158             PyList_SET_ITEM(newlist, j++, PyInt_FromLong(tmp[i]));
00159           }
00160         }
00161         delete [] tmp;
00162       }
00163       break;
00164       case (SymbolTableElement::IS_FLOAT):
00165       {
00166         double *tmp = new double[num_atoms];
00167         elem->keyword_double(&context, num_atoms, tmp, flgs);
00168         int j=0;
00169         for (int i=0; i<num_atoms; i++) {
00170           if (flgs[i])  
00171             PyList_SET_ITEM(newlist, j++, PyFloat_FromDouble(tmp[i]));
00172         }
00173         delete [] tmp;
00174       }
00175       break;
00176     } // end switch
00177   }   // end else
00178   delete [] flgs;
00179   return newlist;
00180 }
00181 
00182 static char set_doc[] = "set(molid, frame, tuple, attribute, values) -> None\nSe attributes for selected atoms using values.";
00183 static PyObject *set(PyObject *self, PyObject *args) {
00184   int i, molid, frame;
00185   PyObject *selected, *val;
00186   char *attr = 0;
00187 
00188   //
00189   // get molid, frame, list, attribute, and value
00190   //
00191   if (!PyArg_ParseTuple(args, (char *)"iiO!sO!", &molid, &frame,
00192                         &PyTuple_Type, &selected, 
00193                         &attr, &PyTuple_Type, &val ))
00194     return NULL;  // bad args
00195 
00196   // 
00197   // check that we have been given either one value or one for each selected
00198   // atom
00199   //
00200   int num_selected = PyTuple_Size(selected);
00201   int tuplesize = PyTuple_Size(val);
00202   if (tuplesize != 1 && tuplesize != num_selected) {
00203     PyErr_SetString(PyExc_ValueError, "wrong number of items");
00204     return NULL; 
00205   }
00206  
00207   //
00208   // check molecule
00209   //
00210   VMDApp *app = get_vmdapp();
00211   Molecule *mol = app->moleculeList->mol_from_id(molid);
00212   if (!mol) {
00213     PyErr_SetString(PyExc_ValueError, "molecule no longer exists");
00214     return NULL;
00215   }
00216   const int num_atoms = mol->nAtoms;
00217 
00218   //
00219   // Check for a valid attribute
00220   //
00221   SymbolTable *table = app->atomSelParser;
00222   int attrib_index = table->find_attribute(attr);
00223   if (attrib_index == -1) {
00224     PyErr_SetString(PyExc_ValueError, "unknown atom attribute");
00225     return NULL;
00226   }
00227   SymbolTableElement *elem = table->fctns.data(attrib_index);
00228   if (elem->is_a != SymbolTableElement::KEYWORD &&
00229       elem->is_a != SymbolTableElement::SINGLEWORD) {
00230     PyErr_SetString(PyExc_ValueError, "attribute is not a keyword or singleword");
00231     return NULL;
00232   }
00233   if (!table->is_changeable(attrib_index)) {
00234     PyErr_SetString(PyExc_ValueError, "attribute is not modifiable");
00235     return NULL; 
00236   }
00237 
00238   // 
00239   // convert the list of selected atoms into an array of integer flags
00240   //
00241   // XXX should check that selected contains valid indices
00242   int *flgs = new int[num_atoms];
00243   memset(flgs,0,num_atoms*sizeof(int));
00244   for (i=0; i<num_selected; i++)
00245     flgs[PyInt_AsLong(PyTuple_GET_ITEM(selected,i))] = 1;
00246  
00247   //  
00248   // set the data
00249   //
00250 
00251   // singlewords can never be set, so macro is NULL.
00252   atomsel_ctxt context(table, mol, frame, NULL);
00253   if (elem->returns_a == SymbolTableElement::IS_INT) {
00254     int *list = new int[num_atoms];
00255     if (tuplesize > 1) {
00256       int j=0;
00257       for (int i=0; i<num_atoms; i++) {
00258         if (flgs[i])
00259           list[i] = PyInt_AsLong(PyTuple_GET_ITEM(val, j++));
00260       }
00261     } else {
00262       for (int i=0; i<num_atoms; i++) {
00263         if (flgs[i])
00264           list[i] = PyInt_AsLong(PyTuple_GET_ITEM(val, 0));
00265       }
00266     }
00267     elem->set_keyword_int(&context, num_atoms, list, flgs);
00268     delete [] list;
00269 
00270   } else if (elem->returns_a == SymbolTableElement::IS_FLOAT) {
00271     double *list = new double[num_atoms];
00272     if (tuplesize > 1) { 
00273       int j=0;
00274       for (int i=0; i<num_atoms; i++) { 
00275         if (flgs[i])
00276           list[i] = PyFloat_AsDouble(PyTuple_GET_ITEM(val, j++));
00277       }
00278     } else {
00279       for (int i=0; i<num_atoms; i++) {
00280         if (flgs[i])
00281           list[i] = PyFloat_AsDouble(PyTuple_GET_ITEM(val, 0));
00282       }
00283     }
00284     elem->set_keyword_double(&context, num_atoms, list, flgs);
00285     delete [] list;
00286 
00287 
00288   } else if (elem->returns_a == SymbolTableElement::IS_STRING) {
00289 
00290     const char **list = new const char *[num_atoms];
00291     if (tuplesize > 1) { 
00292       int j=0;
00293       for (int i=0; i<num_atoms; i++) { 
00294         if (flgs[i])
00295           list[i] = PyString_AsString(PyTuple_GET_ITEM(val, j++));
00296       }
00297     } else {
00298       for (int i=0; i<num_atoms; i++) {
00299         if (flgs[i])
00300           list[i] = PyString_AsString(PyTuple_GET_ITEM(val, 0));
00301       }
00302     }
00303     elem->set_keyword_string(&context, num_atoms, list, flgs);
00304     delete [] list;
00305   }
00306 
00307   // Recompute the color assignments if certain atom attributes are changed.
00308   if (!strcmp(attr, "name") ||
00309       !strcmp(attr, "type") ||
00310       !strcmp(attr, "resname") ||
00311       !strcmp(attr, "chain") ||
00312       !strcmp(attr, "segid") ||
00313       !strcmp(attr, "segname")) 
00314     app->moleculeList->add_color_names(molid);
00315 
00316   mol->force_recalc(DrawMolItem::SEL_REGEN | DrawMolItem::COL_REGEN); 
00317   delete [] flgs;
00318   Py_INCREF(Py_None);
00319   return Py_None;
00320 }
00321 
00322 // getbonds(molid, atomlist)
00323 static char getbonds_doc[] = "getbonds(molid, tuple) -> bondlists\nReturn list of bonds for each atom id in tuple.";
00324 static PyObject *getbonds(PyObject *self, PyObject *args) {
00325   int molid;
00326   PyObject *atomlist;
00327    
00328   if (!PyArg_ParseTuple(args, (char *)"iO!:getbonds", &molid, 
00329                         &PyTuple_Type, &atomlist)) 
00330     return NULL;  // bad args
00331  
00332   Molecule *mol = get_vmdapp()->moleculeList->mol_from_id(molid);
00333   if (!mol) {
00334     PyErr_SetString(PyExc_ValueError, "molecule no longer exists");
00335     return NULL;
00336   }
00337   int num_atoms = mol->nAtoms;
00338   int num_selected = PyTuple_Size(atomlist);
00339   PyObject *newlist = PyList_New(num_selected);
00340   for (int i=0; i< num_selected; i++) {
00341     int id = PyInt_AsLong(PyTuple_GET_ITEM(atomlist, i));
00342     if (PyErr_Occurred()) {
00343       Py_DECREF(newlist);
00344       return NULL;
00345     }
00346     if (id < 0 || id >= num_atoms) {
00347       PyErr_SetString(PyExc_ValueError, (char *)"invalid atom id found");
00348       Py_DECREF(newlist);
00349       return NULL;
00350     }
00351     const MolAtom *atom = mol->atom(id);
00352     PyObject *bondlist = PyList_New(atom->bonds);
00353     for (int j=0; j<atom->bonds; j++) {
00354       PyList_SET_ITEM(bondlist, j, PyInt_FromLong(atom->bondTo[j]));
00355     }
00356     PyList_SET_ITEM(newlist, i, bondlist);
00357   }
00358   return newlist;
00359 }
00360 
00361 // setbonds(molid, atomlist, bondlist)
00362 static char setbonds_doc[] = "setbonds(molid, tuple, bondlist) -> None\nSet bonds for each atom in tuple using bondlist.";
00363 static PyObject *setbonds(PyObject *self, PyObject *args) {
00364   int molid;
00365   PyObject *atomlist, *bondlist; 
00366 
00367   if (!PyArg_ParseTuple(args, (char *)"iO!O!:setbonds", &molid, 
00368                         &PyTuple_Type, &atomlist, &PyList_Type, &bondlist)) 
00369     return NULL;  // bad args
00370  
00371   Molecule *mol = get_vmdapp()->moleculeList->mol_from_id(molid);
00372   if (!mol) {
00373     PyErr_SetString(PyExc_ValueError, "molecule no longer exists");
00374     return NULL;
00375   }
00376   int num_atoms = mol->nAtoms;
00377   int num_selected = PyTuple_Size(atomlist);
00378   if (PyList_Size(bondlist) != num_selected) {
00379     PyErr_SetString(PyExc_ValueError, 
00380       (char *)"setbonds: atomlist and bondlist must have the same size");
00381     return NULL;
00382   }
00383   mol->force_recalc(DrawMolItem::MOL_REGEN); // many reps ignore bonds
00384   for (int i=0; i<num_selected; i++) {
00385     int id = PyInt_AsLong(PyTuple_GET_ITEM(atomlist, i));
00386     if (PyErr_Occurred()) {
00387       return NULL;
00388     }
00389     if (id < 0 || id >= num_atoms) {
00390       PyErr_SetString(PyExc_ValueError, (char *)"invalid atom id found");
00391       return NULL;
00392     }
00393     MolAtom *atom = mol->atom(id);
00394    
00395     PyObject *atomids = PyList_GET_ITEM(bondlist, i);
00396     if (!PyList_Check(atomids)) {
00397       PyErr_SetString(PyExc_TypeError, 
00398         (char *)"bondlist must contain lists");
00399       return NULL;
00400     }
00401     int numbonds = PyList_Size(atomids);
00402     int k=0;
00403     for (int j=0; j<numbonds; j++) {
00404       int bond = PyInt_AsLong(PyList_GET_ITEM(atomids, j));
00405       if (PyErr_Occurred())
00406         return NULL;
00407       if (bond >= 0 && bond < mol->nAtoms) {
00408         atom->bondTo[k++] = bond;
00409       } else {
00410         char buf[40];
00411         sprintf(buf, "Invalid atom id in bondlist: %d", bond);
00412         PyErr_SetString(PyExc_ValueError, buf);
00413         return NULL;
00414       }
00415     }
00416     atom->bonds = k;
00417   }
00418   Py_INCREF(Py_None);
00419   return Py_None;
00420 }
00421 
00422 // macro(name, selection)
00423 static char macro_doc[] = "macro(name=None, selection=None) -> macro information\nIf both name and selection are None, return list of macro names.\nIf selection is None, return definition for name.\nIf both name and selection are given, define new macro.\n";
00424 static PyObject *macro(PyObject *self, PyObject *args, PyObject *keywds) {
00425   char *name = NULL, *selection = NULL;
00426   static char *kwlist[] = {
00427     (char *)"name", (char *)"selection", NULL
00428   };
00429   if (!PyArg_ParseTupleAndKeywords(args, keywds, (char *)"|ss:atomselection.macro", kwlist, &name, &selection))
00430     return NULL;
00431 
00432   if (selection && !name) {
00433     PyErr_SetString(PyExc_ValueError, (char *)"Must specify name for macro");
00434     return NULL;
00435   }
00436   SymbolTable *table = get_vmdapp()->atomSelParser;
00437   if (!name && !selection) {
00438     // return list of defined macros
00439     PyObject *macrolist = PyList_New(0);
00440     for (int i=0; i<table->num_custom_singleword(); i++) {
00441       const char *s = table->custom_singleword_name(i);
00442       if (s && strlen(s))
00443         PyList_Append(macrolist, PyString_FromString(s));
00444     }
00445     return macrolist;
00446   }
00447   if (name && !selection) {
00448     // return definition of macro
00449     const char *s = table->get_custom_singleword(name);
00450     if (!s) {
00451       PyErr_SetString(PyExc_ValueError, (char *)"No macro for given name");
00452       return NULL;
00453     }
00454     return PyString_FromString(s);
00455   }
00456   // must have both and selection.  Define a new macro.
00457   if (!table->add_custom_singleword(name, selection)) {
00458     PyErr_SetString(PyExc_ValueError, (char *)"Unable to create new macro");
00459     return NULL;
00460   }
00461   Py_INCREF(Py_None);
00462   return Py_None;
00463 }
00464 
00465 // delmacro(name)
00466 static char delmacro_doc[] = "delmacro(name) -> None\nDelete macro with given name.";
00467 static PyObject *delmacro(PyObject *self, PyObject *args) {
00468   char *name;
00469   if (!PyArg_ParseTuple(args, (char *)"s:atomselection.delmacro", &name))
00470     return NULL;
00471   if (!get_vmdapp()->atomSelParser->remove_custom_singleword(name)) {
00472     PyErr_SetString(PyExc_ValueError, (char *)"Unable to remove macro");
00473     return NULL;
00474   }
00475   Py_INCREF(Py_None);
00476   return Py_None;
00477 }
00478       
00479 static AtomSel *sel_from_py(int molid, int frame, PyObject *selected, VMDApp *app) {
00480   Molecule *mol = app->moleculeList->mol_from_id(molid);
00481   if (!mol) {
00482     PyErr_SetString(PyExc_ValueError, "molecule no longer exists");
00483     return NULL;
00484   }
00485   if (mol->nAtoms == 0) {
00486     PyErr_SetString(PyExc_ValueError, "Molecule in atom selection contains no atoms");
00487     return NULL;
00488   }
00489   AtomSel *sel = new AtomSel(app->atomSelParser, mol->id());
00490   sel->on = new int[mol->nAtoms];
00491         sel->num_atoms = mol->nAtoms;
00492   if (!selected) {
00493     // turn them all on
00494     for (int i=0; i<sel->num_atoms; i++) sel->on[i] = 1;
00495     sel->selected = sel->num_atoms;
00496   } else {
00497     memset(sel->on, 0, mol->nAtoms*sizeof(int));
00498     const int num_selected = PyTuple_Size(selected);
00499           sel->selected = num_selected;
00500     for (int i=0; i<num_selected; i++) {
00501       int ind = PyInt_AsLong(PyTuple_GET_ITEM(selected,i));
00502       if (ind < 0 || ind >= mol->nAtoms) {
00503         delete sel;
00504         PyErr_SetString(PyExc_ValueError, "Invalid atom id in selection");
00505         return NULL;
00506       }
00507       sel->on[ind] = 1;
00508     }
00509   }
00510   sel->which_frame = frame;
00511   return sel;
00512 }
00513 
00514 // utility routine for parsing weight values.  Uses the sequence protocol
00515 // so that sequence-type structure (list, tuple) will be accepted.
00516 static float *parse_weight(AtomSel *sel, PyObject *wtobj) {
00517   float *weight = new float[sel->selected];
00518   if (!wtobj || wtobj == Py_None) {
00519     for (int i=0; i<sel->selected; i++) weight[i] = 1.0f;
00520     return weight;
00521   }
00522 
00523   PyObject *seq = PySequence_Fast(wtobj, (char *)"weight must be a sequence.");
00524   if (!seq) return NULL;
00525   if (PySequence_Size(seq) != sel->selected) {
00526     Py_DECREF(seq);
00527     PyErr_SetString(PyExc_ValueError, "weight must be same size as selection.");
00528     delete [] weight;
00529     return NULL;
00530   }
00531   for (int i=0; i<sel->selected; i++) {
00532     double tmp = PyFloat_AsDouble(PySequence_Fast_GET_ITEM(seq, i));
00533     if (PyErr_Occurred()) {
00534       PyErr_SetString(PyExc_ValueError, "non-floating point value found in weight.");
00535       Py_DECREF(seq);
00536       delete [] weight;
00537       return NULL;
00538     }
00539     weight[i] = (float)tmp;
00540   }
00541   Py_DECREF(seq);
00542   return weight;
00543 }
00544 
00545 static char minmax_doc[] = "minmax(molid, frame, tuple) -> (min, max)\nReturn minimum and maximum coordinates for selected atoms.";
00546 static PyObject *minmax(PyObject *self, PyObject *args) {
00547   int molid, frame;
00548   PyObject *selected;
00549 
00550   if (!PyArg_ParseTuple(args, (char *)"iiO!", 
00551                         &molid, &frame, &PyTuple_Type, &selected))
00552     return NULL;  // bad args
00553 
00554   VMDApp *app = get_vmdapp();
00555   AtomSel *sel = sel_from_py(molid, frame, selected, app);
00556   if (!sel) return NULL;
00557   const Timestep *ts = app->moleculeList->mol_from_id(molid)->get_frame(frame); 
00558   if (!ts) {
00559     PyErr_SetString(PyExc_ValueError, "No coordinates in selection");
00560     delete sel;
00561     return NULL;
00562   }
00563   float min[3], max[3];
00564   int rc = measure_minmax(sel->num_atoms, sel->on, ts->pos, NULL, min, max);
00565   delete sel;
00566   if (rc < 0) {
00567     PyErr_SetString(PyExc_ValueError, measure_error(rc));
00568     return NULL;
00569   }
00570   PyObject *mintuple, *maxtuple, *result;
00571   mintuple = PyTuple_New(3);
00572   maxtuple = PyTuple_New(3);
00573   result = PyTuple_New(2);
00574   for (int i=0; i<3; i++) {
00575     PyTuple_SET_ITEM(mintuple, i, PyFloat_FromDouble(min[i]));
00576     PyTuple_SET_ITEM(maxtuple, i, PyFloat_FromDouble(max[i]));
00577   }
00578   PyTuple_SET_ITEM(result, 0, mintuple);
00579   PyTuple_SET_ITEM(result, 1, maxtuple);
00580   return result;
00581 }
00582 
00583 static char center_doc[] = "center(molid, frame, tuple, weight) -> (x, y, z)\nReturn a tuple corresponding to the center of the selection,\npossibly weighted by weight.";
00584 static PyObject *center(PyObject *self, PyObject *args) {
00585   int molid, frame;
00586   PyObject *selected;
00587   PyObject *weightobj = NULL;
00588   AtomSel *sel;
00589   // parse arguments
00590   if (!PyArg_ParseTuple(args, (char *)"iiO!|O",
00591         &molid, &frame, &PyTuple_Type, &selected, &weightobj))
00592     return NULL;
00593   VMDApp *app = get_vmdapp();
00594   // get selection
00595   if (!(sel = sel_from_py(molid, frame, selected, app)))
00596     return NULL;
00597   // get weight
00598   float *weight = parse_weight(sel, weightobj);
00599   if (!weight) return NULL;
00600   float cen[3];
00601   // compute center
00602   int ret_val = measure_center(sel, sel->coordinates(app->moleculeList),
00603       weight, cen);
00604   delete [] weight;
00605   delete sel;
00606   if (ret_val < 0) {
00607     PyErr_SetString(PyExc_ValueError, measure_error(ret_val));
00608     return NULL;
00609   }
00610   // return as (x, y, z)
00611   PyObject *cenobj = PyTuple_New(3);
00612   for (int i=0; i<3; i++)
00613     PyTuple_SET_ITEM(cenobj, i, PyFloat_FromDouble(cen[i]));
00614   return cenobj;
00615 }
00616 
00617 static PyObject *py_align(PyObject *self, PyObject *args) {
00618   int selmol, selframe, refmol, refframe, movemol, moveframe;
00619   PyObject *selobj, *refobj, *moveobj, *weightobj = NULL;
00620   if (!PyArg_ParseTuple(args, (char *)"iiO!iiO!iiO!O:atomselection.align",
00621         &selmol, &selframe, &PyTuple_Type, &selobj,
00622         &refmol, &refframe, &PyTuple_Type, &refobj,
00623         &movemol, &moveframe, &PyTuple_Type, &moveobj,
00624         &weightobj))
00625     return NULL;
00626 
00627   // check if movemol is -1.  If so, use the sel molecule and timestep instead
00628   if (movemol == -1) {
00629     movemol = selmol;
00630     moveobj = NULL;
00631   }
00632   VMDApp *app = get_vmdapp();
00633   AtomSel *sel=NULL, *ref=NULL, *move=NULL;
00634   if (!(sel = sel_from_py(selmol, selframe, selobj, app)) ||
00635       !(ref = sel_from_py(refmol, refframe, refobj, app)) ||
00636       !(move = sel_from_py(movemol, moveframe, moveobj, app))) {
00637     delete sel;
00638     delete ref;
00639     delete move;
00640     return NULL;
00641   }
00642   const float *selts, *refts;
00643   float *movets;
00644   if (!(selts = sel->coordinates(app->moleculeList)) ||
00645       !(refts = ref->coordinates(app->moleculeList)) || 
00646       !(movets = move->coordinates(app->moleculeList))) {
00647     delete sel;
00648     delete ref;
00649     delete move;
00650     PyErr_SetString(PyExc_ValueError, "No coordinates in selection");
00651     return NULL;
00652   }
00653   float *weight = parse_weight(sel, weightobj);
00654   if (!weight) {
00655     delete sel;
00656     delete ref;
00657     delete move;
00658     return NULL;
00659   }
00660   // Find the matrix that aligns sel with ref.  Apply the transformation to
00661   // the atoms in move.
00662   // XXX need to add support for the "order" parameter as in Tcl.
00663   Matrix4 mat;
00664   int rc = measure_fit(sel, ref, selts, refts, weight, NULL, &mat);
00665   delete [] weight;
00666   delete sel;
00667   delete ref;
00668   if (rc < 0) {
00669     delete move;
00670     PyErr_SetString(PyExc_ValueError, (char *)measure_error(rc));
00671     return NULL;
00672   }
00673   for (int i=0; i<move->num_atoms; i++) {
00674     if (move->on[i]) {
00675       float *pos = movets+3*i;
00676       mat.multpoint3d(pos, pos);
00677     }
00678   }
00679   Molecule *mol = app->moleculeList->mol_from_id(move->molid());
00680   mol->force_recalc(DrawMolItem::MOL_REGEN);
00681   delete move;
00682   Py_INCREF(Py_None);
00683   return Py_None;
00684 }
00685 
00686 static char rmsd_doc[] = "rmsd(mol1, frame1, tuple1, mol2, frame2, tuple2, weights) -> rms distance.\nWeight must be None or list of same size as tuples.";
00687 static PyObject *py_rmsd(PyObject *self, PyObject *args) {
00688   int mol1, frame1, mol2, frame2;
00689   PyObject *selected1, *selected2, *weightobj = NULL;
00690   if (!PyArg_ParseTuple(args, (char *)"iiO!iiO!O:atomselection.rmsd",
00691         &mol1, &frame1, &PyTuple_Type, &selected1,
00692         &mol2, &frame2, &PyTuple_Type, &selected2,
00693                                 &weightobj))
00694     return NULL;
00695   VMDApp *app = get_vmdapp();
00696   AtomSel *sel1 = sel_from_py(mol1, frame1, selected1, app);
00697   AtomSel *sel2 = sel_from_py(mol2, frame2, selected2, app);
00698   if (!sel1 || !sel2) {
00699     delete sel1;
00700     delete sel2;
00701     return NULL;
00702   }
00703   const Timestep *ts1 =app->moleculeList->mol_from_id(mol1)->get_frame(frame1); 
00704   const Timestep *ts2 =app->moleculeList->mol_from_id(mol2)->get_frame(frame2); 
00705   if (!ts1 || !ts2) {
00706     PyErr_SetString(PyExc_ValueError, "No coordinates in selection");
00707     delete sel1;
00708     delete sel2;
00709     return NULL;
00710   }
00711   float *weight = parse_weight(sel1, weightobj);
00712   if (!weight) {
00713     delete sel1;
00714     delete sel2;
00715     return NULL;
00716   }
00717   float rmsd;
00718   int rc = measure_rmsd(sel1, sel2, sel1->selected, ts1->pos, ts2->pos,
00719       weight, &rmsd);
00720   delete sel1;
00721   delete sel2;
00722   delete [] weight;
00723   if (rc < 0) {
00724     PyErr_SetString(PyExc_ValueError, measure_error(rc));
00725     return NULL;
00726   }
00727   return PyFloat_FromDouble(rmsd);
00728 }
00729 
00730 // Find all atoms p in sel1 and q in sel2 within the cutoff.  
00731 static char contacts_doc[] = "contacts(mol1, frame1, tuple1, mol2, frame2, tuple2, cutoff) -> contact pairs\nReturn pairs of atoms, one from each selection, within cutoff of each other.";
00732 static PyObject *contacts(PyObject *self, PyObject *args) {
00733   
00734   int mol1, frame1, mol2, frame2;
00735   PyObject *selected1, *selected2;
00736   float cutoff;
00737   if (!PyArg_ParseTuple(args, (char *)"iiO!iiO!f:atomselection.contacts",
00738         &mol1, &frame1, &PyTuple_Type, &selected1,
00739         &mol2, &frame2, &PyTuple_Type, &selected2,
00740         &cutoff))
00741     return NULL;
00742   VMDApp *app = get_vmdapp();
00743   AtomSel *sel1 = sel_from_py(mol1, frame1, selected1, app);
00744   AtomSel *sel2 = sel_from_py(mol2, frame2, selected2, app);
00745   if (!sel1 || !sel2) {
00746     delete sel1;
00747     delete sel2;
00748     return NULL;
00749   }
00750   const float *ts1 = sel1->coordinates(app->moleculeList);
00751   const float *ts2 = sel2->coordinates(app->moleculeList);
00752   if (!ts1 || !ts2) {
00753     PyErr_SetString(PyExc_ValueError, "No coordinates in selection");
00754     delete sel1;
00755     delete sel2;
00756     return NULL;
00757   }
00758   Molecule *mol = app->moleculeList->mol_from_id(mol1);
00759 
00760   GridSearchPair *pairlist = vmd_gridsearch3(
00761       ts1, sel1->num_atoms, sel1->on,
00762       ts2, sel2->num_atoms, sel2->on,
00763       cutoff, -1, (sel1->num_atoms + sel2->num_atoms) * 27);
00764 
00765   delete sel1;
00766   delete sel2;
00767   GridSearchPair *p, *tmp;
00768   PyObject *list1 = PyList_New(0);
00769   PyObject *list2 = PyList_New(0);
00770   for (p=pairlist; p != NULL; p=tmp) {
00771     // throw out pairs that are already bonded
00772     MolAtom *a1 = mol->atom(p->ind1);
00773     if (mol1 != mol2 || !a1->bonded(p->ind2)) {
00774       PyList_Append(list1, PyInt_FromLong(p->ind1));
00775       PyList_Append(list2, PyInt_FromLong(p->ind2));
00776     }
00777     tmp = p->next;
00778     free(p);
00779   }
00780   PyObject *result = PyList_New(2);
00781   PyList_SET_ITEM(result, 0, list1);
00782   PyList_SET_ITEM(result, 1, list2);
00783   return result;
00784 }
00785 
00786 static char sasa_doc[] = "sasa(srad, sel, samples=500, points=None, restrict=None)\npoints must be a list; surface points will be appended to the list\nin the order xyzxyzxyz (i.e. a flat list)";
00787 
00788 static PyObject *sasa(PyObject *self, PyObject *args, PyObject *keywds) {
00789   int molid = -1, frame = -1;
00790   float srad = 0;
00791   PyObject *selobj = NULL, *restrictobj = NULL;
00792   int samples = -1;
00793   const int *sampleptr = NULL;
00794   PyObject *pointsobj = NULL;
00795 
00796   static char *kwlist[] = {
00797     (char *)"srad", (char *)"molid", (char *)"frame", (char *)"selected",
00798     (char *)"samples", (char *)"points", (char *)"restrict"
00799   };
00800   if (!PyArg_ParseTupleAndKeywords(args, keywds, 
00801         (char *)"fiiO!|iO!O!:atomselection.sasa", kwlist, 
00802         &srad, &molid, &frame, &PyTuple_Type, &selobj, 
00803         &samples, &PyList_Type, &pointsobj, &PyTuple_Type, &restrictobj))
00804     return NULL;
00805 
00806   // validate srad
00807   if (srad < 0) {
00808     PyErr_SetString(PyExc_ValueError, (char *)"atomselect.sasa: srad must be non-negative.");
00809     return NULL;
00810   }
00811 
00812   // validate selection
00813   VMDApp *app = get_vmdapp();
00814   AtomSel *sel = sel_from_py(molid, frame, selobj, app);
00815   if (!sel) return NULL;
00816 
00817   // fetch the radii and coordinates
00818   const float *radii = 
00819     app->moleculeList->mol_from_id(sel->molid())->extra.data("radius");
00820   const float *coords = sel->coordinates(app->moleculeList);
00821 
00822   // if samples was given and is valid, use it
00823   if (samples > 1) sampleptr = &samples;
00824 
00825   // if restrict is given, validate it
00826   AtomSel *restrictsel = NULL;
00827   if (restrictobj) {
00828     if (!(restrictsel = sel_from_py(molid, frame, restrictobj, app))) {
00829       delete sel;
00830       return NULL;
00831     }
00832   }
00833 
00834   // if points are requested, fetch them
00835   ResizeArray<float> sasapts;
00836   ResizeArray<float> *sasaptsptr = pointsobj ? &sasapts : NULL;
00837  
00838   // go!
00839   float sasa = 0;
00840   int rc = measure_sasa(sel, coords, radii, srad, &sasa, 
00841         sasaptsptr, restrictsel, sampleptr);
00842   delete sel;
00843   delete restrictsel;
00844   if (rc) {
00845     PyErr_SetString(PyExc_ValueError, (char *)measure_error(rc));
00846     return NULL;
00847   }
00848 
00849   // append surface points to the provided list object.
00850   if (pointsobj) {
00851     for (int i=0; i<sasapts.num(); i++) {
00852       PyList_Append(pointsobj, PyFloat_FromDouble(sasapts[i]));
00853     }
00854   }
00855 
00856   // return the total SASA.
00857   return PyFloat_FromDouble(sasa);
00858 }
00859 
00860 static PyMethodDef AtomselectionMethods[] = {
00861   {(char *)"create", (vmdPyMethod)create, METH_VARARGS, create_doc },
00862   {(char *)"get", (vmdPyMethod)get, METH_VARARGS, get_doc },
00863   {(char *)"set", (vmdPyMethod)set, METH_VARARGS, set_doc },
00864   {(char *)"getbonds", (vmdPyMethod)getbonds, METH_VARARGS, getbonds_doc },
00865   {(char *)"setbonds", (vmdPyMethod)setbonds, METH_VARARGS, setbonds_doc },
00866   {(char *)"macro", (PyCFunction)macro, METH_VARARGS | METH_KEYWORDS, macro_doc},
00867   {(char *)"delmacro", (vmdPyMethod)delmacro, METH_VARARGS, delmacro_doc },
00868   {(char *)"minmax", (vmdPyMethod)minmax, METH_VARARGS, minmax_doc },
00869   {(char *)"center", (vmdPyMethod)center, METH_VARARGS, center_doc },
00870   {(char *)"rmsd", (vmdPyMethod)py_rmsd, METH_VARARGS, rmsd_doc },
00871   {(char *)"align", (vmdPyMethod)py_align, METH_VARARGS},
00872   {(char *)"contacts", (vmdPyMethod)contacts, METH_VARARGS, contacts_doc },
00873   {(char *)"sasa", (vmdPyMethod)sasa, METH_VARARGS | METH_KEYWORDS, sasa_doc },
00874   {NULL, NULL}
00875 };
00876 
00877 void initatomselection() {
00878   (void) Py_InitModule((char *)"atomselection", AtomselectionMethods);
00879 }
00880 
00881   
00882   

Generated on Sun Sep 7 01:26:09 2008 for VMD (current) by doxygen1.2.14 written by Dimitri van Heesch, © 1997-2002