Coverage for packages/dqm-ml-core/src/dqm_ml_core/metrics/representativeness.py: 90%
339 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Representativeness metric processor for evaluating distribution fit.
3This module contains the RepresentativenessProcessor class that evaluates
4how well a dataset represents a target statistical distribution using
5various statistical tests.
6"""
8import json
9import logging
10from typing import Any
12import numpy as np
13import pandas as pd
14import pyarrow as pa
15from scipy import stats
17# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed
18from typing_extensions import override
20from dqm_ml_core.api.metrics_processor import MetricsProcessor
21from dqm_ml_core.models.processors import RepresentativenessProcessorConfig
23logger = logging.getLogger(__name__)
26class RepresentativenessProcessor(MetricsProcessor):
27 """
28 Evaluates how well the dataset represents a target statistical distribution.
30 This processor performs on samples discretisation statistical tests to compare the observed
31 distribution of numerical columns against a theoretical target distribution
32 (Normal or Uniform).
34 Supported Metrics:
35 - Chi-square: Goodness-of-fit test for categorical/binned data.
36 - Kolmogorov-Smirnov (KS): Non-parametric test for continuous distributions (approximated via sampling).
37 - Shannon Entropy: Measures the information diversity of the binned data.
38 - GRTE (Geometric Representativeness Trajectory Error): Measures the exponential gap
39 between observed and theoretical entropy.
41 The processor uses a streaming architecture:
42 - Batch level: Computes partial calculus.
43 - Dataset level: Aggregates histograms and performs final statistical tests.
44 """
46 SUPPORTED_METRICS = {
47 "chi-square",
48 "grte",
49 "shannon-entropy",
50 "kolmogorov-smirnov",
51 }
52 SUPPORTED_DISTS = {"normal", "uniform"}
54 # Configuration constants - can be overridden in config
55 DEFAULT_ALPHA = 0.05 # Significance level for statistical tests
56 DEFAULT_SHANNON_ENTROPY_THRESHOLD = 2.0 # Threshold for high/low diversity interpretation
57 DEFAULT_GRTE_THRESHOLD = 0.5 # Threshold for high/low representativeness interpretation
58 DEFAULT_KS_SAMPLE_SIZE = 500 # Maximum sample size for KS test
59 DEFAULT_KS_MIN_SAMPLE_SIZE = 50 # Minimum sample size for KS test
60 DEFAULT_KS_SAMPLE_DIVISOR = 20 # Divisor for calculating sample size per batch
61 DEFAULT_EPSILON = 1e-9 # Small value to avoid division by zero
62 DEFAULT_INTERPRETATION_THRESHOLDS = {
63 "follows_distribution": "follows_distribution",
64 "does_not_follow_distribution": "does_not_follow_distribution",
65 "high_diversity": "high_diversity",
66 "low_diversity": "low_diversity",
67 "high_representativeness": "high_representativeness",
68 "low_representativeness": "low_representativeness",
69 }
71 def __init__(
72 self,
73 name: str = "representativeness",
74 config: dict[str, Any] | None = None,
75 ) -> None:
76 """
77 Initialize the representativeness processor.
79 Args:
80 name: Name of the processor.
81 config: Configuration dictionary containing:
82 - input_columns: List of columns to analyze.
83 - metrics: List of metrics to compute (default: all supported).
84 - bins: Number of bins for histograms (default: 10).
85 - distribution: Target distribution ("normal" or "uniform").
86 - alpha: Significance level (default: 0.05).
87 - distribution_params: Dictionary of params (e.g., mean, std, min, max).
88 """
89 super().__init__(name, config)
90 self.name = name
91 cfg = RepresentativenessProcessorConfig.model_validate({**self.config, "name": self.name})
92 self._init_from_cfg(cfg)
94 # Use device from compute config if available
95 compute_device = getattr(self, "compute_device", None)
96 self.device = self._resolve_device(compute_device) if compute_device else self._resolve_device("cpu")
98 # Use seed from compute config if available
99 compute_seed = getattr(self, "compute_seed", None)
100 if compute_seed is not None: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true
101 self._rng = np.random.default_rng(compute_seed)
102 else:
103 self._rng = np.random.default_rng(42)
105 self._bin_edges: dict[str, np.ndarray] = {}
106 self._bin_params: dict[str, dict[str, float]] = {} # caching for _compute_expected_counts
107 self._initialized: bool = False
109 def _set_optional_subconfig(self, cfg: RepresentativenessProcessorConfig) -> None:
110 self.bins = cfg.histogram.bins if cfg.histogram else 10
111 self.shannon_entropy_threshold = (
112 cfg.shannon.threshold if cfg.shannon else self.DEFAULT_SHANNON_ENTROPY_THRESHOLD
113 )
114 self.grte_threshold = cfg.grte.threshold if cfg.grte else self.DEFAULT_GRTE_THRESHOLD
115 self.grte_scaling_factor = cfg.grte.scaling_factor if cfg.grte else -2.0
116 self.ks_sample_size = cfg.ks.sample_size if cfg.ks else self.DEFAULT_KS_SAMPLE_SIZE
117 self.ks_min_sample_size = cfg.ks.min_sample_size if cfg.ks else self.DEFAULT_KS_MIN_SAMPLE_SIZE
118 self.ks_sample_divisor = cfg.ks.sample_divisor if cfg.ks else self.DEFAULT_KS_SAMPLE_DIVISOR
119 self.interpretation_thresholds = (
120 dict(cfg.interpretation.model_dump()) if cfg.interpretation else self.DEFAULT_INTERPRETATION_THRESHOLDS
121 )
123 @staticmethod
124 def _parse_distribution_params(
125 cfg: RepresentativenessProcessorConfig,
126 ) -> dict[str, dict[str, float]]:
127 dist_params: dict[str, dict[str, float]] = {}
128 if not cfg.distribution_params:
129 return dist_params
130 for p in cfg.distribution_params:
131 dd: dict[str, float] = {}
132 if p.mean is not None:
133 dd["mean"] = p.mean
134 if p.std is not None:
135 dd["std"] = p.std
136 if p.min is not None:
137 dd["min"] = p.min
138 if p.max is not None:
139 dd["max"] = p.max
140 dist_params[p.column] = dd
141 return dist_params
143 def _init_from_cfg(self, cfg: RepresentativenessProcessorConfig) -> None:
144 self.metrics = list(cfg.metrics)
145 self._set_optional_subconfig(cfg)
146 self.distribution = cfg.distribution
147 self.mean_std_estimation = cfg.mean_std_estimation
148 self.expected_counts_method = cfg.expected_counts_method
149 self.alpha = cfg.alpha
150 self.epsilon = cfg.epsilon
151 self.dist_params = self._parse_distribution_params(cfg)
153 @override
154 def generated_metrics(self) -> list[str]:
155 """
156 Return the list of metric columns that will be generated.
158 Returns:
159 List of output metric column names
160 """
161 # TODO : manage output metrics names with configuration
162 # for now we follow a fixed naming convention
163 metrics = []
164 for col in self.input_columns:
165 if "chi-square" in self.metrics: 165 ↛ 169line 165 didn't jump to line 169 because the condition on line 165 was always true
166 metrics.append(f"{col}_chi-square_p_value")
167 metrics.append(f"{col}_chi-square_statistic")
168 metrics.append(f"{col}_chi-square_interpretation")
169 if "kolmogorov-smirnov" in self.metrics: 169 ↛ 173line 169 didn't jump to line 173 because the condition on line 169 was always true
170 metrics.append(f"{col}_kolmogorov-smirnov_p_value")
171 metrics.append(f"{col}_kolmogorov-smirnov_statistic")
172 metrics.append(f"{col}_kolmogorov-smirnov_interpretation")
173 if "shannon-entropy" in self.metrics: 173 ↛ 176line 173 didn't jump to line 176 because the condition on line 173 was always true
174 metrics.append(f"{col}_shannon-entropy_entropy")
175 metrics.append(f"{col}_shannon-entropy_interpretation")
176 if "grte" in self.metrics: 176 ↛ 164line 176 didn't jump to line 164 because the condition on line 176 was always true
177 metrics.append(f"{col}_grte_grte_value")
178 metrics.append(f"{col}_grte_interpretation")
180 return metrics
182 # utils functions
183 @staticmethod
184 def _resolve_device(device: str) -> str:
185 """Resolve ``"auto"`` to CUDA if available, else CPU."""
186 if device == "auto":
187 try:
188 import torch
190 return "cuda" if torch.cuda.is_available() else "cpu"
191 except ImportError:
192 return "cpu"
193 return device
195 @staticmethod
196 def _convert_column_to_numeric(feature_array: pa.Array) -> pd.Series | None:
197 """Convert a PyArrow column array to a numeric pandas Series with NaN handling.
199 Args:
200 feature_array: PyArrow array from the batch.
202 Returns:
203 Numeric pandas Series with NaN dropped, or None if conversion fails.
204 """
205 try:
206 np_col = np.asarray(feature_array.to_numpy(zero_copy_only=False))
207 except Exception:
208 np_col = pd.Series(feature_array.to_pylist()).to_numpy(copy=True)
209 numeric_values = pd.to_numeric(pd.Series(np_col), errors="coerce").dropna()
210 return numeric_values if not numeric_values.empty else None
212 def _compute_batch_ks_sample(self, numeric_values: pd.Series) -> np.ndarray | None:
213 """Compute a random KS sample from numeric values if KS or chi-square is enabled.
215 Args:
216 numeric_values: Numeric pandas Series from a batch.
218 Returns:
219 Sampled numpy array, or None if no sampling is needed.
220 """
221 if "kolmogorov-smirnov" not in self.metrics and "chi-square" not in self.metrics:
222 return None
224 sample_per_batch = min(
225 self.ks_sample_size,
226 max(
227 self.ks_min_sample_size,
228 len(numeric_values) // self.ks_sample_divisor,
229 ),
230 )
231 if len(numeric_values) > sample_per_batch:
232 sample_indices = self._rng.choice(len(numeric_values), sample_per_batch, replace=False)
233 return np.asarray(numeric_values.iloc[sample_indices])
234 return np.asarray(numeric_values)
236 @override
237 def compute_batch_metric(self, features: dict[str, pa.Array]) -> dict[str, pa.Array]:
238 """
239 Compute partial histogram statistics per batch for streaming aggregation.
241 Args:
242 features: Dictionary of column arrays from this batch.
244 Returns:
245 Dictionary containing:
246 - {col}_count: Total valid numeric samples.
247 - {col}_hist: Histogram counts.
248 - {col}_ks_sample: Random subset of data for KS test.
249 """
250 batch_metrics = {}
252 for col in self.input_columns:
253 if col not in features:
254 logger.warning(f"[{self.name}] column '{col}' not found in batch")
255 continue
257 numeric_values = self._convert_column_to_numeric(features[col])
258 if numeric_values is None:
259 logger.warning(f"[{self.name}] column '{col}' has no valid numeric values in this batch")
260 continue
262 if not self._initialized or col not in self._bin_edges:
263 self._initialize_bin_edges(numeric_values.to_numpy(), col)
265 edges = self._bin_edges[col]
266 hist_counts = np.histogram(numeric_values, bins=edges)[0].astype(np.int64)
268 batch_metrics[f"{col}_count"] = pa.array([len(numeric_values)], type=pa.int64())
269 batch_metrics[f"{col}_hist"] = pa.FixedSizeListArray.from_arrays(
270 hist_counts, list_size=hist_counts.shape[0]
271 )
273 ks_sample = self._compute_batch_ks_sample(numeric_values)
274 if ks_sample is not None: 274 ↛ 252line 274 didn't jump to line 252 because the condition on line 274 was always true
275 batch_metrics[f"{col}_ks_sample"] = pa.array(ks_sample.tolist(), type=pa.float64())
277 if not self._initialized and batch_metrics:
278 self._initialized = True
280 return batch_metrics
282 def _get_normal_init_params(self, col: str, sample_data: np.ndarray) -> tuple[float, float]:
283 col_params = self.dist_params.get(col, {})
284 if self.mean_std_estimation == "user_provided":
285 if "mean" not in col_params or "std" not in col_params:
286 raise ValueError(
287 f"[{self.name}] user_provided strategy requires 'mean' and 'std' "
288 f"in distribution_params for column '{col}'"
289 )
290 mean = float(col_params["mean"])
291 std = float(col_params["std"])
292 else:
293 mean = float(col_params.get("mean", np.mean(sample_data)))
294 std = float(col_params.get("std", np.std(sample_data, ddof=0)))
295 std = std if std > 0.0 else self.epsilon
296 return mean, std
298 def _get_uniform_init_params(self, col: str, sample_data: np.ndarray) -> tuple[float, float]:
299 col_params = self.dist_params.get(col, {})
300 if self.mean_std_estimation == "user_provided":
301 if "min" not in col_params or "max" not in col_params: 301 ↛ 302line 301 didn't jump to line 302 because the condition on line 301 was never true
302 raise ValueError(
303 f"[{self.name}] user_provided strategy requires 'min' and 'max' "
304 f"in distribution_params for column '{col}'"
305 )
306 min_val = float(col_params["min"])
307 max_val = float(col_params["max"])
308 else:
309 min_val = float(col_params.get("min", np.min(sample_data)))
310 max_val = float(col_params.get("max", np.max(sample_data)))
311 if max_val <= min_val:
312 max_val = min_val + self.epsilon
313 return min_val, max_val
315 def _initialize_bin_edges(self, sample_data: np.ndarray, col: str) -> None:
316 sample_data = sample_data[np.isfinite(sample_data)]
317 if len(sample_data) == 0: 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true
318 sample_data = np.array([0.0])
319 if self.distribution == "normal":
320 mean, std = self._get_normal_init_params(col, sample_data)
321 edges = self._bin_edges_normal(mean, std, self.bins)
322 self._bin_params[col] = {"mean": mean, "std": std}
323 else:
324 min_val, max_val = self._get_uniform_init_params(col, sample_data)
325 edges = self._bin_edges_uniform(min_val, max_val, self.bins, sample_data)
326 self._bin_edges[col] = edges
328 def _aggregate_column_metrics(
329 self, batch_metrics: dict[str, pa.Array], col: str
330 ) -> tuple[int, np.ndarray, np.ndarray] | None:
331 """Aggregate histogram and count for a single column across all batches.
333 Args:
334 batch_metrics: Dictionary of batch-level metrics.
335 col: Column name.
337 Returns:
338 Tuple of (total_count, obs_counts, edges) or None if aggregation fails.
339 """
340 count_key = f"{col}_count"
341 hist_key = f"{col}_hist"
342 if count_key not in batch_metrics or hist_key not in batch_metrics:
343 logger.warning(f"[{self.name}] no batch metrics for column '{col}'")
344 return None
346 hist_batch_arrays = np.asarray(batch_metrics[hist_key].to_numpy(zero_copy_only=False))
347 if hist_batch_arrays.shape[0] == 0:
348 logger.warning(f"[{self.name}] no histogram batch for '{col}'")
349 return None
351 total_count = int(np.sum(batch_metrics[count_key].to_numpy()))
352 hist_arrays = hist_batch_arrays[0].copy()
353 for batch_hist in hist_batch_arrays[1:]:
354 hist_arrays += batch_hist
355 obs_counts = hist_arrays.astype(float)
357 if total_count <= 0 or obs_counts.sum() <= 0: 357 ↛ 358line 357 didn't jump to line 358 because the condition on line 357 was never true
358 logger.warning(f"[{self.name}] no valid data for column '{col}'")
359 return None
361 if col not in self._bin_edges:
362 logger.warning(f"[{self.name}] no bin edges for column '{col}' - skipping")
363 return None
365 return total_count, obs_counts, self._bin_edges[col]
367 def _resolve_normal_params(self, col: str, batch_metrics: dict[str, pa.Array]) -> tuple[float, float]:
368 if self.mean_std_estimation == "from_first_batch":
369 params = self._bin_params.get(col)
370 if params is not None: 370 ↛ 372line 370 didn't jump to line 372 because the condition on line 370 was always true
371 return params["mean"], params["std"]
372 return self._estimate_normal_params(col, batch_metrics)
373 if self.mean_std_estimation == "per_batch": 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true
374 return self._estimate_normal_params(col, batch_metrics)
375 if self.mean_std_estimation == "user_provided":
376 col_params = self.dist_params.get(col, {})
377 if "mean" not in col_params or "std" not in col_params: 377 ↛ 378line 377 didn't jump to line 378 because the condition on line 377 was never true
378 raise ValueError(
379 f"[{self.name}] user_provided strategy requires 'mean' and 'std' "
380 f"in distribution_params for column '{col}'"
381 )
382 std = float(col_params["std"])
383 return float(col_params["mean"]), std if std > 0.0 else self.epsilon
384 raise NotImplementedError(
385 f"[{self.name}] from_all_data strategy requires two-pass processing and is not yet implemented"
386 )
388 def _compute_expected_counts(
389 self,
390 col: str,
391 batch_metrics: dict[str, pa.Array],
392 total_count: int,
393 edges: np.ndarray,
394 ) -> np.ndarray:
395 if self.distribution == "normal":
396 mean, std = self._resolve_normal_params(col, batch_metrics)
397 if self.expected_counts_method == "cdf": 397 ↛ 400line 397 didn't jump to line 400 because the condition on line 397 was always true
398 expected_probs = stats.norm.cdf(edges[1:], mean, std) - stats.norm.cdf(edges[:-1], mean, std)
399 return np.asarray((expected_probs * total_count).astype(np.float64))
400 expected_values = self._rng.normal(mean, std, total_count)
401 else:
402 col_params = self.dist_params.get(col, {})
403 min_val = float(col_params.get("min", edges[0]))
404 max_val = float(col_params.get("max", edges[-1]))
405 if self.expected_counts_method == "cdf": 405 ↛ 410line 405 didn't jump to line 410 because the condition on line 405 was always true
406 expected_probs = stats.uniform.cdf(edges[1:], min_val, max_val - min_val) - stats.uniform.cdf(
407 edges[:-1], min_val, max_val - min_val
408 )
409 return np.asarray((expected_probs * total_count).astype(np.float64))
410 expected_values = self._rng.uniform(min_val, max_val, total_count)
411 return np.histogram(expected_values, bins=edges)[0].astype(np.float64)
413 def _estimate_from_samples(
414 self,
415 col_params: dict[str, Any],
416 sample_key: str,
417 batch_metrics: dict[str, pa.Array],
418 ) -> tuple[float, float]:
419 """Estimate normal params from KS sample data or fall back to defaults."""
420 if sample_key in batch_metrics: 420 ↛ 427line 420 didn't jump to line 427 because the condition on line 420 was always true
421 sample_arrays = batch_metrics[sample_key].to_numpy()
422 if sample_arrays.ndim > 1: 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 sample_arrays = sample_arrays.flatten()
424 mean = float(col_params.get("mean", np.mean(sample_arrays)))
425 std = float(col_params.get("std", np.std(sample_arrays, ddof=0)))
426 else:
427 mean = float(col_params.get("mean", 0.0))
428 std = float(col_params.get("std", 1.0))
429 return mean, std
431 def _estimate_normal_params(self, col: str, batch_metrics: dict[str, pa.Array]) -> tuple[float, float]:
432 """Estimate normal distribution parameters from config or KS samples.
434 Args:
435 col: Column name.
436 batch_metrics: Batch-level metrics dict.
438 Returns:
439 Tuple of (mean, std).
440 """
441 col_params = self.dist_params.get(col, {})
442 if self.mean_std_estimation == "user_provided":
443 if "mean" not in col_params or "std" not in col_params: 443 ↛ 444line 443 didn't jump to line 444 because the condition on line 443 was never true
444 raise ValueError(
445 f"[{self.name}] user_provided strategy requires 'mean' and 'std' "
446 f"in distribution_params for column '{col}'"
447 )
448 mean = float(col_params["mean"])
449 std = float(col_params["std"])
450 else:
451 sample_key = f"{col}_ks_sample"
452 mean, std = self._estimate_from_samples(col_params, sample_key, batch_metrics)
453 std = std if std > 0.0 else self.epsilon
454 return mean, std
456 def _compute_chi_square_metric(self, obs_counts: np.ndarray, exp_counts: np.ndarray) -> dict[str, Any]:
457 """Compute chi-square goodness-of-fit test between observed and expected counts.
459 Args:
460 obs_counts: Observed frequency counts per bin.
461 exp_counts: Expected frequency counts per bin.
463 Returns:
464 Dict with p_value, statistic, interpretation keys.
465 """
466 mask = exp_counts > 0
467 if mask.sum() < 2:
468 return {
469 "p_value": float("nan"),
470 "statistic": float("nan"),
471 "interpretation": "insufficient_bins",
472 }
474 obs_sum = obs_counts[mask].sum()
475 exp_sum = exp_counts[mask].sum()
477 exp_counts_normalized = exp_counts[mask] * (obs_sum / exp_sum)
478 try:
479 chi = stats.chisquare(f_obs=obs_counts[mask], f_exp=exp_counts_normalized)
480 return {
481 "p_value": float(chi.pvalue),
482 "statistic": float(chi.statistic),
483 "interpretation": self.interpretation_thresholds.get(
484 "follows_distribution" if chi.pvalue >= self.alpha else "does_not_follow_distribution",
485 "follows_distribution",
486 ),
487 }
488 except ValueError as e:
489 return {
490 "p_value": float("nan"),
491 "statistic": float("nan"),
492 "interpretation": f"chi_square_failed: {e!s}",
493 "note": "using observed counts only due to statistical constraints",
494 }
496 def _compute_ks_metric(self, col: str, batch_metrics: dict[str, pa.Array]) -> dict[str, Any]:
497 """Compute Kolmogorov-Smirnov test using sampled data.
499 Args:
500 col: Column name.
501 batch_metrics: Batch-level metrics dict containing KS samples.
503 Returns:
504 Dict with p_value, statistic, interpretation keys.
505 """
506 sample_key = f"{col}_ks_sample"
507 if sample_key not in batch_metrics:
508 return {
509 "p_value": float("nan"),
510 "statistic": float("nan"),
511 "interpretation": "no_sample_data_found",
512 }
514 sample_arrays = batch_metrics[sample_key].to_numpy()
515 ks_samples = sample_arrays if sample_arrays.ndim == 1 else sample_arrays.flatten()
517 if len(ks_samples) == 0:
518 return {
519 "p_value": float("nan"),
520 "statistic": float("nan"),
521 "interpretation": "no_samples_available",
522 }
524 if self.distribution == "normal":
525 mean, std = self._estimate_normal_params(col, batch_metrics)
526 ks = stats.kstest(ks_samples, stats.norm.cdf, args=(mean, std))
527 else:
528 col_params = self.dist_params.get(col, {})
529 min_val = float(col_params.get("min", np.min(ks_samples)))
530 max_val = float(col_params.get("max", np.max(ks_samples)))
531 if max_val <= min_val:
532 max_val = min_val + self.epsilon
533 ks = stats.kstest(ks_samples, stats.uniform.cdf, args=(min_val, max_val - min_val))
535 return {
536 "p_value": float(ks.pvalue),
537 "statistic": float(ks.statistic),
538 "interpretation": self.interpretation_thresholds.get(
539 "follows_distribution" if ks.pvalue >= self.alpha else "does_not_follow_distribution",
540 "follows_distribution",
541 ),
542 "sample_size": len(ks_samples),
543 "note": "approximated_from_random_samples",
544 }
546 def _compute_shannon_entropy_metric(self, exp_counts: np.ndarray) -> dict[str, Any]:
547 """Compute Shannon entropy from expected frequency counts.
549 Args:
550 exp_counts: Expected frequency counts per bin.
552 Returns:
553 Dict with entropy and interpretation.
554 """
555 p_exp = exp_counts / exp_counts.sum()
556 h_exp = float(stats.entropy(p_exp))
557 is_high = h_exp > self.shannon_entropy_threshold
558 return {
559 "entropy": h_exp,
560 "interpretation": self.interpretation_thresholds.get(
561 "high_diversity" if is_high else "low_diversity",
562 "high_diversity",
563 ),
564 }
566 def _compute_grte_metric(self, obs_counts: np.ndarray, exp_counts: np.ndarray) -> dict[str, Any]:
567 """Compute GRTE (exponential gap between observed and theoretical entropy).
569 Args:
570 obs_counts: Observed frequency counts per bin.
571 exp_counts: Expected frequency counts per bin.
573 Returns:
574 Dict with grte_value and interpretation.
575 """
576 p_obs = obs_counts / obs_counts.sum()
577 p_exp = exp_counts / exp_counts.sum()
578 h_obs = float(stats.entropy(p_obs))
579 h_exp = float(stats.entropy(p_exp))
580 grte = float(np.exp(self.grte_scaling_factor * abs(h_exp - h_obs)))
581 is_high = grte > self.grte_threshold
582 return {
583 "grte_value": grte,
584 "interpretation": self.interpretation_thresholds.get(
585 "high_representativeness" if is_high else "low_representativeness",
586 "high_representativeness",
587 ),
588 }
590 @staticmethod
591 def _build_compute_metadata(
592 total_samples: int,
593 batch_metrics: dict[str, pa.Array],
594 input_columns: list[str],
595 distribution: str,
596 metrics: list[str],
597 bins: int,
598 ) -> str:
599 """Build metadata JSON string for compute results.
601 Args:
602 total_samples: Total number of samples processed.
603 batch_metrics: Batch-level metrics dict.
604 input_columns: Input column names.
605 distribution: Target distribution name.
606 metrics: List of computed metric names.
607 bins: Number of histogram bins.
609 Returns:
610 JSON-encoded metadata string.
611 """
612 return json.dumps(
613 {
614 "bins": bins,
615 "distribution": distribution,
616 "metrics_computed": metrics,
617 "total_samples": total_samples,
618 "columns_analyzed": [c for c in input_columns if f"{c}_count" in batch_metrics],
619 "ks_sampling_enabled": "kolmogorov-smirnov" in metrics,
620 "note": "KS test uses random sampling approximation for scalability",
621 }
622 )
624 @staticmethod
625 def _flatten_col_results(col_res: dict[str, Any], col: str, results: dict[str, Any]) -> None:
626 """Flatten nested metric dicts into flat output keys.
628 Args:
629 col_res: Per-column metric results (potentially nested).
630 col: Column name.
631 results: Output dict to populate.
632 """
633 for key, value in col_res.items():
634 if isinstance(value, dict): 634 ↛ 638line 634 didn't jump to line 638 because the condition on line 634 was always true
635 for prop, content in value.items():
636 results[f"{col}_{key}_{prop}"] = content
637 else:
638 results[f"{col}_{key}"] = value
640 def _compute_column_results(
641 self,
642 col: str,
643 batch_metrics: dict[str, pa.Array],
644 results: dict[str, Any],
645 ) -> int:
646 """Compute all metrics for a single column and write them to results.
648 Args:
649 col: Column name.
650 batch_metrics: Batch-level metrics dict.
651 results: Output dict to populate (mutated in place).
653 Returns:
654 Number of samples processed (0 if column has no valid data).
655 """
656 agg = self._aggregate_column_metrics(batch_metrics, col)
657 if agg is None:
658 return 0
660 total_count, obs_counts, edges = agg
661 exp_counts = self._compute_expected_counts(col, batch_metrics, total_count, edges)
663 col_res: dict[str, Any] = {}
665 if "chi-square" in self.metrics: 665 ↛ 668line 665 didn't jump to line 668 because the condition on line 665 was always true
666 col_res["chi-square"] = self._compute_chi_square_metric(obs_counts, exp_counts)
668 if "kolmogorov-smirnov" in self.metrics: 668 ↛ 671line 668 didn't jump to line 671 because the condition on line 668 was always true
669 col_res["kolmogorov-smirnov"] = self._compute_ks_metric(col, batch_metrics)
671 if "shannon-entropy" in self.metrics: 671 ↛ 674line 671 didn't jump to line 674 because the condition on line 671 was always true
672 col_res["shannon-entropy"] = self._compute_shannon_entropy_metric(exp_counts)
674 if "grte" in self.metrics: 674 ↛ 677line 674 didn't jump to line 677 because the condition on line 674 was always true
675 col_res["grte"] = self._compute_grte_metric(obs_counts, exp_counts)
677 if col_res: 677 ↛ 680line 677 didn't jump to line 680 because the condition on line 677 was always true
678 self._flatten_col_results(col_res, col, results)
680 return total_count
682 @override
683 def compute(self, batch_metrics: dict[str, pa.Array] | None = None) -> dict[str, Any]:
684 """
685 Compute final dataset-level metrics by aggregating batch histograms.
687 Args:
688 batch_metrics: Dictionary of batch-level metrics collected during processing.
690 Returns:
691 Dictionary containing final scores and interpretations.
692 """
693 if not batch_metrics:
694 return {"_metadata": {"error": "No batch metrics provided"}}
696 results: dict[str, Any] = {}
697 total_samples = 0
699 for col in self.input_columns:
700 total_samples += self._compute_column_results(col, batch_metrics, results)
702 results["_metadata"] = self._build_compute_metadata(
703 total_samples,
704 batch_metrics,
705 self.input_columns or [],
706 self.distribution,
707 self.metrics,
708 self.bins,
709 )
710 return results
712 @override
713 def reset(self) -> None:
714 """Reset processor state for new processing run."""
715 self._bin_edges = {}
716 self._bin_params = {}
717 self._initialized = False
719 # utils methods for bin edge calculation
721 def _bin_edges_normal(self, mean: float, std: float, bins: int) -> np.ndarray:
722 """Calculate bin edges using the PPF of a Normal distribution.
724 This ensures bins represent equal probability mass under the
725 theoretical distribution. The first and last bins are extended
726 to -inf and +inf respectively.
727 """
728 # logic from dqm-ml v1: use stats.norm.ppf with linspace(1/bins, 1, bins)
729 bin_edges_list = [stats.norm.ppf(i / bins, mean, std) for i in range(1, bins)]
730 return np.array([-np.inf] + bin_edges_list + [np.inf])
732 def _bin_edges_uniform(self, param_min: float, param_max: float, bins: int, data: np.ndarray) -> np.ndarray:
733 """
734 Calculate linearly spaced bin edges for a Uniform distribution.
736 The range is determined by the minimum/maximum of both the configured
737 parameters and the actual observed data.
738 """
739 low_edge = min(param_min, float(np.min(data)))
740 high_edge = max(param_max, float(np.max(data)))
741 # Handle degenerate case where all data is identical
742 if high_edge <= low_edge:
743 high_edge = low_edge + self.epsilon
744 return np.linspace(low_edge, high_edge, bins + 1)