def test_sets_ranges(self):
     cmd.fab('ACDEFGHIKL')
     cmd.alter('resi 9', 'resi="9A"') # insertion code
     self.assertEqual(3, cmd.count_atoms('guide & resi 2-4'))
     self.assertEqual(3, cmd.count_atoms('guide & resi 2:4'))
     self.assertEqual(2, cmd.count_atoms('guide & resi 2+4'))
     self.assertEqual(4, cmd.count_atoms('guide & resi 2-4+6'))
     self.assertEqual(6, cmd.count_atoms('guide & resi 2-4+6-8'))
     self.assertEqual(0, cmd.count_atoms('guide & resi 9'))
     self.assertEqual(1, cmd.count_atoms('guide & resi 9A'))
     self.assertEqual(1, cmd.count_atoms('guide & resi 10'))
     self.assertEqual(0, cmd.count_atoms('guide & resi 10A'))
     self.assertEqual(2, cmd.count_atoms('guide & resi 9-10'))
     self.assertEqual(2, cmd.count_atoms('guide & resi 9A-10A'))
     self.assertEqual(10 + 9, cmd.count_atoms('name CA+CB'))
     self.assertEqual(10 + 9, cmd.count_atoms('name CA+CB+XYZ'))
     self.assertEqual(10, cmd.count_atoms('name C'))
     self.assertEqual(50, cmd.count_atoms('name C*'))
     self.assertEqual(10, cmd.count_atoms('name H'))
     self.assertEqual(24, cmd.count_atoms('name H*'))
     self.assertEqual(10, cmd.count_atoms('name *H'))
     self.assertEqual(74, cmd.count_atoms('name *H*'))
     self.assertEqual(20, cmd.count_atoms('name C+N'))
     self.assertEqual(23, cmd.count_atoms('name C+N*'))
     self.assertEqual(60, cmd.count_atoms('name C*+N'))
     self.assertEqual(63, cmd.count_atoms('name C*+N*'))

def loadBfacts (mol,startaa=1,source="/Users/student/Box Sync/PUBS/Pymol-practice.txt", visual="Y"):
	"""
	Replaces B-factors with a list of values contained in a plain txt file
 
	usage: loadBfacts mol, [startaa, [source, [visual]]]
 
	mol = any object selection (within one single object though)
	startaa = number of first amino acid in 'new B-factors' file (default=1)
	source = name of the file containing new B-factor values (default=newBfactors.txt)
	visual = redraws structure as cartoon_putty and displays bar with min/max values (default=Y)
 
	example: loadBfacts 1LVM and chain A
	"""
	obj=cmd.get_object_list(mol)[0]
	cmd.alter(mol,"b=-1.0")
	inFile = open(source, 'r')
	counter=int(startaa)
	bfacts=[]
	for line in inFile.readlines():	
		bfact=float(line)
		bfacts.append(bfact)
		cmd.alter("%s and resi %s and n. CA"%(mol,counter), "b=%s"%bfact)
		counter=counter+1
	if visual=="Y":
		cmd.show_as("cartoon",mol)
		cmd.cartoon("putty", mol)
		cmd.set("cartoon_putty_scale_min", min(bfacts),obj)
		cmd.set("cartoon_putty_scale_max", max(bfacts),obj)
		cmd.set("cartoon_putty_transform", 0,obj)
		cmd.set("cartoon_putty_radius", 0.2,obj)
		cmd.spectrum("b","rainbow", "%s and n. CA " %mol)
		cmd.ramp_new("count", obj, [min(bfacts), max(bfacts)], "rainbow")
		cmd.recolor()

def useOccRadii(sel="all"):
	for a in cmd.get_model(sel).atom:
		q = a.q
		if q >= 3:
			print "shrik radius"
			q <- 0.1
		cmd.alter("%s and resi %s and name %s"%(sel,a.resi,a.name),"vdw=%f"%(q))

    def testExportStyle(self):
        cmd.fab('ACDEF', 'm1')
        cmd.hide()
        cmd.show('cartoon', 'resi 1-3')
        cmd.show('lines', 'resn CYS')
        cmd.show('sticks', 'resn ASP+PHE')
        cmd.show('spheres', 'resn GLU')
        cmd.set('stick_ball', 1, 'resn PHE')
        cmd.set('stick_ball_ratio', 1.5, 'm1')
        testlabel = 'Hello "World"'
        cmd.label('name SG', repr(testlabel))

        with testing.mktemp('.mae') as filename:
            cmd.save(filename)
            cmd.delete('*')
            cmd.load(filename, 'm2')

        g_labels = []
        cmd.iterate('name SG', 'g_labels.append(label)', space=locals())
        cmd.alter('*', 'b = 1 if s.stick_ball else 0')

        self._assertCountEqual('rep cartoon & guide', 'resi 1-3 & guide')
        self._assertCountEqual('rep lines', 'resn CYS', delta=1)
        self._assertCountEqual('rep sticks', 'resn ASP+PHE')
        self._assertCountEqual('rep spheres', 'resn GLU')
        self._assertCountEqual('b > 0.5', 'resn PHE')
        self.assertTrue(cmd.get_setting_float('stick_ball_ratio', 'm2') > 1.1)
        self.assertEqual(g_labels[0], testlabel)

    def test_wildcard_sets_ranges(self):
        for i in range(10): cmd.pseudoatom('m%d' % i, chain=chr(65 + i))
        cmd.alter('m5', 'chain = "AB"')
        cmd.alter('m6', 'chain = "BA"')
        cmd.alter('m7', 'chain = "CC"')
        cmd.alter('m8', 'chain = "ZA"')
        cmd.alter('m9', 'chain = "ABA"')
        # A patterns
        self.assertEqual(1, cmd.count_atoms('chain A'))
        self.assertEqual(3, cmd.count_atoms('chain A*'))
        self.assertEqual(4, cmd.count_atoms('chain *A'))
        self.assertEqual(5, cmd.count_atoms('chain *A*'))
        self.assertEqual(1, cmd.count_atoms('chain A*A'))
        # B patterns
        self.assertEqual(2, cmd.count_atoms('chain B*'))
        self.assertEqual(2, cmd.count_atoms('chain *B'))
        self.assertEqual(4, cmd.count_atoms('chain *B*'))
        # X patterns (no matches)
        self.assertEqual(0, cmd.count_atoms('chain X*'))
        self.assertEqual(0, cmd.count_atoms('chain *X'))
        self.assertEqual(0, cmd.count_atoms('chain *X*'))
        # list with wildcards
        self.assertEqual(5, cmd.count_atoms('chain B*+A*'))
        self.assertEqual(3, cmd.count_atoms('chain B*+A*A'))
        self.assertEqual(3, cmd.count_atoms('chain B*+A*A+*X'))

        # lexicographical alpha ranges, A:C, will match AB (A <= AB <= C) but not CC (C < CC)
        # no wildcards in alpha ranges possible
        self.assertEqual(6, cmd.count_atoms('chain A:C'))
        self.assertEqual(6, cmd.count_atoms('chain A:CA'))
        self.assertEqual(7, cmd.count_atoms('chain A:CX')) # include CC
        self.assertEqual(6, cmd.count_atoms('chain A:C+Z'))
        self.assertEqual(7, cmd.count_atoms('chain A:C+Z*'))

def sequential_residues(sel1, offset=1):
    """
        PURPOSE: renumbers the selected residues sequentially, regardless of gaps
        USAGE: sequential_residues protName    # first residue is 1
        USAGE: sequential_residues protName, 5 # first residue is 5
        EXAMPLE: sequential_residues *
        """
    offset = int(offset)

    # A counter from offset up
    stored.offset = int(offset) - 1
    stored.curr_res = None

    cmd.alter(
        sel1,
        """
if stored.curr_res != int(resi):
    stored.offset=stored.offset + 1
    stored.curr_res=int(resi)
    resi=stored.offset
else:
    resi=stored.offset
""",
    )
    cmd.sort()

def useRosettaRadii():
	cmd.alter("element C", "vdw=2.00")
	cmd.alter("element N", "vdw=1.75")
	cmd.alter("element O", "vdw=1.55")
	cmd.alter("element H", "vdw=1.00")
	cmd.alter("element P", "vdw=1.90")
	cmd.set("sphere_scale", 1.0)

def corexProtectionFactorUpdateB(obj, infoFile, pfFile):
    """Sets the B-factor column to the protection factors.
    infoFile is used to establish the mapping between pdb and corex numbers
        (*.info).
    pfFile gives the protection factors (*.Protection_Factors).
    """
    pf = CorexProtectionFactors(pfFile)
    info = CorexAtomInfo(infoFile)

    if info is None:
        print "Error parsing %s" % infoFile
        return
    if pf is None:
        print "Error parsing %s" % exchangeRateFile
        return

    assignPDBNums(
            [atom.res for atom in info.getAtoms() if not atom.isHet],
            pf.getResidues()
        )

    cmd.alter(obj, "b = -10")
    for res in pf.getResidues():
        if res.pdbNum is not None and res.modifiedProtectionFactor is not None:
            sel = "%s and i. %s and chain %s" % (
                    obj,  res.pdbNum, res.chain)
            cmd.alter( sel, "b = %f" % res.modifiedProtectionFactor)
    cmd.sort(obj)

def loadFitnessFactors (mol,startaa=1,source="/Users/student/Box Sync/PUBS/Pymol-practice.txt", visual="Y"):
        # adapted from http://www.pymolwiki.org/index.php/Load_new_B-factors
        """
        Replaces B-factors with a list of fitness factor values contained in a plain txt file
 
        usage: loadFitnessFactors mol, [startaa, [source, [visual]]]
 
        mol = any object selection (within one single object though)
        startaa = number of first amino acid in 'new Fitness-factors' file (default=1)
        source = name of the file containing new Fitness-factor values (default=newFitnessFactors.txt)
        visual = redraws structure as cartoon_putty and displays bar with min/max values (default=Y)
 
        example: loadFitnessFactors 1LVM and chain A
        """
        obj=cmd.get_object_list(mol)[0]
        cmd.alter(mol,"b=-1.0")
        inFile = open(source, 'r')
        counter=int(startaa)
        fitnessFacts=[]
        for line in inFile.readlines():
                fitnessFact=float(line)
                fitnessFacts.append(fitnessFact)
                cmd.alter("%s and resi %s and n. CA"%(mol,counter), "b=%s"%fitnessFact)
                counter=counter+1
        if visual=="Y":
                cmd.show_as("cartoon",mol)
                cmd.cartoon("putty", mol)
#                cmd.set("cartoon_putty_scale_min", min(fitnessFacts),obj)
#                cmd.set("cartoon_putty_scale_max", max(fitnessFacts),obj)
                cmd.set("cartoon_putty_transform", 0,obj)
                cmd.set("cartoon_putty_radius", 0.2,obj)
                cmd.spectrum("b","red_white_blue", "%s and n. CA " %mol)
                cmd.ramp_new("count", obj, [min(fitnessFacts), (min(fitnessFacts)+max(fitnessFacts))/2, max(fitnessFacts)], color = ["blue", "white", "red"])
                cmd.recolor()

def corexDiffSAUpdateB(obj, corex1, corex2):
    """Sets the B-factor column to the difference in surface areas between
    files corex1-corex2
    """
    corex1 = CorexAtomInfo(corex1)
    if corex1 is None:
        print "Error parsing %s" % corex1
        return

    corex2 = CorexAtomInfo(corex2)
    if corex2 is None:
        print "Error parsing %s" % corex2
        return

    atoms1 = corex1.getAtoms()
    atoms2 = corex2.getAtoms()

    if len(atoms1) != len(atoms2):
        print "Error: Corex files must have identical atoms"
        #TODO line up atoms1 and atoms2 somehow
    for i in xrange(len(atoms1)):
        sel = "%s and n. %s and i. %s and chain %s" % (
                obj, atoms1[i].name, atoms1[i].res.pdbNum, atoms1[i].res.chain)
        cmd.alter( sel, "b = %f" % (atoms1[i].sa-atoms2[i].sa))
        print "b = %f" % (atoms1[i].sa-atoms2[i].sa)
    cmd.sort(obj)

def corexHDXRateUpdateB(obj, infoFile, exchangeRateFile):
    """Sets the B-factor column to the HDX rate.
    infoFile is used to establish the mapping between pdb and corex numbers
        (*.info).
    exchangeRateFile gives the HDX rates (*.Exchange_Rates).
    """
    hdx = CorexExchangeRates(exchangeRateFile)
    info = CorexAtomInfo(infoFile)

    if info is None:
        print "Error parsing %s" % infoFile
        return
    if hdx is None:
        print "Error parsing %s" % exchangeRateFile
        return

    assignPDBNums(
            [atom.res for atom in info.getAtoms() if not atom.isHet],
            hdx.getResidues()
        )

    cmd.alter(obj, "b = -10")
    for res in hdx.getResidues():
        if res.pdbNum is not None and res.exchangeRate is not None:
            sel = "%s and i. %s and chain %s" % (
                    obj,  res.pdbNum, res.chain)
            cmd.alter( sel, "b = %f" % res.exchangeRate)
    cmd.sort(obj)

def test_msms_surface():
    eps = 1e-3
    cmd.reinitialize()
    cmd.fragment('ala', 'm1')
    # default
    psico.msms.msms_surface(name='surf1')
    extent = cmd.get_extent('surf1')
    assert extent[0] == approx([-2.705, -3.208, -2.413], rel=eps)
    assert extent[1] == approx([3.530, 2.907, 2.676], rel=eps)
    # global solvent_radius
    cmd.set('solvent_radius', 3.5)
    psico.msms.msms_surface(name='surf2')
    extent = cmd.get_extent('surf2')
    assert extent[0] == approx([-2.705, -3.169, -2.436], rel=eps)
    assert extent[1] == approx([3.530, 2.907, 2.676], rel=eps)
    # object-level solvent_radius
    cmd.set('solvent_radius', 2.8, 'm1')
    psico.msms.msms_surface(name='surf3')
    extent = cmd.get_extent('surf3')
    assert extent[0] == approx([-2.705, -3.161, -2.427], rel=eps)
    assert extent[1] == approx([3.530, 2.907, 2.676], rel=eps)
    # modified atom radii
    cmd.alter('m1', 'vdw = 3.0')
    psico.msms.msms_surface(name='surf4')
    extent = cmd.get_extent('surf4')
    assert extent[0] == approx([-4.605, -5.162, -4.418], rel=eps)
    assert extent[1] == approx([5.030, 4.861, 4.681], rel=eps)

    def testCifMissing(self):
        N = 7
        cmd.fragment('gly', 'm1')
        cmd.alter('all', '(chain, segi, resv, alt) = ("?", ".", 5, "")')

        s = cmd.get_str('cif')
        self.assertTrue("'?'" in s or '"?"' in s) # chain
        self.assertTrue("'.'" in s or '"."' in s) # segi
        self.assertTrue(' ? ' in s) # e.g. pdbx_PDB_ins_code
        self.assertTrue(' . ' in s) # e.g. label_alt_id

        cmd.delete('*')
        cmd.set('cif_keepinmemory')
        cmd.load(s, 'm2', format='cifstr')
        self.assertEqual(['?'], cmd.get_chains())
        self.assertEqual(cmd.count_atoms('segi .'), N)
        self.assertEqual(cmd.count_atoms('alt ""'), N)  # no alt
        self.assertEqual(cmd.count_atoms('resi 5'), N)  # no ins_code

        from pymol.querying import cif_get_array
        self.assertEqual(cif_get_array("m2", "_atom_site.type_symbol"), list('NCCOHHH'))
        self.assertEqual(cif_get_array("m2", "_atom_site.id", "i"),     list(range(1, N + 1)))
        self.assertEqual(cif_get_array("m2", "_atom_site.auth_asym_id"),        ['?'] * N)
        self.assertEqual(cif_get_array("m2", "_atom_site.label_asym_id"),       ['.'] * N)
        self.assertEqual(cif_get_array("m2", "_atom_site.pdbx_pdb_ins_code"),   [None] * N)
        self.assertEqual(cif_get_array("m2", "_atom_site.label_alt_id"),        [None] * N)

    def RefreshDisplay(self):
        
        try:
            # Display the atoms spheres
            cmd.show('spheres', self.LigDisplay)
            cmd.refresh()

            cmd.alter(self.LigDisplay,'vdw=0.25')
            cmd.rebuild(self.LigDisplay)
            cmd.refresh()

            util.cbag(self.LigDisplay)
            cmd.refresh()
            
            if self.AnchorAtom != -1:
                AtomSel = self.LigDisplay + ' & id ' + str(self.AnchorAtom)
                cmd.color('white',AtomSel)
                cmd.refresh()

                cmd.alter(AtomSel ,'vdw=0.30')
                cmd.rebuild(AtomSel)
                cmd.refresh()
                        
        except:
            self.ErrorCode = 1

        return self.ErrorCode

def useTempRadii(sel="all"):
	for a in cmd.get_model(sel).atom:
		bfac = a.b
		if bfac >= 3:
			print "shrik radius"
			bfac <- 0.1
		cmd.alter("%s and resi %s and name %s"%(sel,a.resi,a.name),"vdw=%f"%(bfac))

    def updateSS(self):
        if self.ss_asgn_prog is None:
            err_msg = 'Run DSSP or Stride to assign secondary structures first!'
            print 'ERROR: %s' % (err_msg,)
            tkMessageBox.showinfo(title='ERROR', message=err_msg)
        else:
            print 'Update secondary structures for %s' % (self.pymol_sel.get()),
            print 'using secondary structure assignment by %s' % (self.ss_asgn_prog,)
            print 'SSE mapping: (H,G,I) ==> H; (E,B,b) ==> S; (T,S,-,C) ==> L'

            if self.ss_asgn_prog == 'DSSP':     SSE_list = self.DSSP_SSE_list
            elif self.ss_asgn_prog == 'Stride': SSE_list = self.STRIDE_SSE_list

            for sse in SSE_list:
                for sel_obj in self.sel_obj_list:
                    if self.SSE_sel_dict[sel_obj][sse] is not None:
                        cmd.alter(self.SSE_sel_dict[sel_obj][sse], 'ss=\'%s\''% (self.SSE_map[sse],))
                        print 'alter %s, ss=%s' % (self.SSE_sel_dict[sel_obj][sse], self.SSE_map[sse])
                    else:
                        print 'No residue with SSE %s to update ss.' % (sse,)

            # show cartoon for the input selection, and rebuild
            cmd.show('cartoon',self.pymol_sel.get())
            cmd.rebuild(self.pymol_sel.get()) 
            return

def mkc2(sel, a=Vec(0, 0, 1), c=Vec(0, 0, 0)):
    cmd.delete("c1")
    cmd.delete("c2")
    cmd.create("c1", sel)
    cmd.create("c2", sel)
    rot("c2", a, 180, c)
    cmd.alter("c1", 'chain = "A"')
    cmd.alter("c2", 'chain = "B"')

def hilightPolar(sel='all'):
	N = "elem N"# and (not name N)"
	O = "elem O"# and (not name O)" # want to inlude look BB atoms...
	S = "elem S"
	sel = sel + " and " + N+" or "+O+" or "+S
	cmd.alter(sel,"vdw=0.3")
	cmd.rebuild()
	cmd.show('spheres',sel)

 def colorByProp(self, prop, palette="rainbow"):
     stored.propUniqVals = set()
     cmd.iterate(self.objName, "stored.propUniqVals.add(%s)" % prop)
     v = sorted(stored.propUniqVals)
     b = n.arange(1, len(v) + 1, dtype=float)  # / len(v)
     stored.b = dict(zip(v, b))
     cmd.alter(self.objName, "b=stored.b[%s]" % prop)
     cmd.spectrum("b", palette, self.objName)

def renumber(selection='all', start=1, startsele=None, quiet=1):
    '''
DESCRIPTION

    Set residue numbering (resi) based on connectivity.

ARGUMENTS

    selection = string: atom selection to renumber {default: all}

    start = integer: counting start {default: 1}

    startsele = string: residue to start counting from {default: first in
    selection}
    '''
    start, quiet = int(start), int(quiet)
    model = cmd.get_model(selection)
    cmd.iterate(selection, 'atom_it.next().model = model',
                space={'atom_it': iter(model.atom)})
    if startsele is not None:
        startidx = cmd.index('first (' + startsele + ')')[0]
        for atom in model.atom:
            if (atom.model, atom.index) == startidx:
                startatom = atom
                break
        else:
            print(' Error: startsele not in selection')
            raise CmdException
    else:
        startatom = model.atom[0]
    for atom in model.atom:
        atom.adjacent = []
        atom.visited = False
    for bond in model.bond:
        atoms = [model.atom[i] for i in bond.index]
        atoms[0].adjacent.append(atoms[1])
        atoms[1].adjacent.append(atoms[0])
    minmax = [start, start]

    def traverse(atom, resi):
        atom.resi = resi
        atom.visited = True
        for other in atom.adjacent:
            if other.visited:
                continue
            if (atom.name, other.name) in [('C', 'N'), ("O3'", 'P')]:
                minmax[1] = resi + 1
                traverse(other, resi + 1)
            elif (atom.name, other.name) in [('N', 'C'), ('P', "O3'")]:
                minmax[0] = resi - 1
                traverse(other, resi - 1)
            elif (atom.name, other.name) not in [('SG', 'SG')]:
                traverse(other, resi)
    traverse(startatom, start)
    cmd.alter(selection, 'resi = atom_it.next().resi',
              space={'atom_it': iter(model.atom)})
    if not quiet:
        print(' Renumber: range (%d to %d)' % tuple(minmax))