chemprep.cli

 1import argparse
 2import pandas as pd
 3from .fingerprints import generate_fingerprint, get_available_fingerprints
 4from .similarity import calculate_similarity, get_available_similarity_metrics
 5from .descriptors import get_all_descriptors
 6
 7def main():
 8    parser = argparse.ArgumentParser(description="A command-line tool for molecular fingerprinting, similarity calculations, and descriptor generation.")
 9    subparsers = parser.add_subparsers(dest="command")
10
11    # Fingerprint command
12    parser_fp = subparsers.add_parser("fingerprint", help="Generate molecular fingerprints.")
13    parser_fp.add_argument("smiles", nargs="?", help="SMILES string of the molecule.")
14    parser_fp.add_argument("-i", "--input", help="Path to a file with SMILES (one per line, or a CSV with a 'SMILES' column).")
15    parser_fp.add_argument("-t", "--type", default="Morgan", choices=get_available_fingerprints(), help="The type of fingerprint to generate.")
16    parser_fp.add_argument("-o", "--output", help="Path to the output file (CSV).")
17
18    # Similarity command
19    parser_sim = subparsers.add_parser("similarity", help="Calculate molecular similarity.")
20    parser_sim.add_argument("smiles1", help="SMILES string of the first molecule.")
21    parser_sim.add_argument("smiles2", help="SMILES string of the second molecule.")
22    parser_sim.add_argument("-t", "--type", default="Morgan", choices=get_available_fingerprints(), help="The type of fingerprint to use.")
23    parser_sim.add_argument("-m", "--metric", default="Tanimoto", choices=get_available_similarity_metrics(), help="The similarity metric to use.")
24
25    # Descriptor command
26    parser_desc = subparsers.add_parser("descriptors", help="Generate molecular descriptors.")
27    parser_desc.add_argument("smiles", nargs="?", help="SMILES string of the molecule.")
28    parser_desc.add_argument("-i", "--input", help="Path to a file with SMILES (one per line, or a CSV with a 'SMILES' column).")
29    parser_desc.add_argument("-o", "--output", help="Path to the output file (CSV).")
30
31    args = parser.parse_args()
32
33    if args.command == "fingerprint":
34        if args.smiles:
35            smiles_list = [args.smiles]
36        elif args.input:
37            if args.input.endswith(".csv"):
38                df = pd.read_csv(args.input)
39                smiles_list = df["SMILES"].tolist()
40            else:
41                with open(args.input) as f:
42                    smiles_list = [line.strip() for line in f]
43        else:
44            parser_fp.error("Either a SMILES string or an input file must be provided.")
45
46        fingerprints = [generate_fingerprint(s, args.type) for s in smiles_list]
47        df_out = pd.DataFrame({"SMILES": smiles_list, "Fingerprint": fingerprints})
48
49        if args.output:
50            df_out.to_csv(args.output, index=False)
51        else:
52            print(df_out)
53
54    elif args.command == "similarity":
55        similarity = calculate_similarity(args.smiles1, args.smiles2, args.type, args.metric)
56        print(f"The {args.metric} similarity between {args.smiles1} and {args.smiles2} is: {similarity}")
57
58    elif args.command == "descriptors":
59        if args.smiles:
60            smiles_list = [args.smiles]
61        elif args.input:
62            if args.input.endswith(".csv"):
63                df = pd.read_csv(args.input)
64                smiles_list = df["SMILES"].tolist()
65            else:
66                with open(args.input) as f:
67                    smiles_list = [line.strip() for line in f]
68        else:
69            parser_desc.error("Either a SMILES string or an input file must be provided.")
70
71        all_descriptors = [get_all_descriptors(s) for s in smiles_list]
72        df_out = pd.DataFrame(all_descriptors)
73        df_out.insert(0, "SMILES", smiles_list)
74
75
76        if args.output:
77            df_out.to_csv(args.output, index=False)
78        else:
79            print(df_out)
80
81if __name__ == "__main__":
82    main()