Coverage for little_loops / analytics / association.py: 38%

1"""PMI and lift scoring for sequence association analysis.""" 

2 

3from __future__ import annotations 

4 

5import math 

6from dataclasses import dataclass 

7 

8LIFT_THRESHOLD = 1.0 

9"""Lift below this value means the pair co-occurs at or below the frequency-prior rate.""" 

10 

11 

12@dataclass 

13class AssociationScores: 

14 """PMI and lift scores for a token pair. 

15 

16 Attributes: 

17 pmi: Pointwise mutual information = log(P(a,b) / (P(a)*P(b))). 

18 Positive means the pair co-occurs more than chance predicts. 

19 lift: P(b|a) / P(b) = exp(pmi). Values < 1.0 indicate 

20 frequency-prior-equivalent co-occurrence. 

21 """ 

22 

23 pmi: float 

24 lift: float 

25 

26 

27def compute_pmi(count_ab: int, count_a: int, count_b: int, total_unigrams: int) -> float: 

28 """Compute pointwise mutual information for token pair (a, b). 

29 

30 PMI(a,b) = log( count(a,b) * total / (count_a * count_b) ) 

31 

32 Uses integer counts to defer float division until the final log, minimising 

33 accumulation of floating-point error. 

34 

35 Args: 

36 count_ab: Co-occurrence count of the pair. 

37 count_a: Unigram count of token a. 

38 count_b: Unigram count of token b. 

39 total_unigrams: Total unigram count across the corpus. 

40 

41 Returns: 

42 PMI value (float, may be negative). 

43 

44 Raises: 

45 ValueError: If any count is zero or negative. 

46 """ 

47 if count_ab <= 0: 

48 raise ValueError(f"count_ab must be positive, got {count_ab}") 

49 if count_a <= 0: 

50 raise ValueError(f"count_a must be positive, got {count_a}") 

51 if count_b <= 0: 

52 raise ValueError(f"count_b must be positive, got {count_b}") 

53 if total_unigrams <= 0: 

54 raise ValueError(f"total_unigrams must be positive, got {total_unigrams}") 

55 

56 return math.log(count_ab * total_unigrams / (count_a * count_b)) 

57 

58 

59def compute_lift(count_ab: int, count_a: int, count_b: int, total_unigrams: int) -> float: 

60 """Compute lift (confidence ratio) for token pair (a, b). 

61 

62 lift(a,b) = P(b|a) / P(b) = count(a,b) * total / (count_a * count_b) 

63 

64 A lift of 1.0 means the pair co-occurs at exactly the frequency-prior rate. 

65 Values < 1.0 are frequency-prior-equivalent (not worth automating on this 

66 signal alone). Values > 1.0 indicate a genuine non-trivial co-occurrence. 

67 

68 Args: 

69 count_ab: Co-occurrence count of the pair. 

70 count_a: Unigram count of token a. 

71 count_b: Unigram count of token b. 

72 total_unigrams: Total unigram count across the corpus. 

73 

74 Returns: 

75 Lift value (positive float). 

76 

77 Raises: 

78 ValueError: If any count is zero or negative. 

79 """ 

80 if count_ab <= 0: 

81 raise ValueError(f"count_ab must be positive, got {count_ab}") 

82 if count_a <= 0: 

83 raise ValueError(f"count_a must be positive, got {count_a}") 

84 if count_b <= 0: 

85 raise ValueError(f"count_b must be positive, got {count_b}") 

86 if total_unigrams <= 0: 

87 raise ValueError(f"total_unigrams must be positive, got {total_unigrams}") 

88 

89 return count_ab * total_unigrams / (count_a * count_b)