--- a/sucos_max.py	Sat Mar 28 05:16:44 2020 -0400
+++ b/sucos_max.py	Wed Apr 15 09:26:30 2020 -0400
@@ -39,7 +39,7 @@
 from rdkit import Chem
 
 
-def process(inputfilename, clusterfilenames, outputfilename):
+def process(inputfilename, clusterfilenames, outputfilename, filter_value, filter_field):
     all_clusters = {}
     for filename in clusterfilenames:
         cluster = []
@@ -80,6 +80,7 @@
             continue
         scores_max = [0, 0, 0]
         scores_cum = [0, 0, 0]
+        cluster_name = None
         for clusterfilename in all_clusters:
             cluster = all_clusters[clusterfilename]
             index = 0
@@ -104,21 +105,28 @@
                 scores_cum[2] += vol_score
 
 
-        cluster_file_name_only = cluster_name.split(os.sep)[-1]
-
         # utils.log("Max SuCOS:", scores[0], "FM:", scores[1], "P:", scores[2],"File:", cluster_file_name_only, "Index:", cluster_index)
         mol.SetDoubleProp("Max_SuCOS_Score", scores_max[0] if scores_max[0] > 0 else 0)
         mol.SetDoubleProp("Max_SuCOS_FeatureMap_Score", scores_max[1] if scores_max[1] > 0 else 0)
         mol.SetDoubleProp("Max_SuCOS_Protrude_Score", scores_max[2] if scores_max[2] > 0 else 0)
-        mol.SetProp("Max_SuCOS_Cluster", cluster_file_name_only)
-        mol.SetIntProp("Max_SuCOS_Index", cluster_index)
+
+        if cluster_name:
+            cluster_file_name_only = cluster_name.split(os.sep)[-1]
+            mol.SetProp("Max_SuCOS_Cluster", cluster_file_name_only)
+            mol.SetIntProp("Max_SuCOS_Index", cluster_index)
 
         # utils.log("Cum SuCOS:", scores[0], "FM:", scores[1], "P:", scores[2])
         mol.SetDoubleProp("Cum_SuCOS_Score", scores_cum[0] if scores_cum[0] > 0 else 0)
         mol.SetDoubleProp("Cum_SuCOS_FeatureMap_Score", scores_cum[1] if scores_cum[1] > 0 else 0)
         mol.SetDoubleProp("Cum_SuCOS_Protrude_Score", scores_cum[2] if scores_cum[2] > 0 else 0)
 
-        writer.write(mol)
+        if filter_value and filter_field:
+            if mol.HasProp(filter_field):
+                val = mol.GetDoubleProp(filter_field)
+                if val > filter_value:
+                    writer.write(mol)
+        else:
+            writer.write(mol)
 
     input_file.close()
     writer.flush()
@@ -135,11 +143,13 @@
     parser.add_argument('-i', '--input', help='Input file to score in SDF format. Can be gzipped (*.gz).')
     parser.add_argument('-o', '--output', help='Output file in SDF format. Can be gzipped (*.gz).')
     parser.add_argument('clusters', nargs='*', help="One or more SDF files with the clustered hits")
+    parser.add_argument('--filter-value', type=float, help='Filter out values with scores less than this.')
+    parser.add_argument('--filter-field', help='Field to use to filter values.')
 
     args = parser.parse_args()
     utils.log("Max SuCOS Args: ", args)
 
-    process(args.input, args.clusters, args.output)
+    process(args.input, args.clusters, args.output, args.filter_value, args.filter_field)
 
 
 if __name__ == "__main__":