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

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_batches, 

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 prior_count=20.0, 

33 min_weight=0.2, 

34 max_weight=5.0, 

35 njobs=1, 

36 read_buffer_size=10000, 

37 ): 

38 self.subsample = subsample 

39 self.nbins = nbins 

40 self.mappability_file = mappability_file 

41 self.min_insert_size = min_insert_size 

42 self.max_insert_size = max_insert_size 

43 self.min_mappability_threshold = min_mappability_threshold 

44 self.prior_count = prior_count 

45 self.min_weight = min_weight 

46 self.max_weight = max_weight 

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

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

49 self.read_buffer_size = read_buffer_size 

50 self.bin_edges = None 

51 self.weights = None 

52 

53 def fit(self, bam: Union[str, pysam.AlignmentFile]): 

54 close_bam = False 

55 if isinstance(bam, str): 

56 bam = ( 

57 pysam.AlignmentFile(bam, "rb") 

58 if self.njobs == 1 

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

60 ) 

61 close_bam = True 

62 features = collect_fragment_features( 

63 read_info_batches(iter_pysam_reads(bam.fetch()), self.read_buffer_size), 

64 min_insert_size=self.min_insert_size, 

65 max_insert_size=self.max_insert_size, 

66 mappability_path=self.mappability_file, 

67 min_mappability_threshold=self.min_mappability_threshold, 

68 downsample_ratio=self.subsample, 

69 njobs=self.njobs, 

70 ) 

71 if close_bam: 

72 bam.close() 

73 if not features: 

74 self.bin_edges = None 

75 self.weights = None 

76 return self 

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

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

79 observed = histogram > 0 

80 mean_count = np.mean(histogram[observed]) 

81 self.weights[observed] = (mean_count + self.prior_count) / ( 

82 histogram[observed] + self.prior_count 

83 ) 

84 if self.min_weight is not None or self.max_weight is not None: 

85 self.weights[observed] = np.clip( 

86 self.weights[observed], 

87 -np.inf if self.min_weight is None else self.min_weight, 

88 np.inf if self.max_weight is None else self.max_weight, 

89 ) 

90 self.bin_edges = bin_edges 

91 return self 

92 

93 def transform_features(self, features): 

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

95 

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

97 if self.weights is None: 

98 return 1.0 

99 return self.transform_features(fragment_features(read))