Coverage for packages/dqm-ml-core/src/dqm_ml_core/metrics/diversity.py: 98%

1"""Diversity metric processor for evaluating categorical data diversity. 

2 

3This module contains the DiversityProcessor class that computes 

4diversity indices (Simpson, Gini-Simpson, Shannon Entropy, Richness) 

5for categorical or discretized data columns. 

6""" 

7 

8from collections import Counter 

9import logging 

10from typing import Any 

11 

12import numpy as np 

13import pyarrow as pa 

14from typing_extensions import override 

15 

16from dqm_ml_core.api.metrics_processor import MetricsProcessor 

17from dqm_ml_core.models.processors import DiversityProcessorConfig 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22class DiversityProcessor(MetricsProcessor): 

23 """Computes diversity indices for categorical data columns. 

24 

25 Simpson Index (1 - Σn(n-1) / N(N-1)): 

26 Unbiased estimator — probability two random samples belong to 

27 different categories. Ranges [0, 1]; high = diverse. 

28 

29 Gini-Simpson Index (1 - Σp²): 

30 Gini impurity — probability of incorrect classification. 

31 Ranges [0, 1]; high = diverse. 

32 

33 Shannon Entropy (-Σp·log(p)): 

34 Information content / uncertainty. Lower bound 0; no upper bound 

35 (grows with category count and evenness). 

36 

37 Richness: 

38 Simple count of unique categories. 

39 

40 The processor uses a streaming architecture: 

41 - Batch level: Computes value counts per column. 

42 - Dataset level: Merges all batch counts and computes final indices. 

43 """ 

44 

45 SUPPORTED_METRICS = {"simpson", "gini", "shannon", "richness"} 

46 

47 def __init__(self, name: str = "diversity", config: dict[str, Any] | None = None) -> None: 

48 """Initialize the diversity processor. 

49 

50 Args: 

51 name: Name of the processor. 

52 config: Configuration dictionary containing: 

53 - input_columns: List of columns to analyze. 

54 - metrics: List of metrics to compute (default: all supported). 

55 """ 

56 super().__init__(name, config) 

57 

58 cfg = DiversityProcessorConfig.model_validate({**self.config, "name": self.name}) 

59 

60 self.metrics: list[str] = list(cfg.metrics) 

61 

62 @override 

63 def generated_metrics(self) -> list[str]: 

64 """Return the list of metric columns that will be generated. 

65 

66 Returns: 

67 List of output metric column names. 

68 """ 

69 metrics = [] 

70 for col in self.input_columns: 

71 if "simpson" in self.metrics: 71 ↛ 73line 71 didn't jump to line 73 because the condition on line 71 was always true

72 metrics.append(f"{col}_simpson") 

73 if "gini" in self.metrics: 

74 metrics.append(f"{col}_gini") 

75 if "shannon" in self.metrics: 

76 metrics.append(f"{col}_shannon") 

77 if "richness" in self.metrics: 

78 metrics.append(f"{col}_richness") 

79 return metrics 

80 

81 @override 

82 def compute_batch_metric(self, features: dict[str, pa.Array]) -> dict[str, pa.Array]: 

83 """Compute per-batch value counts for streaming aggregation. 

84 

85 Args: 

86 features: Dictionary of column arrays from this batch. 

87 

88 Returns: 

89 Dictionary of batch-level value-count pairs. 

90 """ 

91 batch_metrics: dict[str, pa.Array] = {} 

92 

93 for col in self.input_columns: 

94 if col not in features: 

95 logger.warning("[%s] column '%s' not found in batch", self.name, col) 

96 continue 

97 

98 col_array = features[col] 

99 vc = pa.compute.value_counts(col_array) 

100 if len(vc) == 0: 

101 continue 

102 

103 values = vc.field("values") 

104 counts = vc.field("counts") 

105 

106 values_str = pa.compute.cast(values, pa.string()) 

107 batch_metrics[f"{col}_values"] = values_str 

108 batch_metrics[f"{col}_counts"] = counts 

109 

110 return batch_metrics 

111 

112 @override 

113 def compute(self, batch_metrics: dict[str, pa.Array] | None = None) -> dict[str, Any]: 

114 """Compute final dataset-level diversity indices. 

115 

116 Merges all batch-level value counts and computes diversity 

117 indices for each column. 

118 

119 Args: 

120 batch_metrics: Dictionary of batch-level value-count arrays. 

121 

122 Returns: 

123 Dictionary containing final diversity scores. 

124 """ 

125 if not batch_metrics: 

126 return {"_metadata": {"error": "No batch metrics provided"}} 

127 

128 results: dict[str, Any] = {} 

129 for col in self.input_columns: 

130 col_res = self._compute_column_diversity(col, batch_metrics) 

131 results.update(col_res) 

132 return results 

133 

134 def _compute_column_diversity(self, col: str, batch_metrics: dict[str, pa.Array]) -> dict[str, Any]: 

135 """Compute diversity indices for a single column. 

136 

137 Args: 

138 col: Column name to compute diversity for. 

139 batch_metrics: Dictionary of batch-level value-count arrays. 

140 

141 Returns: 

142 Dictionary mapping column-specific metric names to values. 

143 """ 

144 values_key = f"{col}_values" 

145 counts_key = f"{col}_counts" 

146 

147 if values_key not in batch_metrics or counts_key not in batch_metrics: 

148 logger.warning("[%s] no batch metrics for column '%s'", self.name, col) 

149 return {} 

150 

151 values_arr = batch_metrics[values_key] 

152 counts_arr = batch_metrics[counts_key] 

153 

154 if len(values_arr) == 0: 

155 return {} 

156 

157 counter: Counter[str] = Counter() 

158 total = 0 

159 for i in range(len(values_arr)): 

160 val = values_arr[i].as_py() 

161 cnt = int(counts_arr[i]) 

162 counter[val] += cnt 

163 total += cnt 

164 

165 if total == 0: 

166 return {} 

167 

168 col_res: dict[str, Any] = {} 

169 

170 if "richness" in self.metrics: 

171 col_res["richness"] = len(counter) 

172 

173 if "shannon" in self.metrics: 

174 probs = np.array(list(counter.values()), dtype=np.float64) / total 

175 probs = probs[probs > 0] 

176 entropy = float(-(probs * np.log(probs)).sum()) 

177 col_res["shannon"] = entropy 

178 

179 if "gini" in self.metrics: 

180 probs = np.array(list(counter.values()), dtype=np.float64) / total 

181 gini = float(1.0 - (probs**2).sum()) 

182 col_res["gini"] = gini 

183 

184 if "simpson" in self.metrics: 184 ↛ 189line 184 didn't jump to line 189 because the condition on line 184 was always true

185 counts_np = np.array(list(counter.values()), dtype=np.float64) 

186 simpson = float(1.0 - (counts_np * (counts_np - 1)).sum() / (total * (total - 1))) if total > 1 else 0.0 

187 col_res["simpson"] = simpson 

188 

189 return {f"{col}_{metric_name}": value for metric_name, value in col_res.items()} 

190 

191 @override 

192 def reset(self) -> None: 

193 """Reset processor state for new processing run."""