Adding new cnn

Alexis Morrissey · Alexis Morrissey · commit 34b76b8c904c · 2024-05-09T11:03:02.000-04:00
diff --git a/Allo/allo b/Allo/allo
@@ -14,6 +14,7 @@ parser.add_argument('-seq', type=str, nargs=1, help='Single-end or paired-end se
                    choices=['pe','se'], required=True, dest='seq')
 parser.add_argument('-o', type=str, nargs=1, help='Output file name', dest='outfile', default=None)
 parser.add_argument('--mixed', help='Use CNN trained on a dataset with mixed peaks, narrow by default', action='store_true', default=None)
+parser.add_argument('--rna', help='Use CNN trained on a dataset with mixed peaks, narrow by default', action='store_true', default=None)
 parser.add_argument('-p', type=int, nargs=1, help='Number of processes, 1 by default', dest='processes', default=None)
 parser.add_argument('-max', type=int, nargs=1, help='Maximum value for number of locations a read can map', dest='maxlocations', default=None)
 parser.add_argument('--keep-unmap', help='Keep unmapped reads and reads that include N in their sequence', action='store_true', default=None)
@@ -84,6 +85,10 @@ if __name__ == '__main__':
         d = os.path.dirname(sys.modules["Allo"].__file__)
         m = os.path.join(d, "mixed")
         winSize = 500
+    elif args.rna is not None:
+        d = os.path.dirname(sys.modules["Allo"].__file__)
+        m = os.path.join(d, "rna")
+        winSize = 1000
     else:
         d = os.path.dirname(sys.modules["Allo"].__file__)
         m = os.path.join(d, "narrow")
diff --git a/Allo/allocation.py b/Allo/allocation.py
@@ -292,7 +292,7 @@ def parseUniq(tempFile, winSize, cnn_scores, AS, rc, keep):
     return [cu, cf]
 
 
-def parseMulti(tempFile, winSize, genLand, modelName, cnn_scores, rc, keep, rmz, maxa):
+def parseMulti(tempFile, winSize, genLand, modelName, cnn_scores, rc, keep, rmz, maxa):   
     numLoc = [0,0] #Keep info on average number of places read maps to
     #Getting trained CNN
     try:
@@ -306,6 +306,7 @@ def parseMulti(tempFile, winSize, genLand, modelName, cnn_scores, rc, keep, rmz,
         else:
             modelName = 0
     except:
+        print("Model loading error", flush=True)
         print("Could not load Tensorflow model :( Allo was written with Tensorflow version 2.11")
         sys.exit(0)
 
@@ -705,6 +706,7 @@ def parseMultiPE(tempFile, winSize, genLand, modelName, cnn_scores, rc, keep, rm
         else:
             modelName = 0
     except:
+        print("PE model load error")
         print("Could not load Tensorflow model :( Allo was written with Tensorflow version 2.11")
         sys.exit(0)
 
diff --git a/Allo/predictPeak.py b/Allo/predictPeak.py
@@ -15,7 +15,6 @@
 
 
 def predictNN(counts, winSize, model):
-    winSize = math.floor(winSize/5) #Genome is already binned 5bps
     pic = np.zeros((100, 100), float)
     if sum(counts) == 0:
         pred = model(pic.reshape(-1,100,100), training=False)
@@ -40,6 +39,7 @@ def predictNN(counts, winSize, model):
         return pred.numpy()[0][0]
 
     except:
+        print("Prediction error.", flush=True)
         print("Could not predict with Tensorflow model :( Allo was written with Tensorflow version 2.11")
         sys.exit(0)
 
