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

2

3 from __future__ import annotations

4

5 import math

6 from dataclasses import dataclass

7

8 LIFT_THRESHOLD = 1.0

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

10

11

12@dataclass

13 class 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

27 def 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

59 def 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)