--- a/spring_roc.py	Tue Nov 24 14:02:08 2020 +0000
+++ b/spring_roc.py	Tue Nov 24 17:12:07 2020 +0000
@@ -210,6 +210,9 @@
     # get subcellular locations from UniProt export
     locations = dict()
     if isfile(args.locations):
+        regions = list()
+        if args.regions:
+            regions = args.regions.split(",")
         with open(args.locations) as locFile:
             for line in locFile:
                 searchKey = "SUBCELLULAR LOCATION"
@@ -218,9 +221,11 @@
                     uniId = line.split()[0]
                     locStart = searchPos + len(searchKey) + 1
                     locId = line[locStart:].split()[0]
-                    if locId in ["Nucleus", "Membrane", "Cytoplasm"]:
-                        if uniId in filterA and uniId in filterB:
-                            locations[uniId] = locId
+                    if regions:
+                        if locId not in regions:
+                            continue
+                    if uniId in filterA or uniId in filterB:
+                        locations[uniId] = locId
         print("Found %d subcellular locations." % (len(list(locations.keys()))))
 
     # estimate background noise
@@ -269,6 +274,7 @@
     parser.add_argument('-i', '--input', help='Input prediction file.', required=True)
     parser.add_argument('-b', '--biogrid', help='BioGRID interaction database file', required=True)
     parser.add_argument('-l', '--locations', help='UniProt export table with subcellular locations', required=False)
+    parser.add_argument('-r', '--regions', help='Comma-separated regions', required=False)
     parser.add_argument('-e', '--experiment', help='Type (physical/genetic)', default="", required=False)
     parser.add_argument('-t', '--throughput', help='Throughput (low/high)', default="", required=False)
     parser.add_argument('-m', '--method', help='Method e.g. Two-hybrid', default="", required=False)