'''
assemstats.py

Uses screed to calculate various assembly statistics.

You can obtain screed at github by running
   git clone git://github.com/acr/screed.git

Then, install by running
   python setup.py install
in the newly created screed directory.

Once completed, you should be able to run this script as is.

Author: Jason Pell (pelljaso@cse.msu.edu)
'''

import screed.fasta
import sys
import glob

def trimLens(lens, minLen):
   '''
   Eliminates any reads below a certain threshold.  Function assumes that input
   list lens is sorted smallest to largest.
   '''

   index = 0

   for i in range(len(lens)):
      if lens[i] < minLen:
         index += 1
      else:
         break

   return lens[index:len(lens)]

def getLens(filename):
   '''
   Parses FASTA file using screed to create a sorted list of contig lengths.
   '''
   lens = []
   fd = open(filename, 'r')

   fa_instance = screed.fasta.fasta_iter(fd)
   for record in fa_instance:
      lens.append(len(record['sequence']))

   fd.close()

   return sorted(lens)

def calcNXX(lens, percent):
   '''
   Calculates any NXX (e.g. N50, N90) statistic.
   '''

   lenSum = sum(lens)
   threshold = (float(percent) / 100) * lenSum
   
   runningSum = 0
   nxx = 0
   nxxLen = 0

   for i in range(len(lens)-1, -1, -1):
      myLen = lens[i]
      nxx += 1
      runningSum += myLen

      if runningSum >= threshold:
         nxxLen = myLen
         break

   return nxx, nxxLen

def main():
   '''
   Outputs assembly statistics for provided FASTA files.
   '''
   if len(sys.argv) < 3:
      print "Usage: python assemstats.py <min contig length> [ FASTA files ]"
      return

   try:
      minLen = int(sys.argv[1])
   except ValueError:
      print "Minimum contig length must be an integer."
      return

   print "N\tsum\tmax\tfilename"

   for expr in sys.argv[2:len(sys.argv)]:
      for filename in glob.glob(expr):
         lens = getLens(filename)
         trimmedLens = trimLens(lens, minLen)

         if len(trimmedLens) == 0:
            print filename + " - no sequences longer than threshold"
            continue

         statN = len(lens)
         statTrimmedN = len(trimmedLens)
         statSum = sum(trimmedLens)
         statMin = min(trimmedLens)
         statMax = max(trimmedLens)
         statMed = trimmedLens[(len(trimmedLens)-1)/2]
         statMean = int(statSum / float(statTrimmedN))
         statN50, statN50Len = calcNXX(trimmedLens, 50)
         statN90, statN90Len = calcNXX(trimmedLens, 90)

         print "%d\t%d\t%d\t%s" % (statTrimmedN, statSum, statMax, filename)

main()

@@ -1527,16 +1527,47 @@
 %0A   '''%0A
+   totalN = 0%0A   totalSum = 0%0A%0A
    if le
@@ -1848,20 +1848,24 @@
 %0A   for 
-expr
+filename
  in sys.
@@ -1875,66 +1875,11 @@
 v%5B2:
-len(sys.argv)%5D:%0A      for filename in glob.glob(expr):%0A   
+%5D:%0A
     
@@ -1911,19 +1911,16 @@
 )%0A      
-   
 trimmedL
@@ -1955,26 +1955,19 @@
 %0A%0A      
-   if len(
+if 
 trimmedL
@@ -1973,134 +1973,10 @@
 Lens
-) == 0
 :%0A
-            print filename + %22 - no sequences longer than threshold%22%0A            continue%0A%0A         statN = len(lens)%0A
     
@@ -2066,16 +2066,16 @@
 tatM
-in
+ax
  = m
-in
+ax
 (tri
@@ -2094,336 +2094,313 @@
     
-   statMax = max(trimmedLens)%0A         statMed = trimmedLens%5B(len(trimmedLens)-1)/2%5D%0A         statMean = int(statSum / float(statTrimmedN))%0A         statN50, statN50Len = calcNXX(trimmedLens, 50)%0A         statN90, statN90Len = calcNXX(trimmedLens, 90)%0A%0A         print %22%25d%5Ct%25d%5Ct%25d%5Ct%25s%22 %25 (statTrimmedN, statSum, statMax, filename
+else:%0A         statTrimmedN = 0%0A         statSum = 0%0A         statMax = 0%0A%0A      totalN += statTrimmedN%0A      totalSum += statSum%0A%0A      print %22%25d%5Ct%25d%5Ct%25d%5Ct%25s%22 %25 (statTrimmedN, statSum, statMax, filename)%0A%0A   if len(sys.argv) %3E 3:%0A      print '--'%0A      print 'TOTAL: %25g in %25d contigs' %25 (totalSum, totalN
 )%0A%0Am

cleaned up no-contigs output; provided sum