wehub-resource-sync
208 lines
7.0 KiB
Python
208 lines
7.0 KiB
Python
import math
|
|
from dataclasses import dataclass
|
|
|
|
# Recommended smoothing parameter for NPMI calculation
|
|
# Using Jeffreys prior (alpha=0.5) to minimize bias while providing robust estimates
|
|
JEFFREYS_PRIOR = 0.5
|
|
|
|
|
|
@dataclass
|
|
class TraceCorrelationCounts:
|
|
"""
|
|
Count statistics for trace correlation analysis.
|
|
|
|
This dataclass encapsulates the four fundamental counts needed
|
|
for correlation analysis between two trace filters.
|
|
|
|
Attributes:
|
|
total_count: Total number of traces in the experiment(s)
|
|
filter1_count: Number of traces matching filter 1
|
|
filter2_count: Number of traces matching filter 2
|
|
joint_count: Number of traces matching both filters
|
|
"""
|
|
|
|
total_count: int
|
|
filter1_count: int
|
|
filter2_count: int
|
|
joint_count: int
|
|
|
|
|
|
@dataclass
|
|
class NPMIResult:
|
|
"""
|
|
Result of NPMI calculation containing both unsmoothed and smoothed values.
|
|
|
|
Attributes:
|
|
npmi: Unsmoothed NPMI value with explicit -1.0 rule for zero joint count.
|
|
Returns NaN when undefined (e.g., when filter1_count=0 or filter2_count=0).
|
|
npmi_smoothed: NPMI calculated with Jeffreys prior smoothing (alpha=0.5).
|
|
More robust for small sample sizes and confidence interval estimation.
|
|
"""
|
|
|
|
npmi: float
|
|
npmi_smoothed: float | None
|
|
|
|
|
|
def calculate_npmi_from_counts(
|
|
joint_count: int,
|
|
filter1_count: int,
|
|
filter2_count: int,
|
|
total_count: int,
|
|
) -> NPMIResult:
|
|
"""
|
|
Calculate both unsmoothed and smoothed NPMI from count data.
|
|
|
|
Implements the recommended policy for NPMI calculation:
|
|
- Returns NaN (undefined) when either filter has zero support (n1=0 or n2=0)
|
|
- Returns -1.0 for unsmoothed when filters never co-occur despite both having support
|
|
- Calculates smoothed version using Jeffreys prior for robustness
|
|
|
|
NPMI measures the association between two events, normalized to [-1, 1]:
|
|
- -1: Perfect negative correlation (events never co-occur)
|
|
- 0: Independence (events occur independently)
|
|
- 1: Perfect positive correlation (events always co-occur)
|
|
- NaN: Undefined (when one or both events have zero count)
|
|
|
|
Args:
|
|
joint_count: Number of times both events occur together
|
|
filter1_count: Number of times event 1 occurs
|
|
filter2_count: Number of times event 2 occurs
|
|
total_count: Total number of observations
|
|
|
|
Returns:
|
|
NPMIResult containing both unsmoothed and smoothed NPMI values.
|
|
|
|
Examples:
|
|
>>> result = calculate_npmi_from_counts(10, 20, 15, 100)
|
|
>>> result.npmi # Unsmoothed value
|
|
>>> result.npmi_smoothed # Smoothed value
|
|
"""
|
|
# No population
|
|
if total_count <= 0:
|
|
return NPMIResult(npmi=float("nan"), npmi_smoothed=float("nan"))
|
|
|
|
# Return NaN if either filter has zero support
|
|
if filter1_count == 0 or filter2_count == 0:
|
|
return NPMIResult(npmi=float("nan"), npmi_smoothed=float("nan"))
|
|
|
|
n11 = joint_count # Both occur
|
|
n10 = filter1_count - joint_count # Only filter1
|
|
n01 = filter2_count - joint_count # Only filter2
|
|
n00 = total_count - filter1_count - filter2_count + joint_count # Neither
|
|
|
|
if min(n11, n10, n01, n00) < 0:
|
|
# Inconsistent counts, return undefined
|
|
return NPMIResult(npmi=float("nan"), npmi_smoothed=float("nan"))
|
|
|
|
# Calculate unsmoothed NPMI with explicit -1.0 rule
|
|
if joint_count == 0 and filter1_count > 0 and filter2_count > 0:
|
|
npmi_unsmoothed = -1.0
|
|
else:
|
|
npmi_unsmoothed = _calculate_npmi_core(n11, n10, n01, n00, smoothing=0)
|
|
|
|
# Calculate smoothed NPMI for robustness
|
|
npmi_smoothed = _calculate_npmi_core(n11, n10, n01, n00, smoothing=JEFFREYS_PRIOR)
|
|
|
|
return NPMIResult(npmi=npmi_unsmoothed, npmi_smoothed=npmi_smoothed)
|
|
|
|
|
|
def _calculate_npmi_core(
|
|
n11: float,
|
|
n10: float,
|
|
n01: float,
|
|
n00: float,
|
|
smoothing: float = 0,
|
|
) -> float:
|
|
"""
|
|
Core NPMI calculation with optional smoothing.
|
|
|
|
Internal function that performs the actual NPMI calculation
|
|
on a 2x2 contingency table with optional additive smoothing.
|
|
|
|
Args:
|
|
n11: Count of both events occurring
|
|
n10: Count of only event 1 occurring
|
|
n01: Count of only event 2 occurring
|
|
n00: Count of neither event occurring
|
|
smoothing: Additive smoothing parameter (0 for no smoothing)
|
|
|
|
Returns:
|
|
NPMI value in [-1, 1], or NaN if undefined.
|
|
"""
|
|
n11_s = n11 + smoothing
|
|
n10_s = n10 + smoothing
|
|
n01_s = n01 + smoothing
|
|
n00_s = n00 + smoothing
|
|
|
|
N = n11_s + n10_s + n01_s + n00_s
|
|
n1 = n11_s + n10_s # Total event 1 count
|
|
n2 = n11_s + n01_s # Total event 2 count
|
|
|
|
# NB: When marginals are zero (degenerate cases where no events occur), we return NaN
|
|
# rather than forcing a sentinel value like -1. This is mathematically correct since
|
|
# PMI is undefined when P(x)=0 or P(y)=0 (division by zero). NaN properly represents
|
|
# this undefined state and can be handled by our RPC layer, providing a more accurate
|
|
# signal than an arbitrary sentinel value.
|
|
if n1 <= 0 or n2 <= 0 or n11_s <= 0:
|
|
return float("nan")
|
|
|
|
# Handle perfect co-occurrence - check pre-smoothing values
|
|
# With smoothing, n11_s == N is never true since smoothing adds mass to other cells
|
|
if n10 == 0 and n01 == 0 and n00 == 0:
|
|
# Perfect co-occurrence: both events always occur together
|
|
return 1.0
|
|
|
|
# Calculate PMI using log-space arithmetic for numerical stability
|
|
# PMI = log(P(x,y) / (P(x) * P(y))) = log(n11*N / (n1*n2))
|
|
log_n11 = math.log(n11_s)
|
|
log_N = math.log(N)
|
|
log_n1 = math.log(n1)
|
|
log_n2 = math.log(n2)
|
|
|
|
pmi = (log_n11 + log_N) - (log_n1 + log_n2)
|
|
|
|
# Normalize by -log(P(x,y)) to get NPMI
|
|
denominator = -(log_n11 - log_N) # -log(n11/N)
|
|
|
|
npmi = pmi / denominator
|
|
|
|
# Clamp to [-1, 1] to handle floating point errors
|
|
return max(-1.0, min(1.0, npmi))
|
|
|
|
|
|
def calculate_smoothed_npmi(
|
|
joint_count: int,
|
|
filter1_count: int,
|
|
filter2_count: int,
|
|
total_count: int,
|
|
smoothing: float = JEFFREYS_PRIOR,
|
|
) -> float:
|
|
"""
|
|
Calculate smoothed NPMI for confidence interval estimation.
|
|
|
|
This function applies additive smoothing (Jeffreys prior by default) to all cells
|
|
of the contingency table. Used for uncertainty quantification via Dirichlet sampling.
|
|
|
|
Args:
|
|
joint_count: Number of times both events occur together
|
|
filter1_count: Number of times event 1 occurs
|
|
filter2_count: Number of times event 2 occurs
|
|
total_count: Total number of observations
|
|
smoothing: Additive smoothing parameter (default: JEFFREYS_PRIOR = 0.5)
|
|
|
|
Returns:
|
|
Smoothed NPMI value in [-1, 1], or NaN if undefined.
|
|
"""
|
|
if total_count <= 0:
|
|
return float("nan")
|
|
|
|
n11 = joint_count
|
|
n10 = filter1_count - joint_count
|
|
n01 = filter2_count - joint_count
|
|
n00 = total_count - filter1_count - filter2_count + joint_count
|
|
|
|
if min(n11, n10, n01, n00) < 0:
|
|
return float("nan")
|
|
|
|
return _calculate_npmi_core(n11, n10, n01, n00, smoothing)
|