add option to specify chromosomes to run imputation for

atgu · Feb 5, 2025 · a2e3364 · a2e3364
1 parent 0ff3e08
commit a2e3364
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/gwaspy/imputation/glimpse2_impute.py b/gwaspy/imputation/glimpse2_impute.py
@@ -23,6 +23,7 @@ def glimpse_phase_impute(
         batch: hb.Batch = None,
         bam_files: str = None,
         reference_path: str = None,
+        chromosomes: str = "all",
         output_filename: str = None,
         output_path: str = None):
 
@@ -286,7 +287,10 @@ def ligate_chunks(
     bam_list = [(sample_id, bam_path) for sample_id, bam_path in zip(bams['sample'], bams['path'])]
     bam_sizes = [size(bam_list[s][1]) for s in range(len(bam_list))]
 
-    for i in range(1, 23):
+    chroms = chromosomes.replace(" ", "") # remove spaces if there are any
+    chroms = [i for i in range(1, 23)] if chroms == "all" else chroms.split(",")
+
+    for i in chroms:
         ref_chrom_path = reference_path.replace('CNUMBER', str(i))
         ref_idx = f'{ref_chrom_path}.tbi' if hfs.exists(f'{ref_chrom_path}.tbi') else f'{ref_chrom_path}.csi'
         ref_vcf = batch.read_input_group(**{'vcf': ref_chrom_path,

diff --git a/gwaspy/imputation/impute.py b/gwaspy/imputation/impute.py
@@ -8,8 +8,9 @@
 
 
 def run_impute(backend: Union[hb.ServiceBackend, hb.LocalBackend] = None,
-               input_vcf: str = None,
+               input_file: str = None,
                vcf_ref: str = None,
+               chromosomes: str = "all",
                software: str = 'impute5',
                output_filename: str = None,
                n_samples: int = None,
@@ -34,8 +35,9 @@ def run_impute(backend: Union[hb.ServiceBackend, hb.LocalBackend] = None,
         print(f'\nIMPUTING GENOTYPES USING IMPUTE5\n')
         impute5_imputation(
             batch=b,
-            input_path=input_vcf,
+            input_path=input_file,
             reference_path=ref_path,
+            chromosomes=chromosomes,
             output_filename=output_filename,
             n_samples=n_samples,
             n_panel_samples=n_panel_samples,
@@ -44,8 +46,9 @@ def run_impute(backend: Union[hb.ServiceBackend, hb.LocalBackend] = None,
     elif software == 'glimpse2':
         glimpse_phase_impute(
             batch=b,
-            bam_files=input_vcf,
+            bam_files=input_file,
             reference_path=ref_path,
+            chromosomes = chromosomes,
             output_filename=output_filename,
             output_path=out_dir
         )
@@ -54,8 +57,9 @@ def run_impute(backend: Union[hb.ServiceBackend, hb.LocalBackend] = None,
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--input-vcf', type=str, required=True)
+    parser.add_argument('--input-file', type=str, required=True)
     parser.add_argument('--vcf-ref', type=str, default='hgdp1kgp')
+    parser.add_argument('--chromosomes', type=str, default='all')
     parser.add_argument('--local', action='store_true')
     parser.add_argument('--billing-project', required=True)
     parser.add_argument('--n-samples', type=int, required=True)
@@ -73,8 +77,9 @@ def main():
                                     remote_tmpdir=f'{args.out_dir}/tmp/')
 
     run_impute(backend=backend,
-               input_vcf=args.input_vcf,
+               input_file=args.input_file,
                vcf_ref=args.vcf_ref,
+               chromosomes=args.chromosomes,
                software=args.software,
                output_filename=args.output_filename,
                n_samples=args.n_samples,

diff --git a/gwaspy/imputation/impute5_impute.py b/gwaspy/imputation/impute5_impute.py
@@ -24,6 +24,7 @@ def impute5_imputation(
         batch: hb.Batch = None,
         input_path: str = None,
         reference_path: str = None,
+        chromosomes: str = "all",
         output_filename: str = None,
         n_samples: int = None,
         n_panel_samples: int = 4091,
@@ -117,7 +118,10 @@ def concatenate_imputed_chunks(
 
         return j
 
-    for i in range(1, 23):
+    chroms = chromosomes.replace(" ", "")  # remove spaces if there are any
+    chroms = [i for i in range(1, 23)] if chroms == "all" else chroms.split(",")
+
+    for i in chroms:
         # read chrom input files
         if "CNUMBER" in input_path:     # input VCF is already split by chromosome
             vcf_path = input_path.replace('CNUMBER', str(i))