def cafit_orientation(selection, visualize=1, quiet=0): ''' DESCRIPTION Get the center and direction of a peptide by least squares linear fit on CA atoms. USAGE cafit_orientation selection [, visualize] NOTES Requires python module "numpy". SEE ALSO helix_orientation ''' visualize, quiet = int(visualize), int(quiet) import numpy stored.x = list() cmd.iterate_state(STATE, '(%s) and name CA' % (selection), 'stored.x.append([x,y,z])') x = numpy.array(stored.x) U, s, Vh = numpy.linalg.svd(x - x.mean(0)) vec = cpv.normalize(Vh[0]) if cpv.dot_product(vec, x[-1] - x[0]) < 0: vec = cpv.negate(vec) return _common_orientation(selection, vec, visualize, quiet)
def get_alignment_coords(name, active_only=0, state=-1, quiet=0): ''' DESCRIPTION API only function. Returns a dictionary with items (object name, Nx3 coords list) N is the number of alignment columns without gaps. EXAMPLE import numpy from psico.multistuff import * from psico.querying import * extra_fit('name CA', cycles=0, object='aln') x = get_alignment_coords('aln') m = numpy.array(x.values()) ''' active_only, state, quiet = int(active_only), int(state), int(quiet) aln = cmd.get_raw_alignment(name, active_only) object_list = cmd.get_object_list(name) idx2coords = dict() cmd.iterate_state(state, name, 'idx2coords[model,index] = (x,y,z)', space={'idx2coords': idx2coords}) allcoords = dict((model, []) for model in object_list) for pos in aln: if len(pos) != len(object_list): continue for model,index in pos: allcoords[model].append(idx2coords[model,index]) return allcoords
def pymol_selection_to_atom_list(sele: str) -> t.List[u.Atom]: """ pymol_selection_to_atom_list(...) converts a interface_Pymol selection and returns a list of restraintmaker_Logic.utilities.Atom :return: A list of Atoms :rtype: t.List[u.Atom] """ atom_info: t.List[t.Dict] = [] # Get all the information about the Atoms as a dict. I can not directly put it into an atom, because the interface_Pymol interpreter will not know what an u.Atom is. cmd.iterate_state( -1, selection=sele, expression= "atom_info.append({\"elem\":elem,\"id\":ID, \"name\":name, \"x\":x, \"y\":y, \"z\":z, \"model\": model, \"chain\":chain, \"resn\":resn, \"resi\":resi, \"alt\":alt, \"label\":label, \"b\":b})", space=locals()) # Convert dict to Atoms atoms: t.List[u.Atom] = [] for a in atom_info: x = u.Atom(elem=a['elem'], id=a['id'], name=a['name'], x=a['x'], y=a['y'], z=a['z'], chain=a['chain'], resn=a['resn'], resi=a['resi'], alt=a['alt'], label=a['label'], b=a['b']) atoms.append(x) return (atoms)
def save_xyzr(filename, selection='all', state=1, _colorsout=None): ''' DESCRIPTION Write the given selection to an xyzr or xyzrn (determined by extension) file for MSMS. ''' if filename.endswith('xyzrn'): expr = 'callback(x, y, z, vdw, name, resn, resi)' fmt = '%.3f %.3f %.3f %.2f 1 %s_%s_%s\n' else: expr = 'callback(x, y, z, vdw)' fmt = '%.3f %.3f %.3f %.2f\n' if isinstance(_colorsout, list): expr += ';_colorsout.append(color)' handle = open(filename, 'w') try: cmd.iterate_state(state, selection, expr, space={ 'callback': lambda *args: handle.write(fmt % args), '_colorsout': _colorsout }) finally: handle.close()
def visualize( *args, **kwargs ): if len(args) < 1: print("No molecules selected") return stored.b_factors = [] for selection in args: cmd.iterate_state(1, selector.process(selection), "stored.b_factors.append(b)") max = 0.00 for value in stored.b_factors: if abs(value) > max: max = abs(value) print("Maximum absolute b-factor value: %s") % (max) #Uses same magnitude for maximum and minimum to stay symmetrical around zero for selection_item in args: cmd.spectrum("b", "red_white_green", minimum = (0-max), maximum = max, selection = selection_item)
def testPseudoatom(self): cmd.set('retain_order') cmd.pseudoatom('m1', '', 'A1', 'B2', 1, 'D', 'E5', 'F', 7.0, 0, 9.0, 0.1, '11', 'foo12', (13, 14, 15), 1) cmd.pseudoatom('m2', '', 'A1', 'B2', 2, 'A', 'E5', 'F', 7.0, 0, 9.0, 0.1, '12', 'bar12', (13, 10, 15), 1) # append to m1, pos and vdw vom selection cmd.pseudoatom('m1', 'all', 'A2', 'B2', 3, 'D', 'E5', 'F', -1, 1, 19.0, 0.2, '13', 'com12', None, 1, 'extent') r_list = [] f_indices = [7, 9, 10, 13, 14, 15] cmd.iterate_state(1, 'all', 'r_list.append([model, name, resn, resi, chain, segi, elem, ' 'vdw, type, b, q, color, label, x, y, z])', space=locals()) for ref, values in zip(r_list, [ ['m1', 'A1', 'B2', '1', 'D', 'E5', 'F', 7.0, 'ATOM', 9.0, 0.1, 11, 'foo12', 13.0, 14.0, 15.0], ['m1', 'A2', 'B2', '3', 'D', 'E5', 'F', 2.0, 'HETATM', 19.0, 0.2, 13, 'com12', 13.0, 12.0, 15.0], ['m2', 'A1', 'B2', '2', 'A', 'E5', 'F', 7.0, 'ATOM', 9.0, 0.1, 12, 'bar12', 13.0, 10.0, 15.0], ]): for i in f_indices: values[i] = round(values[i], 1) ref[i] = round(ref[i], 1) self.assertEqual(ref, values)
def superpose(self): #get the position of the selected residue's O atom stored.xyz = [] if self.currentLabel.modifiedAA: cmd.iterate_state(1, "%s & name O" % self.residue1Name, "stored.xyz.append([x,y,z])") args = [] i = 0 while i < len(self.currentLabel.atomsForSuperposition): args.append( "%s & name %s" % ("currentLabel", self.currentLabel.atomsForSuperposition[i])) args.append("%s & name %s" % (self.residue1Name, self.currentLabel.atomsForSuperposition[i])) i += 1 #print args if apply(cmd.pair_fit, args): #set the label's O atom to the stored position if self.currentLabel.modifiedAA: cmd.alter_state(1, "%s & name O" % "currentLabel", "(x,y,z)=stored.xyz.pop(0)") return True else: return False
def loop_orientation(selection, visualize=1, quiet=0): ''' DESCRIPTION Get the center and approximate direction of a peptide. Works for any secondary structure. Averages direction of N(i)->C(i) pseudo bonds. USAGE loop_orientation selection [, visualize] SEE ALSO helix_orientation ''' visualize, quiet = int(visualize), int(quiet) stored.x = dict() cmd.iterate_state( STATE, '(%s) and name N+C' % (selection), 'stored.x.setdefault(chain + resi, dict())[name] = x,y,z') vec = cpv.get_null() count = 0 for x in stored.x.itervalues(): if 'C' in x and 'N' in x: vec = cpv.add(vec, cpv.sub(x['C'], x['N'])) count += 1 if count == 0: print 'warning: count == 0' raise CmdException vec = cpv.normalize(vec) return _common_orientation(selection, vec, visualize, quiet)
def principalellipsoid(selection, name="ellipsoid", scale_factor=1, state=1, color="red red green green blue blue",segments=40,**kwargs): """Fit an ellipsoid to the selection based on the principal axes """ scale_factor = float(scale_factor) xyz = [] cmd.iterate_state(state,selection,"xyz.append( (x,y,z) )", space={'xyz':xyz} ) #create coordinates array coord = numpy.array(xyz, float) axis1,axis2,axis3,center = computeprincipalaxes(coord) a = numpy.linalg.norm(axis1) b = numpy.linalg.norm(axis2) c = numpy.linalg.norm(axis3) # Convert axes into transformation matrix M = numpy.vstack((axis1/a,axis2/b,axis3/c,center)).transpose() M = numpy.vstack((M,numpy.array([[0,0,0,1]]))) # Convert to row-major array transf = M.reshape(-1).tolist() ellipsoid.ellipsoid(name,0,0,0, a*scale_factor, b*scale_factor, c*scale_factor, color=color, segs=segments, transformation=transf, **kwargs)
def hs_resize(meta_file, selection): if not is_sele(selection): raise RuntimeError( "Selection \"{}\" does not exists.".format(selection)) # Find free sele name temp_sele = _random_string() while is_sele(temp_sele): temp_sele = _random_string() states = cmd.count_states(selection=selection) for state in range(1, states + 1): stored.info = [] cmd.iterate_state(state, selection, "stored.info.append((ID, partial_charge))") for id_, partial_charge in stored.info: size = log(partial_charge / ref + 1) print(ref) print(size) cmd.select(temp_sele, "{} and id {}".format(selection, id_), state=state) cmd.set("sphere_scale", value=size, selection=temp_sele) cmd.alter(temp_sele, "b={}".format(partial_charge)) cmd.delete(temp_sele)
def rao_dihedrals(selection, quiet='yes'): """ rao_dihedrals Calculates the virtual torsion angles a1, a2 and a3 of a hexopyranose ring, using the definitions of Rao and co-workers, as in the book. """ stored.names=[]; states=cmd.count_states(selection); print(states); a1=[]; a2=[]; a3=[]; for i in range(1, states+1, 1): cmd.iterate_state(i, selection, 'stored.names.append(name)'); a1.append(dihe(selection, stored.names[3], stored.names[4], stored.names[1], stored.names[0], state=i)); a2.append(dihe(selection, stored.names[5], stored.names[1], stored.names[3], stored.names[4], state=i)); a3.append(dihe(selection, stored.names[1], stored.names[3], stored.names[5], stored.names[4], state=i)); av1=numpy.mean(a1); stdErr1=numpy.std(a1)/math.sqrt(states); av2=numpy.mean(a2); stdErr2=numpy.std(a2)/math.sqrt(states); av3=numpy.mean(a3); stdErr3=numpy.std(a3)/math.sqrt(states); if quiet=="no": print("a1 a2 a3"); for i in range(len(a1)): print("%.3f %.3f %.3f" %(a1[i],a2[i],a3[i])); print("Averages(Mean +/- Standard Error):\na1 a2 a3\n %.3f+/-%.3f %.3f+/-%.3f %.3f+/-%.3f" %(av1,stdErr1, av2,stdErr2, av3,stdErr3));
def test_atom_state_settings_3f(self, f, data): m1 = "pseudo01" lp1, lp2 = lp = map(f, data) cmd.pseudoatom(m1) cmd.pseudoatom(m1) cmd.create(m1, m1, 1, 2) stored.pos = {} stored.lp1 = lp1 stored.lp2 = lp2 stored.origp = None # get origp (should be 0.,0.,0. cmd.alter("(index 1)", "stored.origp = s['label_placement_offset']") # change atom-state setting to lp1 for atom in both states cmd.alter_state(0, "(index 1)", "s['label_placement_offset']=stored.lp1") # change atom-level setting for all atoms to something else cmd.alter("all", "s['label_placement_offset']=stored.lp2") # get atom-state settings cmd.iterate_state(0, "all", "stored.pos['%s-%s-%s' % (model, state, index)]=s['label_placement_offset']") # atom-state setting should be lp1 for atom 1 self.assertEqual(tuple(lp1), stored.pos['%s-%s-%s' % (m1, 1, 1)]) self.assertEqual(tuple(lp1), stored.pos['%s-%s-%s' % (m1, 2, 1)]) # atom setting should override to lp2 for atom 2 in both states self.assertEqual(tuple(lp2), stored.pos['%s-%s-%s' % (m1, 1, 2)]) self.assertEqual(tuple(lp2), stored.pos['%s-%s-%s' % (m1, 2, 2)]) # unset all atom-level settings, atom-state settings should still exist cmd.alter("all", "s['label_placement_offset']=None") stored.pos = {} cmd.iterate_state(0, "all", "stored.pos['%s-%s-%s' % (model, state, index)]=s['label_placement_offset']") self.assertEqual(tuple(lp1), stored.pos['%s-%s-%s' % (m1, 1, 1)]) self.assertEqual(tuple(lp1), stored.pos['%s-%s-%s' % (m1, 2, 1)]) self.assertEqual(tuple(stored.origp), stored.pos['%s-%s-%s' % (m1, 1, 2)]) self.assertEqual(tuple(stored.origp), stored.pos['%s-%s-%s' % (m1, 2, 2)])
def loop_orientation(selection, visualize=1, quiet=0): ''' DESCRIPTION Get the center and approximate direction of a peptide. Works for any secondary structure. Averages direction of N(i)->C(i) pseudo bonds. USAGE loop_orientation selection [, visualize] SEE ALSO helix_orientation ''' visualize, quiet = int(visualize), int(quiet) stored.x = dict() cmd.iterate_state(STATE, '(%s) and name N+C' % (selection), 'stored.x.setdefault(chain + resi, dict())[name] = x,y,z') vec = cpv.get_null() count = 0 for x in stored.x.itervalues(): if 'C' in x and 'N' in x: vec = cpv.add(vec, cpv.sub(x['C'], x['N'])) count += 1 if count == 0: print 'warning: count == 0' raise CmdException vec = cpv.normalize(vec) return _common_orientation(selection, vec, visualize, quiet)
def plane_orientation(selection, state=-1, visualize=1, quiet=1): ''' DESCRIPTION Fit plane (for example beta-sheet). Can also be used with angle_between_helices (even though this does not fit helices). Returns center and normal vector of plane. ''' try: import numpy except ImportError: print ' Error: numpy not available' raise CmdException state, visualize, quiet = int(state), int(visualize), int(quiet) coords = list() cmd.iterate_state(state, '(%s) and guide' % (selection), 'coords.append([x,y,z])', space=locals()) if len(coords) < 3: print 'not enough guide atoms in selection' raise CmdException x = numpy.array(coords) U,s,Vh = numpy.linalg.svd(x - x.mean(0)) # normal vector of plane is 3rd principle component vec = cpv.normalize(Vh[2]) if cpv.dot_product(vec, x[-1] - x[0]) < 0: vec = cpv.negate(vec) center = x.mean(0).tolist() _common_orientation(selection, center, vec, visualize, 4.0, quiet) # plane visualize if visualize: from pymol import cgo dir1 = cpv.normalize(Vh[0]) dir2 = cpv.normalize(Vh[1]) sx = [max(i/4.0, 2.0) for i in s] obj = [ cgo.BEGIN, cgo.TRIANGLES, cgo.COLOR, 0.5, 0.5, 0.5 ] for vertex in [ cpv.scale(dir1, sx[0]), cpv.scale(dir2, sx[1]), cpv.scale(dir2, -sx[1]), cpv.scale(dir1, -sx[0]), cpv.scale(dir2, -sx[1]), cpv.scale(dir2, sx[1]), ]: obj.append(cgo.VERTEX) obj.extend(cpv.add(center, vertex)) obj.append(cgo.END) cmd.load_cgo(obj, cmd.get_unused_name('planeFit')) return center, vec
def _inner_colorize_4(): time_step = const_time_max / 100 cols = [] cmd.iterate_state(state=-1, selection=_atom_list_to_pymol_selection(atoms), expression='cols.append(ID)', space=locals) print(cols, mv=1)
def create_pymol_objects_for_molecules(exclude: t.List[str] = []): """ creates separate selecetions for all molecules in interface_Pymol TODO; WRITE THE NAME OF THE MOLECULE IN TO THE ATOM SOMEHOW TODO: make more stable if protein is present! Parameters ---------- exclude: t.List[str] A list of molecules that should not be considered, e.g. 'solv;, 'protein', etc Returns ------- t.List[u.Atom] The list of the molecule object names """ molecules: str = [] def _append_molecule(m: str): if not m in molecules: molecules.append(m) """ Deapreacitated, as chains are maybe more general! #try chain wise cmd.iterate_state(-1, selection='all', expression='_append_molecule(resi)', space=locals()) molecule_names = [] for m in molecules: if not m in exclude: name = 'mol_'+ m cmd.create(name=name, selection='resi ' + m) molecule_names.append(name) """ cmd.iterate_state(-1, selection='all', expression='_append_molecule(resn)', space=locals()) molecule_names = [] i = 1 for m in molecules: if not m in exclude: name = 'mol_' + str(i) cmd.create(name=name, selection='resn ' + m) molecule_names.append(name) i += 1 return molecule_names # The selection containing all the Molecules must be deleted. If it is not, iterate state will duplicate each atom! cmd.delete(cmd.get_object_list() [0]) # Delete the object containing all molecules at once
def get_atom_cords(sele): stored.xyz = [] cmd.iterate_state(1, sele, "stored.xyz.append([x,y,z])") list_xyz = [] for atom in stored.xyz: list_xyz.append(atom) return list_xyz
def cubes(selection='all', name='', state=0, scale=0.5, atomcolors=1, _func=cgo_cube): ''' DESCRIPTION Create a cube representation CGO for all atoms in selection. ARGUMENTS selection = string: atom selection {default: all} name = string: name of CGO object to create state = int: object state {default: 0 = all states} scale = float: scaling factor. If scale=1.0, the corners of the cube will be on the VDW surface of the atom {default: 0.5} atomcolors = 0/1: use atom colors (cannot be changed), otherwise apply one color to the object (can be changed with color command) {default: 1} SEE ALSO tetrahedra ''' if not name: name = cmd.get_unused_name('cubes') state, scale, atomcolors = int(state), float(scale), int(atomcolors) if state < 0: state = cmd.get_setting_int('state') states = [state] if state else list( range(1, cmd.count_states(selection) + 1)) def callback(x, y, z, vdw, color): if atomcolors: obj.append(cgo.COLOR) obj.extend(cmd.get_color_tuple(color)) obj.extend(_func(x, y, z, vdw * scale)) space = {'xcb': callback} for state in states: obj = [] cmd.iterate_state(state, selection, 'xcb(x, y, z, vdw, color)', space=space) cmd.load_cgo(obj, name, state) if not atomcolors: cmd.color('auto', name)
def get_ring_coords(resn_list, matrix): """obtain coordinates of sugar rings""" matrix_coords = [] for state in range(1, cmd.count_states()+1): coords = [] for i, resi in enumerate(resn_list): stored.coords = [] cmd.iterate_state(state, 'resi %s and name %s' % (resi, '+'.join(matrix[i])), 'stored.coords.append([x,y,z])') coords.append(stored.coords) matrix_coords.append(coords) return matrix_coords
def helix_orientation(selection, visualize=1, sigma_cutoff=1.5, quiet=0): ''' DESCRIPTION Get the center and direction of a helix as vectors. Will only work for helices and gives slightly different results than loop_orientation. Averages direction of C(i)->O(i) bonds. USAGE helix_orientation selection [, visualize [, sigma_cutoff]] ARGUMENTS selection = string: atom selection of helix visualize = 0 or 1: show fitted vector as arrow {default: 1} sigma_cutoff = float: drop outliers outside (standard_deviation * sigma_cutoff) {default: 1.5} SEE ALSO angle_between_helices, helix_orientation_hbond, loop_orientation, cafit_orientation ''' visualize, quiet, sigma_cutoff = int(visualize), int(quiet), float( sigma_cutoff) stored.x = dict() cmd.iterate_state( STATE, '(%s) and name C+O' % (selection), 'stored.x.setdefault(chain + resi, dict())[name] = x,y,z') vec_list = [] count = 0 for x in stored.x.itervalues(): if 'C' in x and 'O' in x: vec_list.append(cpv.sub(x['O'], x['C'])) count += 1 if count == 0: print 'warning: count == 0' raise CmdException vec = _vec_sum(vec_list) if count > 2 and sigma_cutoff > 0: angle_list = [cpv.get_angle(vec, x) for x in vec_list] angle_mu, angle_sigma = _mean_and_std(angle_list) vec_list = [ vec_list[i] for i in range(len(vec_list)) if abs(angle_list[i] - angle_mu) < angle_sigma * sigma_cutoff ] if not quiet: print 'Dropping %d outlier(s)' % (len(angle_list) - len(vec_list)) vec = _vec_sum(vec_list) vec = cpv.normalize(vec) return _common_orientation(selection, vec, visualize, quiet)
def mutXYZposition(residues): cmd.hide("everything") cmd.show("sticks") for resi in residues: for caa in residues[resi]: print(resi, caa) filename = "residue" + str(resi) + "to" + caa + ".csv" outfile = open(filename, "w") row = ("rotamer" + "," + "X" + "," + "Y" + "," + "Z" + "\n") cmd.wizard("mutagenesis") cmd.do("refresh_wizard") cmd.get_wizard().set_mode(caa) sel = "/test//A/" + str(resi) cmd.get_wizard().do_select(sel) fram = cmd.count_states("mutation") stored.pos = [] cmd.iterate_state(0, 'mutation', 'stored.pos.append((x,y,z))', atomic=0) counter = 1 state = 1 natom = int(len(stored.pos) / fram) for po in stored.pos: if counter < natom: row = (str(state) + "," + str(po[0]) + "," + str(po[1]) + "," + str(po[2]) + "\n") outfile.write(row) counter += 1 elif counter == natom: row = (str(state) + "," + str(po[0]) + "," + str(po[1]) + "," + str(po[2]) + "\n") outfile.write(row) counter = 1 state += 1 outfile.close() cmd.frame(fram) cmd.get_wizard().apply() cmd.set_wizard("done") cmd.save("residue" + str(resi) + "to" + caa + "rotamer" + str(fram) + ".pdb") for f in range(1, fram): fra = fram - f cmd.wizard("mutagenesis") cmd.do("refresh_wizard") cmd.get_wizard().do_select(sel) cmd.frame(fra) cmd.get_wizard().apply() cmd.set_wizard("done") cmd.save("residue" + str(resi) + "to" + caa + "rotamer" + str(fra) + ".pdb")
def helix_orientation(selection, visualize=1, sigma_cutoff=1.5, quiet=0): """ DESCRIPTION Get the center and direction of a helix as vectors. Will only work for helices and gives slightly different results than loop_orientation. Averages direction of C(i)->O(i) bonds. USAGE helix_orientation selection [, visualize [, sigma_cutoff]] ARGUMENTS selection = string: atom selection of helix visualize = 0 or 1: show fitted vector as arrow {default: 1} sigma_cutoff = float: drop outliers outside (standard_deviation * sigma_cutoff) {default: 1.5} SEE ALSO angle_between_helices, helix_orientation_hbond, loop_orientation, cafit_orientation """ visualize, quiet, sigma_cutoff = int(visualize), int(quiet), float(sigma_cutoff) stored.x = dict() cmd.iterate_state( STATE, "(%s) and name C+O" % (selection), "stored.x.setdefault(chain + resi, dict())[name] = x,y,z" ) vec_list = [] count = 0 for x in stored.x.values(): if "C" in x and "O" in x: vec_list.append(cpv.sub(x["O"], x["C"])) count += 1 if count == 0: print("warning: count == 0") raise CmdException vec = _vec_sum(vec_list) if count > 2 and sigma_cutoff > 0: angle_list = [cpv.get_angle(vec, x) for x in vec_list] angle_mu, angle_sigma = _mean_and_std(angle_list) vec_list = [ vec_list[i] for i in range(len(vec_list)) if abs(angle_list[i] - angle_mu) < angle_sigma * sigma_cutoff ] if not quiet: print("Dropping %d outlier(s)" % (len(angle_list) - len(vec_list))) vec = _vec_sum(vec_list) vec = cpv.normalize(vec) return _common_orientation(selection, vec, visualize, quiet)
def testIterateState(self): cmd.fragment('ala') cmd.create('ala', 'ala', 1, 2) self.assertEqual(2, cmd.count_states()) v_count = cmd.count_atoms('all') expr = 'v_xyz.append(((model,index), (x,y,z)))' # current state v_xyz = [] cmd.iterate_state(-1, 'all', expr, space=locals()) self.assertEqual(len(v_xyz), v_count) # all states v_xyz = [] cmd.iterate_state(0, 'all', expr, space=locals()) self.assertEqual(len(v_xyz), v_count * 2) # atomic=0 stored.v_xyz = [] cmd.iterate_state(1, 'all', 'stored.v_xyz.append((x,y,z))', atomic=0) self.assertEqual(len(stored.v_xyz), v_count) space = {'self': self, 'NameError': NameError, 'v_list': []} cmd.iterate_state( 1, 'all', 'v_list.append(self.assertRaises(NameError, lambda: (model, index)))', space=space) self.assertEqual(len(space['v_list']), v_count)
def find_flips(self): """ Yield the indices of any positions where the wildtype and mut sidechains have at least one atom further apart than the distance threshold. """ from math import sqrt for i in range(len(self.aligned_resis[0])): wt_seq = self.aligned_seqs[0][i] mut_seq = self.aligned_seqs[1][i] if wt_seq != mut_seq: continue if wt_seq in 'X-' or mut_seq in 'X-': continue # I wanted to use `cmd.rms_cur()` to calculate this, but it refuses # to work on residues with different numbers, which I often have. wt_xyzs = {} mut_xyzs = {} sele = '({}) and resi {} and chain {} and not hydro' cmd.iterate_state(1, sele.format(self.wildtype_obj, *self.aligned_resis[0][i]), 'wt_xyzs[name] = x,y,z', space=locals(), ) cmd.iterate_state(1, sele.format(self.mutant_obj, *self.aligned_resis[1][i]), 'mut_xyzs[name] = x,y,z', space=locals(), ) # Even though the two residue are the same type, they might have # different sets of atoms. Some ways this could happen: (i) some # atoms weren't resolved in one structure, (ii) Rosetta adds an # oxygen to the C-terminus (OXT). keys = set(wt_xyzs.keys()) & set(mut_xyzs.keys()) max_dist = 0 for k in keys: dist = sqrt(sum( (wt_xyzs[k][j] - mut_xyzs[k][j])**2 for j in range(3) )) max_dist = max(dist, max_dist) if max_dist > self.flip_dist_cutoff: yield i
def testIterateState(self): cmd.fragment('ala') cmd.create('ala', 'ala', 1, 2) self.assertEqual(2, cmd.count_states()) v_count = cmd.count_atoms('all') expr = 'v_xyz.append(((model,index), (x,y,z)))' # current state v_xyz = [] cmd.iterate_state(-1, 'all', expr, space=locals()) self.assertEqual(len(v_xyz), v_count) # all states v_xyz = [] cmd.iterate_state(0, 'all', expr, space=locals()) self.assertEqual(len(v_xyz), v_count * 2) # atomic=0 stored.v_xyz = [] cmd.iterate_state(1, 'all', 'stored.v_xyz.append((x,y,z))', atomic=0) self.assertEqual(len(stored.v_xyz), v_count) space = {'self': self, 'NameError': NameError, 'v_list': []} cmd.iterate_state(1, 'all', 'v_list.append(self.assertRaises(NameError, lambda: (model, index)))', space=space) self.assertEqual(len(space['v_list']), v_count)
def testAtomLevelSettingsOnRemove(self): if cmd.get_setting_int('suspend_undo'): self.skipTest("need suspend_undo=0") cmd.load(self.datafile("1molecule.mae")) cmd.load(self.datafile("1molecule.mae")) cmd.label("index 10", "'Test Label'") plv = [ 5., 0., 0.] cmd.alter("index 10", "s.label_placement_offset = %s" % plv) cmd.remove("index 10") cmd.undo() stored.offset = None cmd.iterate_state(1, "index 10", "stored.offset = list(s.label_placement_offset)") self.assertEqual(stored.offset, plv)
def sphere(name, model_and_center_atom, radius, color, tr): ''' DESCRIPTION "sphere" creates a sphere with the given center-coordinates and radius USAGE sphere name, x_center, y_center, z_center, radius, color, tr name = name of sphere center_atom = center of sphere radius = radius of sphere color = color name tr = transparency value ''' color_rgb = cmd.get_color_tuple(cmd.get_color_index(color)) r = color_rgb[0] g = color_rgb[1] b = color_rgb[2] str_list = [] #str_list.append(str(center_atom)) #res_str = str(center_atom) #str_list.append(str(model)) #str_list.append(str("and resi")) #str_list.append(str(res_str)) #str_list.append(str("and name Ca")) sel_str = model_and_center_atom #string.join(str_list, ' ') print sel_str stored.xyz = [] #stored.xyz.append([x_center,y_center,z_center]) cmd.create("sphere", sel_str) cmd.iterate_state(1,"sphere","stored.xyz.append([x,y,z])") cmd.delete("sphere") print stored.xyz obj = [] obj.extend([cgo.ALPHA, tr]) obj.extend([ BEGIN, SPHERE, COLOR, r, g, b, SPHERE, stored.xyz[0][0], stored.xyz[0][1], stored.xyz[0][2], radius, END ]) cmd.load_cgo(obj, name)
def do_pick(self, picked_bond): self.reset() cmd.iterate("pk1", "setattr(cmd.get_wizard(),'pk1_st'," "'%s/%s/%s/%s/%s'%(model,segi,chain,resi,name))") if picked_bond: cmd.iterate("pk1", "setattr(cmd.get_wizard(),'pk2_st'," "'%s/%s/%s/%s/%s'%(model,segi,chain,resi,name))") else: # for single atom, also get 3D coordinates (EXAMPLE) cmd.iterate_state(cmd.get_state(), "pk1", "setattr(cmd.get_wizard(),'pk1_xyz',(x,y,z))") cmd.unpick() cmd.refresh_wizard()
def testOrigin(self): from chempy import cpv cmd.pseudoatom('m1') cmd.pseudoatom('m2') cmd.pseudoatom('m3', pos=[1,0,0]) # by selection cmd.origin('m3') cmd.rotate('y', 90, 'm1') # by position cmd.origin(position=[-1,0,0]) cmd.rotate('y', 90, 'm2') coords = [] cmd.iterate_state(1, 'm1 m2', 'coords.append([x,y,z])', space=locals()) d = cpv.distance(*coords) self.assertAlmostEqual(d, 2 * 2**0.5)
def get_bonds_colors(resn_list, matrix, color): """obtain colors for the bonds""" matrix_colors = [] if color == 'auto': for state in range(1, cmd.count_states()+1): colors = [] for bond in matrix: stored.colors = [] cmd.iterate_state(state, 'resi %s and name c1 or resi %s and name c1' % (bond[0], bond[2]), 'stored.colors.append(color)') colors.append((stored.colors[0], stored.colors[1])) matrix_colors.append(colors) else: for state in range(1, cmd.count_states()+1): matrix_colors.append([(color, color)] * (len(resn_list)-1)) return matrix_colors
def testOrigin(self): from chempy import cpv cmd.pseudoatom('m1') cmd.pseudoatom('m2') cmd.pseudoatom('m3', pos=[1, 0, 0]) # by selection cmd.origin('m3') cmd.rotate('y', 90, 'm1') # by position cmd.origin(position=[-1, 0, 0]) cmd.rotate('y', 90, 'm2') coords = [] cmd.iterate_state(1, 'm1 m2', 'coords.append([x,y,z])', space=locals()) d = cpv.distance(*coords) self.assertAlmostEqual(d, 2 * 2**0.5)
def get_colors_c1(resn_list, color): """obtain colors of c1 atoms""" matrix_colors = [] if color == 'auto': for state in range(1, cmd.count_states()+1): colors = [] for i, resi in enumerate(resn_list): stored.colors = [] cmd.iterate_state(state, 'resi %s and name c1' % resi, 'stored.colors.append(color)') colors.extend(stored.colors) matrix_colors.append(colors) else: for state in range(1, cmd.count_states()+1): matrix_colors.append([color] * len(resn_list)) return matrix_colors
def calculate_target_grid_center(self): if self.target_grid_center_selection_mode.get( ) == GRID_CENTER_FROM_SELECTION: sel = self.target_grid_center_selection_user.get() if sel: stored.xyz_target = [] cmd.iterate_state(1, sel, "stored.xyz_target.append([x,y,z])") xx = average([a[0] for a in stored.xyz_target]) yy = average([a[1] for a in stored.xyz_target]) zz = average([a[2] for a in stored.xyz_target]) self.target_grid_center[0].set(round(xx, 2)) self.target_grid_center[1].set(round(yy, 2)) self.target_grid_center[2].set(round(zz, 2)) else: self.target_grid_center_selection_mode.set( GRID_CENTER_FROM_COORDINATES)
def calcCogFromStr(selection: str) -> Tuple[float, float, float]: """ calculates the center of geometry of a given PyMOL selection """ stored.cogX, stored.cogY, stored.cogZ = 0, 0, 0 stored.i = 1 # has to be in an if statement since otherwise there have to be multiple for loops (pyMOL) cmd.iterate_state(-1, selection, """\ if(True): stored.cogX += x stored.cogY += y stored.cogZ += z stored.i += 1 """) return(stored.cogX/stored.i, stored.cogY/stored.i, stored.cogZ/stored.i)
def computecenterRA(self, selection="(all)"): stored.xyz = [] cmd.iterate_state(1, selection, "stored.xyz.append([x, y, z])") centx = 0 centy = 0 centz = 0 cnt = 0 for atom in stored.xyz: centx += atom[0] centy += atom[1] centz += atom[2] cnt += 1 centx /= cnt centy /= cnt centz /= cnt return (centx, centy, centz)
def get_atoms(sel, attrs, state=1): """Get the atoms and attributes of a selection.""" coords = None if "coords" in attrs: coords = pm.get_coords(sel, state) attrs.remove("coords") atoms = pd.DataFrame(coords, columns=["x", "y", "z"]) if attrs: fields_str = ", ".join(attrs) stored.atoms = [] pm.iterate_state(state, sel, f"stored.atoms.append(({fields_str}))") atoms = pd.concat( [atoms, pd.DataFrame(stored.atoms, columns=attrs)], axis=1) del stored.atoms return atoms
def calculate_offtarget_grid_center(self): if self.offtarget_grid_center_selection_mode.get( ) == GRID_CENTER_FROM_SELECTION_B: sel = self.offtarget_grid_center_selection_user.get() if sel: stored.xyz_offtarget = [] cmd.iterate_state(1, sel, "stored.xyz_offtarget.append([x,y,z])") xx = average(map(lambda a: a[0], stored.xyz_offtarget)) yy = average(map(lambda a: a[1], stored.xyz_offtarget)) zz = average(map(lambda a: a[2], stored.xyz_offtarget)) self.offtarget_grid_center[0].set(round(xx, 2)) self.offtarget_grid_center[1].set(round(yy, 2)) self.offtarget_grid_center[2].set(round(zz, 2)) else: self.offtarget_grid_center_selection_mode.set( GRID_CENTER_FROM_COORDINATES_B)
def cubes(selection='all', name='', state=0, scale=0.5, atomcolors=1, _func=cgo_cube): ''' DESCRIPTION Create a cube representation CGO for all atoms in selection. ARGUMENTS selection = string: atom selection {default: all} name = string: name of CGO object to create state = int: object state {default: 0 = all states} scale = float: scaling factor. If scale=1.0, the corners of the cube will be on the VDW surface of the atom {default: 0.5} atomcolors = 0/1: use atom colors (cannot be changed), otherwise apply one color to the object (can be changed with color command) {default: 1} SEE ALSO tetrahedra ''' if not name: name = cmd.get_unused_name('cubes') state, scale, atomcolors = int(state), float(scale), int(atomcolors) if state < 0: state = cmd.get_setting_int('state') states = [state] if state else range(1, cmd.count_states(selection) + 1) def callback(x, y, z, vdw, color): if atomcolors: obj.append(cgo.COLOR) obj.extend(cmd.get_color_tuple(color)) obj.extend(_func(x, y, z, vdw * scale)) space = {'xcb': callback} for state in states: obj = [] cmd.iterate_state(state, selection, 'xcb(x, y, z, vdw, color)', space=space) cmd.load_cgo(obj, name, state) if not atomcolors: cmd.color('auto', name)
def cafit_orientation(selection, state=STATE, visualize=1, guide=1, quiet=1): ''' DESCRIPTION Get the center and direction of a peptide by least squares linear fit on CA atoms. USAGE cafit_orientation selection [, visualize ] NOTES Requires python module "numpy". SEE ALSO helix_orientation ''' try: import numpy except ImportError: print(' Error: numpy not available') raise CmdException state, visualize, quiet = int(state), int(visualize), int(quiet) if int(guide): selection = '(%s) and guide' % (selection) coords = [] cmd.iterate_state(state, selection, 'coords.append([x,y,z])', space=locals()) x = numpy.array(coords) center = x.mean(0).tolist() U, s, Vh = numpy.linalg.svd(x - center) vec = cpv.normalize(Vh[0]) if cpv.dot_product(vec, x[-1] - x[0]) < 0: vec = cpv.negate(vec) _common_orientation(selection, center, vec, visualize, s[0], quiet) return center, vec
def sphere(name, model_and_center_atom, radius, color, tr): ''' DESCRIPTION "sphere" creates a sphere with the given center-coordinates and radius USAGE sphere name, x_center, y_center, z_center, radius, color, tr name = name of sphere center_atom = center of sphere radius = radius of sphere color = color name tr = transparency value ''' color_rgb = cmd.get_color_tuple(cmd.get_color_index(color)) r = color_rgb[0] g = color_rgb[1] b = color_rgb[2] str_list = [] #str_list.append(str(center_atom)) #res_str = str(center_atom) #str_list.append(str(model)) #str_list.append(str("and resi")) #str_list.append(str(res_str)) #str_list.append(str("and name Ca")) sel_str = model_and_center_atom #string.join(str_list, ' ') print sel_str stored.xyz = [] #stored.xyz.append([x_center,y_center,z_center]) cmd.create("sphere", sel_str) cmd.iterate_state(1, "sphere", "stored.xyz.append([x,y,z])") cmd.delete("sphere") print stored.xyz obj = [] obj.extend([cgo.ALPHA, tr]) obj.extend([ BEGIN, SPHERE, COLOR, r, g, b, SPHERE, stored.xyz[0][0], stored.xyz[0][1], stored.xyz[0][2], radius, END ]) cmd.load_cgo(obj, name)
def _get_solvent_radius(selection, state): '''Get object-state level solvent_radius''' radius = [None] try: cmd.iterate_state(state, 'first ({})'.format(selection), 'radius[0] = s.solvent_radius', space=locals()) if radius[0] is not None: # Note: One of my test cases failed with radius 2.75, rounding # to one digit worked return '{:.1}'.format(radius[0]) except: print('Using global solvent_radius') return cmd.get('solvent_radius')
def _common_orientation(selection, vec, visualize=1, quiet=0): """ Common part of different helix orientation functions. Does calculate the center of mass and does the visual feedback. """ stored.x = [] cmd.iterate_state(STATE, "(%s) and name CA" % (selection), "stored.x.append([x,y,z])") if len(stored.x) < 2: print("warning: count(CA) < 2") raise CmdException center = cpv.scale(_vec_sum(stored.x), 1.0 / len(stored.x)) if visualize: scale = cpv.distance(stored.x[0], stored.x[-1]) visualize_orientation(vec, center, scale, True) cmd.zoom(selection, buffer=2) if not quiet: print("Center: (%.2f, %.2f, %.2f) Direction: (%.2f, %.2f, %.2f)" % tuple(center + vec)) return center, vec
def testAlterState(self): self.load_big_example_multistate() v_count = cmd.count_atoms() * cmd.count_states() assert v_count > 10**5 xyz = [] with self.timing('i', 5.0): cmd.iterate_state(0, 'all', 'xyz.append((x,y,z))', space=locals()) self.assertEqual(v_count, len(xyz)) with self.timing('a', 5.0): cmd.alter_state(0, 'all', '(x,y,z) = next(xyz_rev)', space={'xyz_rev': reversed(xyz), 'next': next}) cmd.iterate_state(0, 'last all', 'stored.xyz = (x,y,z)') self.assertEqual(stored.xyz, xyz[0])
def cafit_orientation(selection, state=STATE, visualize=1, guide=1, quiet=1): ''' DESCRIPTION Get the center and direction of a peptide by least squares linear fit on CA atoms. USAGE cafit_orientation selection [, visualize ] NOTES Requires python module "numpy". SEE ALSO helix_orientation ''' try: import numpy except ImportError: print ' Error: numpy not available' raise CmdException state, visualize, quiet = int(state), int(visualize), int(quiet) if int(guide): selection = '(%s) and guide' % (selection) coords = [] cmd.iterate_state(state, selection, 'coords.append([x,y,z])', space=locals()) x = numpy.array(coords) center = x.mean(0).tolist() U,s,Vh = numpy.linalg.svd(x - center) vec = cpv.normalize(Vh[0]) if cpv.dot_product(vec, x[-1] - x[0]) < 0: vec = cpv.negate(vec) _common_orientation(selection, center, vec, visualize, s[0], quiet) return center, vec
def data( userSelection ): # this array will be used to hold the coordinates. It # has access to PyMOL objects and, we have access to it. stored.alphaCarbons = [] # let's just get the alpha carbons, so make the # selection just for them userSelection = userSelection + " and n. CA" # iterate over state 1, or the userSelection -- this just means # for each item in the selection do what the next parameter says. # And, that is to append the (x,y,z) coordinates to the stored.alphaCarbon # array. cmd.iterate_state(1, selector.process(userSelection), "stored.alphaCarbons.append([x,y,z])") return stored.alphaCarbons #userSelection = [] #cmd.extend( "data", data(userSelection) ); # stored.alphaCarbons now has the data you want.
def save_mopac(filename, dist, selection='all', zero='none', state=-1, quiet=1): ''' DESCRIPTION Save to MOPAC format ARGUMENTS filename = string: file path to be written dist = string: beyond which distance to fix atoms selection = string: atoms to save {default: all} zero = string: atoms to save with zero flag {default: none} state = integer: state to save {default: -1 (current state)} ''' #cmd.select("rest", "(byres (sub expand %s))" % dist) #cmd.select("rest", "((byobj sub) and not rest)") state, quiet = int(state), int(quiet) fmt = '%5s(%6i %3s%4i) %12.8f +%i %12.8f +%i %12.8f +%i %26.4f\n' zero_idx = set() cmd.iterate(zero, 'zero_idx.add((model,index))', space=locals()) serial = [0] def callback(model, index, e, resn, resv, x, y, z, c): if (model, index) in zero_idx: flag = 0 else: flag = 1 serial[0] += 1 # print "Saving: %s" % save_string handle.write(fmt % (e, serial[0], resn, resv, x, flag, y, flag, z, flag, c)) save_string = filename.split('-')[0] + "-con-" + filename.split('-')[2] +\ "-%02d" % float(dist) + "-ste-ini.pdb" handle = open(save_string, 'w') cmd.iterate_state(state, selection, 'callback(model, index, elem, resn, resv,' ' x, y, z, partial_charge)', space=locals()) handle.close() if not quiet: print ' Save-MOPAC: Wrote %i atoms to file' % (serial[0])
def save_mopac(filename, selection='all', zero='none', state=-1, quiet=1): ''' DESCRIPTION Save to MOPAC format ARGUMENTS filename = string: file path to be written selection = string: atoms to save {default: all} zero = string: atoms to save with zero flag {default: none} state = integer: state to save {default: -1 (current state)} ''' state, quiet = int(state), int(quiet) fmt = '%5s(%6i %3s%4i) %12.8f +%i %12.8f +%i %12.8f +%i %26.4f\n' zero_idx = set() cmd.iterate(zero, 'zero_idx.add((model,index))', space=locals()) serial = [0] def callback(model, index, e, resn, resv, x, y, z, c): flag = (model, index) not in zero_idx and 1 or 0 serial[0] += 1 handle.write(fmt % (e, serial[0], resn, resv, x, flag, y, flag, z, flag, c)) handle = open(filename, 'w') handle.write('PM6\n\nGenerated by PyMOL\n') cmd.iterate_state(state, selection, 'callback(model, index, elem, resn, resv, x, y, z, partial_charge)', space=locals()) handle.close() if not quiet: print ' Save-MOPAC: Wrote %i atoms to file' % (serial[0])
def loop_orientation(selection, state=STATE, visualize=1, quiet=1): ''' DESCRIPTION Get the center and approximate direction of a peptide. Works for any secondary structure. Averages direction of N(i)->C(i) pseudo bonds. USAGE loop_orientation selection [, visualize ] SEE ALSO helix_orientation ''' state, visualize, quiet = int(state), int(visualize), int(quiet) coords = dict() cmd.iterate_state(state, '(%s) and name N+C' % (selection), 'coords.setdefault(chain + resi, {})[name] = x,y,z', space=locals()) vec = cpv.get_null() center = cpv.get_null() count = 0 for x in coords.itervalues(): if 'C' in x and 'N' in x: vec = cpv.add(vec, cpv.sub(x['C'], x['N'])) for coord in x.itervalues(): center = cpv.add(center, coord) count += 1 if count == 0: print 'warning: count == 0' raise CmdException vec = cpv.normalize(vec) center = cpv.scale(center, 1./count) _common_orientation(selection, center, vec, visualize, 2.0*len(coords), quiet) return center, vec
def contacts(selection, filename='contacts.png', cutoff=8): # get residue numbers stored.resis = [] cmd.iterate(selection, 'stored.resis.append(int(resi))') # get coordinates stored.positions = [] cmd.iterate_state(1, selection, 'stored.positions.append((x,y,z))') # a little check for duplicate and missing residues prev = stored.resis[0]-1 for resi in sorted(stored.resis): if resi == prev: print 'duplicate',resi if resi == prev+2: print 'missing',prev+1 elif resi > prev+2: print 'missing',prev+1,'-',resi-1 prev = resi # plot the contacts on the canvas and save it out canvas = PNGCanvas(cellsize*len(stored.resis), cellsize*len(stored.resis)) plot_contacts(stored.resis, stored.positions, float(cutoff), canvas) with open(filename, 'wb') as f: f.write(canvas.dump())
def test2667(self): cmd.fragment('ala', 'm1') cmd.alter_state(1, '%m1', 's.label_placement_offset = [1., 0., 0.]') with testing.mktemp('.pse') as filename: cmd.save(filename) cmd.set_name('m1', 'm2') cmd.load(filename, partial=1) # now we have two copies (m1 m2). If unique ids are converted upon # loading, then there will be no settings cross-leaking. Otherwise # changing m1 settings will affect m2. cmd.alter_state(1, '%m1', 's.label_placement_offset = [0., 1., 0.]') m2_setting = [] cmd.iterate_state(1, 'first %m2', 'm2_setting.append(s.label_placement_offset)', space=locals()) self.assertArrayEqual([1., 0., 0.], m2_setting[0], 0.001)
def getCylinderCoords(cylinderCenterName, cylinderEdgeNames): """Queries pymol for the specified coordinates representing the cylinders. cylinderCenterName is a string specifying the pdb atom name field of the cylinder centers cylinderEdgeNames is a two element list specifying the pdb atom names of the cylinder top and bottom returns nothing, stores the results in pymol.stored.centCoords, pymol.stored.cylCoords0, pymol.stored.cylCoords1 """ stored.centCoords = []; stored.cylCoords0 = []; stored.cylCoords1 = []; cyl0Sel = "cylAtoms0"; cyl1Sel = "cylAtoms1"; centerSel = "centers"; cmd.select(centerSel, "(name " + cylinderCenterName + ")"); cmd.select(cyl0Sel, "(name " + cylinderEdgeNames[0] + ")"); cmd.select(cyl1Sel, "(name " + cylinderEdgeNames[1] + ")"); cmd.iterate_state(1, selector.process(centerSel), "stored.centCoords.append([x,y,z])"); cmd.iterate_state(1, selector.process(cyl0Sel), "stored.cylCoords0.append([x,y,z])"); cmd.iterate_state(1, selector.process(cyl1Sel), "stored.cylCoords1.append([x,y,z])");
def helix_orientation_hbond(selection, visualize=1, cutoff=3.5, quiet=0): ''' DESCRIPTION Get the center and direction of a helix as vectors. Will only work for alpha helices and gives slightly different results than helix_orientation. Averages direction of O(i)->N(i+4) hydrogen bonds. USAGE helix_orientation selection [, visualize [, cutoff]] ARGUMENTS cutoff = float: maximal hydrogen bond distance {default: 3.5} SEE ALSO helix_orientation ''' visualize, quiet, cutoff = int(visualize), int(quiet), float(cutoff) stored.x = dict() cmd.iterate_state(STATE, '(%s) and name N+O' % (selection), 'stored.x.setdefault(resv, dict())[name] = x,y,z') vec_list = [] for resi in stored.x: resi_other = resi + 4 if 'O' in stored.x[resi] and resi_other in stored.x: if 'N' in stored.x[resi_other]: vec = cpv.sub(stored.x[resi_other]['N'], stored.x[resi]['O']) if cpv.length(vec) < cutoff: vec_list.append(vec) if len(vec_list) == 0: print 'warning: count == 0' raise CmdException vec = _vec_sum(vec_list) vec = cpv.normalize(vec) return _common_orientation(selection, vec, visualize, quiet)
def rmsf2b(selection='all', linearscale=1.0, var='b', quiet=1): ''' DESCRIPTION Determine the root mean square fluctuation (RMSF) per atom for a multi-state object and assign b-factor ARGUMENTS selection = string: atom selection {default: name CA} linearscale = float: if linearscale <= 0, then use real b-factor equation, else use b=(rmsf*linearscale) {default: 1.0} SEE ALSO spheroid, rmsf_states.py from Robert Campbell ''' from numpy import asfarray, sqrt, pi linearscale = float(linearscale) n_atoms = cmd.count_atoms(selection) n_states = cmd.count_states(selection) if n_atoms == 0 or n_states < 2: print(' Error: not enough atoms or states') raise CmdException coords = [] cmd.iterate_state(0, selection, 'coords.append((x,y,z))', atomic=0, space={'coords': coords}) coords = asfarray(coords).reshape((cmd.count_states(selection), -1, 3)) u_sq = coords.var(0).sum(1) # var over states, sum over x,y,z b_array = sqrt(u_sq) * linearscale if linearscale > 0.0 \ else 8 * pi**2 * u_sq cmd.alter(selection, var + ' = b_iter.next()', space={'b_iter': iter(b_array)}) if not int(quiet): print(' Average RMSF: %.2f' % (sqrt(u_sq).mean())) return b_array