| 11 | |
| 12 | |
| 13 | def parse(): |
| 14 | parser = ArgumentParser(description='split train / valid / test') |
| 15 | parser.add_argument('--data', type=str, required=True, help='Path to the data file') |
| 16 | parser.add_argument('--out_dir', type=str, default=None, help='Directory to save results. Default the same as input data.') |
| 17 | parser.add_argument('--valid_ratio', type=float, default=0.1, |
| 18 | help='Ratio of validation set') |
| 19 | parser.add_argument('--test_ratio', type=float, default=0.1, |
| 20 | help='Ratio of test set') |
| 21 | parser.add_argument('--cdr', type=str, choices=[f'cdrh{i}' for i in range(1, 4)] + [f'cdrl{i}' for i in range(1, 4)], |
| 22 | default='cdrh3', help='Cluster according to which cdr') |
| 23 | parser.add_argument('--filter', type=str, default='1*1', help='Filter out complex with heavy / light / antigen.' + \ |
| 24 | 'The code refers to heavy / light / antigen sequentially. 1 for has, 0 for not has, * for either.' + \ |
| 25 | 'e.g default 1*1 means has heavy chain and antigen, either has light chain or not.') |
| 26 | parser.add_argument('--k_fold', type=int, default=-1, help='K fold dataset. -1 for not do k-fold.' + \ |
| 27 | 'Note that if this is enabled, the test/valid ratio will be automatically calculated.') |
| 28 | parser.add_argument('--seed', type=int, default=2022, help='seed') |
| 29 | parser.add_argument('--rabd', type=str, default=None, help='Path to rabd json file. If this is enabled, '+ \ |
| 30 | 'RAbD complexes will be used as test set and complexes from data will be used as train/valid.' + \ |
| 31 | 'Note that complexes sharing clusters with RAbD will be dropped.' + \ |
| 32 | 'K fold will also be turned off.') |
| 33 | return parser.parse_args() |
| 34 | |
| 35 | |
| 36 | def load_file(fpath): |