Add CASP-CAPRI support to feature generation pipeline

Author: amorehead

project/datasets/builder/generate_hhsuite_features.py

@@ -23,7 +23,7 @@
 @click.option('--num_cpu_jobs', '-j', default=cpu_count() // 2, type=int)
 @click.option('--num_cpus_per_job', '-c', default=2, type=int)
 @click.option('--num_iter', '-i', default=2, type=int)
-@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling']))
+@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling', 'casp_capri']))
 @click.option('--write_file/--read_file', '-w', default=True)
 def main(pkl_dataset: str, pruned_dataset: str, hhsuite_db: str, output_dir: str, rank: int,
          size: int, num_cpu_jobs: int, num_cpus_per_job: int, num_iter: int, source_type: str, write_file: bool):

project/datasets/builder/generate_psaia_features.py

@@ -18,7 +18,7 @@
 @click.argument('pkl_dataset', type=click.Path(exists=True))
 @click.argument('pruned_dataset', type=click.Path(exists=True))
 @click.argument('output_dir', type=click.Path())
-@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling']))
+@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling', 'casp_capri']))
 @click.option('--rank', '-r', default=0)
 @click.option('--size', '-s', default=1)
 def main(psaia_dir: str, psaia_config: str, pdb_dataset: str, pkl_dataset: str,

project/datasets/builder/make_dataset.py

@@ -22,7 +22,7 @@
 @click.option('--neighbor_def', default='non_heavy_res',
               type=click.Choice(['non_heavy_res', 'non_heavy_atom', 'ca_res', 'ca_atom']))
 @click.option('--cutoff', default=6)
-@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling']))
+@click.option('--source_type', default='rcsb', type=click.Choice(['rcsb', 'db5', 'evcoupling', 'casp_capri']))
 @click.option('--unbound/--bound', default=False)
 def main(input_dir: str, output_dir: str, num_cpus: int, rank: int, size: int,
          neighbor_def: str, cutoff: int, source_type: str, unbound: bool):
@@ -44,7 +44,7 @@ def main(input_dir: str, output_dir: str, num_cpus: int, rank: int, size: int,
         get_neighbors = nb.build_get_neighbors(neighbor_def, cutoff)
         get_pairs = pair.build_get_pairs(neighbor_def, source_type, unbound, get_neighbors, False)
         complexes = comp.read_complexes(complexes_dill)
-        pair.all_complex_to_pairs(complexes, get_pairs, pairs_dir, num_cpus)
+        pair.all_complex_to_pairs(complexes, source_type, get_pairs, pairs_dir, num_cpus)
 
 
 if __name__ == '__main__':

requirements.txt

@@ -4,6 +4,6 @@ dill==0.3.3
 tqdm==4.49.0
 Sphinx==4.0.1
 easy-parallel-py3==0.1.6.4
-atom3-py3==0.1.9.1
+atom3-py3==0.1.9.3
 click==7.0.0
 # mpi4py==3.0.3  # On Andes, do 'source venv/bin/activate', 'module load gcc/10.3.0', and 'pip install mpi4py --no-cache-dir --no-binary :all:'