Coverage for src / optwps / weighting.py: 80%

1import pysam 

2from typing import Union 

3import numpy as np 

4import joblib 

5 

6from .read_processing import ( 

7 collect_fragment_features, 

8 fragment_features, 

9 iter_pysam_reads, 

10 read_info, 

11 weight_from_features, 

12) 

13 

14 

15class WeightsCalculator: 

16 """Calculates correction weights for WPS based on fragment features. 

17 Weights are computed based on the distribution of fragment features in the dataset, 

18 allowing for correction of biases in WPS calculations. 

19 The weights are determined by binning the fragment features and calculating the inverse 

20 of the frequency of fragments in each bin, which can then be applied to adjust WPS 

21 scores accordingly. 

22 """ 

23 

24 def __init__( 

25 self, 

26 mappability_file: str = None, 

27 nbins=10, 

28 subsample=0.05, 

29 min_insert_size=None, 

30 max_insert_size=None, 

31 min_mappability_threshold=0.9, 

32 njobs=1, 

33 read_buffer_size=10000, 

34 ): 

35 self.subsample = subsample 

36 self.nbins = nbins 

37 self.mappability_file = mappability_file 

38 self.min_insert_size = min_insert_size 

39 self.max_insert_size = max_insert_size 

40 self.min_mappability_threshold = min_mappability_threshold 

41 self.njobs = joblib.cpu_count() + njobs if njobs < 0 else njobs 

42 self.njobs = max(1, self.njobs) 

43 self.read_buffer_size = read_buffer_size 

44 self.bin_edges = None 

45 self.weights = None 

46 

47 def fit(self, bam: Union[str, pysam.AlignmentFile], y=None): 

48 close_bam = False 

49 if isinstance(bam, str): 

50 bam = ( 

51 pysam.AlignmentFile(bam, "rb") 

52 if self.njobs == 1 

53 else pysam.AlignmentFile(bam, "rb", threads=self.njobs) 

54 ) 

55 close_bam = True 

56 read_batches = [] 

57 batch = [] 

58 for read in iter_pysam_reads(bam.fetch()): 

59 batch.append(read_info(read)) 

60 if len(batch) >= self.read_buffer_size: 

61 read_batches.append(batch) 

62 batch = [] 

63 if batch: 

64 read_batches.append(batch) 

65 features = collect_fragment_features( 

66 read_batches, 

67 min_insert_size=self.min_insert_size, 

68 max_insert_size=self.max_insert_size, 

69 mappability_path=self.mappability_file, 

70 min_mappability_threshold=self.min_mappability_threshold, 

71 downsample_ratio=self.subsample, 

72 njobs=self.njobs, 

73 read_buffer_size=self.read_buffer_size, 

74 ) 

75 if close_bam: 

76 bam.close() 

77 if not features: 

78 self.bin_edges = None 

79 self.weights = None 

80 return self 

81 histogram, bin_edges = np.histogramdd(np.array(features), bins=self.nbins) 

82 self.weights = np.ones_like(histogram, dtype=float) 

83 observed = histogram > 0 

84 self.weights[observed] = np.mean(histogram[observed]) / histogram[observed] 

85 self.bin_edges = bin_edges 

86 return self 

87 

88 def transform_features(self, features): 

89 return weight_from_features(features, self.bin_edges, self.weights) 

90 

91 def transform(self, read: pysam.AlignedSegment): 

92 if self.weights is None: 

93 return 1.0 

94 return self.transform_features(fragment_features(read))