add function to sort for species

test_for_cells.py

@@ -4,11 +4,61 @@ from DataParserFactory import DatParser
 import numpy as np
 import os
 import matplotlib.pyplot as plt
+import pyrelacs.DataLoader as Dl
 
 data_save_path = "test_routines/test_files/"
 read = False
 
-if read:
+
+def main():
+    # test_kraken_files()
+    test_for_species()
+
+
+def test_for_species():
+
+    directory = "invivo_data/"
+    sorted_cells = {}
+    error_cells = []
+    for cell in os.listdir(directory):
+        if "thresh" in cell:
+            continue
+        cell_path = os.path.join(directory, cell)
+        # print(cell_path)
+        info_file = os.path.join(cell_path, "info.dat")
+
+        for metadata in Dl.load(info_file):
+            if "Species" in metadata.keys():
+                species = metadata["Species"]
+            else:
+                species = metadata["Subject"]["Species"]
+            if species not in sorted_cells.keys():
+                sorted_cells[species] = []
+            sorted_cells[species].append(cell_path)
+
+
+    print("Errors:", len(error_cells))
+    for species in sorted_cells.keys():
+        print("{}: {}".format(species, len(sorted_cells[species])))
+
+
+    print()
+    print("errors:")
+    for cell in error_cells:
+        print(cell)
+
+    # print()
+    # print("eigen:")
+    # for cell in sorted_cells["Eigenmannia virescens"]:
+    #     print(cell)
+    #
+    # print()
+    # print("albi:")
+    # for cell in sorted_cells["Apteronotus albifrons"]:
+    #     print(cell)
+
+def test_kraken_files():
+    if read:
         directory = "/mnt/invivo_data/"
         fi_curve_min_contrasts = 7
         fi_curve_min_trials = 7
@@ -66,7 +116,7 @@ if read:
         np.save(data_save_path + "accepted", accepted)
         print("Total good:", count)
 
-else:
+    else:
         files = np.load(data_save_path + "files.npy", allow_pickle=True)
         baseline = np.load(data_save_path + "baseline.npy", allow_pickle=True)
         ficurve = np.load(data_save_path + "ficurve.npy", allow_pickle=True)
@@ -128,3 +178,7 @@ else:
             count += 1
 
         print("Fullfilled by {:} of {:} test cells".format(count, all_cells))
+
+
+if __name__ == '__main__':
+    main()