parsed arguments correctly

inference_VS_2.py

@@ -49,6 +49,8 @@ def parse_arguments(arglist = None):
     p.add_argument('-i', '--inference_path', type = str, default = None, help='path to some pdb files for which you want to run inference')
     p.add_argument('--checkpoint', type=str, help='path to .pt file in a checkpoint directory')
     p.add_argument('-o','--output_directory', type=str, default=None, help='path where to put the predicted results')
+    p.add_argument('--protein', type=str, default=None, help='file name of the receptor for multiligand docking')
+    p.add_argument('--small_molecule_library', type=str, default=None, help='file name of the ligands for multiligand docking')
     p.add_argument('--run_corrections', type=bool, default=True,
                    help='whether or not to run the fast point cloud ligand fitting')
     #p.add_argument('--run_dirs', type=list, default=[], help='path directory with saved runs')

main.py

@@ -4,50 +4,29 @@
 import yaml
 # import fastapi # a package we use to receive and return results via an API
 
-def predict(
-    protein: str, # a PDB protein structure file
-    small_molecule_library: str, # "an SDF file containing >=2 small molecule ligands")
-    ):
-    # custom changes
-    args = inference_VS_2.parse_arguments()
-    args = args[0]
-    args.inference_path = '/src/tmp' # copy the input files to this directory for renaming and processing
-    args.output_directory = '/outputs'
-
-    # formatting input
-    protein = str(protein)
-    small_molecule_library = str(small_molecule_library) # '/inputs/test.sdf'
-    print("received input: " + protein + " " + small_molecule_library)     
-
+def predict(args):
     # isolate the argument path basenames
-    protein_base = os.path.basename(protein)
-    small_molecule_library_base = os.path.basename(small_molecule_library)
-
+    protein_base = os.path.basename(args.protein)
+    small_molecule_library_base = os.path.basename(args.small_molecule_library)
     # defining file name
     protein_destination = args.inference_path + '/dummy/' + protein_base
     small_molecule_library_destination = args.inference_path + '/dummy/' + small_molecule_library_base
-
     # moving files from the paths defined in the arguments to the input directory for processing
     print("moving files to " + args.inference_path + " for processing")
     os.system('mkdir -p ' + args.inference_path + '/dummy')
-    os.system('cp ' + protein + ' ' + protein_destination)
-    os.system('cp ' + small_molecule_library + ' ' + small_molecule_library_destination)
-
+    os.system('cp ' + args.protein + ' ' + protein_destination)
+    os.system('cp ' + args.small_molecule_library + ' ' + small_molecule_library_destination)
     # renaming file names to match the expected input pattern
     print("renaming file names to match the expected input pattern")
     os.system('mv ' + protein_destination + ' ' + args.inference_path + '/dummy/protein_' + protein_base)
-    os.system('mv ' + small_molecule_library + ' ' + args.inference_path + '/dummy/ligands_' + small_molecule_library_base)
-
-
+    os.system('mv ' + small_molecule_library_destination + ' ' + args.inference_path + '/dummy/ligands_' + small_molecule_library_base)
     # the dataurl go package does not like .sdf files, the input should be given in .txt - something to add to petri
     #small_molecule_library_sdf = os.path.splitext(small_molecule_library_destination)[0]+'.sdf'
     #print("converting library to sdf: " + small_molecule_library_sdf)
     #os.system('mv ' + small_molecule_library_destination + ' ' + small_molecule_library_sdf)
-
     # adding missings args, only works for one run_dir
     args.multi_ligand = True
     # args.model_parameters['noise_initial'] = 0
-
     # running the inference - this is the main function - lifted from __main__ in inference_VS_2.py
     if args.config:
         config_dict = yaml.load(args.config, Loader=yaml.FullLoader)
@@ -64,8 +43,6 @@ def predict(
         args.config = args.config.name
     else:
         config_dict = {}
-
-
     for run_dir in args.run_dirs:
         args.checkpoint = f'runs/{run_dir}/best_checkpoint.pt'
         config_dict['checkpoint'] = f'runs/{run_dir}/best_checkpoint.pt'
@@ -89,25 +66,36 @@ def predict(
             inference_VS_2.multi_lig_inference(args)
         else:
             inference_VS_2.inference_from_files(args)
-
     # moving the output file to the output directory
     output_name = os.listdir(args.output_directory + '/dummy')[0]
     output_path_sdf = args.output_directory + '/dummy/' + output_name
     print(output_path_sdf)
-
     # the dataurl go package does not like .sdf files, the input should be given in .txt - something to add to petri
     # output_path = os.path.splitext(output_path_sdf)[0]+'.txt'
     # print("converting library to txt: " + output_path)
     output_path = args.output_directory + '/' + output_name
     os.system('mv ' + output_path_sdf + ' ' + output_path)
-
     # output
     # output = Path(output_path)
     print(output_path)
     return(output_path)
 
+def parse_arguments_main(): 
+    print("parsing arguments")
+    # parsing arguments
+    args = inference_VS_2.parse_arguments()
+    args = args[0]
+    # setting a custom default
+    args.inference_path = '/src/tmp' # copy the input files to this directory for renaming and processing
+    args.output_directory = '/outputs'
+    print("received input: " + args.protein + " " + args.small_molecule_library + " for protein and small molecule library, respectively")
+    return(args)
+
 if __name__ == "__main__":
-    #sys.argv[1]
-    print(sys.argv[1])
-    print(sys.argv[2])
-    print(predict(protein = sys.argv[1], small_molecule_library= sys.argv[2]))
+    args = parse_arguments_main()
+    print("printing arguments:")
+    print("####################")
+    print(args)
+    print("####################")
+    print("running prediction:")
+    predict(args)