#!/usr/bin/env python

import tm3
import string
import sys

def deleteLeaves(tmData, attributeName, threshold):
  column = tmData.attributeTitles.index(attributeName)
  for node in tmData.bottomUpTraversal():
    if node in tmData.tree:  # may have already been merged
      if tmData.isLeaf(node.getId()):
        if node.attributes[column] < threshold:
          tmData.deleteLeaf(node.getId())

#main is only run for testing import of tm3 from commandline
if -1 != string.find(sys.argv[0], "tm3filterVolumeJustThreshold.py"):
  try:
    threshold = float(sys.argv[1])
    if 2 == len(sys.argv):
      print "no input files specified"
    else:
      for filename in sys.argv[2:]:
        tmData = tm3.tmTreeFromFile(filename)
        deleteLeaves(tmData, "Volume", threshold)
        tmData.write(filename+".v"+str(threshold)+".jt.tm3")
  except (TypeError, IndexError):
    print "tm3filterVolumeJustThreshold.py threshold file [more files]"
#!/usr/bin/env python

# compares the roots of the trees passed in

import tm3
import string
import sys
import dot

if -1 != string.find(sys.argv[0], "tm3rootcomp.py"):
    tmDataList = []
    for filename in sys.argv[1:]:
        tmDataList.append(tm3.tmTreeFromFile(filename))
    dotData = dot.dotRoot(tmDataList)
    # print dotData.treeToBestNode
    # dotData.write("temp.dot")
    # dotData.writeGdl("temp.gdl", force=False)
    # print "dot -Tpng temp.dot > temp.png"
    print "output file also written to temp.tab.txt and temp.animation.py"
    matchList = tm3.findSimilarNodes(
        tmDataList,
        [
            "Surface Area",
            "Volume",
            "height",
            # "mean absolute Curvature",
            "mean Curvature",
            "mouths",
            "longest dimension",
            "middle dimension",
            "short dimension",
import tm3
import string
import sys

def animateScript(tmData, tstData, sortCol, outName):
  '''sort by threshold, make a script to open and display each pocket in turn'''
  nodes = tmData.idNode.values()
  nodes.sort(
      lambda xList, yList: cmp(
          xList.attributes[sortCol], yList.attributes[sortCol]))
  nodes.reverse()  # sorted high to low now
  print "from pymol import cmd"
  print 'cmd.do("delete all")'
  print 'cmd.do("tstOpen ' + tstData + '")'
  print 'cmd.do("tstDraw surf=pdb")'
  print 'cmd.do("color grey")'
  for index, node in enumerate(nodes):
    print 'cmd.do("turn x,1")'
    print 'cmd.do("tstDraw group=', str(node.getId()), '")'
    if index % 5 == 0:
      print 'cmd.do("ray")'
      print 'cmd.do("png ', str(outName + string.zfill(index, 6)+'.png') + '")'
  print 'cmd.do("ray")'
  print 'cmd.do("png ', str(outName + string.zfill(index+1, 6)+'.png') + '")'


if -1 != string.find(sys.argv[0], "tm3animatepockets.py"):
  tmData = tm3.tmTreeFromFile(sys.argv[1])
  tstData = sys.argv[1].replace(".tree.tm3", "")
  animateScript(tmData, tstData, 3, tstData+".movie.turn.")