Single variant TabixLD fetching & range bugfix

Scripts/data_access/linkage.py

@@ -108,9 +108,10 @@ def get_range(self, variant: Variant, bp_range: int, ld_threshold_:Optional[floa
         return ld_data
 
 class TabixLD(LDAccess):
-    def __init__(self,path_template:str):
+    def __init__(self,path_template:str,single_variant_ld:bool):
         self.token_refresh_time = time.time()
         self.path_template = path_template
+        self.single_variant_ld = single_variant_ld
         ## NOTE: we assume that data is in chromosomes 1..22,X, and that the files are named so
         chroms = [str(a).replace("23","X") for a in range(1,24)]
         paths={str(a):path_template.replace("{CHROM}",f"{a}") for a in chroms}
@@ -140,8 +141,15 @@ def get_range(self, variant: Variant, bp_range: int, ld_threshold_:Optional[floa
                 self.refresh_token()
                 self.token_refresh_time = current_time
         variant_id=f"chr{variant.chrom.replace('chr','')}_{variant.pos}_{variant.ref}_{variant.alt}"
+
         start = max(0,variant.pos-bp_range)
         end = variant.pos+bp_range
+        if self.single_variant_ld:
+            fetch_start = variant.pos-1
+            fetch_end = variant.pos+1
+        else:
+            fetch_start = start
+            fetch_end = end
         sequence = variant.chrom.replace("23","X")
         if sequence not in self.sequences:
             raise Exception(f"Error in fetching LD for variant {variant}: File for chromosome {sequence} not in available files.")
@@ -151,20 +159,29 @@ def get_range(self, variant: Variant, bp_range: int, ld_threshold_:Optional[floa
             ld_threshold = ld_threshold_
         data = []
         tries = 0
+        # Pysam does not always error on tabix fileobject corruption.
+        # Errors logged by tbx_itr_next and bgzf_read_block do not seem to trigger exceptions.
+        # This check triggers once if the query was empty, to try and find the problematic cases.
+        restart_fileobject_fetch_failure = False
         while True:
             try:
                 data = []
-                iter = self.fileobjects[sequence].fetch(sequence, start,end)
+                iter = self.fileobjects[sequence].fetch(sequence, fetch_start,fetch_end)
+                iter_len = 0
                 for l in iter:
                     cols = l.split("\t")
                     if cols[self.hdi["variant1"]].replace("chrX","chr23") == variant_id:
                         data.append(cols)
+                    iter_len +=1
+                if iter_len == 0 and not restart_fileobject_fetch_failure:
+                    restart_fileobject_fetch_failure = True
+                    raise Exception("LD Iterator was empty, restart fileobject and redo")
                 break
             except:
-                print(f"Error loading data from region {sequence}:{start}-{end}",file=sys.stderr)
+                print(f"Error loading data from region {sequence}:{fetch_start}-{fetch_end}",file=sys.stderr)
                 #assume error is in accessing over gcp, so we retry
                 if tries > MAX_RETRIES:
-                    raise Exception(f"Accessing LD region {sequence}:{start}-{end} from file {self.paths[sequence]} failed after {tries} tries!")
+                    raise Exception(f"Accessing LD region {sequence}:{fetch_start}-{fetch_end} from file {self.paths[sequence]} failed after {tries} tries!")
 
                 print(f"Error when accessing data from LD tabix file, assuming error is with access over GCP. Waiting {2**tries} seconds, recycling fileobject.",file=sys.stderr)
                 time.sleep(2**tries)

Scripts/group_annotation.py

@@ -56,14 +56,14 @@ def generate_chrom_ranges(variants:set[Variant],maximum_range_length:Optional[in
                 #if at beginning, change start to be first pos minus 1 or 0 whichever is bigger
                 if region_start < 0:
                     region_start=max(v.pos-1,0)
+                    region_end = v.pos
                 #if we have started, but the region including this variant is not too large
                 elif v.pos-region_start <= maximum_range_length-1:
                     region_end = v.pos
                 # if we have started, and the region is too long, we wrap up range.
                 else:
                     chr_d[c].append(Region(c,region_start,region_end))
-                    #if we reset both to -1, then it starts with the next variant correctly.
-                    region_start = v.pos
+                    region_start = max(v.pos-1,0)
                     region_end = v.pos
             if region_end != -1 and region_start != -1:
                 chr_d[c].append(Region(c,region_start,region_end))
@@ -73,7 +73,7 @@ class TabixAnnotation(AnnotationSource):
     def __init__(self,opts: TabixOptions):
         self.cpra = [opts.c,opts.p,opts.r,opts.a]
         self.fname = opts.fname
-        self.variant_switch = 100_000
+        self.variant_switch = 5000
         with tb_resource_manager(self.fname,opts.c,opts.p,opts.r,opts.a) as tb_resource:
             self.sequences= tb_resource.sequences
             if any([a for a in (self.cpra) if a not in tb_resource.header ]):
@@ -166,7 +166,6 @@ def __init__(self, opts: PreviousReleaseOptions) -> None:
         self.beta_col = opts.beta_col
         self.other_cols = opts.other_cols
         self.cols = [self.pval_col,self.beta_col]+self.other_cols
-        self.variant_switch = 50_000
         #validate source file
         with tb_resource_manager(self.fname,opts.c,opts.p,opts.r,opts.a) as tb_resource:
 
@@ -196,7 +195,6 @@ class ExtraColAnnotation(TabixAnnotation):
     def __init__(self, opts: TabixOptions,extra_columns:List[str]) -> None:
         super().__init__(opts)
         self.extra_columns = extra_columns
-        self.variant_switch = 50_000
         with tb_resource_manager(self.fname,opts.c,opts.p,opts.r,opts.a) as tb_resource:
             if any([a for a in extra_columns if a not in tb_resource.header ]):
                 raise Exception(f"Summary statistic file did not contain all extra columns! Missing columns: {[a for a in self.extra_columns if a not in tb_resource.header ]}. Supplied columns:{self.extra_columns}")
@@ -315,7 +313,6 @@ def __init__(self,fname:str):
             "nfsee.AF":tryfloat,
             "nfsee.homozygote_count":tryint
         }
-        self.variant_switch = 1_000_000
         #make sure all columns are in 
         with tb_resource_manager(self.fname,opts.c,opts.p,opts.r,opts.a) as tb_resource:
             if any([a for a in [b for b in self.columntypes.keys()] if a not in tb_resource.header ]):
@@ -379,7 +376,6 @@ def __init__(self,fname:str):
             "rsids":"rsid",
             "FG_AF":"AF"
         }
-        self.variant_switch = 1_000_000
 
     def _create_annotation(self, cols: List[str], hdi: Dict[str, int])->Annotation:
         most_severe_gene = cols[hdi[self.colnames["most_severe_gene"]]]
@@ -432,7 +428,6 @@ def calculate_enrichment(nfe_AC:List[int],nfe_AN:List[int],fi_af:float):
 class GnomadGenomeAnnotation(TabixAnnotation):
     def __init__(self,fname: str):
         super().__init__(TabixOptions(fname,"#CHROM","POS","REF","ALT"))
-        self.variant_switch = 1_000_000
         self.columns=["AF_fin",
         "AF_nfe",
         "AF_nfe_est",
@@ -483,7 +478,6 @@ def get_output_columns() -> List[str]:
 class GnomadExomeAnnotation(TabixAnnotation):
     def __init__(self,fname: str):
         super().__init__(TabixOptions(fname,"#CHROM","POS","REF","ALT"))
-        self.variant_switch = 50_000
         self.columns=["AF_nfe_bgr",
         "AF_fin",
         "AF_nfe",
@@ -549,7 +543,6 @@ def get_output_columns() -> List[str]:
 class Gnomad4Annotation(TabixAnnotation):
     def __init__(self,fname: str):
         super().__init__(TabixOptions(fname,"#chr","pos","ref","alt"))
-        self.variant_switch = 500_000
 
 
     def _create_annotation(self, cols: List[str], hdi: Dict[str, int]) -> Annotation:

Scripts/main.py

@@ -52,7 +52,7 @@ def main(args):
         elif args.ld_api_choice == "online":
             ld_api = OnlineLD(url="http://api.finngen.fi/api/ld")
         elif args.ld_api_choice == "tabix":
-            ld_api = TabixLD(args.ld_panel_path)
+            ld_api = TabixLD(args.ld_panel_path,args.tabixld_assume_narrow_ld)
         else:
             raise ValueError("Wrong argument for --ld-api:{}".format(args.ld_api_choice))
     args.sig_treshold_2 = max(args.sig_treshold, args.sig_treshold_2)
@@ -193,6 +193,7 @@ def create_fname(path:str,prefix:str)->str:
     parser.add_argument("--locus-width-kb",dest="loc_width",type=int,default=250,help="locus width to include for each SNP, in kb")
     parser.add_argument("--alt-sign-treshold",dest="sig_treshold_2",type=float, default=5e-8,help="optional group treshold")
     parser.add_argument("--ld-panel-path",dest="ld_panel_path",type=str,help="Filename to the genotype data for ld calculation, without suffix")
+    parser.add_argument("--tabix-ld-wide-fetch",dest ="tabixld_assume_narrow_ld",default=True,action="store_false",help = "Only on TabixLD: When using FinnGen tabix panel, we assume by default that all LD pairs are keyed on first variant. That greatly reduces runtime. Pass this flag to disable this assumption when using LD files not compatible with that assumption." )
 
     #r2 static bound or dynamic per peak
     r2_group = parser.add_mutually_exclusive_group()

testing/test_annotation.py

@@ -21,14 +21,17 @@ def test_generate_ranges(self):
         out2 = generate_chrom_ranges(set(test_data),maximum_range_length=4)
         validation2 = {'2': [
                 Region(chrom='2', start=14, end=17),
-                Region(chrom='2', start=18, end=21),
-                Region(chrom='2', start=22, end=25),
+                Region(chrom='2', start=17, end=20),
+                Region(chrom='2', start=20, end=23),
+                Region(chrom='2', start=23, end=26),
                 Region(chrom='2', start=26, end=29)],
             '1': [
                 Region(chrom='1', start=0, end=3),
-                Region(chrom='1', start=4, end=7),
-                Region(chrom='1', start=8, end=11),
-                Region(chrom='1', start=12, end=14)]}
+                Region(chrom='1', start=3, end=6),
+                Region(chrom='1', start=6, end=9),
+                Region(chrom='1', start=9, end=12),
+                Region(chrom='1', start=12, end=14)],
+            '3' :[Region(chrom='3', start=29, end=30)]}
         self.assertEqual(out2,validation2)