add orthogroup gene count script

davidemms · Nov 6, 2020 · f488aed · f488aed
1 parent 0ce67f7
commit f488aed
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/scripts_of/test_consensus_tree.py → tests/test_consensus_tree.py b/scripts_of/test_consensus_tree.py → tests/test_consensus_tree.py
diff --git a/tools/orthogroup_gene_count.py b/tools/orthogroup_gene_count.py
@@ -0,0 +1,20 @@
+import os
+import sys
+import csv
+
+if len(sys.argv) != 2 or sys.argv[1] == "-h" or sys.argv[1] == "-help" or sys.argv[1] == "--help":
+    print("Usage: orthogroup_gene_count.py Orthogroups.csv")
+    sys.exit()
+
+inFN = sys.argv[1]
+outFN = os.path.splitext(inFN)[0] + ".GeneCount.csv"
+
+with open(inFN, 'r') as infile, open(outFN, 'w') as outfile:
+    reader = csv.reader(infile, delimiter="\t")
+    writer = csv.writer(outfile, delimiter="\t")
+    header = next(reader)
+    n_col_skip = 3 if header[0] == "HOG" else 1
+    writer.writerow(header)
+    for line in reader:
+        writer.writerow(line[:n_col_skip] + [0 if "" == cell else len(cell.split(", ")) for cell in line[n_col_skip:]])
+print("Orthogroup gene count table has been written to %s" % outFN)