Coverage for packages/dqm-ml-pytorch/src/dqm_ml_pytorch/domain_gap.py: 94%
453 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Domain gap processor for measuring distribution distance between datasets.
3This module contains the DomainGapProcessor class that computes statistical
4distances (KL divergence, MMD variants, FID, Wasserstein, PAD, CMD) between
5source and target datasets using image embeddings.
6"""
8from __future__ import annotations
10import logging
11from math import comb
12import os
13from pathlib import Path
14import tempfile
15from typing import Any
17from dqm_ml_core import GapProcessor
18from dqm_ml_core.models.processors import DomainGapProcessorConfig
19from dqm_ml_core.utils.matching import has_pattern, resolve_include_exclude
20import numpy as np
21import pyarrow as pa
22import torch
24# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed
25from typing_extensions import override
27_MISSING_EMB_MSG = "missing __emb__ — set summary.store_embeddings=true"
29logger = logging.getLogger(__name__)
32def _debug_enabled() -> bool:
33 """Check whether debug data generation is enabled via environment variable.
35 Returns:
36 True if ``DQM_ML_DEBUG`` is set to a truthy value ("1", "true", "yes").
37 """
38 return os.environ.get("DQM_ML_DEBUG", "").strip().lower() in {
39 "1",
40 "true",
41 "yes",
42 }
45# Known ResNet-18 layer embedding dimensions (C*H*W) mapped to channel counts (C).
46# Used in _compute_batch_metric_cmd to spatially pool flattened features
47# so CMD moments match the per-channel computation of v1.
48_CMD_RESNET18_EMBDIM_CHANNELS: dict[int, int] = {
49 200704: 64, # maxpool / layer1.1.relu_1: 64 * 56 * 56
50 100352: 128, # layer2.1.relu_1: 128 * 28 * 28
51 50176: 256, # layer3.1.relu_1: 256 * 14 * 14
52 25088: 512, # layer4.1.relu_1: 512 * 7 * 7
53}
56def _fixed_to_matrix(arr: pa.FixedSizeListArray) -> np.ndarray:
57 """Convert a FixedSizeListArray to a (N, D) numpy float64 matrix.
59 Args:
60 arr: Input FixedSizeListArray with entries of equal length.
62 Returns:
63 A 2D numpy array of shape (N, D).
64 """
65 vals = np.asarray(arr.values.to_numpy(), dtype=np.float64)
66 dim = len(arr[0])
67 return vals.reshape(-1, dim)
70def _sum_fixed(
71 fixed_list_array: pa.FixedSizeListArray,
72) -> tuple[np.ndarray, int]:
73 """Sum all FixedSizeList entries into a single numpy vector.
75 Args:
76 fixed_list_array: Input FixedSizeListArray.
78 Returns:
79 Tuple of (sum_vector, list_size).
80 """
81 vals = np.asarray(fixed_list_array.values.to_numpy(), dtype=np.float64)
82 list_size = len(fixed_list_array[0])
83 return vals.reshape(-1, list_size).sum(axis=0), list_size
86def _sum_scalar(arr: pa.Array) -> int:
87 """Sum all values in a pyarrow Array and return as int.
89 Args:
90 arr: Input pyarrow Array.
92 Returns:
93 Integer sum of all elements.
94 """
95 return int(np.asarray(arr.to_numpy()).sum())
98def _mmd_rbf(src_emb: np.ndarray, tgt_emb: np.ndarray, gamma: float) -> float:
99 """Compute Maximum Mean Discrepancy with an RBF kernel.
101 Uses the biased estimator matching the legacy implementation:
102 MMD^2 = mean(K_xx) + mean(K_yy) - 2 * mean(K_xy)
104 The RBF kernel uses non-squared Euclidean distance (matching legacy):
105 K(x, y) = exp(-gamma * ||x - y||)
107 Args:
108 src_emb: Source embeddings, shape (N, D).
109 tgt_emb: Target embeddings, shape (M, D).
110 gamma: RBF kernel coefficient.
112 Returns:
113 Scalar MMD^2 value.
114 """
115 m, n = len(src_emb), len(tgt_emb)
116 if m <= 1 or n <= 1:
117 return 0.0
119 # pairwise Euclidean distances (non-squared, matching legacy torch.cdist)
120 sq_src = np.sum(src_emb**2, axis=1, keepdims=True)
121 sq_tgt = np.sum(tgt_emb**2, axis=1, keepdims=True)
122 cross = src_emb @ tgt_emb.T
123 dist_xy = np.sqrt(np.maximum(0.0, sq_src - 2 * cross + sq_tgt.T))
125 k_xy = np.exp(-gamma * dist_xy)
127 # Within-source
128 sq_src_src = np.sum(src_emb**2, axis=1, keepdims=True)
129 cross_xx = src_emb @ src_emb.T
130 dist_xx = np.sqrt(np.maximum(0.0, sq_src_src - 2 * cross_xx + sq_src_src.T))
131 k_xx = np.exp(-gamma * dist_xx)
133 # Within-target
134 sq_tgt_tgt = np.sum(tgt_emb**2, axis=1, keepdims=True)
135 cross_yy = tgt_emb @ tgt_emb.T
136 dist_yy = np.sqrt(np.maximum(0.0, sq_tgt_tgt - 2 * cross_yy + sq_tgt_tgt.T))
137 k_yy = np.exp(-gamma * dist_yy)
139 mmd2 = k_xx.mean() + k_yy.mean() - 2 * k_xy.mean()
140 return float(max(mmd2, 0.0))
143def _mmd_poly(
144 src_emb: np.ndarray,
145 tgt_emb: np.ndarray,
146 degree: float,
147 gamma: float,
148 coefficient0: float,
149) -> float:
150 """Compute Maximum Mean Discrepancy with a polynomial kernel.
152 Uses the biased estimator matching the legacy implementation:
153 MMD^2 = mean(K_xx) + mean(K_yy) - 2 * mean(K_xy)
155 Polynomial kernel: K(a, b) = (gamma * <a, b> + coefficient0)^degree
157 Args:
158 src_emb: Source embeddings, shape (N, D).
159 tgt_emb: Target embeddings, shape (M, D).
160 degree: Polynomial degree.
161 gamma: Scaling factor for the dot product.
162 coefficient0: Bias term.
164 Returns:
165 Scalar MMD^2 value.
166 """
167 m, n = len(src_emb), len(tgt_emb)
168 if m <= 1 or n <= 1:
169 return 0.0
171 def _poly_kernel(a: np.ndarray, b: np.ndarray) -> np.ndarray:
172 return (gamma * (a @ b.T) + coefficient0) ** degree # type: ignore[no-any-return]
174 k_xx = _poly_kernel(src_emb, src_emb)
175 k_yy = _poly_kernel(tgt_emb, tgt_emb)
176 k_xy = _poly_kernel(src_emb, tgt_emb)
178 mmd2 = k_xx.mean() + k_yy.mean() - 2 * k_xy.mean()
179 return float(max(mmd2, 0.0))
182def _pad_distance(src_emb: np.ndarray, tgt_emb: np.ndarray, evaluator: str) -> float:
183 """Compute Proxy A-Distance (PAD) using a linear SVM.
185 Trains an SVM to discriminate source vs target, then returns
186 2 * (1 - 2 * error) where error is MSE or MAE of the classifier.
188 Args:
189 src_emb: Source embeddings, shape (N, D).
190 tgt_emb: Target embeddings, shape (M, D).
191 evaluator: Error metric, "mse" or "mae".
193 Returns:
194 PAD scalar value.
196 Raises:
197 ImportError: If scikit-learn is not installed.
198 """
199 from sklearn.calibration import CalibratedClassifierCV
200 from sklearn.svm import SVC
202 x_svm = np.vstack([src_emb, tgt_emb])
203 y = np.hstack([np.zeros(len(src_emb)), np.ones(len(tgt_emb))])
205 svm = CalibratedClassifierCV(
206 SVC(C=1, kernel="linear", random_state=42, verbose=0, gamma="auto"),
207 ensemble=False,
208 )
209 svm.fit(x_svm, y)
210 pred = svm.predict_proba(x_svm)
212 y_onehot = np.zeros_like(pred)
213 y_onehot[np.arange(len(y)), y.astype(int)] = 1
215 error = float(np.mean((pred - y_onehot) ** 2)) if evaluator == "mse" else float(np.mean(np.abs(pred - y_onehot)))
217 return 2.0 * (1.0 - 2.0 * error)
220class DomainGapProcessor(GapProcessor):
221 """Computes statistical distances between source and target
222 dataselections using image embeddings.
224 This processor works in two stages:
225 1. Dataset Summary: Aggregates high-dimensional embeddings into
226 compact statistics (mean, variance, outer products, histograms).
227 2. Delta Computation: Uses these summaries to calculate distance
228 metrics between a source and a target dataset.
230 Supported Delta Metrics:
231 - ``klmvn_diag``: KL divergence assuming a multivariate Normal
232 distribution with a diagonal covariance matrix.
233 - ``mmd_linear``: Maximum Mean Discrepancy with a linear kernel.
234 - ``mmd_rbf``: Maximum Mean Discrepancy with an RBF kernel.
235 - ``mmd_poly``: Maximum Mean Discrepancy with a polynomial kernel.
236 - ``fid``: Frechet Inception Distance.
237 - ``wasserstein_1d``: Average 1D Wasserstein distance across
238 embedding dimensions, approximated via histograms.
239 - ``pad``: Proxy A-Distance via linear SVM.
240 - ``cmd``: Central Moment Discrepancy (multi-layer only).
241 """
243 def __init__(
244 self,
245 name: str = "domain_gap",
246 config: dict[str, Any] | None = None,
247 ):
248 """Initialize the domain gap processor.
250 Args:
251 name: Unique name of the processor instance.
252 config: Configuration dictionary containing:
253 - input:
254 - embedding_col: Column name containing embeddings (default: "embedding").
255 - embedding_cols: List of column names for multi-layer metrics (CMD).
256 - summary:
257 - collect_sum_outer: Whether to compute outer products (needed for FID).
258 - collect_hist_1d: Whether to compute histograms (needed for Wasserstein).
259 - hist_dims: Number of dimensions to histogram.
260 - hist_bins: Number of bins per histogram.
261 - store_embeddings: Whether to store raw embeddings for full-data metrics.
262 - delta:
263 - metric: Target metric name.
264 - k: Number of moments (CMD only, default 5).
265 - feature_weights: Per-layer weights (CMD only).
266 - kernel_params: Kernel parameters (MMD-RBF/Poly).
267 - method:
268 - evaluator: Error metric for PAD ("mse" or "mae").
269 """
270 super().__init__(name, config)
272 cfg = DomainGapProcessorConfig.model_validate({**self.config, "name": self.name})
273 self._validate_and_set_columns(cfg)
274 self.delta_metric = cfg.distance.metric.lower()
275 self.is_cmd = self.delta_metric == "cmd"
276 self._configure_summary(cfg)
277 self._configure_cmd(cfg)
278 self._configure_kernel_and_pad(cfg)
280 def _validate_and_set_columns(self, cfg: DomainGapProcessorConfig) -> None:
281 """Validate and set embedding column configuration."""
282 if not cfg.columns.input: 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true
283 raise ValueError("columns.input is required for domain_gap processor")
284 self.embedding_col = cfg.columns.input[0]
285 self.embedding_cols = list(cfg.columns.input)
287 def _resolve_summary_bool(self, cfg: DomainGapProcessorConfig, attr: str, default: bool) -> bool:
288 """Resolve a summary boolean config value with a fallback default."""
289 if cfg.summary and getattr(cfg.summary, attr, None) is not None:
290 return bool(getattr(cfg.summary, attr))
291 return default
293 def _configure_summary(self, cfg: DomainGapProcessorConfig) -> None:
294 """Configure summary collection flags and histogram parameters."""
295 full_data_metrics = {"mmd_rbf", "mmd_poly", "pad", "cmd"}
296 auto_store_emb = self.delta_metric in full_data_metrics
297 auto_sum_outer = self.delta_metric == "fid"
299 self.collect_sum_outer = self._resolve_summary_bool(cfg, "collect_sum_outer", auto_sum_outer)
300 self.store_embeddings = self._resolve_summary_bool(cfg, "store_embeddings", auto_store_emb)
302 self.hist_dims = 64
303 self.hist_bins = 32
304 self.hist_range = (-3.0, 3.0)
305 if cfg.summary and cfg.summary.histogram:
306 self.collect_hist_1d = True
307 self.hist_dims = cfg.summary.histogram.dims
308 self.hist_bins = cfg.summary.histogram.bins
309 self.hist_range = (
310 float(cfg.summary.histogram.range[0]),
311 float(cfg.summary.histogram.range[1]),
312 )
313 else:
314 self.collect_hist_1d = self.delta_metric == "wasserstein_1d"
316 def _configure_cmd(self, cfg: DomainGapProcessorConfig) -> None:
317 """Configure CMD-specific parameters."""
318 if not self.is_cmd:
319 return
320 self.cmd_k = cfg.distance.k or 5
321 self.cmd_embedding_cols = cfg.columns.input if cfg.columns and cfg.columns.input else [self.embedding_col]
322 self.cmd_feature_weights = list(cfg.distance.feature_weights or [1.0] * len(self.cmd_embedding_cols))
324 def _configure_kernel_and_pad(self, cfg: DomainGapProcessorConfig) -> None:
325 """Configure kernel parameters and PAD evaluator."""
326 self.kernel_params = dict(cfg.distance.kernel_params) if cfg.distance.kernel_params else {}
327 self.pad_evaluator = cfg.distance.evaluator or "mse"
328 self.epsilon = cfg.distance.epsilon
329 self.klmvn_var_eps = cfg.distance.klmvn_var_eps
331 def check_config(self) -> None:
332 """Validate configuration.
334 Kept for backward compatibility. All config is already
335 parsed in ``__init__``.
336 """
338 def _embedding_cols(self) -> list[str]:
339 """Get the embedding columns based on metric type.
341 For CMD, returns all configured embedding columns.
342 For other metrics, returns the single primary embedding column.
344 Returns:
345 List of embedding column names.
346 """
347 if self.is_cmd:
348 return self.cmd_embedding_cols
349 return [self.embedding_col]
351 @override
352 def needed_columns(self) -> list[str]:
353 """Return the list of columns required for domain gap computation.
355 Returns:
356 List of embedding column names needed for the configured metric.
357 """
358 return self._embedding_cols()
360 # utils functions
361 @staticmethod
362 def _resolve_device(device: str) -> str:
363 """Resolve ``"auto"`` to CUDA if available, else CPU."""
364 if device == "auto":
365 return "cuda" if torch.cuda.is_available() else "cpu"
366 return device
368 def _resolve_embedding_patterns(self, available: list[str]) -> None:
369 """Resolve wildcard patterns in embedding column config against available columns.
370 Updates ``embedding_col`` and ``embedding_cols`` / ``cmd_embedding_cols`` in place.
371 """
372 if has_pattern(self.embedding_col):
373 matched = resolve_include_exclude([self.embedding_col], None, available)
374 if matched: 374 ↛ exitline 374 didn't return from function '_resolve_embedding_patterns' because the condition on line 374 was always true
375 self.embedding_col = matched[0]
376 self.embedding_cols = matched
377 if self.is_cmd: 377 ↛ 378line 377 didn't jump to line 378 because the condition on line 377 was never true
378 self.cmd_embedding_cols = list(matched)
379 self.cmd_feature_weights = [1.0] * len(matched)
381 @override
382 def compute_batch_metric(self, features: dict[str, pa.Array]) -> dict[str, pa.Array]:
383 """Reduce a batch of embeddings into summary statistics.
385 For single-column metrics, computes count, sum, sum_sq, and
386 optionally sum_outer, hist_counts, and raw embeddings.
388 For CMD, computes raw moments up to order k for each embedding
389 column.
391 Args:
392 features: Dictionary of feature arrays from the batch.
394 Returns:
395 Dictionary of aggregated statistics per batch.
396 """
397 self._resolve_embedding_patterns(list(features.keys()))
398 if self.is_cmd:
399 return self._compute_batch_metric_cmd(features)
401 emb = features.get(self.embedding_col)
402 if emb is None or not isinstance(emb, pa.FixedSizeListArray):
403 return {}
405 num_samples = len(emb)
406 embed_dim = len(emb[0])
407 flat_values = emb.values
408 emb_matrix = np.asarray(flat_values.to_numpy()).reshape(num_samples, embed_dim)
410 out: dict[str, pa.Array] = {}
411 out["count"] = pa.array([num_samples], type=pa.int64())
412 sum_vec = emb_matrix.sum(axis=0).astype(np.float64)
413 out["sum"] = pa.FixedSizeListArray.from_arrays(pa.array(sum_vec), embed_dim)
414 sum_sq_vec = (emb_matrix * emb_matrix).sum(axis=0).astype(np.float64)
415 out["sum_sq"] = pa.FixedSizeListArray.from_arrays(pa.array(sum_sq_vec), embed_dim)
417 # optional: sum_outer for FID
418 if self.collect_sum_outer:
419 sum_outer_product = (emb_matrix.T @ emb_matrix).reshape(-1).astype(np.float64)
420 outer_dim = embed_dim * embed_dim
421 out["sum_outer"] = pa.FixedSizeListArray.from_arrays(pa.array(sum_outer_product), outer_dim)
423 # optional: histograms for Wasserstein-1D
424 if self.collect_hist_1d:
425 use_dims = min(embed_dim, self.hist_dims)
426 low, high = self.hist_range
427 hist_list: list[np.ndarray] = []
428 for j in range(use_dims):
429 hist_1d, _ = np.histogram(emb_matrix[:, j], bins=self.hist_bins, range=(low, high))
430 hist_list.append(hist_1d.astype(np.int64))
431 hist_all = np.stack(hist_list, axis=0).reshape(-1)
432 out["hist_counts"] = pa.FixedSizeListArray.from_arrays(pa.array(hist_all), self.hist_bins * use_dims)
434 # optional: raw embeddings for full-data metrics (MMD-RBF, MMD-Poly, PAD)
435 if self.store_embeddings:
436 out["__emb__"] = emb
438 return out
440 def _compute_batch_metric_cmd(self, features: dict[str, pa.Array]) -> dict[str, pa.Array]:
441 """Compute per-batch raw moment power sums for CMD.
443 For each CMD column, applies sigmoid and accumulates sum(x^j)
444 for j=1..k (raw moment sums). These are aggregated across batches
445 in _compute_cmd_aggregate and converted to central moments in
446 _compute_delta_cmd.
448 Args:
449 features: Dictionary of feature arrays from the batch.
451 Returns:
452 Dictionary of power sums and counts per batch.
453 """
454 out: dict[str, pa.Array] = {}
455 for col in self.cmd_embedding_cols:
456 emb = features.get(col)
457 if emb is None or not isinstance(emb, pa.FixedSizeListArray): 457 ↛ 458line 457 didn't jump to line 458 because the condition on line 457 was never true
458 continue
459 mat = _fixed_to_matrix(emb)
460 batch_n = len(mat)
461 if batch_n == 0: 461 ↛ 462line 461 didn't jump to line 462 because the condition on line 461 was never true
462 continue
464 # Apply sigmoid (matching v1 behavior)
465 mat = 1.0 / (1.0 + np.exp(-mat))
467 # Per-channel spatial reshaping to match v1's moment computation.
468 # v1 treats each individual spatial element as a sample, computing
469 # moments over all C x H x W values per channel across all images.
470 # Reshape flattened (N, C*H*W) → (N, C, H*W) so we can sum over
471 # both N and H*W, matching v1's element-wise treatment.
472 channels = self._resolve_cmd_channels(col, mat.shape[1], features)
473 hw = mat.shape[1] // channels
474 mat = mat.reshape(-1, channels, hw)
476 out[f"cmd_{col}_n"] = pa.array([batch_n * hw], type=pa.int64())
478 # Raw moment sums: sum(x^j) for j=1..k over all spatial elements
479 for j in range(1, self.cmd_k + 1):
480 power_sum = np.power(mat, j).sum(axis=(0, 2)).astype(np.float64)
481 out[f"cmd_{col}_sum_{j}"] = pa.FixedSizeListArray.from_arrays(pa.array(power_sum), len(power_sum))
483 return out
485 def _resolve_cmd_channels(self, col: str, flattened_dim: int, features: dict[str, pa.Array]) -> int:
486 """Determine number of channels for CMD spatial moment computation.
488 Tries to read channel count from metadata column '{col}_channels'.
489 Falls back to legacy ResNet-18 dimension lookup if metadata unavailable.
491 Args:
492 col: Embedding column name.
493 flattened_dim: Total flattened dimension of embeddings.
494 features: Dictionary of feature arrays (may contain channels column).
496 Returns:
497 Number of channels (C) for reshaping (N, C*H*W) -> (N, C, H*W).
499 Raises:
500 ValueError: If channels cannot be determined from metadata or lookup.
501 """
502 channels_col = f"{col}_channels"
503 channels_arr = features.get(channels_col)
504 if channels_arr is not None and len(channels_arr) > 0:
505 c = int(channels_arr[0].as_py())
506 if flattened_dim % c == 0: 506 ↛ 509line 506 didn't jump to line 509 because the condition on line 506 was always true
507 return c
509 _c = _CMD_RESNET18_EMBDIM_CHANNELS.get(flattened_dim)
510 if _c is not None and flattened_dim % _c == 0:
511 return _c
513 raise ValueError(
514 f"Cannot determine channels for embedding column '{col}' "
515 f"(flattened_dim={flattened_dim}). "
516 f"The image_embedding processor did not produce a "
517 f"'{channels_col}' metadata column, and the dimension "
518 f"is not in the legacy lookup table. "
519 f"Provide 'cmd_channels' in the domain_gap delta config "
520 f"or ensure the image_embedding processor outputs "
521 f"'{channels_col}'."
522 )
524 @override
525 def compute(self, batch_metrics: dict[str, pa.Array]) -> dict[str, pa.Array]:
526 """Aggregate batch-level summary statistics into global dataselection statistics.
528 For summary-based metrics, aggregates count, sum, sum_sq, etc.
529 For CMD, aggregates per-batch power sums for later central moment
530 computation in compute_delta.
531 For store_embeddings, concatenates raw embedding arrays.
533 Args:
534 batch_metrics: Dictionary containing batch-level statistics.
536 Returns:
537 Dictionary containing aggregated dataset-level statistics.
538 """
539 if not batch_metrics:
540 return {}
542 if self.is_cmd:
543 return self._compute_cmd_aggregate(batch_metrics)
545 out: dict[str, pa.Array] = {}
547 # count
548 if "count" not in batch_metrics:
549 return {}
550 total_n = _sum_scalar(batch_metrics["count"])
551 out["count"] = pa.array([total_n], type=pa.int64())
553 # sum / sum_sq
554 if "sum" in batch_metrics: 554 ↛ 557line 554 didn't jump to line 557 because the condition on line 554 was always true
555 sum_vec, list_size = _sum_fixed(batch_metrics["sum"])
556 out["sum"] = pa.FixedSizeListArray.from_arrays(pa.array(sum_vec), list_size)
557 if "sum_sq" in batch_metrics: 557 ↛ 562line 557 didn't jump to line 562 because the condition on line 557 was always true
558 sum_sq_vec, list_size2 = _sum_fixed(batch_metrics["sum_sq"])
559 out["sum_sq"] = pa.FixedSizeListArray.from_arrays(pa.array(sum_sq_vec), list_size2)
561 # optional sum_outer
562 if "sum_outer" in batch_metrics:
563 so_vals = np.asarray(batch_metrics["sum_outer"].values.to_numpy(), dtype=np.float64)
564 outer_dim = len(batch_metrics["sum_outer"][0])
565 out["sum_outer"] = pa.FixedSizeListArray.from_arrays(
566 pa.array(so_vals.reshape(-1, outer_dim).sum(axis=0)), outer_dim
567 )
569 # optional hist_counts
570 if "hist_counts" in batch_metrics:
571 h_vals = np.asarray(batch_metrics["hist_counts"].values.to_numpy(), dtype=np.int64)
572 h_len = len(batch_metrics["hist_counts"][0])
573 out["hist_counts"] = pa.FixedSizeListArray.from_arrays(
574 pa.array(h_vals.reshape(-1, h_len).sum(axis=0)), h_len
575 )
577 # raw embeddings for full-data metrics
578 if self.store_embeddings and "__emb__" in batch_metrics:
579 vals = np.asarray(batch_metrics["__emb__"].values.to_numpy(), dtype=np.float64)
580 dim = len(batch_metrics["__emb__"][0])
581 out["__emb__"] = pa.FixedSizeListArray.from_arrays(pa.array(vals), dim)
583 return out
585 def _compute_cmd_aggregate(self, batch_metrics: dict[str, pa.Array]) -> dict[str, pa.Array]:
586 """Aggregate CMD power sums across batches.
588 Args:
589 batch_metrics: Dictionary containing per-batch power sums.
591 Returns:
592 Dictionary with aggregated power sums and total count per layer.
593 """
594 out: dict[str, pa.Array] = {}
595 for col in self.cmd_embedding_cols:
596 n_key = f"cmd_{col}_n"
597 if n_key not in batch_metrics:
598 continue
600 total_n = _sum_scalar(batch_metrics[n_key])
601 if total_n == 0:
602 continue
603 out[n_key] = pa.array([total_n], type=pa.int64())
605 for j in range(1, self.cmd_k + 1):
606 sum_key = f"cmd_{col}_sum_{j}"
607 if sum_key in batch_metrics: 607 ↛ 605line 607 didn't jump to line 605 because the condition on line 607 was always true
608 sum_vec, dim = _sum_fixed(batch_metrics[sum_key])
609 out[sum_key] = pa.FixedSizeListArray.from_arrays(pa.array(sum_vec), dim)
611 return out
613 @override
614 def compute_delta(self, source: dict[str, pa.Array], target: dict[str, pa.Array]) -> dict[str, pa.Array]:
615 """Calculate the domain gap metric between source and target statistics.
617 Args:
618 source: Dataselection statistics from the source dataset.
619 target: Dataselection statistics from the target dataset.
621 Returns:
622 Dictionary containing the calculated metric value.
623 """
624 metric = self.delta_metric
626 if self.is_cmd:
627 return self._compute_delta_cmd(source, target)
629 if metric in {"klmvn_diag", "mmd_linear", "fid"}:
630 return self._compute_delta_summary(source, target, metric)
632 if metric == "wasserstein_1d":
633 return self._compute_delta_wasserstein(source, target)
635 if metric == "mmd_rbf":
636 return self._compute_delta_mmd_rbf(source, target)
638 if metric == "mmd_poly":
639 return self._compute_delta_mmd_poly(source, target)
641 if metric == "pad":
642 return self._compute_delta_pad(source, target)
644 return {
645 "metric": pa.array([metric]),
646 "note": pa.array(["unsupported metric or invalid inputs"]),
647 }
649 @staticmethod
650 def _compute_mmd_linear(mean_src: np.ndarray, mean_tgt: np.ndarray) -> dict[str, pa.Array]:
651 diff = mean_src - mean_tgt
652 val = float(np.dot(diff, diff))
653 return {"mmd_linear": pa.array([val], type=pa.float64())}
655 def _compute_klmvn_diag(
656 self,
657 mean_src: np.ndarray,
658 mean_tgt: np.ndarray,
659 var_src: np.ndarray,
660 var_tgt: np.ndarray,
661 ) -> dict[str, pa.Array]:
662 if self.klmvn_var_eps > 0:
663 mean_var = 0.5 * (var_src.mean() + var_tgt.mean())
664 var_src = var_src + self.klmvn_var_eps * mean_var
665 var_tgt = var_tgt + self.klmvn_var_eps * mean_var
666 term_var = np.sum(var_src / var_tgt - 1.0 - np.log(var_src / var_tgt))
667 term_mean = np.sum((mean_tgt - mean_src) ** 2 / var_tgt)
668 val = 0.5 * (term_var + term_mean)
669 return {"klmvn_diag": pa.array([float(val)], type=pa.float64())}
671 @staticmethod
672 def _compute_fid(
673 mean_src: np.ndarray,
674 mean_tgt: np.ndarray,
675 source: dict[str, pa.Array],
676 target: dict[str, pa.Array],
677 n_src: int,
678 n_tgt: int,
679 eps: float,
680 ) -> dict[str, pa.Array]:
681 from scipy.linalg import sqrtm
683 sum_outer_src = _sum_fixed(source["sum_outer"])[0]
684 sum_outer_tgt = _sum_fixed(target["sum_outer"])[0]
685 embed_dim = int(np.sqrt(sum_outer_src.size))
686 cov_src = (sum_outer_src.reshape(embed_dim, embed_dim) / n_src) - np.outer(mean_src, mean_src)
687 cov_tgt = (sum_outer_tgt.reshape(embed_dim, embed_dim) / n_tgt) - np.outer(mean_tgt, mean_tgt)
689 cov_src += eps * np.eye(embed_dim)
690 cov_tgt += eps * np.eye(embed_dim)
691 covmean = sqrtm(cov_src.dot(cov_tgt))
692 if np.iscomplexobj(covmean):
693 covmean = covmean.real
695 diff = mean_src - mean_tgt
696 fid = diff.dot(diff) + np.trace(cov_src) + np.trace(cov_tgt) - 2 * np.trace(covmean)
697 return {"fid": pa.array([float(abs(fid))], type=pa.float64())}
699 def _compute_delta_summary(
700 self,
701 source: dict[str, pa.Array],
702 target: dict[str, pa.Array],
703 metric: str,
704 ) -> dict[str, pa.Array]:
705 """Compute KLMVN, MMD-Linear, or FID from summary statistics.
707 Args:
708 source: Source dataset statistics.
709 target: Target dataset statistics.
710 metric: One of "klmvn_diag", "mmd_linear", "fid".
712 Returns:
713 Dictionary with the metric value.
714 """
715 need: set[str] = {"count", "sum"}
716 if metric in {"klmvn_diag", "fid"}:
717 need |= {"sum_sq"}
718 if metric == "fid":
719 need |= {"sum_outer"}
720 for dataset_stats, name in ((source, "source"), (target, "target")):
721 if not need.issubset(dataset_stats.keys()):
722 return {
723 "metric": pa.array([metric]),
724 "note": pa.array([f"missing keys in {name}: {sorted(need)}"]),
725 }
727 n_src = _sum_scalar(source["count"])
728 n_tgt = _sum_scalar(target["count"])
729 if n_src <= 0 or n_tgt <= 0:
730 return {
731 "metric": pa.array([metric]),
732 "note": pa.array(["empty summaries"]),
733 }
735 mean_src = _sum_fixed(source["sum"])[0] / n_src
736 mean_tgt = _sum_fixed(target["sum"])[0] / n_tgt
738 if metric == "mmd_linear":
739 return self._compute_mmd_linear(mean_src, mean_tgt)
741 var_src = np.maximum(_sum_fixed(source["sum_sq"])[0] / n_src - mean_src * mean_src, 1e-9)
742 var_tgt = np.maximum(_sum_fixed(target["sum_sq"])[0] / n_tgt - mean_tgt * mean_tgt, 1e-9)
744 if metric == "klmvn_diag":
745 return self._compute_klmvn_diag(mean_src, mean_tgt, var_src, var_tgt)
747 if metric == "fid": 747 ↛ 750line 747 didn't jump to line 750 because the condition on line 747 was always true
748 return self._compute_fid(mean_src, mean_tgt, source, target, n_src, n_tgt, self.epsilon)
750 return {"metric": pa.array([metric]), "note": pa.array(["unreachable"])}
752 def _compute_delta_wasserstein(
753 self, source: dict[str, pa.Array], target: dict[str, pa.Array]
754 ) -> dict[str, pa.Array]:
755 """Compute 1D Wasserstein distance from histogram summaries.
757 Args:
758 source: Source dataset statistics.
759 target: Target dataset statistics.
761 Returns:
762 Dictionary with wasserstein_1d value.
763 """
764 if "hist_counts" not in source or "hist_counts" not in target:
765 return {
766 "metric": pa.array(["wasserstein_1d"]),
767 "note": pa.array(["missing hist_counts"]),
768 }
769 h_src = np.asarray(source["hist_counts"].values.to_numpy(), dtype=np.int64)
770 h_tgt = np.asarray(target["hist_counts"].values.to_numpy(), dtype=np.int64)
771 use_dims = self.hist_dims
772 bins = self.hist_bins
773 if h_src.size != h_tgt.size or h_src.size != bins * use_dims: 773 ↛ 774line 773 didn't jump to line 774 because the condition on line 773 was never true
774 return {
775 "metric": pa.array(["wasserstein_1d"]),
776 "note": pa.array(["hist_counts length mismatch"]),
777 }
778 width = (self.hist_range[1] - self.hist_range[0]) / bins
779 total = 0.0
780 used = 0
781 for j in range(use_dims):
782 h_src_slice = h_src[j * bins : (j + 1) * bins].astype(np.float64)
783 h_tgt_slice = h_tgt[j * bins : (j + 1) * bins].astype(np.float64)
784 if h_src_slice.sum() == 0 and h_tgt_slice.sum() == 0:
785 continue
786 prob_src = h_src_slice / max(1.0, h_src_slice.sum())
787 prob_tgt = h_tgt_slice / max(1.0, h_tgt_slice.sum())
788 cdf_src = np.cumsum(prob_src)
789 cdf_tgt = np.cumsum(prob_tgt)
790 total += float(np.sum(np.abs(cdf_src - cdf_tgt)) * width)
791 used += 1
792 val = total / max(1, used)
793 return {"wasserstein_1d": pa.array([val], type=pa.float64())}
795 def _compute_delta_mmd_rbf(self, source: dict[str, pa.Array], target: dict[str, pa.Array]) -> dict[str, pa.Array]:
796 """Compute MMD with RBF kernel from stored embeddings.
798 Args:
799 source: Source dataset statistics including "__emb__".
800 target: Target dataset statistics including "__emb__".
802 Returns:
803 Dictionary with mmd_rbf value.
804 """
805 if "__emb__" not in source or "__emb__" not in target:
806 return {
807 "metric": pa.array(["mmd_rbf"]),
808 "note": pa.array([_MISSING_EMB_MSG]),
809 }
810 src = _fixed_to_matrix(source["__emb__"])
811 tgt = _fixed_to_matrix(target["__emb__"])
812 gamma = float(self.kernel_params.get("gamma", 1.0))
813 val = _mmd_rbf(src, tgt, gamma)
814 return {"mmd_rbf": pa.array([val], type=pa.float64())}
816 def _compute_delta_mmd_poly(self, source: dict[str, pa.Array], target: dict[str, pa.Array]) -> dict[str, pa.Array]:
817 """Compute MMD with polynomial kernel from stored embeddings.
819 Args:
820 source: Source dataset statistics including "__emb__".
821 target: Target dataset statistics including "__emb__".
823 Returns:
824 Dictionary with mmd_poly value.
825 """
826 if "__emb__" not in source or "__emb__" not in target:
827 return {
828 "metric": pa.array(["mmd_poly"]),
829 "note": pa.array([_MISSING_EMB_MSG]),
830 }
831 src = _fixed_to_matrix(source["__emb__"])
832 tgt = _fixed_to_matrix(target["__emb__"])
833 degree = float(self.kernel_params.get("degree", 3.0))
834 gamma = float(self.kernel_params.get("gamma", 1.0))
835 coefficient0 = float(self.kernel_params.get("coefficient0", 1.0))
836 val = _mmd_poly(src, tgt, degree, gamma, coefficient0)
837 return {"mmd_poly": pa.array([val], type=pa.float64())}
839 def _compute_delta_pad(self, source: dict[str, pa.Array], target: dict[str, pa.Array]) -> dict[str, pa.Array]:
840 """Compute Proxy A-Distance from stored embeddings.
842 Args:
843 source: Source dataset statistics including "__emb__".
844 target: Target dataset statistics including "__emb__".
846 Returns:
847 Dictionary with pad value.
848 """
849 if "__emb__" not in source or "__emb__" not in target:
850 return {
851 "metric": pa.array(["pad"]),
852 "note": pa.array([_MISSING_EMB_MSG]),
853 }
854 src = _fixed_to_matrix(source["__emb__"])
855 tgt = _fixed_to_matrix(target["__emb__"])
856 val = _pad_distance(src, tgt, self.pad_evaluator)
857 return {"pad": pa.array([val], type=pa.float64())}
859 def _compute_delta_cmd(self, source: dict[str, pa.Array], target: dict[str, pa.Array]) -> dict[str, pa.Array]:
860 """Compute Central Moment Discrepancy between source and target.
862 Computes per-layer raw moments from power sums, converts to central
863 moments, and compares them using Euclidean distance (matching v1's
864 RMSELoss). Weighted averaging follows v1's formula:
865 layer_loss = (1/k) * sum(rmse(moment) for moment in 0..k-1)
866 total_loss = sum(weight * layer_loss for each layer)
867 cmd = total_loss / sum(weights)
869 Args:
870 source: Source dataset statistics including cmd_{col}_n and
871 cmd_{col}_sum_{j} for j=1..k.
872 target: Target dataset statistics (same keys as source).
874 Returns:
875 Dictionary with cmd value.
876 """
877 total_loss = 0.0
878 total_weight = 0.0
879 debug_data: dict[str, np.ndarray] | None = {} if _debug_enabled() else None
881 for col, weight in zip(self.cmd_embedding_cols, self.cmd_feature_weights, strict=True):
882 if weight == 0:
883 continue
885 layer_result = self._compute_layer_cmd(col, source, target, debug_data)
886 if layer_result is None:
887 continue
889 layer_loss, layer_debug = layer_result
890 total_weight += weight
891 total_loss += weight * layer_loss
893 if debug_data is not None and layer_debug is not None: 893 ↛ 894line 893 didn't jump to line 894 because the condition on line 893 was never true
894 debug_data.update(layer_debug)
896 if debug_data is not None: 896 ↛ 897line 896 didn't jump to line 897 because the condition on line 896 was never true
897 tmp_path = str(Path(tempfile.gettempdir()) / f"debug_moments_{os.getpid()}.npz")
898 np.savez_compressed(tmp_path, **debug_data) # type: ignore[arg-type]
900 if total_weight == 0:
901 return {
902 "metric": pa.array(["cmd"]),
903 "note": pa.array(["no valid layers"]),
904 }
906 final_loss = total_loss / total_weight
907 return {"cmd": pa.array([final_loss], type=pa.float64())}
909 def _collect_raw_moments(
910 self,
911 col: str,
912 source: dict[str, pa.Array],
913 target: dict[str, pa.Array],
914 all_j: list[int],
915 n_src: int,
916 n_tgt: int,
917 ) -> tuple[list[np.ndarray], list[np.ndarray]]:
918 """Collect raw moments from power sums for a single layer.
920 Args:
921 col: Layer column name.
922 source: Source statistics.
923 target: Target statistics.
924 all_j: List of moment orders.
925 n_src: Number of source samples.
926 n_tgt: Number of target samples.
928 Returns:
929 Tuple of (src_raw, tgt_raw) moment lists.
930 """
931 src_raw: list[np.ndarray] = []
932 tgt_raw: list[np.ndarray] = []
933 for j in all_j:
934 src_sum, _ = _sum_fixed(source[f"cmd_{col}_sum_{j}"])
935 tgt_sum, _ = _sum_fixed(target[f"cmd_{col}_sum_{j}"])
936 src_raw.append(src_sum / n_src)
937 tgt_raw.append(tgt_sum / n_tgt)
938 return src_raw, tgt_raw
940 def _compute_cmd_loss(
941 self,
942 src_raw: list[np.ndarray],
943 tgt_raw: list[np.ndarray],
944 mu_src: np.ndarray,
945 mu_tgt: np.ndarray,
946 ) -> float:
947 """Convert raw moments to central moments and compute CMD distance.
949 Args:
950 src_raw: Source raw moments.
951 tgt_raw: Target raw moments.
952 mu_src: Source mean.
953 mu_tgt: Target mean.
955 Returns:
956 Layer CMD loss value.
957 """
958 src_cm: list[np.ndarray] = [mu_src]
959 tgt_cm: list[np.ndarray] = [mu_tgt]
960 for order in range(2, self.cmd_k + 1):
961 cm_src = np.zeros_like(mu_src)
962 cm_tgt = np.zeros_like(mu_tgt)
963 for i in range(order + 1):
964 coeff = float(comb(order, i))
965 if i == 0:
966 raw_src = np.array(1.0)
967 raw_tgt = np.array(1.0)
968 else:
969 raw_src = src_raw[i - 1]
970 raw_tgt = tgt_raw[i - 1]
971 cm_src += coeff * raw_src * ((-mu_src) ** (order - i))
972 cm_tgt += coeff * raw_tgt * ((-mu_tgt) ** (order - i))
973 src_cm.append(cm_src)
974 tgt_cm.append(cm_tgt)
975 layer_loss = 0.0
976 for t in range(self.cmd_k):
977 diff = src_cm[t] - tgt_cm[t]
978 dist = float(np.sqrt(np.sum(diff**2)))
979 layer_loss += dist
980 layer_loss /= self.cmd_k
981 return layer_loss
983 def _compute_layer_cmd(
984 self,
985 col: str,
986 source: dict[str, pa.Array],
987 target: dict[str, pa.Array],
988 debug_data: dict[str, np.ndarray] | None = None,
989 ) -> tuple[float, dict[str, np.ndarray]] | None:
990 """Compute CMD loss for a single embedding layer.
992 Args:
993 col: Layer column name.
994 source: Source statistics.
995 target: Target statistics.
996 debug_data: Optional debug dict to populate.
998 Returns:
999 Tuple of (layer_loss, debug_entries) or None if layer is invalid.
1000 """
1001 n_src_key = f"cmd_{col}_n"
1002 n_tgt_key = f"cmd_{col}_n"
1003 if n_src_key not in source or n_tgt_key not in target:
1004 return None
1006 n_src = int(source[n_src_key].to_numpy()[0])
1007 n_tgt = int(target[n_tgt_key].to_numpy()[0])
1008 if n_src <= 0 or n_tgt <= 0:
1009 return None
1011 all_j = list(range(1, self.cmd_k + 1))
1012 if not all(f"cmd_{col}_sum_{j}" in source and f"cmd_{col}_sum_{j}" in target for j in all_j):
1013 return None
1015 src_raw, tgt_raw = self._collect_raw_moments(col, source, target, all_j, n_src, n_tgt)
1017 layer_debug: dict[str, np.ndarray] = {}
1018 if debug_data is not None: 1018 ↛ 1019line 1018 didn't jump to line 1019 because the condition on line 1018 was never true
1019 layer_key = col
1020 for prefix in ["image_embedding_cmd_", "image_embedding_"]:
1021 if col.startswith(prefix):
1022 layer_key = col[len(prefix) :]
1023 break
1024 layer_debug[f"{layer_key}/mean_src"] = src_raw[0]
1025 layer_debug[f"{layer_key}/mean_tgt"] = tgt_raw[0]
1026 layer_debug[f"{layer_key}/raw_moment2_src"] = src_raw[1]
1027 layer_debug[f"{layer_key}/raw_moment2_tgt"] = tgt_raw[1]
1028 layer_debug[f"{layer_key}/n_src"] = np.array([n_src], dtype=np.int64)
1029 layer_debug[f"{layer_key}/n_tgt"] = np.array([n_tgt], dtype=np.int64)
1031 mu_src = src_raw[0]
1032 mu_tgt = tgt_raw[0]
1033 layer_loss = self._compute_cmd_loss(src_raw, tgt_raw, mu_src, mu_tgt)
1034 return (layer_loss, layer_debug)