init repo

.gitignore

@@ -0,0 +1,2 @@
+data/
+tests/

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.10-slim
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    libgl1 \
+    libglx-mesa0 \
+    libglib2.0-0 \
+    libxcb1 \
+    libxkbcommon-x11-0 \
+    libxrender1 \
+    && rm -rf /var/lib/apt/lists/*
+
+
+RUN useradd -m -u 1001 dvarfe
+USER dvarfe
+ENV PATH="/home/dvarfe/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=dvarfe ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=dvarfe . /app
+CMD ["python", "dashboard_app.py"]

analysis/cache_utils.py

@@ -0,0 +1,787 @@
+"""
+Утилиты для кэширования и загрузки активаций SAE.
+"""
+
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from analysis.models import _iqa_activations
+
+
+def _df_mb(df: pd.DataFrame) -> float:
+    """Объём памяти DataFrame в МБ."""
+    return df.memory_usage(deep=True).sum() / 1024 ** 2
+
+
+def _sparse_mb(mat: sp.csr_matrix) -> float:
+    """Объём памяти CSR-матрицы (данные + индексы) в МБ."""
+    return (mat.data.nbytes + mat.indices.nbytes + mat.indptr.nbytes) / 1024 ** 2
+
+
+def _cache_paths(base_path: str) -> Tuple[str, str, str]:
+    """
+    Возвращает пути к трём файлам кэша из базового пути.
+
+    Пример:
+        'cache/kadid_acts.feather'
+          → ('cache/kadid_acts_meta.feather', 'cache/kadid_acts_codes.npz',
+              'cache/kadid_acts_steps.npz')
+    """
+    p = Path(base_path)
+    stem = p.stem.removesuffix('.feather')
+    return (
+        str(p.parent / f'{stem}_meta.feather'),
+        str(p.parent / f'{stem}_codes.npz'),
+        str(p.parent / f'{stem}_steps.npz'),
+    )
+
+
+def _pristine_cache_paths(base_path: str) -> Tuple[str, str, str]:
+    """Return paths for pristine cache files derived from base cache path."""
+    meta_path, codes_path, steps_path = _cache_paths(base_path)
+    return (
+        meta_path.replace('.feather', '_pristine.feather'),
+        codes_path.replace('.npz', '_pristine.npz'),
+        steps_path.replace('.npz', '_pristine.npz'),
+    )
+
+
+def load_parquet_cache(cache_path: Optional[str], *, label: str = 'cache') -> Optional[pd.DataFrame]:
+    """Load cached parquet table if present."""
+    if cache_path is None:
+        return None
+
+    cache = Path(cache_path)
+    if not cache.exists():
+        return None
+
+    print(f'[cache] Loading {label} from {cache}')
+    return pd.read_parquet(cache)
+
+
+def save_parquet_cache(df: pd.DataFrame, cache_path: Optional[str], *, label: str = 'cache') -> None:
+    """Persist a dataframe to parquet cache if path is provided."""
+    if cache_path is None:
+        return
+
+    cache = Path(cache_path)
+    cache.parent.mkdir(parents=True, exist_ok=True)
+    df.to_parquet(cache)
+    print(f'[cache] Saved {label} to {cache}')
+
+
+_PATCH_LABEL_META_KEYS = frozenset({'dist_type', 'dist_group'})
+
+
+def _patch_labels_to_dist_meta(
+    patch_labels: np.ndarray,
+    *,
+    label_to_dist_type: Dict[int, str],
+    label_to_dist_group: Dict[str, str],
+) -> Tuple[List[str], List[str]]:
+    flat_labels = patch_labels.reshape(-1)
+    dist_types = [label_to_dist_type.get(int(label_id), 'background') for label_id in flat_labels]
+    dist_groups = [label_to_dist_group.get(dist_type, dist_type) for dist_type in dist_types]
+    return dist_types, dist_groups
+
+
+def _process_dataloader(
+    dataloader,
+    iqa,
+    sae,
+    layer_name,
+    scaling_factor,
+    device,
+    patches_per_image,
+    patch_grid_shape,
+    meta_keys,
+    max_batches,
+    max_memory_gb,
+    add_patch_mask_stats,
+    show_progress_bars: bool = True,
+    *,
+    label_to_dist_type: Optional[Dict[int, str]] = None,
+    label_to_dist_group: Optional[Dict[str, str]] = None,
+):
+    all_sparse_codes = []
+    all_meta = []
+    all_sparse_steps = []
+
+    n_patches_known = patches_per_image
+    patch_grid_known = patch_grid_shape
+    image_offset = 0
+
+    for batch_i, batch in enumerate(tqdm(dataloader, desc='Caching activations', disable=not show_progress_bars)):
+        if max_batches is not None and batch_i >= max_batches:
+            break
+
+        imgs = batch['images'].to(device)
+        B = imgs.shape[0]
+
+        with torch.no_grad():
+            iqa(imgs)
+            acts = _iqa_activations[layer_name].to(device)
+            acts = acts * scaling_factor
+            enc_out = sae.get_acts(acts)
+
+            if isinstance(enc_out, tuple):
+                codes, activation_steps = enc_out
+            else:
+                codes, activation_steps = enc_out, None
+
+        codes_np = codes.cpu().float().numpy()
+
+        if activation_steps is None:
+            steps_np = np.zeros_like(codes_np, dtype=np.int32)
+        else:
+            steps_np = activation_steps.cpu().numpy().astype(np.int32)
+
+        if n_patches_known is None:
+            n_patches_known = codes_np.shape[0] // B
+            print(f'  Detected {n_patches_known} patches per image')
+
+        P = n_patches_known
+
+        use_patch_label_meta = (
+            label_to_dist_type is not None
+            and label_to_dist_group is not None
+        )
+        meta = {}
+        for k in meta_keys:
+            if k in batch:
+                if use_patch_label_meta and k in _PATCH_LABEL_META_KEYS:
+                    continue
+                vals = batch[k]
+                meta[k] = [v for v in vals for _ in range(P)]
+
+        meta['patch_idx'] = list(range(P)) * B
+        meta['image_idx'] = [image_offset + i for i in range(B) for _ in range(P)]
+
+        if add_patch_mask_stats and 'masks' in batch:
+            masks = batch['masks'].to(device=device, dtype=torch.float32)
+            if patch_grid_known is not None:
+                grid_h, grid_w = patch_grid_known
+                if grid_h * grid_w == P:
+                    mask_labels = masks.to(dtype=torch.int64)
+                    max_label = int(mask_labels.max().item())
+                    max_cov = None
+
+                    if max_label <= 0:
+                        patch_labels = torch.zeros((B, grid_h, grid_w), device=device, dtype=torch.int64)
+                    else:
+                        class_coverages = []
+                        for label_id in range(1, max_label + 1):
+                            label_cov = F.adaptive_avg_pool2d(
+                                (mask_labels == label_id).to(dtype=torch.float32),
+                                (grid_h, grid_w),
+                            )
+                            class_coverages.append(label_cov)
+
+                        coverages = torch.cat(class_coverages, dim=1)  # (B, classes, H, W)
+                        max_cov, max_idx = coverages.max(dim=1)
+                        patch_labels = torch.where(
+                            max_cov > 0,
+                            max_idx.to(dtype=torch.int64) + 1,
+                            torch.zeros_like(max_idx, dtype=torch.int64),
+                        )
+
+                    patch_labels_np = patch_labels.reshape(B, P).cpu().numpy().astype(np.int16)
+                    patch_is_dist = (patch_labels_np > 0).astype(np.int8)
+                    patch_coverage_np = max_cov.reshape(B, P).cpu().numpy() if max_label > 0 else patch_is_dist.astype(np.float32)
+
+                    meta['patch_mask_label'] = patch_labels_np.reshape(-1).tolist()
+                    meta['patch_mask_coverage'] = patch_coverage_np.reshape(-1).tolist()
+                    meta['patch_is_distorted'] = patch_is_dist.reshape(-1).tolist()
+
+                    if use_patch_label_meta:
+                        dist_types, dist_groups = _patch_labels_to_dist_meta(
+                            patch_labels_np,
+                            label_to_dist_type=label_to_dist_type,
+                            label_to_dist_group=label_to_dist_group,
+                        )
+                        if 'dist_type' in meta_keys:
+                            meta['dist_type'] = dist_types
+                        if 'dist_group' in meta_keys:
+                            meta['dist_group'] = dist_groups
+
+        all_meta.append(pd.DataFrame(meta))
+        all_sparse_codes.append(sp.csr_matrix(codes_np))
+        all_sparse_steps.append(sp.csr_matrix(steps_np))
+
+        image_offset += B
+
+    meta_df = pd.concat(all_meta, ignore_index=True)
+    codes_csr = sp.vstack(all_sparse_codes, format='csr')
+    steps_csr = sp.vstack(all_sparse_steps, format='csr')
+
+    return meta_df, codes_csr, steps_csr
+
+
+def collect_and_cache(
+    dataloader: DataLoader,
+    iqa: torch.nn.Module,
+    sae,
+    layer_name: str,
+    output_path: str,
+    scaling_factor: float = 1.0,
+    patches_per_image: Optional[int] = None,
+    patch_grid_shape: Optional[Tuple[int, int]] = None,
+    meta_keys: Sequence[str] = (
+        'dist_type',
+        'dist_group',
+        'dist_level',
+        'mos',
+        'distorted_img_path',
+        'original_img_path',
+        'sample_id',
+    ),
+    device: str = 'cuda',
+    max_batches: Optional[int] = None,
+    max_memory_gb: Optional[float] = None,
+    add_patch_mask_stats: bool = True,
+    pristine_dataloader: Optional[DataLoader] = None,
+    show_progress_bars: bool = True,
+    label_to_dist_type: Optional[Dict[int, str]] = None,
+    label_to_dist_group: Optional[Dict[str, str]] = None,
+) -> Tuple[pd.DataFrame, sp.csr_matrix]:
+
+    meta_df, codes_csr, steps_csr = _process_dataloader(
+        dataloader=dataloader,
+        iqa=iqa,
+        sae=sae,
+        layer_name=layer_name,
+        scaling_factor=scaling_factor,
+        device=device,
+        patches_per_image=patches_per_image,
+        patch_grid_shape=patch_grid_shape,
+        meta_keys=meta_keys,
+        max_batches=max_batches,
+        max_memory_gb=max_memory_gb,
+        add_patch_mask_stats=add_patch_mask_stats,
+        show_progress_bars=show_progress_bars,
+        label_to_dist_type=label_to_dist_type,
+        label_to_dist_group=label_to_dist_group,
+    )
+    sparse_mb = _sparse_mb(codes_csr)
+    steps_mb = _sparse_mb(steps_csr)
+
+    print(f'Activations: shape={codes_csr.shape}, {sparse_mb:.1f} МБ (sparse)')
+    print(f'Activation steps: shape={steps_csr.shape}, {steps_mb:.1f} МБ (sparse)')
+
+    meta_path, codes_path, steps_path = _cache_paths(output_path)
+
+    meta_df.to_feather(meta_path)
+    sp.save_npz(codes_path, codes_csr)
+    sp.save_npz(steps_path, steps_csr)
+
+    print(f'Saved metadata ({len(meta_df)} rows) -> {meta_path}')
+    print(f'Saved activations -> {codes_path}')
+    print(f'Saved activation steps -> {steps_path}')
+    print(f'  Metadata:    {_df_mb(meta_df):.1f} МБ')
+    print(f'  Activations: {sparse_mb:.1f} МБ')
+    print(f'  Steps:       {steps_mb:.1f} МБ')
+
+    if pristine_dataloader is not None:
+        print('\nProcessing pristine dataset...')
+
+        pristine_meta, pristine_codes, pristine_steps = _process_dataloader(
+            dataloader=pristine_dataloader,
+            iqa=iqa,
+            sae=sae,
+            layer_name=layer_name,
+            scaling_factor=scaling_factor,
+            device=device,
+            patches_per_image=patches_per_image,
+            patch_grid_shape=patch_grid_shape,
+            meta_keys=meta_keys,
+            max_batches=max_batches,
+            max_memory_gb=max_memory_gb,
+            add_patch_mask_stats=False,
+            show_progress_bars=show_progress_bars,
+        )
+        pristine_sparse_mb = _sparse_mb(pristine_codes)
+        pristine_steps_mb = _sparse_mb(pristine_steps)
+
+        pristine_meta_path = meta_path.replace(".feather", "_pristine.feather")
+        pristine_codes_path = codes_path.replace(".npz", "_pristine.npz")
+        pristine_steps_path = steps_path.replace(".npz", "_pristine.npz")
+
+        pristine_meta.to_feather(pristine_meta_path)
+        sp.save_npz(pristine_codes_path, pristine_codes)
+        sp.save_npz(pristine_steps_path, pristine_steps)
+
+        print(f'Pristine activations: shape={pristine_codes.shape}, {pristine_sparse_mb:.1f} МБ')
+        print(f'Pristine steps: shape={pristine_steps.shape}, {pristine_steps_mb:.1f} МБ')
+
+        print(f'Saved pristine metadata ({len(pristine_meta)} rows) -> {pristine_meta_path}')
+        print(f'Saved pristine activations -> {pristine_codes_path}')
+        print(f'Saved pristine activation steps -> {pristine_steps_path}')
+
+    return meta_df, codes_csr
+
+
+def build_activation_cache(
+    *,
+    dataset: str,
+    cache_path: str,
+    checkpoint_path: str,
+    kadid_path: str,
+    layer_num: int,
+    iqa_metric: str,
+    swin_num: int,
+    device: str,
+    batch_size: int,
+    num_workers: int,
+    crop_size: int,
+    scaling_factor: float = 1.0,
+    min_distortion_level: Optional[int] = None,
+    max_batches: Optional[int] = None,
+    max_memory_gb: Optional[float] = None,
+    add_patch_mask_stats: bool = True,
+    include_pristine: bool = True,
+    show_progress_bars: bool = True,
+) -> Dict[str, Any]:
+    """Build activation cache end-to-end for KADID/local-KADID datasets."""
+    from .datasets import (
+        Kadid10kDataset,
+        KadidPristineDataset,
+        LocalKadidPresavedDataset,
+        LocalKadidPristineDataset,
+        QGroundDataset,
+        SRGroundSmallDataset,
+        available_distortions_qground,
+        available_distortions_srground,
+        distortion_types_mapping_qground,
+        distortion_types_mapping_srground,
+        kadid_collate_fn,
+        kadid_pristine_collate_fn,
+        local_kadid_collate_fn,
+        local_kadid_pristine_collate_fn,
+        qground_collate_fn,
+        srground_collate_fn,
+    )
+    from .models import _iqa_activation_grids, load_iqa_model, load_sae, read_sae_config
+
+    if min_distortion_level is not None and not (1 <= min_distortion_level <= 5):
+        raise ValueError('min_distortion_level must be in [1, 5]')
+
+    label_to_dist_type = None
+    label_to_dist_group = None
+
+    if dataset == 'local_kadid':
+        data = LocalKadidPresavedDataset(root=kadid_path, crop_size=crop_size)
+        collate_fn = local_kadid_collate_fn
+        meta_keys = [
+            'dist_type',
+            'dist_group',
+            'dist_level',
+            'mos',
+            'local_dist_type',
+            'local_dist_level',
+            'mask_shape',
+            'mask_coverage',
+            'sample_id',
+            'distorted_img_path',
+            'original_img_path',
+        ]
+        pristine_data = LocalKadidPristineDataset(root=kadid_path, crop_size=crop_size) if include_pristine else None
+        pristine_collate = local_kadid_pristine_collate_fn
+    elif dataset in {'kadid10k', 'kadid'}:
+        data = Kadid10kDataset(
+            root=kadid_path,
+            crop_size=crop_size,
+            min_distortion_level=min_distortion_level or 1,
+        )
+        collate_fn = kadid_collate_fn
+        meta_keys = [
+            'dist_type',
+            'dist_group',
+            'dist_level',
+            'mos',
+            'distorted_img_path',
+            'original_img_path',
+        ]
+        pristine_data = KadidPristineDataset(root=kadid_path, crop_size=crop_size) if include_pristine else None
+        pristine_collate = kadid_pristine_collate_fn
+    elif dataset == 'QGround':
+        data = QGroundDataset(
+            root=kadid_path,
+            split='test',
+            crop_size=crop_size,
+        )
+        collate_fn = qground_collate_fn
+        meta_keys = [
+            'dist_type',
+            'dist_group',
+            'dist_level',
+            'mos',
+            'mask_coverage',
+            'qground_ann_id',
+            'sample_id',
+            'distorted_img_path',
+            'original_img_path',
+            'image_path',
+            'mask_path',
+            'split',
+        ]
+        label_to_dist_type = distortion_types_mapping_qground
+        label_to_dist_group = available_distortions_qground
+        pristine_data = None
+        pristine_collate = None
+    elif dataset == 'SRGround':
+        data = SRGroundSmallDataset(
+            root=kadid_path,
+            json_path='srground_train.json',
+            crop_size=crop_size,
+        )
+        collate_fn = srground_collate_fn
+        meta_keys = [
+            'dist_type',
+            'dist_group',
+            'dist_level',
+            'mos',
+            'mask_coverage',
+            'sample_id',
+            'distorted_img_path',
+            'image_path',
+            'real_distortions_ann_path',
+            'sr_artifacts_ann_path',
+        ]
+        label_to_dist_type = distortion_types_mapping_srground
+        label_to_dist_group = available_distortions_srground
+        pristine_data = None
+        pristine_collate = None
+    else:
+        raise ValueError(f'Unsupported dataset: {dataset}')
+
+    loader = DataLoader(
+        data,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+    )
+
+    pristine_loader = None
+    if pristine_data is not None:
+        pristine_loader = DataLoader(
+            pristine_data,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            collate_fn=pristine_collate,
+        )
+
+    iqa_model, layer_name = load_iqa_model(
+        layer_num=layer_num,
+        device=device,
+        iqa_metric=iqa_metric,
+        swin_num=swin_num,
+    )
+    dtype = torch.float16 if iqa_metric == 'qalign' else torch.float32
+    sae_cfg = read_sae_config(checkpoint_path)
+    sae_model = load_sae(checkpoint_path, device=device, dtype=dtype, sae_config=sae_cfg)
+
+    with torch.no_grad():
+        dummy = torch.rand(1, 3, crop_size, crop_size, device=device).clamp(0, 1)
+        iqa_model(dummy)
+
+    if layer_name not in _iqa_activations or layer_name not in _iqa_activation_grids:
+        raise RuntimeError(f'Cannot infer activation grid for layer {layer_name}')
+
+    patch_grid_shape = _iqa_activation_grids[layer_name]
+    patches_per_image = patch_grid_shape[0] * patch_grid_shape[1]
+
+    Path(cache_path).parent.mkdir(parents=True, exist_ok=True)
+    collect_and_cache(
+        dataloader=loader,
+        iqa=iqa_model,
+        sae=sae_model,
+        layer_name=layer_name,
+        output_path=cache_path,
+        scaling_factor=scaling_factor,
+        patches_per_image=patches_per_image,
+        patch_grid_shape=patch_grid_shape,
+        meta_keys=meta_keys,
+        device=device,
+        max_batches=max_batches,
+        max_memory_gb=max_memory_gb,
+        add_patch_mask_stats=add_patch_mask_stats,
+        pristine_dataloader=pristine_loader,
+        show_progress_bars=show_progress_bars,
+        label_to_dist_type=label_to_dist_type,
+        label_to_dist_group=label_to_dist_group,
+    )
+
+    return {
+        'layer_name': layer_name,
+        'patch_grid_shape': patch_grid_shape,
+        'patches_per_image': patches_per_image,
+        'sae_config': sae_cfg,
+    }
+
+
+def load_cache(
+    path: str,
+    return_activation_steps: bool = False,
+    min_distortion_level: Optional[int] = None,
+    max_distortion_level: Optional[int] = None,
+) -> Union[
+    Tuple[pd.DataFrame, sp.csr_matrix],
+    Tuple[pd.DataFrame, sp.csr_matrix, sp.csr_matrix],
+]:
+    """Загружает кэш активаций SAE из раздельных файлов.
+
+    Если ``return_activation_steps=True``, дополнительно возвращает CSR-матрицу
+    порядка активаций, где значение ``0`` означает отсутствие активации,
+    а ``k>0`` соответствует шагу ``k`` в pursuit.
+    """
+    meta_path, codes_path, steps_path = _cache_paths(path)
+    meta = pd.read_feather(meta_path)
+    codes = sp.load_npz(codes_path)
+
+    print(f'Loaded from {meta_path}: {meta.shape[0]} rows × {meta.shape[1]} cols')
+    print(f'Loaded from {codes_path}: shape={codes.shape}, dtype={codes.dtype}')
+    print(f'  Metadata:    {_df_mb(meta):.1f} МБ')
+    print(f'  Activations: {_sparse_mb(codes):.1f} МБ (sparse)')
+
+    keep_idx: Optional[np.ndarray] = None
+    if min_distortion_level is not None or max_distortion_level is not None:
+        if 'dist_level' not in meta.columns:
+            raise ValueError('Cannot filter by distortion level: metadata has no "dist_level" column')
+
+        min_level = 1 if min_distortion_level is None else int(min_distortion_level)
+        max_level = 5 if max_distortion_level is None else int(max_distortion_level)
+        if min_level > max_level:
+            raise ValueError(
+                f'Invalid distortion-level range: min_distortion_level={min_level} > max_distortion_level={max_level}'
+            )
+        if 'Ground' not in path:
+            keep_mask = (meta['dist_level'] >= min_level) & (meta['dist_level'] <= max_level)
+            keep_idx = np.flatnonzero(keep_mask.to_numpy())
+        else: 
+            keep_mask = (meta['dist_level'] >= -1000) 
+            keep_idx = np.flatnonzero(keep_mask.to_numpy()) # Костыль -- поправить
+ 
+    if return_activation_steps:
+        if Path(steps_path).exists():
+            steps = sp.load_npz(steps_path)
+            if steps.shape != codes.shape:
+                raise ValueError(
+                    f'Steps cache shape mismatch: expected {codes.shape}, got {steps.shape}'
+                )
+            print(f'Loaded from {steps_path}: shape={steps.shape}, dtype={steps.dtype}')
+        else:
+            print(f'No steps cache found. Using all-zero activation steps.')
+            steps = sp.csr_matrix(codes.shape, dtype=np.int32)
+
+        if keep_idx is not None:
+            meta = meta.iloc[keep_idx].reset_index(drop=True)
+            codes = codes[keep_idx]
+            steps = steps[keep_idx]
+            print(
+                f'Applied dist_level filter [{min_level}, {max_level}] -> '
+                f'{meta.shape[0]} rows kept'
+            )
+
+        print(f'  Steps:       {_sparse_mb(steps):.1f} МБ (sparse)')
+        return meta, codes, steps
+
+    if keep_idx is not None:
+        meta = meta.iloc[keep_idx].reset_index(drop=True)
+        codes = codes[keep_idx]
+        print(
+            f'Applied dist_level filter [{min_level}, {max_level}] -> '
+            f'{meta.shape[0]} rows kept'
+        )
+
+    return meta, codes
+
+
+def load_pristine_cache(
+    path: str,
+    return_activation_steps: bool = False,
+) -> Union[
+    Tuple[pd.DataFrame, sp.csr_matrix],
+    Tuple[pd.DataFrame, sp.csr_matrix, sp.csr_matrix],
+]:
+    """Load pristine (original-image) activation cache saved by collect_and_cache."""
+    meta_path, codes_path, steps_path = _pristine_cache_paths(path)
+    meta = pd.read_feather(meta_path)
+    codes = sp.load_npz(codes_path)
+
+    print(f'Loaded pristine from {meta_path}: {meta.shape[0]} rows × {meta.shape[1]} cols')
+    print(f'Loaded pristine from {codes_path}: shape={codes.shape}, dtype={codes.dtype}')
+    print(f'  Metadata:    {_df_mb(meta):.1f} МБ')
+    print(f'  Activations: {_sparse_mb(codes):.1f} МБ (sparse)')
+
+    if return_activation_steps:
+        if Path(steps_path).exists():
+            steps = sp.load_npz(steps_path)
+            if steps.shape != codes.shape:
+                raise ValueError(
+                    f'Pristine steps cache shape mismatch: expected {codes.shape}, got {steps.shape}'
+                )
+            print(f'Loaded pristine from {steps_path}: shape={steps.shape}, dtype={steps.dtype}')
+        else:
+            print('No pristine steps cache found. Using all-zero activation steps.')
+            steps = sp.csr_matrix(codes.shape, dtype=np.int32)
+
+        print(f'  Steps:       {_sparse_mb(steps):.1f} МБ (sparse)')
+        return meta, codes, steps
+
+    return meta, codes
+
+
+def ensure_cache_ready(
+    cache_path: str,
+    *,
+    force_recache: bool = False,
+    build_cache_if_missing: bool = True,
+    load_cache_kwargs: Optional[Dict[str, Any]] = None,
+    build_cache_fn: Optional[Callable[[], None]] = None,
+) -> None:
+    """Проверяет доступность кэша и при необходимости собирает его.
+
+    Поведение:
+    - пытается загрузить кэш через ``load_cache``;
+    - если кэш отсутствует или выставлен ``force_recache=True``, запускает сборку;
+    - если сборка отключена, пробрасывает ``FileNotFoundError``.
+    """
+    needs_rebuild = force_recache
+    if not needs_rebuild:
+        try:
+            load_cache(cache_path, **(load_cache_kwargs or {}))
+            return
+        except FileNotFoundError:
+            needs_rebuild = True
+
+    if not needs_rebuild:
+        return
+
+    if not build_cache_if_missing:
+        raise FileNotFoundError(
+            f'Activation cache not found at {cache_path}, and build is disabled. '
+            'Use --build-cache-if-missing or provide existing cache files.'
+        )
+
+    if build_cache_fn is None:
+        raise ValueError(
+            'build_cache_fn must be provided when cache rebuild is required '
+            '(missing cache or force_recache=True).'
+        )
+
+    print('[cache] Building activation cache...')
+    build_cache_fn()
+
+
+def zero_codes_outside_activation_steps(
+    codes_csr: sp.csr_matrix,
+    activation_steps_csr: sp.csr_matrix,
+    activation_steps_to_keep: List[int],
+) -> sp.csr_matrix:
+    """Обнуляет активации, ��аг появления которых не входит в allow-list.
+
+    Параметры
+    ----------
+    codes_csr : CSR-матрица активаций SAE.
+    activation_steps_csr : CSR-матрица шагов активаций (0 = не активирован).
+    activation_steps_to_keep : список шагов, которые нужно сохранить.
+
+    Возвращает
+    ----------
+    CSR-матрицу той же формы, где вне указанных шагов значения занулены.
+    Если список шагов пуст, возвращается исходная матрица без изменений.
+    """
+    if not activation_steps_to_keep:
+        return codes_csr
+
+    if codes_csr.shape != activation_steps_csr.shape:
+        raise ValueError(
+            f'Codes/steps shape mismatch: {codes_csr.shape} vs {activation_steps_csr.shape}'
+        )
+
+    keep_steps = sorted({int(step) for step in activation_steps_to_keep})
+    if any(step <= 0 for step in keep_steps):
+        raise ValueError('activation_steps_to_keep must contain only positive integers')
+
+    codes_coo = codes_csr.tocoo(copy=False)
+    steps_coo = activation_steps_csr.tocoo(copy=False)
+
+    # Steps matrix stores indices for nonzero entries of codes, so coordinates must match.
+    if (
+        codes_coo.nnz != steps_coo.nnz
+        or not np.array_equal(codes_coo.row, steps_coo.row)
+        or not np.array_equal(codes_coo.col, steps_coo.col)
+    ):
+        raise ValueError('Codes and steps matrices must have the same sparsity pattern. Something weird is going on.')
+    else:
+        steps_for_codes = steps_coo.data
+
+    keep_mask = np.isin(np.asarray(steps_for_codes), np.asarray(keep_steps, dtype=np.int32))
+
+    filtered = sp.coo_matrix(
+        (codes_coo.data[keep_mask], (codes_coo.row[keep_mask], codes_coo.col[keep_mask])),
+        shape=codes_csr.shape,
+        dtype=codes_csr.dtype,
+    )
+    return filtered.tocsr()
+
+
+def ensure_activation_cache(
+    dataset: str,
+    acts_cache_path: str,
+    kadid_images_path: str,
+    min_distortion_level: int,
+    params: dict,
+    include_pristine_cache: Optional[bool] = None,
+) -> None:
+    """Build distorted+pristine activation cache if missing."""
+    cache_filter_min = int(min_distortion_level) if dataset == 'kadid10k' else None
+    if include_pristine_cache is None:
+        needs_pristine_cache = dataset in {'kadid10k', 'local_kadid'}
+    else:
+        needs_pristine_cache = bool(include_pristine_cache)
+
+    try:
+        load_cache(
+            acts_cache_path,
+            return_activation_steps=True,
+            min_distortion_level=cache_filter_min,
+            max_distortion_level=params.get('KADID_MAX_DISTORTION_LEVEL') if dataset == 'kadid10k' else None,
+        )
+        if needs_pristine_cache:
+            load_pristine_cache(acts_cache_path, return_activation_steps=True)
+        return
+    except FileNotFoundError:
+        pass
+
+    print(f'[run] Activation cache not found for {acts_cache_path}. Building cache...')
+    build_activation_cache(
+        dataset=dataset,
+        cache_path=acts_cache_path,
+        checkpoint_path=params.get('SAE_CHECKPOINT_PATH'),
+        kadid_path=kadid_images_path,
+        layer_num=params.get('LAYER_NUM'),
+        iqa_metric=params.get('IQA_METRIC'),
+        swin_num=params.get('SWIN_NUM'),
+        device=params.get('DEVICE'),
+        batch_size=params.get('BATCH_SIZE'),
+        num_workers=params.get('NUM_WORKERS'),
+        crop_size=params.get('CROP_SIZE'),
+        scaling_factor=params.get('SCALING_FACTOR'),
+        min_distortion_level=min_distortion_level,
+        max_batches=None,
+        max_memory_gb=30.0,
+        add_patch_mask_stats=True,
+        include_pristine=needs_pristine_cache,
+    )
+    print('[run] Activation cache build completed.')

analysis/config.py

@@ -0,0 +1,332 @@
+"""
+Конфигурация SAE-визуализаций.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Dict, List, Mapping, Sequence, Tuple
+
+DEFAULT_SAE_VIS_CONFIG_PATH = os.path.join(
+    os.path.dirname(__file__), '..', 'default_configs', 'sae_vis_config.json'
+)
+SUPPORTED_DATASETS: Tuple[str, ...] = ('kadid10k', 'local_kadid', 'QGround', 'SRGround')
+
+_DATASET_IMAGE_SUBDIRS = {
+    'kadid10k': 'kadid10k',
+    'local_kadid': 'local_kadid',
+    'QGround': 'QGround',
+    'SRGround': 'SRGround',
+}
+
+
+def dataset_images_root(datasets_root: str, dataset: str) -> str:
+    """Filesystem root for a dataset's image files (under DATASETS_ROOT)."""
+    subdir = _DATASET_IMAGE_SUBDIRS.get(str(dataset), str(dataset))
+    return os.path.join(datasets_root, subdir)
+
+
+@dataclass(frozen=True)
+class DatasetCachePaths:
+    dataset: str
+    cache_dir: str
+    dataset_cache_dir: str
+    acts_cache_path: str
+    corr_group_cache_path: str
+    corr_type_cache_path: str
+    corr_group_patch_cache_path: str
+    corr_type_patch_cache_path: str
+    mi_group_cache_path: str
+    mi_type_cache_path: str
+    mi_group_patch_cache_path: str
+    mi_type_patch_cache_path: str
+    auc_group_cache_path: str
+    auc_type_cache_path: str
+    auc_group_patch_cache_path: str
+    auc_type_patch_cache_path: str
+    precision_group_cache_path: str
+    precision_type_cache_path: str
+    precision_group_patch_cache_path: str
+    precision_type_patch_cache_path: str
+    recall_group_cache_path: str
+    recall_type_cache_path: str
+    recall_group_patch_cache_path: str
+    recall_type_patch_cache_path: str
+    iou_group_cache_path: str
+    iou_type_cache_path: str
+
+
+@dataclass
+class SaeVisConfig:
+    """Runtime configuration loaded from SAE vis JSON (+ optional sae_config.json)."""
+
+    SAE_CHECKPOINT_PATH: str
+    SAE_CONFIG_PATH: str
+    DATASETS_ROOT: str
+    IN_DIM: int
+    INNER_DIM: int
+    LAYER_NUM: int
+    IQA_METRIC: str
+    SWIN_NUM: int
+    CROP_SIZE: int
+    DOWNSCALE_FACTOR: int
+    KADID_MIN_DISTORTION_LEVEL: int
+    KADID_MAX_DISTORTION_LEVEL: int
+    SCALING_FACTOR: float
+    BATCH_SIZE: int
+    NUM_WORKERS: int
+    DEVICE: str
+    ACTIVATION_STEPS_TO_KEEP: List[int]
+    CORR_TOP_K: int
+    HEATMAP_CRITERION: str
+    N_TOP_FEATURES_HEATMAPS: int
+    N_IMAGES_PER_FEATURE: int
+    HEATMAP_AGGREGATIONS: List[str]
+    FEATURE_FILTERS: List[dict]
+    SELECTOR_CONFIGS: List[dict]
+    SCATTER_TOP_K_PATCHES: int
+    SCATTER_GROUP_NAME: str
+    SCATTER_BACKEND: str
+    SCATTER_METRICS: List[str]
+    SCATTER_ACTIVATION_THRESHOLD: float
+    SCATTER_METRIC_LEVEL: str
+    BUILD_CACHE_IF_MISSING: bool
+    SAVE_TABULAR_ARTIFACTS: bool
+    SHOW_PROGRESS_BARS: bool
+    CACHE_DIR: str
+    DATASET: str
+    SUPPORTED_DATASETS: Tuple[str, ...]
+    DATASET_CACHE_CONFIGS: Dict[str, DatasetCachePaths]
+    CACHE_PATHS: DatasetCachePaths
+    KADID_IMAGES_PATH: str
+    config_path: str = field(repr=False, compare=False, default='')
+
+    def as_cache_params(self) -> Dict[str, object]:
+        """Dict for ensure_activation_cache / build_activation_cache call sites."""
+        return {
+            'SAE_CHECKPOINT_PATH': self.SAE_CHECKPOINT_PATH,
+            'LAYER_NUM': self.LAYER_NUM,
+            'IQA_METRIC': self.IQA_METRIC,
+            'SWIN_NUM': self.SWIN_NUM,
+            'DEVICE': self.DEVICE,
+            'BATCH_SIZE': self.BATCH_SIZE,
+            'NUM_WORKERS': self.NUM_WORKERS,
+            'CROP_SIZE': self.CROP_SIZE,
+            'SCALING_FACTOR': self.SCALING_FACTOR,
+            'KADID_MAX_DISTORTION_LEVEL': self.KADID_MAX_DISTORTION_LEVEL,
+        }
+
+
+def _read_json_object(path: str) -> Dict[str, object]:
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f'SAE visualization config not found: {path}. '
+            f'Please create JSON config or set SAE_VIS_CONFIG_PATH.'
+        )
+    with open(path, 'r', encoding='utf-8') as f:
+        payload = json.load(f)
+    if not isinstance(payload, dict):
+        raise ValueError(f'SAE visualization config must be a JSON object: {path}')
+    return payload
+
+
+def _cfg_get(cfg: Mapping[str, object], key: str, default: object = None) -> object:
+    return cfg.get(key, default)
+
+
+def resolve_sae_config_path(checkpoint_path: str) -> str:
+    if os.path.isdir(checkpoint_path):
+        return os.path.join(os.path.dirname(checkpoint_path), 'sae_config.json')
+    return os.path.join(os.path.dirname(os.path.dirname(checkpoint_path)), 'sae_config.json')
+
+
+def build_dataset_cache_paths(
+    cache_dir: str,
+    datasets: Sequence[str],
+) -> Dict[str, DatasetCachePaths]:
+    dataset_cache_configs: Dict[str, DatasetCachePaths] = {}
+    acts_filenames = {
+        'kadid10k': 'kadid_acts.feather',
+        'local_kadid': 'local_kadid_acts.feather',
+        'QGround': 'QGround_acts.feather',
+        'SRGround': 'SRGround_acts.feather',
+    }
+
+    common_cache_filenames = {
+        'corr_group': 'corr_group.parquet',
+        'corr_type': 'corr_type.parquet',
+        'corr_group_patch': 'corr_group_patch.parquet',
+        'corr_type_patch': 'corr_type_patch.parquet',
+        'mi_group': 'mi_group.parquet',
+        'mi_type': 'mi_type.parquet',
+        'mi_group_patch': 'mi_group_patch.parquet',
+        'mi_type_patch': 'mi_type_patch.parquet',
+        'auc_group': 'auc_group.parquet',
+        'auc_type': 'auc_type.parquet',
+        'auc_group_patch': 'auc_group_patch.parquet',
+        'auc_type_patch': 'auc_type_patch.parquet',
+        'precision_group': 'precision_group.parquet',
+        'precision_type': 'precision_type.parquet',
+        'precision_group_patch': 'precision_group_patch.parquet',
+        'precision_type_patch': 'precision_type_patch.parquet',
+        'recall_group': 'recall_group.parquet',
+        'recall_type': 'recall_type.parquet',
+        'recall_group_patch': 'recall_group_patch.parquet',
+        'recall_type_patch': 'recall_type_patch.parquet',
+        'iou_group': 'iou_group.parquet',
+        'iou_type': 'iou_type.parquet',
+    }
+
+    for dataset in datasets:
+        dataset_cache_dir = os.path.join(cache_dir, dataset)
+        acts_filename = acts_filenames.get(dataset, f'{dataset}_acts.feather')
+
+        dataset_paths: Dict[str, str] = {
+            'acts': os.path.join(dataset_cache_dir, acts_filename),
+        }
+        for key, filename in common_cache_filenames.items():
+            dataset_paths[key] = os.path.join(dataset_cache_dir, filename)
+
+        dataset_cache_configs[dataset] = DatasetCachePaths(
+            dataset=dataset,
+            cache_dir=cache_dir,
+            dataset_cache_dir=dataset_cache_dir,
+            acts_cache_path=dataset_paths['acts'],
+            corr_group_cache_path=dataset_paths['corr_group'],
+            corr_type_cache_path=dataset_paths['corr_type'],
+            corr_group_patch_cache_path=dataset_paths['corr_group_patch'],
+            corr_type_patch_cache_path=dataset_paths['corr_type_patch'],
+            mi_group_cache_path=dataset_paths['mi_group'],
+            mi_type_cache_path=dataset_paths['mi_type'],
+            mi_group_patch_cache_path=dataset_paths['mi_group_patch'],
+            mi_type_patch_cache_path=dataset_paths['mi_type_patch'],
+            auc_group_cache_path=dataset_paths['auc_group'],
+            auc_type_cache_path=dataset_paths['auc_type'],
+            auc_group_patch_cache_path=dataset_paths['auc_group_patch'],
+            auc_type_patch_cache_path=dataset_paths['auc_type_patch'],
+            precision_group_cache_path=dataset_paths['precision_group'],
+            precision_type_cache_path=dataset_paths['precision_type'],
+            precision_group_patch_cache_path=dataset_paths['precision_group_patch'],
+            precision_type_patch_cache_path=dataset_paths['precision_type_patch'],
+            recall_group_cache_path=dataset_paths['recall_group'],
+            recall_type_cache_path=dataset_paths['recall_type'],
+            recall_group_patch_cache_path=dataset_paths['recall_group_patch'],
+            recall_type_patch_cache_path=dataset_paths['recall_type_patch'],
+            iou_group_cache_path=dataset_paths['iou_group'],
+            iou_type_cache_path=dataset_paths['iou_type'],
+        )
+
+    return dataset_cache_configs
+
+
+def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
+    """Load and validate SAE visualization config from JSON."""
+    config_path = path or os.environ.get('SAE_VIS_CONFIG_PATH', DEFAULT_SAE_VIS_CONFIG_PATH)
+    # config_path = "/home/28m_gov@lab.graphicon.ru/SAE/xIQA/logs/arniqa_logs/ARNIQA_layer5_lambda5e-4_scale1_exp37/vis/vis_srground/config.json"
+    vis_cfg = _read_json_object(config_path)
+
+    sae_checkpoint_path = str(_cfg_get(vis_cfg, 'SAE_CHECKPOINT_PATH', '')).strip()
+    if not sae_checkpoint_path:
+        raise ValueError('SAE_CHECKPOINT_PATH must be set in SAE vis config JSON')
+
+    sae_config_path = resolve_sae_config_path(sae_checkpoint_path)
+    sae_cfg: Dict[str, object] = _read_json_object(sae_config_path) if os.path.exists(sae_config_path) else {}
+
+    datasets_root = str(_cfg_get(vis_cfg, 'DATASETS_ROOT', '')).strip()
+    if not datasets_root:
+        legacy_kadid_images_path = str(_cfg_get(vis_cfg, 'KADID_IMAGES_PATH', '')).strip()
+        if legacy_kadid_images_path:
+            datasets_root = os.path.dirname(legacy_kadid_images_path)
+    if not datasets_root:
+        raise ValueError('DATASETS_ROOT must be set in SAE vis config JSON')
+
+    kadid_min_distortion_level = int(_cfg_get(vis_cfg, 'KADID_MIN_DISTORTION_LEVEL', 1))
+    kadid_max_distortion_level = int(_cfg_get(vis_cfg, 'KADID_MAX_DISTORTION_LEVEL', 5))
+    if not (1 <= kadid_min_distortion_level <= 5):
+        raise ValueError('KADID_MIN_DISTORTION_LEVEL must be in [1, 5]')
+    if not (1 <= kadid_max_distortion_level <= 5):
+        raise ValueError('KADID_MAX_DISTORTION_LEVEL must be in [1, 5]')
+    if kadid_min_distortion_level > kadid_max_distortion_level:
+        raise ValueError('KADID_MIN_DISTORTION_LEVEL must be <= KADID_MAX_DISTORTION_LEVEL')
+
+    activation_steps_to_keep = [int(step) for step in _cfg_get(vis_cfg, 'ACTIVATION_STEPS_TO_KEEP', [])]
+    if any(step <= 0 for step in activation_steps_to_keep):
+        raise ValueError('ACTIVATION_STEPS_TO_KEEP must contain only positive integers')
+
+    dataset = str(_cfg_get(vis_cfg, 'DATASET', 'kadid10k')).strip()
+    supported_raw = _cfg_get(vis_cfg, 'SUPPORTED_DATASETS', list(SUPPORTED_DATASETS))
+    supported_datasets = tuple(str(value) for value in supported_raw)
+    if dataset not in supported_datasets:
+        raise ValueError(f'Unsupported DATASET={dataset!r}; expected one of {supported_datasets}')
+
+    cache_dir = str(
+        _cfg_get(
+            vis_cfg,
+            'CACHE_DIR',
+            os.path.join(os.path.dirname(os.path.dirname(sae_checkpoint_path)), 'cache/'),
+        )
+    )
+    dataset_cache_configs = build_dataset_cache_paths(cache_dir, supported_datasets)
+
+    raw_selector_configs = _cfg_get(vis_cfg, 'SELECTOR_CONFIGS', [])
+    if raw_selector_configs is None:
+        raw_selector_configs = []
+
+    feature_filters = _cfg_get(
+        vis_cfg,
+        'FEATURE_FILTERS',
+        [
+            {'name': 'nonzero_max', 'params': {}},
+            {
+                'name': 'kruskal_wallis',
+                'params': {'alpha': 0.05, 'group_col': 'dist_type', 'min_group_size': 3},
+            },
+        ],
+    )
+
+    dataset_images_subdir = _DATASET_IMAGE_SUBDIRS[dataset]
+
+    return SaeVisConfig(
+        SAE_CHECKPOINT_PATH=sae_checkpoint_path,
+        SAE_CONFIG_PATH=sae_config_path,
+        DATASETS_ROOT=datasets_root,
+        IN_DIM=int(_cfg_get(sae_cfg, 'sae_input_dim', 64)),
+        INNER_DIM=int(_cfg_get(sae_cfg, 'inner_dim', 6400)),
+        LAYER_NUM=int(_cfg_get(sae_cfg, 'layer_num', 3)),
+        IQA_METRIC=str(_cfg_get(sae_cfg, 'iqa_metric', 'arniqa-kadid')),
+        SWIN_NUM=int(_cfg_get(sae_cfg, 'swin_num', 2)),
+        CROP_SIZE=int(_cfg_get(vis_cfg, 'CROP_SIZE', 224)),
+        DOWNSCALE_FACTOR=int(_cfg_get(vis_cfg, 'DOWNSCALE_FACTOR', 2)),
+        KADID_MIN_DISTORTION_LEVEL=kadid_min_distortion_level,
+        KADID_MAX_DISTORTION_LEVEL=kadid_max_distortion_level,
+        SCALING_FACTOR=float(_cfg_get(sae_cfg, 'scaling_factor', 1.0)),
+        BATCH_SIZE=int(_cfg_get(vis_cfg, 'BATCH_SIZE', 32)),
+        NUM_WORKERS=int(_cfg_get(vis_cfg, 'NUM_WORKERS', 4)),
+        DEVICE=str(_cfg_get(vis_cfg, 'DEVICE', 'cuda')),
+        ACTIVATION_STEPS_TO_KEEP=activation_steps_to_keep,
+        CORR_TOP_K=int(_cfg_get(vis_cfg, 'CORR_TOP_K', 30)),
+        HEATMAP_CRITERION=str(_cfg_get(vis_cfg, 'HEATMAP_CRITERION', 'max')),
+        N_TOP_FEATURES_HEATMAPS=int(_cfg_get(vis_cfg, 'N_TOP_FEATURES_HEATMAPS', 3)),
+        N_IMAGES_PER_FEATURE=int(_cfg_get(vis_cfg, 'N_IMAGES_PER_FEATURE', 5)),
+        HEATMAP_AGGREGATIONS=[str(value) for value in _cfg_get(vis_cfg, 'HEATMAP_AGGREGATIONS', ['max', 'mean_acts', 'sum'])],
+        FEATURE_FILTERS=list(feature_filters),
+        SELECTOR_CONFIGS=list(raw_selector_configs),
+        SCATTER_TOP_K_PATCHES=int(_cfg_get(vis_cfg, 'SCATTER_TOP_K_PATCHES', 1000)),
+        SCATTER_GROUP_NAME=str(_cfg_get(vis_cfg, 'SCATTER_GROUP_NAME', 'blur')),
+        SCATTER_BACKEND=str(_cfg_get(vis_cfg, 'SCATTER_BACKEND', 'matplotlib')),
+        SCATTER_METRICS=[str(value) for value in _cfg_get(vis_cfg, 'SCATTER_METRICS', ['entropy', 'iou', 'roc_auc', 'precision', 'recall'])],
+        SCATTER_ACTIVATION_THRESHOLD=float(_cfg_get(vis_cfg, 'SCATTER_ACTIVATION_THRESHOLD', 0.0)),
+        SCATTER_METRIC_LEVEL=str(_cfg_get(vis_cfg, 'SCATTER_METRIC_LEVEL', 'patch')),
+        BUILD_CACHE_IF_MISSING=bool(_cfg_get(vis_cfg, 'BUILD_CACHE_IF_MISSING', True)),
+        SAVE_TABULAR_ARTIFACTS=bool(_cfg_get(vis_cfg, 'SAVE_TABULAR_ARTIFACTS', False)),
+        SHOW_PROGRESS_BARS=bool(_cfg_get(vis_cfg, 'SHOW_PROGRESS_BARS', True)),
+        CACHE_DIR=cache_dir,
+        DATASET=dataset,
+        SUPPORTED_DATASETS=supported_datasets,
+        DATASET_CACHE_CONFIGS=dataset_cache_configs,
+        CACHE_PATHS=dataset_cache_configs[dataset],
+        KADID_IMAGES_PATH=os.path.join(datasets_root, dataset_images_subdir),
+        config_path=config_path,
+    )

analysis/datasets.py

@@ -0,0 +1,1148 @@
+"""
+Датасеты и вспомогательные структуры данных для KADID-10k. Взяты из репозитория PatchSAE.
+"""
+
+import gzip
+import re
+import json
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import List, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image, ImageColor
+from scipy.special import softmax
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+distortion_types_mapping = {
+    1:  "gaussian_blur",
+    2:  "lens_blur",
+    3:  "motion_blur",
+    4:  "color_diffusion",
+    5:  "color_shift",
+    6:  "color_quantization",
+    7:  "color_saturation_1",
+    8:  "color_saturation_2",
+    9:  "jpeg2000",
+    10: "jpeg",
+    11: "white_noise",
+    12: "white_noise_color_component",
+    13: "impulse_noise",
+    14: "multiplicative_noise",
+    15: "denoise",
+    16: "brighten",
+    17: "darken",
+    18: "mean_shift",
+    19: "jitter",
+    20: "non_eccentricity_patch",
+    21: "pixelate",
+    22: "quantization",
+    23: "color_block",
+    24: "high_sharpen",
+    25: "contrast_change",
+}
+
+available_distortions = {
+    "gaussian_blur":              "blur",
+    "lens_blur":                  "blur",
+    "motion_blur":                "blur",
+    "color_diffusion":            "color_distortion",
+    "color_shift":                "color_distortion",
+    "color_quantization":         "color_distortion",
+    "color_saturation_1":         "color_distortion",
+    "color_saturation_2":         "color_distortion",
+    "jpeg2000":                   "jpeg",
+    "jpeg":                       "jpeg",
+    "white_noise":                "noise",
+    "white_noise_color_component": "noise",
+    "impulse_noise":              "noise",
+    "multiplicative_noise":       "noise",
+    "denoise":                    "noise",
+    "brighten":                   "brightness_change",
+    "darken":                     "brightness_change",
+    "mean_shift":                 "brightness_change",
+    "jitter":                     "spatial_distortion",
+    "non_eccentricity_patch":     "spatial_distortion",
+    "pixelate":                   "spatial_distortion",
+    "quantization":               "spatial_distortion",
+    "color_block":                "spatial_distortion",
+    "high_sharpen":               "sharpness_contrast",
+    "contrast_change":            "sharpness_contrast",
+}
+
+distortion_groups = {
+    "blur":              ["gaussian_blur", "lens_blur", "motion_blur"],
+    "color_distortion":   ["color_diffusion", "color_shift", "color_quantization",
+                           "color_saturation_1", "color_saturation_2"],
+    "jpeg":               ["jpeg2000", "jpeg"],
+    "noise":              ["white_noise", "white_noise_color_component", "impulse_noise",
+                           "multiplicative_noise", "denoise"],
+    "brightness_change":  ["brighten", "darken", "mean_shift"],
+    "spatial_distortion": ["jitter", "non_eccentricity_patch", "pixelate",
+                           "quantization", "color_block"],
+    "sharpness_contrast":  ["high_sharpen", "contrast_change"],
+}
+
+QGROUND_DISTORTION_TYPES = {
+    'jitter': np.array(ImageColor.getrgb('#4b54e1')),
+    'noise': np.array(ImageColor.getrgb('#93fff0')),
+    'overexposure': np.array(ImageColor.getrgb('#cde55d')),
+    'blur': np.array(ImageColor.getrgb('#e45c5c')),
+    'low light': np.array(ImageColor.getrgb('#35e344')),
+}
+
+distortion_types_mapping_qground = {
+    0: 'background',
+    1: 'jitter',
+    2: 'noise',
+    3: 'overexposure',
+    4: 'blur',
+    5: 'low light',
+}
+
+available_distortions_qground = {
+    'background': 'background',
+    'jitter': 'jitter',
+    'noise': 'noise',
+    'overexposure': 'overexposure',
+    'blur': 'blur',
+    'low light': 'low light',
+}
+
+SRGROUND_CLASS_ORDER = (
+    'no_distortion',
+    'blur',
+    'jitter',
+    'lowlight',
+    'noise',
+    'overexposure',
+    'sr_artifact',
+)
+
+SRGROUND_DISTORTION_TYPES = {
+    'blur': np.array(ImageColor.getrgb('#e45c5c')),
+    'jitter': np.array(ImageColor.getrgb('#4b54e1')),
+    'lowlight': np.array(ImageColor.getrgb('#35e344')),
+    'noise': np.array(ImageColor.getrgb('#93fff0')),
+    'overexposure': np.array(ImageColor.getrgb('#cde55d')),
+    'sr_artifact': np.array(ImageColor.getrgb('#c000a0')),
+}
+
+distortion_types_mapping_srground = {
+    0: 'background',
+    1: 'blur',
+    2: 'jitter',
+    3: 'lowlight',
+    4: 'noise',
+    5: 'overexposure',
+    6: 'sr_artifact',
+}
+
+available_distortions_srground = {
+    'background': 'background',
+    'blur': 'blur',
+    'jitter': 'jitter',
+    'lowlight': 'lowlight',
+    'noise': 'noise',
+    'overexposure': 'overexposure',
+    'sr_artifact': 'sr_artifact',
+}
+
+SRGROUND_SR_ARTIFACT_THRESHOLD = 0.3
+
+
+def _load_npy_gz(path: Path) -> np.ndarray:
+    with gzip.open(path, 'rb') as handle:
+        return np.load(handle)
+
+
+def _real_distortion_labels(
+    real_maps: np.ndarray,
+    prominences: Optional[object] = None,
+) -> np.ndarray:
+    real_maps = np.asarray(real_maps, dtype=np.float64)
+    real_prom = np.array(prominences)[:-1, None, None]
+    real_prob = softmax(real_maps, axis=0) * real_prom
+    return np.argmax(real_prob, axis=0).astype(np.uint8)
+
+
+def _sr_artifact_labels(
+    sr_maps: np.ndarray,
+    prominences: Optional[object] = None,
+    threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+) -> np.ndarray:
+    sr_maps = np.asarray(sr_maps, dtype=np.float64)
+    if sr_maps.ndim == 3:
+        if sr_maps.shape[0] == 1:
+            sr_maps = sr_maps[0]
+        else:
+            raise ValueError(f'Expected SR artifact maps with shape (1, H, W) or (H, W), got {sr_maps.shape}')
+
+    sr_prom = prominences[-1]
+    return np.where(sr_maps * sr_prom >= threshold, 6, 0).astype(np.uint8)
+
+
+def merge_srground_masks(annot_rd: np.ndarray, annot_sr: np.ndarray) -> np.ndarray:
+    return np.where(annot_sr == 6, 6, annot_rd).astype(np.uint8)
+
+
+def label2rgb_srground(mask_label: np.ndarray) -> np.ndarray:
+    mask_rgb = np.zeros(mask_label.shape + (3,), dtype=np.uint8)
+    for label_id, dist_name in distortion_types_mapping_srground.items():
+        if label_id == 0:
+            continue
+        mask_rgb[mask_label == label_id] = SRGROUND_DISTORTION_TYPES[dist_name]
+    return mask_rgb
+
+
+SRGROUND_LEGEND_LABELS = {
+    'blur': 'blur',
+    'jitter': 'jitter',
+    'lowlight': 'low light',
+    'noise': 'noise',
+    'overexposure': 'overexposure',
+    'sr_artifact': 'SR artifact',
+}
+
+
+def _parse_prominences(value: object) -> np.ndarray | None:
+    """Convert a prominences field from JSON/CSV into a 1D float array.
+
+    Does not read any files — callers load ``prominences`` from ``srground_train.json``
+    (or another source) and pass the cell value here.
+    """
+    if value is None:
+        return None
+    if isinstance(value, float) and np.isnan(value):
+        return None
+    if isinstance(value, str):
+        text = value.strip()
+        if not text:
+            return None
+        try:
+            value = json.loads(text)
+        except json.JSONDecodeError:
+            return None
+    try:
+        return np.asarray(value, dtype=np.float64)
+    except (TypeError, ValueError):
+        return None
+
+
+def _center_crop_label_map(mask_label: np.ndarray, crop_size: int, reference_size: tuple[int, int]) -> np.ndarray:
+    """Center-crop a label map to match heatmap preprocessing (reference_size is W×H)."""
+    width, height = reference_size
+    mask_img = Image.fromarray(mask_label.astype(np.uint8), mode='L')
+    if mask_label.shape[:2] != (height, width):
+        mask_img = mask_img.resize((width, height), resample=Image.NEAREST)
+    return np.asarray(transforms.CenterCrop(int(crop_size))(mask_img), dtype=np.uint8)
+
+
+@lru_cache(maxsize=8)
+def _srground_train_dataframe(datasets_root: str) -> pd.DataFrame:
+    from analysis.config import dataset_images_root as _dataset_images_root
+
+    root = Path(_dataset_images_root(datasets_root, 'SRGround'))
+    return pd.read_json(root / 'srground_train.json')
+
+
+def _resolve_datasets_root(datasets_root: str | None) -> str:
+    if datasets_root is not None:
+        return str(datasets_root)
+    from analysis.config import load_sae_vis_config
+
+    return str(load_sae_vis_config().DATASETS_ROOT)
+
+
+def srground_image_key(path: str | Path, *, datasets_root: str | None = None) -> str:
+    """Normalize a path to ``image_path`` keys used in ``srground_train.json`` (relative to SRGround root)."""
+    from analysis.config import dataset_images_root
+
+    root = _resolve_datasets_root(datasets_root)
+    path_obj = Path(path)
+    sr_root = Path(dataset_images_root(root, 'SRGround'))
+    if path_obj.is_absolute():
+        return _to_relative_dataset_path(path_obj, sr_root)
+    return path_obj.as_posix()
+
+
+@lru_cache(maxsize=8)
+def srground_prominences_index(datasets_root: str | None = None) -> dict[str, np.ndarray]:
+    """Cached ``image_path`` → prominences array from ``srground_train.json``."""
+    root = _resolve_datasets_root(datasets_root)
+    df = _srground_train_dataframe(root)
+    out: dict[str, np.ndarray] = {}
+    for image_path, raw_prom in zip(df['image_path'].astype(str), df['prominences']):
+        prom = _parse_prominences(raw_prom)
+        if prom is not None:
+            out[str(image_path)] = prom
+    return out
+
+
+def _image_rel_from_meta_row(meta_row: pd.Series | None) -> str | None:
+    if meta_row is None:
+        return None
+    for column in ('image_path', 'distorted_img_path'):
+        if column not in meta_row:
+            continue
+        value = meta_row.get(column)
+        if value is None or (isinstance(value, float) and np.isnan(value)):
+            continue
+        text = str(value).strip()
+        if text:
+            return text
+    return None
+
+
+def srground_rows_for_image_paths(
+    image_paths: Sequence[str],
+    *,
+    datasets_root: str | None = None,
+) -> pd.DataFrame:
+    """Subset of ``srground_train.json`` for the given ``image_path`` keys."""
+    if not image_paths:
+        return pd.DataFrame()
+
+    root = _resolve_datasets_root(datasets_root)
+    keys = {srground_image_key(path, datasets_root=root) for path in image_paths if path}
+    if not keys:
+        return pd.DataFrame()
+
+    df = _srground_train_dataframe(root)
+    return df[df['image_path'].astype(str).isin(keys)].copy()
+
+
+def srground_prominences_by_image_paths(
+    image_paths: Sequence[str],
+    *,
+    datasets_root: str | None = None,
+) -> dict[str, np.ndarray]:
+    """Look up prominences for paths (absolute or SRGround-relative) via cached index."""
+    if not image_paths:
+        return {}
+
+    root = _resolve_datasets_root(datasets_root)
+    index = srground_prominences_index(root)
+    keys = {srground_image_key(path, datasets_root=root) for path in image_paths if path}
+    return {key: index[key] for key in keys if key in index}
+
+
+def srground_mask_rgb_for_meta_row(
+    meta_row: pd.Series | None,
+    *,
+    datasets_root: str | None = None,
+    include_sr_artifact: bool = True,
+    crop_size: int = 224,
+    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+    prominences: np.ndarray | None = None,
+) -> np.ndarray | None:
+    """Build an RGB annotation mask for one SRGround image (matches dashboard center crop)."""
+    if meta_row is None:
+        return None
+
+    root = _resolve_datasets_root(datasets_root)
+
+    def _row_path(column: str) -> str | None:
+        if column not in meta_row:
+            return None
+        value = meta_row.get(column)
+        if value is None or (isinstance(value, float) and np.isnan(value)):
+            return None
+        text = str(value).strip()
+        return text or None
+
+    image_rel = _image_rel_from_meta_row(meta_row)
+    real_rel = _row_path('real_distortions_ann_path')
+    sr_rel = _row_path('sr_artifacts_ann_path')
+    if not real_rel and not (include_sr_artifact and sr_rel):
+        return None
+
+    if prominences is None and image_rel:
+        prominences = srground_prominences_by_image_paths([image_rel], datasets_root=root).get(image_rel)
+
+    annot_rd = None
+    annot_sr = None
+    if real_rel:
+        real_path = resolve_dataset_image_path('SRGround', real_rel, datasets_root=root)
+        if real_path.is_file():
+            real_maps = _load_npy_gz(real_path)
+            prom = prominences
+            if prom is None:
+                prom = np.ones(int(real_maps.shape[0]) + 1, dtype=np.float64)
+            annot_rd = _real_distortion_labels(real_maps, prom)
+
+    if include_sr_artifact and sr_rel:
+        sr_path = resolve_dataset_image_path('SRGround', sr_rel, datasets_root=root)
+        if sr_path.is_file():
+            sr_maps = _load_npy_gz(sr_path)
+            prom = prominences
+            if prom is None:
+                prom = np.ones(6, dtype=np.float64)
+            annot_sr = _sr_artifact_labels(
+                sr_maps,
+                prom,
+                threshold=sr_artifact_threshold,
+            )
+
+    if annot_rd is None and annot_sr is None:
+        return None
+
+    if annot_rd is None:
+        annot_rd = np.zeros_like(annot_sr, dtype=np.uint8)
+    if include_sr_artifact and annot_sr is not None:
+        merged = merge_srground_masks(annot_rd, annot_sr)
+    else:
+        merged = annot_rd.astype(np.uint8, copy=False)
+
+    if image_rel:
+        image_path = resolve_dataset_image_path('SRGround', image_rel, datasets_root=root)
+        if image_path.is_file():
+            with Image.open(image_path) as image:
+                reference_size = image.size
+            merged = _center_crop_label_map(merged, crop_size, reference_size)
+        else:
+            merged = _center_crop_label_map(merged, crop_size, (merged.shape[1], merged.shape[0]))
+    else:
+        merged = _center_crop_label_map(merged, crop_size, (merged.shape[1], merged.shape[0]))
+
+    return label2rgb_srground(merged)
+
+
+def _rgb2label_qground(mask_rgb: np.ndarray) -> np.ndarray:
+    mask_label = np.zeros(mask_rgb.shape[:2], dtype=np.uint8)
+    for label, rgb_code in enumerate(QGROUND_DISTORTION_TYPES.values(), start=1):
+        matches = np.isclose(mask_rgb, rgb_code, rtol=0.2, atol=20).all(axis=-1)
+        mask_label[matches] = label
+    return mask_label
+
+
+def _label2rgb_qground(mask_label: np.ndarray) -> np.ndarray:
+    mask_rgb = np.zeros(mask_label.shape + (3,), dtype=np.uint8)
+    for label, rgb_code in enumerate(QGROUND_DISTORTION_TYPES.values(), start=1):
+        mask_rgb[mask_label == label] = rgb_code
+    return mask_rgb
+
+
+def _infer_kadid_original_path(distorted_path: Path) -> Path | None:
+    match = re.search(r'(I\d+)_\d+_\d+\.png$', distorted_path.name)
+    if match:
+        return distorted_path.with_name(f'{match.group(1)}.png')
+    return None
+
+
+def _to_relative_dataset_path(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
+class Kadid10kDataset(Dataset):
+    """
+    KADID-10k dataset. При семплинге применяется RandomCrop.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        crop_size: int = 224,
+        min_distortion_level: int = 1,
+        transform=None,
+    ):
+        self.root = Path(root)
+        self.crop_size = crop_size
+        self.mos_range = (1, 5)
+        self.min_distortion_level = int(min_distortion_level)
+
+        if not (1 <= self.min_distortion_level <= 5):
+            raise ValueError(
+                f"min_distortion_level must be in [1, 5], got {self.min_distortion_level}"
+            )
+
+        if transform is None:
+            self.transform = transforms.Compose([
+                transforms.RandomCrop(self.crop_size),
+                transforms.ToTensor(),
+            ])
+        else:
+            self.transform = transform
+
+        scores_csv = pd.read_csv(self.root / "dmos.csv")
+        scores_csv = scores_csv[["dist_img", "dmos"]]
+
+        self.images = np.array([
+            self.root / "images" / el
+            for el in scores_csv["dist_img"].values.tolist()
+        ])
+        self.mos = np.array(scores_csv["dmos"].values.tolist())
+
+        self.distortion_types = []
+        self.distortion_groups = []
+        self.distortion_levels = []
+
+        for image in self.images:
+            match = re.search(r'I\d+_(\d+)_(\d+)\.png$', str(image))
+            dist_type = distortion_types_mapping[int(match.group(1))]
+            self.distortion_types.append(dist_type)
+            self.distortion_groups.append(available_distortions[dist_type])
+            self.distortion_levels.append(int(match.group(2)))
+
+        self.distortion_types = np.array(self.distortion_types)
+        self.distortion_groups = np.array(self.distortion_groups)
+        self.distortion_levels = np.array(self.distortion_levels)
+
+        if self.min_distortion_level > 1:
+            keep_mask = self.distortion_levels >= self.min_distortion_level
+            self.images = self.images[keep_mask]
+            self.mos = self.mos[keep_mask]
+            self.distortion_types = self.distortion_types[keep_mask]
+            self.distortion_groups = self.distortion_groups[keep_mask]
+            self.distortion_levels = self.distortion_levels[keep_mask]
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def __getitem__(self, index: int) -> dict:
+        img = Image.open(self.images[index]).convert("RGB")
+        img = self.transform(img)
+        original_path = _infer_kadid_original_path(Path(self.images[index]))
+        return {
+            "img":        img,
+            "mos":        float(self.mos[index]),
+            "dist_type":  self.distortion_types[index],
+            "dist_group": self.distortion_groups[index],
+            "dist_level": int(self.distortion_levels[index]),
+            "distorted_img_path": _to_relative_dataset_path(Path(self.images[index]), self.root),
+            "original_img_path": _to_relative_dataset_path(original_path, self.root) if original_path is not None else None,
+        }
+
+
+class LocalKadidPresavedDataset(Dataset):
+    """Presaved KADID dataset with local distortions and binary masks.
+
+    Expects a dataset root directory produced by the local-distortion presave script.
+    The root must contain index.csv.
+    Required columns:
+        distorted_img_path, mask_path
+    Optional metadata columns are returned as-is in each sample.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        crop_size: Optional[int] = 224,
+    ):
+        self.root = Path(root)
+        self.index_path = self.root / "index.csv"
+        self.index_dir = self.root
+        if not self.index_path.exists():
+            raise FileNotFoundError(f"index.csv not found in local_kadid root: {self.root}")
+        self.df = pd.read_csv(self.index_path)
+
+        required_cols = {"distorted_img_path", "mask_path"}
+        missing = required_cols - set(self.df.columns)
+        if missing:
+            raise ValueError(f"Index is missing columns: {sorted(missing)}")
+
+        self.image_to_tensor = transforms.ToTensor()
+        self.mask_to_tensor = transforms.ToTensor()
+        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def _resolve_path(self, value: str) -> Path:
+        p = Path(value)
+        return p if p.is_absolute() else (self.index_dir / p)
+
+    def __getitem__(self, index: int) -> dict:
+        row = self.df.iloc[index]
+
+        img_path = self._resolve_path(str(row["distorted_img_path"]))
+        mask_path = self._resolve_path(str(row["mask_path"]))
+
+        img = Image.open(img_path).convert("RGB")
+        mask = Image.open(mask_path).convert("L")
+
+        img = self.image_to_tensor(img)
+        mask = self.mask_to_tensor(mask)
+        mask = (mask > 0.5).to(img.dtype)
+
+        if self.crop is not None:
+            img = self.crop(img)
+            mask = self.crop(mask)
+
+        sample = {
+            "img": img,
+            "mask": mask,
+        }
+
+        for key, value in row.items():
+            if key == "mask_path":
+                continue
+            if isinstance(value, np.generic):
+                value = value.item()
+            if key in ("distorted_img_path", "original_img_path"):
+                value = _to_relative_dataset_path(Path(value), self.root)
+            sample[key] = value
+
+        return sample
+
+
+def kadid_collate_fn(batch: List[dict]) -> dict:
+    """
+    Collate для Kadid10kDataset.
+
+    Возвращает:
+        images: Tensor (B, C, H, W)
+        + все остальные ключи как списки длины B
+    """
+    images = torch.stack([item["img"] for item in batch], dim=0)
+    collated: dict = {"images": images}
+    for key in batch[0]:
+        if key == "img":
+            continue
+        collated[key] = [item[key] for item in batch]
+    return collated
+
+
+def local_kadid_collate_fn(batch: List[dict]) -> dict:
+    """Collate for LocalKadidPresavedDataset.
+
+    Returns:
+        images: Tensor (B, C, H, W)
+        masks:  Tensor (B, 1, H, W)
+        + remaining keys as lists with length B
+    """
+    images = torch.stack([item["img"] for item in batch], dim=0)
+    masks = torch.stack([item["mask"] for item in batch], dim=0)
+
+    collated: dict = {
+        "images": images,
+        "masks": masks,
+    }
+    for key in batch[0]:
+        if key in ("img", "mask"):
+            continue
+        collated[key] = [item[key] for item in batch]
+    return collated
+
+
+class QGroundDataset(Dataset):
+    """QGround dataset stored locally as JSON index files plus image/mask folders.
+
+    Expected layout:
+        root/
+          qground_train.json
+          qground_test.json
+          images/
+          masks/
+    """
+
+    def __init__(
+        self,
+        root: str,
+        split: str = 'test',
+        json_path: Optional[str] = None,
+        crop_size: Optional[int] = 224,
+        annotation_index: int = 0,
+        transform=None,
+    ):
+        self.root = Path(root)
+        self.images_root = self.root / 'images'
+        self.masks_root = self.root / 'masks'
+        self.split = str(split).strip().lower()
+        self.annotation_index = int(annotation_index)
+
+        if json_path is None:
+            candidates = [
+                self.root / f'qground_{self.split}.json',
+                self.root / f'QGround_{self.split}.json',
+            ]
+            self.json_path = next((path for path in candidates if path.exists()), candidates[0])
+        else:
+            self.json_path = Path(json_path)
+            if not self.json_path.is_absolute():
+                self.json_path = self.root / self.json_path
+
+        if not self.json_path.exists():
+            raise FileNotFoundError(f'QGround split file not found: {self.json_path}')
+
+        with self.json_path.open('r', encoding='utf-8') as handle:
+            raw_samples = json.load(handle)
+
+        if not isinstance(raw_samples, list):
+            raise ValueError(f'QGround JSON must contain a list of samples: {self.json_path}')
+
+        self.samples = []
+        for sample in raw_samples:
+            if not isinstance(sample, dict):
+                continue
+
+            ann_list = sample.get('ann_list') or []
+            if isinstance(ann_list, dict):
+                ann_list = [ann_list]
+            if not ann_list:
+                continue
+
+            ann = ann_list[min(self.annotation_index, len(ann_list) - 1)]
+            image_rel = sample.get('image')
+            mask_rel = ann.get('segmentation_mask')
+            if not image_rel or not mask_rel:
+                continue
+
+            self.samples.append({
+                'sample_id': sample.get('id'),
+                'image_rel': image_rel,
+                'mask_rel': mask_rel,
+                'ann_id': ann.get('id'),
+                'quality_description': ann.get('quality_description'),
+            })
+
+        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
+        self.image_to_tensor = transforms.ToTensor() if transform is None else transform
+        self.image_paths = [self._resolve_path(self.images_root, sample['image_rel']) for sample in self.samples]
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def _resolve_path(self, base_dir: Path, rel_path: str) -> Path:
+        rel = Path(str(rel_path))
+        candidates = [
+            base_dir / rel,
+            self.root / rel,
+            base_dir / rel.name,
+            self.root / rel.name,
+        ]
+        for candidate in candidates:
+            if candidate.exists():
+                return candidate
+        return candidates[0]
+
+    def __getitem__(self, index: int) -> dict:
+        sample = self.samples[index]
+
+        img_path = self._resolve_path(self.images_root, sample['image_rel'])
+        mask_path = self._resolve_path(self.masks_root, sample['mask_rel'])
+
+        image = Image.open(img_path).convert('RGB')
+        mask_image = Image.open(mask_path).convert('RGB')
+
+        if self.crop is not None:
+            image = self.crop(image)
+            mask_image = self.crop(mask_image)
+
+        image = self.image_to_tensor(image)
+        if isinstance(image, Image.Image):
+            image = transforms.ToTensor()(image)
+
+        mask_rgb = np.asarray(mask_image, dtype=np.uint8)
+        mask_label = _rgb2label_qground(mask_rgb)
+        mask = torch.from_numpy(mask_label.astype(np.float32)).unsqueeze(0)
+
+        return {
+            'img': image,
+            'mask': mask,
+            'mos': float('nan'),
+            'dist_level': 5,
+            'mask_coverage': float((mask_label > 0).mean()),
+            'qground_ann_id': sample['ann_id'],
+            'sample_id': str(sample['sample_id'] or Path(sample['image_rel']).stem),
+            'distorted_img_path': _to_relative_dataset_path(img_path, self.root),
+            'original_img_path': '',  # no reference images in QGround
+            'image_path': _to_relative_dataset_path(img_path, self.root),
+            'mask_path': _to_relative_dataset_path(mask_path, self.root),
+            'split': self.split,
+        }
+
+
+def qground_collate_fn(batch: List[dict]) -> dict:
+    """Collate for QGroundDataset.
+
+    Returns:
+        images: Tensor (B, C, H, W)
+        masks: Tensor (B, 1, H, W)
+        + remaining keys as lists with length B
+    """
+    images = torch.stack([item['img'] for item in batch], dim=0)
+    masks = torch.stack([item['mask'] for item in batch], dim=0)
+
+    collated: dict = {
+        'images': images,
+        'masks': masks,
+    }
+    for key in batch[0]:
+        if key in ('img', 'mask'):
+            continue
+        collated[key] = [item[key] for item in batch]
+    return collated
+
+
+class SRGroundSmallDataset(Dataset):
+    """
+    Dataset wrapper for SRGround JSON indexes such as `srground_train.json`.
+
+    Expects entries with fields like `image_path`, `real_distortions_ann_path`,
+    `sr_artifacts_ann_path`, `has_markup`, `prominences`.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        json_path: Optional[str] = None,
+        require_markup: bool = True,
+        require_sr: bool = True,
+        allowed_methods: Optional[List[str]] = ['DiT4SR_x2'],
+        crop_size: Optional[int] = None,
+        sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+        transform=None,
+    ):
+        self.root = Path(root)
+        self.sr_artifact_threshold = float(sr_artifact_threshold)
+        self.json_path = self.root / 'srground_train.json'
+        df = pd.read_json(self.json_path)
+
+        if require_markup:
+            df = df[df['has_markup'].fillna(False).astype(bool)]
+        
+        df = df[~df['image_path'].str.contains('blur')]
+
+        df = df[df['image_path'].notna() & (df['image_path'].astype(str) != '')]
+
+        if require_sr:
+            df = df[df['image_path'].astype(str).str.contains('@SR@', na=False)]
+
+        if allowed_methods is not None:
+            method_pattern = '|'.join(re.escape(method) for method in allowed_methods)
+            df = df[df['image_path'].astype(str).str.contains(method_pattern, na=False)]
+
+        df = df.assign(sample_id=df['image_path'].astype(str).map(lambda path: Path(path).stem))
+        df = df.reset_index(drop=True)
+
+
+        self.df = df.copy()
+        self.image_to_tensor = transforms.ToTensor() if transform is None else transform
+        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
+        self.image_paths = [self.root / Path(path) for path in self.df['image_path'].tolist()]
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def _resolve_path(self, rel_path: Optional[str]) -> Path:
+        return self.root / Path(str(rel_path))
+
+    def _load_mask_labels(self, sample: dict) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
+        prominences = sample.get('prominences')
+        annot_rd = None
+        annot_sr = None
+
+        real_path = sample.get('real_distortions_ann_path')
+        real_ann_path = self._resolve_path(real_path)
+        if real_ann_path.exists():
+            annot_rd = _real_distortion_labels(_load_npy_gz(real_ann_path), prominences)
+
+        sr_path = sample.get('sr_artifacts_ann_path')
+        sr_ann_path = self._resolve_path(sr_path)
+        if sr_ann_path.exists():
+            annot_sr = _sr_artifact_labels(_load_npy_gz(sr_ann_path), prominences, threshold=self.sr_artifact_threshold)
+
+        if annot_rd is None and annot_sr is None:
+            return None, None
+        return annot_rd, annot_sr
+
+    def _resize_mask(self, mask_label: Optional[np.ndarray], image: Image.Image) -> Optional[np.ndarray]:
+        if mask_label is None:
+            return None
+        if mask_label.shape[:2] == image.size[::-1]:
+            return mask_label
+
+        mask_image = Image.fromarray(mask_label.astype(np.uint8), mode='L')
+        mask_image = mask_image.resize(image.size, resample=Image.NEAREST)
+        return np.asarray(mask_image, dtype=np.uint8)
+
+    def __getitem__(self, index: int) -> dict:
+        sample = self.df.iloc[index].to_dict()
+        img_path = self.image_paths[index]
+
+        image = Image.open(img_path).convert('RGB')
+        annot_rd, annot_sr = self._load_mask_labels(sample)
+        annot_rd = self._resize_mask(annot_rd, image)
+        annot_sr = self._resize_mask(annot_sr, image)
+
+        if self.crop is not None:
+            image = self.crop(image)
+            if annot_rd is not None:
+                mask_image = Image.fromarray(annot_rd.astype(np.uint8), mode='L')
+                annot_rd = np.asarray(self.crop(mask_image), dtype=np.uint8)
+            if annot_sr is not None:
+                mask_image = Image.fromarray(annot_sr.astype(np.uint8), mode='L')
+                annot_sr = np.asarray(self.crop(mask_image), dtype=np.uint8)
+
+        img_tensor = self.image_to_tensor(image)
+
+        mask_rd = torch.from_numpy(annot_rd.astype(np.float32)).unsqueeze(0)
+        mask_sr = torch.from_numpy(annot_sr.astype(np.float32)).unsqueeze(0)
+        
+        mask = torch.from_numpy(merge_srground_masks(annot_rd, annot_sr).astype(np.float32)).unsqueeze(0)
+        mask_coverage = float((mask > 0).float().mean().item())
+
+        real_ann_path = self._resolve_path(sample.get('real_distortions_ann_path'))
+        sr_ann_path = self._resolve_path(sample.get('sr_artifacts_ann_path'))
+
+        return {
+            'img': img_tensor,
+            'mask': mask,
+            'mask_rd': mask_rd,
+            'mask_sr': mask_sr,
+            'mos': float('nan'),
+            'dist_level': 5,
+            'mask_coverage': mask_coverage,
+            'prominences': sample.get('prominences'),
+            'has_markup': sample.get('has_markup', False),
+            'sr_artifacts_ann_path': _to_relative_dataset_path(sr_ann_path, self.root),
+            'real_distortions_ann_path': _to_relative_dataset_path(real_ann_path, self.root),
+            'sample_id': sample.get('sample_id'),
+            'distorted_img_path': _to_relative_dataset_path(img_path, self.root),
+            'image_path': _to_relative_dataset_path(img_path, self.root),
+            'mask_path': _to_relative_dataset_path(real_ann_path, self.root),
+        }
+
+
+def srground_collate_fn(batch: List[dict]) -> dict:
+    """Collate for SRGroundSmallDataset.
+
+    Returns:
+        images: Tensor (B, C, H, W)
+        masks: Tensor (B, 1, H, W)
+        + remaining keys as lists length B
+    """
+    images = torch.stack([item['img'] for item in batch], dim=0)
+    masks = torch.stack([item['mask'] for item in batch], dim=0)
+
+    collated: dict = {
+        'images': images,
+        'masks': masks,
+    }
+    for key in batch[0]:
+        if key in ('img', 'mask'):
+            continue
+        collated[key] = [item[key] for item in batch]
+    return collated
+
+
+class KadidPristineDataset(Dataset):
+    """
+    KADID-10k pristine (reference) images dataset. 
+    Возвращает только оригинальные изображения без искажений с RandomCrop.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        crop_size: int = 224,
+        transform=None,
+    ):
+        self.root = Path(root)
+        self.crop_size = crop_size
+
+        if transform is None:
+            self.transform = transforms.Compose([
+                transforms.RandomCrop(self.crop_size),
+                transforms.ToTensor(),
+            ])
+        else:
+            self.transform = transform
+
+        # Находим все оригинальные изображения (формат I{number}.png без суффиксов)
+        images_dir = self.root / "images"
+        pristine_pattern = re.compile(r'^I\d+\.png$')
+
+        self.images = sorted([
+            images_dir / fname
+            for fname in os.listdir(images_dir)
+            if pristine_pattern.match(fname)
+        ])
+
+        if len(self.images) == 0:
+            raise ValueError(f"No pristine images found in {images_dir}")
+
+        print(f"Found {len(self.images)} pristine images")
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def __getitem__(self, index: int) -> dict:
+        img_path = self.images[index]
+        img = Image.open(img_path).convert("RGB")
+        img = self.transform(img)
+        img_rel_path = _to_relative_dataset_path(Path(img_path), self.root)
+
+        return {
+            "img": img,
+            "mos": 5.0,  # максимальное качество для pristine
+            "dist_type": "pristine",
+            "dist_group": "pristine",
+            "dist_level": 0,  # уровень искажений = 0
+            "distorted_img_path": img_rel_path,
+            "original_img_path": img_rel_path,  # сам себе референс
+            "sample_id": img_path.stem,  # например "I01"
+        }
+
+
+class LocalKadidPristineDataset(Dataset):
+    """
+    Pristine KADID dataset.
+
+    Ожидает корневую директорию с index.csv.
+    Обязательные колонки: original_img_path
+    Опциональные: mask_path (если есть маска области без искажений)
+    """
+
+    def __init__(
+        self,
+        root: str,
+        crop_size: Optional[int] = 224,
+    ):
+        self.root = Path(root)
+        self.index_path = self.root / "index.csv"
+        self.index_dir = self.root
+
+        if not self.index_path.exists():
+            raise FileNotFoundError(f"index.csv not found in pristine root: {self.root}")
+
+        df = pd.read_csv(self.index_path)
+
+        required_cols = {"original_img_path"}
+
+        missing = required_cols - set(df.columns)
+        if missing:
+            raise ValueError(f"Index is missing columns: {sorted(missing)}")
+
+        self.images = sorted(df['original_img_path'].unique())
+        self.image_to_tensor = transforms.ToTensor()
+        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def _resolve_path(self, value: str) -> Path:
+        p = Path(value)
+        return p if p.is_absolute() else (self.index_dir / p)
+
+    def __getitem__(self, index: int) -> dict:
+        img_path = self._resolve_path(str(self.images[index]))
+
+        img = Image.open(img_path).convert("RGB")
+        img = self.image_to_tensor(img)
+        img_rel_path = _to_relative_dataset_path(Path(img_path), self.root)
+
+        sample = {
+            "img": img,
+            "dist_type": "pristine",
+            "dist_group": "pristine",
+            "dist_level": 0,
+            "distorted_img_path": img_rel_path,
+            "original_img_path": img_rel_path,
+            "sample_id": img_path.stem,
+        }
+
+        if self.crop is not None:
+            sample["img"] = self.crop(sample["img"])
+
+        return sample
+
+
+def kadid_pristine_collate_fn(batch: List[dict]) -> dict:
+    """
+    Collate для KadidPristineDataset.
+
+    Возвращает:
+        images: Tensor (B, C, H, W)
+        + все остальные ключи как списки длины B
+    """
+    images = torch.stack([item["img"] for item in batch], dim=0)
+    collated: dict = {"images": images}
+
+    for key in batch[0]:
+        if key == "img":
+            continue
+        collated[key] = [item[key] for item in batch]
+
+    return collated
+
+
+def local_kadid_pristine_collate_fn(batch: List[dict]) -> dict:
+    """
+    Collate для LocalKadidPristinePresavedDataset.
+
+    Возвращает:
+        images: Tensor (B, C, H, W)
+        masks: Tensor (B, 1, H, W) если has_masks=True
+        + остальные ключи как списки длины B
+    """
+    images = torch.stack([item["img"] for item in batch], dim=0)
+
+    collated: dict = {"images": images}
+
+    # Если есть маски в батче
+    if "mask" in batch[0]:
+        masks = torch.stack([item["mask"] for item in batch], dim=0)
+        collated["masks"] = masks
+
+    for key in batch[0]:
+        if key in ("img", "mask"):
+            continue
+        collated[key] = [item[key] for item in batch]
+
+    return collated
+
+
+def resolve_dataset_image_path(
+    dataset: str,
+    path_from_meta: str,
+    datasets_root: str | None = None,
+) -> Path:
+    """Resolve a path stored in activation-cache metadata to an absolute file path."""
+    from analysis.config import dataset_images_root, load_sae_vis_config
+
+    root = Path(
+        datasets_root
+        if datasets_root is not None
+        else load_sae_vis_config().DATASETS_ROOT
+    )
+    dataset_root = Path(dataset_images_root(str(root), dataset))
+    path = Path(path_from_meta)
+    return path if path.is_absolute() else (dataset_root / path)
+
+
+def dataset_image_paths(
+    dataset: str,
+    kadid_images_path: str,
+    crop_size: int,
+    min_distortion_level: int,
+) -> List[str]:
+    if dataset == 'kadid10k':
+        ds = Kadid10kDataset(
+            kadid_images_path,
+            crop_size=crop_size,
+            min_distortion_level=min_distortion_level,
+        )
+        return [_to_relative_dataset_path(Path(p), ds.root) for p in ds.images]
+
+    if dataset == 'local_kadid':
+        ds = LocalKadidPresavedDataset(kadid_images_path, crop_size=crop_size)
+        return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['distorted_img_path'].tolist()]
+
+    if dataset == 'QGround':
+        ds = QGroundDataset(kadid_images_path, split='test', crop_size=crop_size)
+        return [str(sample['image_rel']) for sample in ds.samples]
+
+    if dataset == 'SRGround':
+        ds = SRGroundSmallDataset(kadid_images_path, json_path='srground_train.json', allowed_methods=['DiT4SR_x2'])
+        return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['image_path'].tolist()]
+
+    raise ValueError(f'Unsupported dataset: {dataset}')

analysis/features/feature_filters.py

@@ -0,0 +1,437 @@
+"""Feature filtering strategies applied before feature selection."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import hashlib
+import json
+from pathlib import Path
+from typing import Callable, Dict, List, Mapping, Sequence
+
+from tqdm import tqdm
+import numpy as np
+import pandas as pd
+from scipy import stats
+from scipy.sparse import csc_matrix
+from statsmodels.stats.multitest import multipletests
+
+
+def _json_default(value: object) -> object:
+    """JSON serializer for numpy/pandas scalar-like values."""
+    if isinstance(value, (np.integer, np.floating)):
+        return value.item()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return str(value)
+
+
+def _build_filter_cache_key(
+    *,
+    filter_name: str,
+    params: Mapping[str, object],
+    meta_df: pd.DataFrame,
+    codes_shape: tuple[int, int],
+    group_col: str,
+) -> str:
+    """Build a stable cache key for feature-filter indices."""
+    if group_col not in meta_df.columns:
+        group_signature = {'missing_group_col': group_col}
+    else:
+        counts = meta_df[group_col].value_counts(dropna=False)
+        group_signature = {
+            'group_counts': [(str(idx), int(cnt)) for idx, cnt in counts.items()],
+        }
+
+    payload = {
+        'filter': filter_name,
+        'params': dict(params),
+        'n_samples': int(codes_shape[0]),
+        'n_features': int(codes_shape[1]),
+        'group_col': group_col,
+        'meta_rows': int(meta_df.shape[0]),
+        **group_signature,
+    }
+    raw = json.dumps(payload, sort_keys=True, default=_json_default).encode('utf-8')
+    return hashlib.sha256(raw).hexdigest()[:8]
+
+
+def _resolve_filter_cache_path(
+    *,
+    context: Mapping[str, object],
+    filter_name: str,
+    cache_key: str,
+) -> Path | None:
+    """Resolve cache file path for feature-filter indices from context."""
+    cache_dir = context.get('feature_filter_cache_dir')
+    if cache_dir is None:
+        return None
+    return Path(str(cache_dir)) / f'{filter_name}_{cache_key}.npz'
+
+
+def _load_feature_ids_from_cache(cache_path: Path, n_features: int) -> List[int]:
+    """Load and validate cached feature indices."""
+    with np.load(cache_path, allow_pickle=False) as data:
+        keep_feature_ids = np.asarray(data['keep_feature_ids'])
+    return [int(fid) for fid in keep_feature_ids.tolist()]
+
+
+def _save_feature_ids_to_cache(cache_path: Path, keep_feature_ids: Sequence[int]) -> None:
+    """Persist feature indices to cache as compressed NPZ."""
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(cache_path, keep_feature_ids=np.asarray(keep_feature_ids, dtype=np.int64))
+
+
+def my_kruskal(
+    col_data: np.ndarray,
+    col_indices: np.ndarray,
+    n_samples: int,
+    group_indices: list[np.ndarray],  # список массивов индексов для каждой группы
+) -> float:
+    """
+    Kruskal-Wallis H-test для одной колонки CSC-матрицы.
+
+    Параметры
+    ----------
+    col_data     : ненулевые значения колонки (csc.data[csc.indptr[j]:csc.indptr[j+1]])
+    col_indices  : строковые индексы ненулевых элементов (csc.indices[...])
+    n_samples    : полное число строк (n)
+    group_indices: group_indices[g] — массив строковых индексов группы g
+
+    Возвращает
+    ----------
+    p-value (float), или np.nan если тест неприменим
+    """
+    n = n_samples
+    n_zeros = n - len(col_data)
+
+    # --- 1. Ранги ненулевых элементов среди ВСЕХ n значений ----------------
+    # Нули занимают позиции 1..n_zeros в общем порядке.
+    # Ненулевые элементы начинаются с позиции n_zeros + 1.
+
+    # Сортируем ненулевые значения, получаем ранги внутри них
+    nonzero_vals = col_data.astype(np.float64)
+    order = np.argsort(nonzero_vals, kind='stable')
+    # Обрабатываем ties: каждому уникальному значению — средний ранг
+    ranks_local = np.empty(len(nonzero_vals), dtype=np.float64)
+    # временно назначаем 1-based ранги среди ненулевых
+    ranks_local[order] = np.arange(1, len(nonzero_vals) + 1, dtype=np.float64)
+    # Разрешаем ties усреднением
+    sorted_vals = nonzero_vals[order]
+    unique_vals, inverse, counts = np.unique(sorted_vals, return_inverse=True, return_counts=True)
+    # средний ранг для каждого уникального значения (1-based среди ненулевых)
+    cum = np.concatenate([[0], np.cumsum(counts)])
+    mean_local_ranks = (cum[:-1] + cum[1:] + 1) / 2.0  # (first + last + 1) / 2
+    ranks_local = mean_local_ranks[inverse]
+
+    # Сдвигаем ранги ненулевых элементов на n_zeros (они все правее нулей)
+    # Но нужно учесть ties между нулями и ненулевыми.
+    global_ranks_nonzero = ranks_local + n_zeros  # сдвиг на количество нулей
+
+    # Средний ранг нулей: нули занимают ранги 1..n_zeros
+    zero_mean_rank = (n_zeros + 1) / 2.0 if n_zeros > 0 else 0.0
+
+    # --- 2. Для каждой группы считаем сумму рангов -------------------------
+    n_groups = len(group_indices)
+    if n_groups < 2:
+        return np.nan
+
+    # Маска ненулевых: строим dict index->global_rank для быстрого lookup
+    rank_map = dict(zip(col_indices, global_ranks_nonzero))
+
+    H_num = 0.0
+    n_total_valid = 0
+
+    group_sizes = []
+    group_rank_sums = []
+
+    for g_idx in group_indices:
+        ng = len(g_idx)
+        if ng == 0:
+            continue
+
+        # Ранги элементов группы
+        r_sum = 0.0
+        for i in g_idx:
+            r_sum += rank_map.get(int(i), zero_mean_rank)
+
+        group_sizes.append(ng)
+        group_rank_sums.append(r_sum)
+        n_total_valid += ng
+
+    if len(group_sizes) < 2:
+        return np.nan
+
+    n_t = n_total_valid  # должно равняться n
+
+    # --- 3. H-статистика (стандартная формула) ------------------------------
+    # H = 12 / (n*(n+1)) * sum(R_i^2 / n_i) - 3*(n+1)
+    H = 0.0
+    for ng, rs in zip(group_sizes, group_rank_sums):
+        H += (rs ** 2) / ng
+
+    H = (12.0 / (n_t * (n_t + 1))) * H - 3.0 * (n_t + 1)
+
+    # --- 4. Поправка на ties ------------------------------------------------
+    # T = sum(t^3 - t) для каждой группы ties, C = 1 - T / (n^3 - n)
+    # Ties среди ненулевых:
+    tie_correction = 1.0
+    if len(unique_vals) < len(nonzero_vals):
+        T = np.sum(counts ** 3 - counts, dtype=np.float64)
+        # Добавляем ties нулей:
+        if n_zeros > 1:
+            T += n_zeros ** 3 - n_zeros
+        denom = float(n_t) ** 3 - n_t
+        if denom > 0:
+            tie_correction = 1.0 - T / denom
+
+    if tie_correction > 0:
+        H /= tie_correction
+
+    if H < 0:
+        H = 0.0
+
+    # --- 5. p-value из chi2 с df = n_groups - 1 ----------------------------
+    df = len(group_sizes) - 1
+    p_value = float(stats.chi2.sf(H, df))
+    return p_value
+
+
+class BaseFeatureFilter(ABC):
+    """Common interface for pre-selection feature filters."""
+
+    name: str
+    description: str
+
+    @abstractmethod
+    def apply_filter(
+        self,
+        context: Dict[str, object],
+    ) -> Dict[str, object]:
+        """Return filtered context with feature tables/matrix subsetted by feature ids."""
+
+    def filter_dataset(
+        self,
+        context: Dict[str, object],
+    ) -> Dict[str, object]:
+        before_count = int(context['codes_csr'].shape[1])
+        if before_count <= 0:
+            raise ValueError('Feature filter received empty feature space')
+
+        base_feature_ids = context.get('feature_ids')
+        if base_feature_ids is None:
+            base_feature_ids_arr = np.arange(before_count, dtype=np.int64)
+        else:
+            base_feature_ids_arr = np.asarray(base_feature_ids, dtype=np.int64)
+            if base_feature_ids_arr.shape[0] != before_count:
+                raise ValueError(
+                    f'Feature provenance length mismatch for filter {self.name!r}: '
+                    f'expected {before_count}, got {base_feature_ids_arr.shape[0]}'
+                )
+
+        filtered_context = self.apply_filter(context)
+        after_count = int(filtered_context['codes_csr'].shape[1])
+        if after_count > before_count:
+            raise ValueError(
+                f'Filter {self.name!r} increased feature count from {before_count} to {after_count}')
+        if after_count <= 0:
+            raise ValueError(f'Filter {self.name!r} removed all candidate features')
+
+        keep_feature_ids = filtered_context.get('filtered_feature_ids')
+        if keep_feature_ids is None:
+            raise ValueError(
+                f'Filter {self.name!r} must return filtered_feature_ids to preserve feature provenance'
+            )
+
+        keep_feature_ids_arr = np.asarray(keep_feature_ids, dtype=np.int64)
+        filtered_context['filtered_feature_ids'] = keep_feature_ids_arr.tolist()
+        filtered_context['feature_ids'] = base_feature_ids_arr[keep_feature_ids_arr].astype(np.int64).tolist()
+        return filtered_context
+
+
+class KruskalWallisFeatureFilter(BaseFeatureFilter):
+    """Filter features by Kruskal-Wallis significance across distortion groups."""
+
+    name = 'kruskal_wallis'
+    description = 'Keep features with p-value <= alpha in Kruskal-Wallis test across groups'
+
+    def __init__(
+        self,
+        alpha: float = 0.05,
+        group_col: str = 'dist_type',
+        min_group_size: int = 3,
+        batch_size: int = 1,
+        show_progress: bool = True,
+    ):
+        self.alpha = float(alpha)
+        self.group_col = group_col
+        self.min_group_size = int(min_group_size)
+        self.batch_size = batch_size
+        self.show_progress = bool(show_progress)
+
+    def apply_filter(
+        self,
+        context: Dict[str, object],
+    ) -> Dict[str, object]:
+        codes_csr = context['codes_csr']
+        meta_df = context['meta_df']
+
+        alpha = float(context.get('kruskal_alpha', self.alpha))
+        min_group_size = int(context.get('kruskal_min_group_size', self.min_group_size))
+        group_col = str(context.get('kruskal_group_col', self.group_col))
+        show_progress = bool(context.get('kruskal_show_progress', self.show_progress))
+
+        group_values = meta_df[group_col].to_numpy()
+        unique_groups = pd.unique(group_values)
+        n_samples = int(codes_csr.shape[0])
+        n_features = int(codes_csr.shape[1])
+
+        cache_params = {
+            'alpha': alpha,
+            'group_col': group_col,
+            'min_group_size': min_group_size,
+        }
+        cache_key = _build_filter_cache_key(
+            filter_name=self.name,
+            params=cache_params,
+            meta_df=meta_df,
+            codes_shape=(n_samples, n_features),
+            group_col=group_col,
+        )
+        cache_path = _resolve_filter_cache_path(
+            context=context,
+            filter_name=self.name,
+            cache_key=cache_key,
+        )
+        if cache_path is not None and cache_path.exists():
+            try:
+                keep_feature_ids = _load_feature_ids_from_cache(cache_path, n_features=n_features)
+                print(f'[cache][feature_filter] hit: {cache_path}')
+                filtered_context = dict(context)
+                filtered_context['codes_csr'] = codes_csr[:, keep_feature_ids]
+                filtered_context['filtered_feature_ids'] = keep_feature_ids
+                filtered_context['feature_filter_cache_path'] = str(cache_path)
+                return filtered_context
+            except Exception as exc:
+                print(f'[cache][feature_filter] invalid cache {cache_path}: {exc}. Recomputing...')
+
+        # Предвычисляем индексы строк для каждой группы — один раз
+        group_indices = []
+        for group in unique_groups:
+            idx = np.where(group_values == group)[0]
+            if idx.size >= min_group_size:
+                group_indices.append(idx)
+
+        codes_csc = codes_csr.tocsc()
+
+        valid_feature_ids = []
+        p_values = []
+
+        for fid in tqdm(range(n_features), disable=not show_progress):
+            start, stop = codes_csc.indptr[fid], codes_csc.indptr[fid + 1]
+            col_data = codes_csc.data[start:stop]
+            col_indices = codes_csc.indices[start:stop]
+
+            # Пропускаем константные колонки
+            n_nonzero = stop - start
+            all_same = (
+                n_nonzero == 0
+                or (n_nonzero == n_samples and np.all(col_data == col_data[0]))
+                or (n_nonzero < n_samples and n_nonzero == 0)
+            )
+            if all_same and n_nonzero == 0:
+                continue  # все нули — константа
+            if n_nonzero > 0 and n_nonzero == n_samples and np.all(col_data == col_data[0]):
+                continue  # все одинаковые ненулевые
+
+            p_value = my_kruskal(col_data, col_indices, n_samples, group_indices)
+
+            if np.isfinite(p_value):
+                valid_feature_ids.append(fid)
+                p_values.append(p_value)
+
+        if len(p_values) == 0:
+            keep_feature_ids = []
+        else:
+            rejected, _, _, _ = multipletests(p_values, alpha=alpha, method='fdr_bh')
+            keep_feature_ids = [
+                fid
+                for fid, r in zip(valid_feature_ids, rejected)
+                if r
+            ]
+
+        if cache_path is not None:
+            try:
+                _save_feature_ids_to_cache(cache_path, keep_feature_ids)
+                print(f'[cache][feature_filter] saved: {cache_path}')
+            except Exception as exc:
+                print(f'[cache][feature_filter] failed to save {cache_path}: {exc}')
+
+        filtered_context = dict(context)
+        filtered_context['codes_csr'] = codes_csr[:, keep_feature_ids]
+        filtered_context['filtered_feature_ids'] = keep_feature_ids
+        if cache_path is not None:
+            filtered_context['feature_filter_cache_path'] = str(cache_path)
+        return filtered_context
+
+
+class NonZeroMaxFeatureFilter(BaseFeatureFilter):
+    """Keep features whose column-wise maximum activation is not zero."""
+
+    name = 'nonzero_max'
+    description = 'Keep features with max activation != 0'
+
+
+    def __init__(
+        self,
+        show_progress: bool = False,
+    ):
+        pass
+
+    def apply_filter(
+        self,
+        context: Dict[str, object],
+    ) -> Dict[str, object]:
+        codes_csr = context['codes_csr']
+        keep_feature_ids = np.flatnonzero(np.asarray(codes_csr.getnnz(axis=0)).ravel() > 0).tolist()
+
+        filtered_context = dict(context)
+        filtered_context['codes_csr'] = codes_csr[:, keep_feature_ids]
+        filtered_context['filtered_feature_ids'] = keep_feature_ids
+        return filtered_context
+
+
+FILTER_REGISTRY: Mapping[str, BaseFeatureFilter] = {}
+
+FILTER_BUILDERS: Mapping[str, Callable[..., BaseFeatureFilter]] = {
+    'nonzero_max': NonZeroMaxFeatureFilter,
+    'kruskal_wallis': KruskalWallisFeatureFilter,
+}
+
+
+def get_filter_registry() -> Dict[str, BaseFeatureFilter]:
+    """Return a mutable copy of built-in filter registry."""
+    return dict(FILTER_REGISTRY)
+
+
+def load_filter(filter_name: str) -> BaseFeatureFilter:
+    """Load built-in feature filter by name."""
+    filter_obj = FILTER_REGISTRY.get(filter_name)
+    if filter_obj is None:
+        available = ', '.join(sorted(FILTER_REGISTRY.keys()))
+        raise ValueError(f'Unknown filter {filter_name!r}. Available: {available}')
+    return filter_obj
+
+
+def load_filters(filter_names: Sequence[str]) -> List[BaseFeatureFilter]:
+    """Load an ordered sequence of filters."""
+    return [load_filter(name) for name in filter_names]
+
+
+def build_filter(
+    filter_name: str,
+    params: Mapping[str, object] | None = None,
+) -> BaseFeatureFilter:
+    """Build a feature filter instance from name and kwargs-like params."""
+    builder = FILTER_BUILDERS.get(filter_name)
+    kwargs = dict(params or {})
+    return builder(**kwargs)

analysis/features/feature_indexing.py

@@ -0,0 +1,116 @@
+"""CSR activations with global SAE id per column: ``FeatureMatrix``."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+import scipy.sparse as sp
+
+
+def global_to_column(column_feature_ids: Sequence[int], global_feature_id: int) -> int | None:
+    try:
+        return column_feature_ids.index(int(global_feature_id))
+    except ValueError:
+        return None
+
+
+def column_for_global(column_feature_ids: Sequence[int], global_feature_id: int) -> int:
+    col = global_to_column(column_feature_ids, global_feature_id)
+    if col is None:
+        raise KeyError(
+            f'Global feature id {int(global_feature_id)!r} is not in column_feature_ids'
+        )
+    return col
+
+
+def global_id_set(column_feature_ids: Sequence[int]) -> frozenset[int]:
+    return frozenset(int(fid) for fid in column_feature_ids)
+
+
+@dataclass(frozen=True)
+class FeatureMatrix:
+    """Sparse activations; ``column_feature_ids[col]`` is the global SAE id for column ``col``."""
+
+    codes: sp.csr_matrix
+    column_feature_ids: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        n_cols = int(self.codes.shape[1])
+        if len(self.column_feature_ids) != n_cols:
+            raise ValueError(
+                f'column_feature_ids length ({len(self.column_feature_ids)}) '
+                f'must match codes.shape[1] ({n_cols})'
+            )
+
+    @classmethod
+    def from_codes(
+        cls,
+        codes: sp.csr_matrix,
+        column_feature_ids: Sequence[int] | None = None,
+    ) -> FeatureMatrix:
+        if column_feature_ids is None:
+            labels = tuple(range(int(codes.shape[1])))
+        else:
+            labels = tuple(int(fid) for fid in column_feature_ids)
+        return cls(codes=codes, column_feature_ids=labels)
+
+    @property
+    def n_features(self) -> int:
+        return len(self.column_feature_ids)
+
+    def column_for(self, global_feature_id: int) -> int:
+        return column_for_global(self.column_feature_ids, global_feature_id)
+
+    def global_to_column(self, global_feature_id: int) -> int | None:
+        return global_to_column(self.column_feature_ids, global_feature_id)
+
+    def global_id_set(self) -> frozenset[int]:
+        return global_id_set(self.column_feature_ids)
+
+    def subset(self, global_feature_ids: Sequence[int] | None = None) -> FeatureMatrix:
+        """Keep only the given global ids (order preserved). ``None`` = all columns."""
+        if global_feature_ids is None:
+            return self
+        cols: list[int] = []
+        labels: list[int] = []
+        for gid in global_feature_ids:
+            col = column_for_global(self.column_feature_ids, int(gid))
+            cols.append(col)
+            labels.append(int(gid))
+        return FeatureMatrix(
+            codes=self.codes[:, cols],
+            column_feature_ids=tuple(labels),
+        )
+
+    def slice_columns(self, column_indices: Sequence[int]) -> FeatureMatrix:
+        """Slice by matrix column index; global ids follow the selected columns."""
+        idx = [int(i) for i in column_indices]
+        return FeatureMatrix(
+            codes=self.codes[:, idx],
+            column_feature_ids=tuple(self.column_feature_ids[i] for i in idx),
+        )
+
+    def with_row_mask(self, mask) -> FeatureMatrix:
+        """Return a view with rows filtered (e.g. one distortion group)."""
+        return FeatureMatrix(
+            codes=self.codes[mask],
+            column_feature_ids=self.column_feature_ids,
+        )
+
+
+def columns_for_global_subset(
+    column_feature_ids: Sequence[int],
+    global_feature_ids: Sequence[int] | None = None,
+) -> tuple[list[int] | None, list[int]]:
+    """Deprecated: use ``FeatureMatrix.subset``."""
+    col_ids = [int(fid) for fid in column_feature_ids]
+    if global_feature_ids is None:
+        return None, col_ids
+    cols: list[int] = []
+    names: list[int] = []
+    for gid in global_feature_ids:
+        col = column_for_global(col_ids, int(gid))
+        cols.append(col)
+        names.append(int(gid))
+    return cols, names

analysis/features/feature_matrix.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import List, Sequence
+
+import numpy as np
+
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+def _zscore_columns(x: np.ndarray) -> np.ndarray:
+    mu = x.mean(axis=0, keepdims=True)
+    sigma = x.std(axis=0, keepdims=True)
+    sigma[sigma < 1e-8] = 1.0
+    return (x - mu) / sigma
+
+
+def build_feature_matrix(
+    feature_id: int,
+    features: FeatureMatrix,
+    rows_img: Sequence[int],
+    rows_patch: Sequence[int],
+    n_images_used: int,
+    patches_per_image: int,
+) -> np.ndarray:
+    """Build image-level matrix (n_images_used, patches_per_image) for one global SAE feature."""
+    col = features.column_for(feature_id)
+    vals = features.codes[:, col].toarray().ravel().astype(np.float32)
+    x = np.zeros((int(n_images_used), int(patches_per_image)), dtype=np.float32)
+    x[rows_img, rows_patch] = vals
+    return _zscore_columns(x)
+
+
+def build_image_feature_matrix(
+    features: FeatureMatrix,
+    image_row_indices: Sequence[Sequence[int]],
+    n_images_used: int,
+    aggregation_mode: str = 'max',
+) -> np.ndarray:
+    """Build image-level feature matrix (n_images_used, n_features_total).
+
+    `image_row_indices` is a sequence where each element is an array of row indices
+    (patch indices) belonging to that image.
+    Aggregation modes: 'max', 'sum', 'mean_acts' (mean over positive activations) or 'mean'.
+    """
+    return _zscore_columns(
+        build_image_feature_matrix_raw(
+            features,
+            image_row_indices,
+            n_images_used,
+            aggregation_mode,
+        )
+    )
+
+
+def build_image_feature_matrix_raw(
+    features: FeatureMatrix,
+    image_row_indices: Sequence[Sequence[int]],
+    n_images_used: int,
+    aggregation_mode: str = 'max',
+) -> np.ndarray:
+    """Build an unnormalized image-level feature matrix.
+
+    This keeps the same aggregation logic as ``build_image_feature_matrix`` but
+    skips the final z-score normalization so it can be used for paired deltas.
+    """
+    mode = str(aggregation_mode)
+    if mode not in {'max', 'sum', 'mean_acts', 'mean'}:
+        raise ValueError(f'Unknown aggregation mode: {mode!r}')
+
+    n_features_total = features.n_features
+    codes_subset = features.codes
+    x = np.zeros((int(n_images_used), int(n_features_total)), dtype=np.float32)
+
+    for img_i, row_ids in enumerate(image_row_indices):
+        if len(row_ids) == 0:
+            continue
+        chunk = codes_subset[row_ids, :]
+        if mode == 'max':
+            x[img_i] = np.asarray(chunk.max(axis=0)).ravel().astype(np.float32)
+        elif mode == 'sum':
+            x[img_i] = np.asarray(chunk.sum(axis=0)).ravel().astype(np.float32)
+        elif mode == 'mean':
+            x[img_i] = np.asarray(chunk.mean(axis=0)).ravel().astype(np.float32)
+        else:  # mean_acts: mean over positive activations
+            if hasattr(chunk, 'multiply'):
+                positive_mask = chunk > 0
+                positive_chunk = chunk.multiply(positive_mask)
+                pos_sum = np.asarray(positive_chunk.sum(axis=0)).ravel().astype(np.float32)
+                pos_count = np.asarray(positive_mask.sum(axis=0)).ravel().astype(np.float32)
+            else:
+                chunk_arr = np.asarray(chunk, dtype=np.float32)
+                positive_mask = chunk_arr > 0
+                pos_sum = np.asarray((chunk_arr * positive_mask).sum(axis=0)).ravel().astype(np.float32)
+                pos_count = np.asarray(positive_mask.sum(axis=0)).ravel().astype(np.float32)
+            mean_pos = np.zeros(int(n_features_total), dtype=np.float32)
+            nonzero = pos_count > 0
+            mean_pos[nonzero] = pos_sum[nonzero] / pos_count[nonzero]
+            x[img_i] = mean_pos
+
+    return x

analysis/features/feature_selectors.py

@@ -0,0 +1,918 @@
+"""Feature selection strategies for automated SAE analysis reports."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+
+from analysis.features.feature_indexing import FeatureMatrix
+from analysis.metrics.correlations import compute_distortion_correlations
+from analysis.metrics.mutual_information import compute_distortion_mutual_information
+from analysis.metrics.paired_deltas import build_paired_delta_tables
+from analysis.metrics.roc_auc import compute_distortion_roc_auc
+
+
+@dataclass(frozen=True)
+class SelectorResult:
+    """Output bundle produced by selector execution."""
+
+    selected_features: pd.DataFrame
+    metric_tables: Dict[str, pd.DataFrame]
+    selection_stats: Dict[str, object]
+    selector_name: str
+    selector_description: str
+    relation_label: str
+
+
+@dataclass(frozen=True)
+class SelectorInputs:
+    """Typed selector inputs with required core data and optional pristine data."""
+
+    meta_df: pd.DataFrame
+    features: FeatureMatrix
+    dataset: str
+    selector_top_k: int = 20
+    selector_params: Mapping[str, object] = field(default_factory=dict)
+    cache_paths: Any | None = None
+    pristine_meta_df: pd.DataFrame | None = None
+    pristine_features: FeatureMatrix | None = None
+    feature_weight_norms: Mapping[int, float] | None = None
+
+
+def _match_metric_df_columns(corr_df: pd.DataFrame, feature_ids: Sequence[int]) -> List[object]:
+    id_set = {int(x) for x in feature_ids}
+    matched_columns: List[object] = []
+    for col in corr_df.columns:
+        try:
+            col_id = int(col)
+        except (TypeError, ValueError):
+            continue
+        if col_id in id_set:
+            matched_columns.append(col)
+    return matched_columns
+
+
+def _build_feature_importance_table(
+    importance_series: pd.Series,
+    top_k: int,
+    metric_name: str,
+) -> pd.DataFrame:
+    top_series = importance_series.nlargest(top_k)
+    result = pd.DataFrame({
+        'feature_id': [int(fid) for fid in top_series.index.tolist()],
+        'importance_score': [float(x) for x in top_series.values.tolist()],
+    })
+    result['importance_rank'] = range(1, len(result) + 1)
+    result['importance_metric'] = metric_name
+    return result[['feature_id', 'importance_score', 'importance_rank', 'importance_metric']]
+
+
+def _collect_table_stats(
+    stats: Dict[str, object],
+    table_key: str,
+    table_df: pd.DataFrame,
+    feature_ids: Sequence[int],
+) -> None:
+    matched_cols = _match_metric_df_columns(table_df, feature_ids)
+    if not matched_cols:
+        return
+
+    table_values = table_df.loc[:, matched_cols]
+    stats[f'{table_key}_matched_cols'] = len(matched_cols)
+    stats[f'{table_key}_mean_abs'] = float(table_values.abs().values.mean())
+    stats[f'{table_key}_max_abs'] = float(table_values.abs().values.max())
+
+
+class BaseFeatureSelector(ABC):
+    """Common interface for feature selectors."""
+
+    name: str
+    description: str
+    relation_label: str = 'feature relevance'
+
+    @abstractmethod
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        """Compute metric tables required by the selector."""
+
+    @abstractmethod
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        """Return per-feature importance scores and metric name."""
+
+    def execute(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> SelectorResult:
+        """Run full selector pipeline and return selected features + computed measures."""
+        metric_tables = self.compute_metric_tables(inputs, feature_ids=feature_ids)
+        importance_series, metric_name = self.compute_importance(
+            metric_tables=metric_tables,
+            inputs=inputs,
+            feature_ids=feature_ids,
+        )
+        top_k = int(inputs.selector_top_k)
+        selected_features = _build_feature_importance_table(
+            importance_series=importance_series,
+            top_k=top_k,
+            metric_name=metric_name,
+        )
+
+        selected_features = selected_features.copy()
+        selected_features['feature_id'] = selected_features['feature_id'].astype(int)
+        selected_features['importance_score'] = selected_features['importance_score'].astype(float)
+        selected_features['importance_rank'] = selected_features['importance_rank'].astype(int)
+
+        selected_features = selected_features.drop_duplicates(
+            subset=['feature_id'], keep='first').reset_index(drop=True)
+        if selected_features.empty:
+            raise ValueError('Selector returned empty feature list')
+
+        selected_features['importance_rank'] = range(1, len(selected_features) + 1)
+        stats = self.selection_stats(selected_features, metric_tables)
+        return SelectorResult(
+            selected_features=selected_features,
+            metric_tables=metric_tables,
+            selection_stats=stats,
+            selector_name=self.name,
+            selector_description=self.description,
+            relation_label=self.relation_label,
+        )
+
+    def select_features(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> pd.DataFrame:
+        """Template method: normalize and return top-k with importance stats."""
+        metric_tables = self.compute_metric_tables(inputs, feature_ids=feature_ids)
+        importance_series, metric_name = self.compute_importance(
+            metric_tables=metric_tables,
+            inputs=inputs,
+            feature_ids=feature_ids,
+        )
+
+        top_k = int(inputs.selector_top_k)
+        selected_features = _build_feature_importance_table(
+            importance_series=importance_series,
+            top_k=top_k,
+            metric_name=metric_name,
+        )
+
+        selected_features = selected_features.copy()
+        selected_features['feature_id'] = selected_features['feature_id'].astype(int)
+        selected_features['importance_score'] = selected_features['importance_score'].astype(float)
+        selected_features['importance_rank'] = selected_features['importance_rank'].astype(int)
+
+        selected_features = selected_features.drop_duplicates(
+            subset=['feature_id'], keep='first').reset_index(drop=True)
+        if selected_features.empty:
+            raise ValueError('Selector returned empty feature list')
+
+        selected_features['importance_rank'] = range(1, len(selected_features) + 1)
+        return selected_features
+
+    def selection_stats(
+        self,
+        selected_features: pd.DataFrame,
+        metric_tables: Mapping[str, pd.DataFrame],
+    ) -> Dict[str, object]:
+        feature_ids = selected_features['feature_id'].tolist()
+        stats: Dict[str, object] = {
+            'selector_name': self.name,
+            'selector_description': self.description,
+            'n_selected': len(feature_ids),
+        }
+        if not selected_features.empty:
+            stats['mean_importance_score'] = float(selected_features['importance_score'].mean())
+            stats['max_importance_score'] = float(selected_features['importance_score'].max())
+            stats['min_importance_score'] = float(selected_features['importance_score'].min())
+        if feature_ids:
+            numeric_ids = [int(x) for x in feature_ids]
+            stats['min_feature_id'] = min(numeric_ids)
+            stats['max_feature_id'] = max(numeric_ids)
+
+        for corr_key in ('corr_type_df', 'corr_group_df'):
+            corr_df = metric_tables.get(corr_key)
+            if isinstance(corr_df, pd.DataFrame) and len(feature_ids) > 0:
+                matched_cols = _match_metric_df_columns(corr_df, feature_ids)
+                if matched_cols:
+                    abs_corr = corr_df.loc[:, matched_cols].abs()
+                    stats[f'{corr_key}_matched_cols'] = len(matched_cols)
+                    stats[f'{corr_key}_mean_abs_corr'] = float(abs_corr.values.mean())
+                    stats[f'{corr_key}_max_abs_corr'] = float(abs_corr.values.max())
+
+        for mi_key in ('mi_type_df', 'mi_group_df'):
+            mi_df = metric_tables.get(mi_key)
+            if isinstance(mi_df, pd.DataFrame) and len(feature_ids) > 0:
+                matched_cols = _match_metric_df_columns(mi_df, feature_ids)
+                if matched_cols:
+                    mi_values = mi_df.loc[:, matched_cols]
+                    stats[f'{mi_key}_matched_cols'] = len(matched_cols)
+                    stats[f'{mi_key}_mean'] = float(mi_values.values.mean())
+                    stats[f'{mi_key}_max'] = float(mi_values.values.max())
+
+        for auc_key in ('auc_type_df', 'auc_group_df'):
+            auc_df = metric_tables.get(auc_key)
+            if isinstance(auc_df, pd.DataFrame) and len(feature_ids) > 0:
+                matched_cols = _match_metric_df_columns(auc_df, feature_ids)
+                if matched_cols:
+                    auc_values = auc_df.loc[:, matched_cols]
+                    stats[f'{auc_key}_matched_cols'] = len(matched_cols)
+                    stats[f'{auc_key}_mean'] = float(auc_values.values.mean())
+                    stats[f'{auc_key}_max'] = float(auc_values.values.max())
+                    stats[f'{auc_key}_min'] = float(auc_values.values.min())
+
+        for key, value in metric_tables.items():
+            if not (isinstance(key, str) and key.startswith('paired_') and key.endswith('_df')):
+                continue
+            if isinstance(value, pd.DataFrame) and len(feature_ids) > 0:
+                _collect_table_stats(stats, key, value, feature_ids)
+
+        return stats
+
+
+@dataclass(frozen=True)
+class TopKAbsCorrSelector(BaseFeatureSelector):
+    """Top-k by max absolute correlation across category rows."""
+
+    name: str = 'topk_abs_corr'
+    description: str = 'Top-k features by max absolute dist_type correlation.'
+    corr_key: str = 'corr_type_df'
+    relation_label: str = 'высокая корреляция'
+
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        selector_params = dict(inputs.selector_params)
+        level = str(selector_params.get('level', 'patch'))
+        binarize = bool(selector_params.get('binarize', False))
+        group_col = 'dist_group' if self.corr_key == 'corr_group_df' else 'dist_type'
+
+        cache_path = selector_params.get('cache_path')
+        if cache_path is None:
+            cache_paths = inputs.cache_paths
+            if cache_paths is not None:
+                attr_map = {
+                    ('corr_type_df', 'patch'): 'corr_type_patch_cache_path',
+                    ('corr_group_df', 'patch'): 'corr_group_patch_cache_path',
+                    ('corr_type_df', 'image'): 'corr_type_cache_path',
+                    ('corr_group_df', 'image'): 'corr_group_cache_path',
+                }
+                cache_attr = attr_map.get((self.corr_key, level))
+                if cache_attr is not None and hasattr(cache_paths, cache_attr):
+                    cache_path = getattr(cache_paths, cache_attr)
+
+        corr_df = compute_distortion_correlations(
+            meta=inputs.meta_df,
+            features=inputs.features,
+            group_col=group_col,
+            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
+            cache_path=cache_path,
+        )
+        return {self.corr_key: corr_df}
+
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        corr_df = metric_tables[self.corr_key]
+        if feature_ids is not None:
+            corr_df = corr_df[[int(fid) for fid in feature_ids]]
+        importance_series = corr_df.abs().max(axis=0)
+        return importance_series, f'max_abs_corr:{self.corr_key}'
+
+
+@dataclass(frozen=True)
+class TopKAbsCorrGroupSelector(TopKAbsCorrSelector):
+    """Top-k by max absolute correlation at distortion-group level."""
+
+    name: str = 'topk_abs_corr_group'
+    description: str = 'Top-k features by max absolute dist_group correlation.'
+    corr_key: str = 'corr_group_df'
+
+
+@dataclass(frozen=True)
+class TopKMutualInfoSelector(BaseFeatureSelector):
+    """Top-k by max mutual information across category rows."""
+
+    name: str = 'topk_mi'
+    description: str = 'Top-k features by max dist_type mutual information.'
+    mi_key: str = 'mi_type_df'
+    relation_label: str = 'высокая взаимная информация'
+
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        selector_params = dict(inputs.selector_params)
+        level = str(selector_params.get('level', 'patch'))
+        binarize = bool(selector_params.get('binarize', False))
+        n_bins = int(selector_params.get('n_bins', 16))
+        feature_chunk_size = int(selector_params.get('feature_chunk_size', 512))
+        group_col = 'dist_group' if self.mi_key == 'mi_group_df' else 'dist_type'
+
+        cache_path = selector_params.get('cache_path')
+        if cache_path is None:
+            cache_paths = inputs.cache_paths
+            if cache_paths is not None:
+                attr_map = {
+                    ('mi_type_df', 'patch'): 'mi_type_patch_cache_path',
+                    ('mi_group_df', 'patch'): 'mi_group_patch_cache_path',
+                    ('mi_type_df', 'image'): 'mi_type_cache_path',
+                    ('mi_group_df', 'image'): 'mi_group_cache_path',
+                }
+                cache_attr = attr_map.get((self.mi_key, level))
+                if cache_attr is not None and hasattr(cache_paths, cache_attr):
+                    cache_path = getattr(cache_paths, cache_attr)
+
+        mi_df = compute_distortion_mutual_information(
+            meta=inputs.meta_df,
+            features=inputs.features,
+            group_col=group_col,
+            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
+            binarize=binarize,
+            level=level,
+            n_bins=n_bins,
+            feature_chunk_size=feature_chunk_size,
+            cache_path=cache_path,
+        )
+        return {self.mi_key: mi_df}
+
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        mi_df = metric_tables[self.mi_key]
+        if feature_ids is not None:
+            mi_df = mi_df[[int(fid) for fid in feature_ids]]
+        importance_series = mi_df.max(axis=0)
+        return importance_series, f'max_mi:{self.mi_key}'
+
+
+@dataclass(frozen=True)
+class TopKMutualInfoGroupSelector(TopKMutualInfoSelector):
+    """Top-k by max mutual information at distortion-group level."""
+
+    name: str = 'topk_mi_group'
+    description: str = 'Top-k features by max dist_group mutual information.'
+    mi_key: str = 'mi_group_df'
+
+
+@dataclass(frozen=True)
+class TopKRocAucSelector(BaseFeatureSelector):
+    """Top-k by one-vs-rest ROC-AUC across distortion categories."""
+
+    name: str = 'topk_auc'
+    description: str = 'Top-k features by one-vs-rest ROC-AUC at dist_type level.'
+    auc_key: str = 'auc_type_df'
+    relation_label: str = 'высокая ROC-AUC различимость'
+
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        selector_params = dict(inputs.selector_params)
+        level = str(selector_params.get('level', 'patch'))
+        feature_chunk_size = int(selector_params.get('feature_chunk_size', 512))
+        group_col = 'dist_group' if self.auc_key == 'auc_group_df' else 'dist_type'
+
+        cache_path = selector_params.get('cache_path')
+        if cache_path is None:
+            cache_paths = inputs.cache_paths
+            if cache_paths is not None:
+                attr_map = {
+                    ('auc_type_df', 'patch'): 'auc_type_patch_cache_path',
+                    ('auc_group_df', 'patch'): 'auc_group_patch_cache_path',
+                    ('auc_type_df', 'image'): 'auc_type_cache_path',
+                    ('auc_group_df', 'image'): 'auc_group_cache_path',
+                }
+                cache_attr = attr_map.get((self.auc_key, level))
+                if cache_attr is not None and hasattr(cache_paths, cache_attr):
+                    cache_path = getattr(cache_paths, cache_attr)
+
+        auc_df = compute_distortion_roc_auc(
+            meta=inputs.meta_df,
+            features=inputs.features,
+            group_col=group_col,
+            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
+            level=level,
+            feature_chunk_size=feature_chunk_size,
+            cache_path=cache_path,
+        )
+        return {self.auc_key: auc_df}
+
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        auc_df = metric_tables[self.auc_key]
+        if feature_ids is not None:
+            auc_df = auc_df[[int(fid) for fid in feature_ids]]
+
+        # Symmetric separability score: high both for direct and inverse predictors.
+        direct_auc = auc_df.max(axis=0)
+        inverse_auc = 1.0 - auc_df.min(axis=0)
+        importance_series = pd.concat([direct_auc, inverse_auc], axis=1).max(axis=1)
+        return importance_series, f'symmetric_auc:{self.auc_key}'
+
+
+@dataclass(frozen=True)
+class TopKRocAucGroupSelector(TopKRocAucSelector):
+    """Top-k by one-vs-rest ROC-AUC at distortion-group level."""
+
+    name: str = 'topk_auc_group'
+    description: str = 'Top-k features by one-vs-rest ROC-AUC at dist_group level.'
+    auc_key: str = 'auc_group_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaSelector(BaseFeatureSelector):
+    """Top-k by original-vs-distorted paired activation deltas."""
+
+    name: str = 'paired_delta'
+    description: str = 'Top-k features by max relative paired activation delta.'
+    delta_key: str = 'paired_relative_dist_type_df'
+    relation_label: str = 'высокая парная дельта активаций'
+
+    @staticmethod
+    def _delta_mode_from_key(delta_key: str) -> str:
+        if '_signed_' in delta_key:
+            return 'signed'
+        if '_relative_' in delta_key:
+            return 'relative'
+        return 'abs'
+
+    @staticmethod
+    def _group_col_from_key(delta_key: str) -> str:
+        if 'dist_group' in delta_key:
+            return 'dist_group'
+        return 'dist_type'
+
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        if inputs.pristine_meta_df is None or inputs.pristine_features is None:
+            raise ValueError(
+                f"Selector '{self.name}' requires 'pristine_meta_df' and 'pristine_features' for paired deltas."
+            )
+
+        selector_params = dict(inputs.selector_params)
+        level = str(selector_params.get('level', 'patch'))
+        ranking_mode = str(selector_params.get('ranking_mode', 'default')).strip().lower()
+        if ranking_mode not in {'default', 'weighted_average'}:
+            raise ValueError(
+                f"Selector '{self.name}': unsupported ranking_mode={ranking_mode!r}. "
+                "Allowed: {'default', 'weighted_average'}."
+            )
+        delta_mode = self._delta_mode_from_key(self.delta_key)
+        group_col = self._group_col_from_key(self.delta_key)
+
+        cache_prefix = selector_params.get('cache_prefix')
+        if cache_prefix is None:
+            cache_paths = inputs.cache_paths
+            dataset_cache_dir = getattr(cache_paths, 'dataset_cache_dir', None) if cache_paths is not None else None
+            if dataset_cache_dir is not None:
+                cache_prefix = str(Path(dataset_cache_dir) / f'paired_delta_{level}')
+
+        tables = build_paired_delta_tables(
+            distorted_meta=inputs.meta_df,
+            distorted=inputs.features,
+            original_meta=inputs.pristine_meta_df,
+            original=inputs.pristine_features,
+            group_cols=(group_col,),
+            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
+            delta_mode=delta_mode,
+            level=level,
+            cache_prefix=cache_prefix,
+        )
+        if self.delta_key not in tables:
+            available = ', '.join(sorted(tables.keys()))
+            raise ValueError(
+                f"Selector '{self.name}' expected table '{self.delta_key}', available: {available}"
+            )
+
+        metric_table = tables[self.delta_key]
+        if ranking_mode == 'weighted_average':
+            if not inputs.feature_weight_norms:
+                raise ValueError(
+                    f"Selector '{self.name}' (ranking_mode='weighted_average') requires per-feature decoder weight norms."
+                )
+
+            norm_by_feature_id: Dict[int, float] = {}
+            for raw_feature_id, raw_norm in inputs.feature_weight_norms.items():
+                try:
+                    feature_id = int(raw_feature_id)
+                    norm_by_feature_id[feature_id] = float(raw_norm)
+                except (TypeError, ValueError):
+                    continue
+
+            if not norm_by_feature_id:
+                raise ValueError(
+                    f"Selector '{self.name}' (ranking_mode='weighted_average') received empty/invalid decoder weight norms."
+                )
+
+            weight_series = pd.Series(index=metric_table.columns, dtype='float64')
+            missing_columns: List[object] = []
+            for column in metric_table.columns:
+                try:
+                    weight_series.loc[column] = norm_by_feature_id[int(column)]
+                except (TypeError, ValueError, KeyError):
+                    missing_columns.append(column)
+
+            if missing_columns:
+                missing_preview = ', '.join(str(x) for x in missing_columns[:8])
+                raise ValueError(
+                    f"Selector '{self.name}' has no decoder norms for {len(missing_columns)} features "
+                    f"(first: {missing_preview})."
+                )
+
+            weighted_abs = metric_table.abs().mul(weight_series, axis=1)
+            col_min = weighted_abs.min(axis=0)
+            col_max = weighted_abs.max(axis=0)
+            denom = col_max - col_min
+            normalized = weighted_abs.subtract(col_min, axis=1)
+            nonzero_denom = denom > 0
+            if nonzero_denom.any():
+                normalized.loc[:, nonzero_denom] = normalized.loc[:, nonzero_denom].div(denom[nonzero_denom], axis=1)
+            if (~nonzero_denom).any():
+                normalized.loc[:, ~nonzero_denom] = 0.0
+            metric_table = normalized.astype('float32', copy=False)
+
+        return {self.delta_key: metric_table}
+
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        delta_df = metric_tables[self.delta_key]
+        if feature_ids is not None:
+            delta_df = delta_df[[int(fid) for fid in feature_ids]]
+
+        ranking_mode = 'default'
+        if inputs is not None:
+            ranking_mode = str(inputs.selector_params.get('ranking_mode', 'default')).strip().lower()
+
+        if ranking_mode == 'weighted_average':
+            importance_series = delta_df.mean(axis=0)
+            return importance_series, f'mean_weighted_delta_norm01:{self.delta_key}'
+
+        delta_abs_mean = delta_df.abs().mean(axis=0)
+        return delta_abs_mean, f'mean_abs_delta:{self.delta_key}'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaAbsTypeSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_abs_type'
+    description: str = 'Top-k features by absolute paired delta at dist_type level.'
+    delta_key: str = 'paired_abs_dist_type_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaAbsGroupSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_abs_group'
+    description: str = 'Top-k features by absolute paired delta at dist_group level.'
+    delta_key: str = 'paired_abs_dist_group_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaSignedTypeSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_signed_type'
+    description: str = 'Top-k features by signed paired delta at dist_type level.'
+    delta_key: str = 'paired_signed_dist_type_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaSignedGroupSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_signed_group'
+    description: str = 'Top-k features by signed paired delta at dist_group level.'
+    delta_key: str = 'paired_signed_dist_group_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaRelativeTypeSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_relative_type'
+    description: str = 'Top-k features by relative paired delta at dist_type level.'
+    delta_key: str = 'paired_relative_dist_type_df'
+
+
+@dataclass(frozen=True)
+class TopKPairedDeltaRelativeGroupSelector(TopKPairedDeltaSelector):
+    name: str = 'paired_delta_relative_group'
+    description: str = 'Top-k features by relative paired delta at dist_group level.'
+    delta_key: str = 'paired_relative_dist_group_df'
+
+
+@dataclass(frozen=True)
+class TopKIoUSelector(BaseFeatureSelector):
+    """Top-k by median IoU (Intersection over Union) with distortion masks at dist_type level."""
+
+    name: str = 'topk_iou_type'
+    description: str = 'Top-k features by median IoU with distortion regions at dist_type level.'
+    relation_label: str = 'высокая пространственная локализация'
+
+    def compute_metric_tables(
+        self,
+        inputs: SelectorInputs,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Dict[str, pd.DataFrame]:
+        from analysis.metrics.iou_utils import compute_iou_per_distortion_type_and_feature
+
+        selector_params = dict(inputs.selector_params)
+        iou_key = inputs.selector_params.get('iou_key', 'iou_type_df')
+        group_col = 'dist_group' if iou_key == 'iou_group_df' else 'dist_type'
+
+        cache_path = selector_params.get('cache_path')
+        if cache_path is None:
+            cache_paths = inputs.cache_paths
+            if cache_paths is not None:
+                attr_map = {
+                    ('iou_type_df',): 'iou_type_cache_path',
+                    ('iou_group_df',): 'iou_group_cache_path',
+                }
+                cache_attr = attr_map.get((iou_key,))
+                if cache_attr is not None and hasattr(cache_paths, cache_attr):
+                    cache_path = getattr(cache_paths, cache_attr)
+
+        # Infer spatial shape from patch counts
+        # Assume square grid: n_patches = height * width
+        n_samples_total = inputs.features.codes.shape[0]
+        n_unique_images = inputs.meta_df['image_idx'].nunique()
+        patches_per_image = n_samples_total // n_unique_images if n_unique_images > 0 else 49
+        spatial_dim = int(np.sqrt(patches_per_image))
+        
+        # Allow override via selector_params
+        if 'spatial_shape' in selector_params:
+            spatial_shape = tuple(selector_params['spatial_shape'])
+        else:
+            spatial_shape = (spatial_dim, spatial_dim)
+
+        feature_batch_size = int(selector_params.get('feature_batch_size', 1024))
+
+        iou_df = compute_iou_per_distortion_type_and_feature(
+            features=inputs.features,
+            meta_df=inputs.meta_df,
+            spatial_shape=spatial_shape,
+            group_col=group_col,
+            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
+            feature_batch_size=feature_batch_size,
+            cache_path=cache_path,
+            dataset=inputs.dataset,
+        )
+
+        return {iou_key: iou_df}
+
+    def compute_importance(
+        self,
+        metric_tables: Dict[str, pd.DataFrame],
+        inputs: SelectorInputs | None = None,
+        feature_ids: Sequence[int] | None = None,
+    ) -> Tuple[pd.Series, str]:
+        iou_key = inputs.selector_params.get('iou_key', 'iou_type_df')
+        iou_df = metric_tables[iou_key]
+        if iou_df.empty:
+            # Return empty series if no IoU data
+            if feature_ids is not None:
+                feature_ids_int = [int(fid) for fid in feature_ids]
+                importance_series = pd.Series(0.0, index=feature_ids_int)
+            else:
+                importance_series = pd.Series(dtype=float)
+            return importance_series, f'max_median_iou:{iou_key}'
+        
+        if feature_ids is not None:
+            iou_df = iou_df[[int(fid) for fid in feature_ids]]
+
+        # Take max IoU across all distortion types (each row is a distortion type)
+        importance_series = iou_df.max(axis=0)
+        return importance_series, f'max_median_iou:{iou_key}'
+
+
+@dataclass(frozen=True)
+class TopKIoUDistGroupSelector(TopKIoUSelector):
+    """Top-k by median IoU (Intersection over Union) with distortion masks at dist_group level."""
+
+    name: str = 'topk_iou_group'
+    description: str = 'Top-k features by median IoU with distortion regions at dist_group level.'
+    iou_key: str = 'iou_group_df'
+
+
+SELECTOR_REGISTRY: Mapping[str, BaseFeatureSelector] = {
+    'topk_abs_corr': TopKAbsCorrSelector(),
+    'topk_abs_corr_group': TopKAbsCorrGroupSelector(),
+    'topk_mi': TopKMutualInfoSelector(),
+    'topk_mi_group': TopKMutualInfoGroupSelector(),
+    'topk_auc_type_image': TopKRocAucSelector(),
+    'topk_auc_group_image': TopKRocAucGroupSelector(),
+    'topk_auc_type_patch': TopKRocAucSelector(),
+    'topk_auc_group_patch': TopKRocAucGroupSelector(),
+    'paired_delta_abs_type': TopKPairedDeltaAbsTypeSelector(),
+    'paired_delta_abs_group': TopKPairedDeltaAbsGroupSelector(),
+    'paired_delta_signed_type': TopKPairedDeltaSignedTypeSelector(),
+    'paired_delta_signed_group': TopKPairedDeltaSignedGroupSelector(),
+    'paired_delta_relative_type': TopKPairedDeltaRelativeTypeSelector(),
+    'paired_delta_relative_group': TopKPairedDeltaRelativeGroupSelector(),
+    'topk_iou_type': TopKIoUSelector(),
+    'topk_iou_group': TopKIoUDistGroupSelector(),
+}
+
+
+def get_selector_registry() -> Dict[str, BaseFeatureSelector]:
+    """Return a mutable copy of built-in selector registry."""
+    return dict(SELECTOR_REGISTRY)
+
+
+def load_selector(selector_name: str) -> BaseFeatureSelector:
+    """Load built-in selector by name."""
+    normalized = selector_name or 'topk_abs_corr'
+    selector = SELECTOR_REGISTRY.get(normalized)
+    if selector is None:
+        available = ', '.join(sorted(SELECTOR_REGISTRY.keys()))
+        raise ValueError(f'Unknown selector {selector_name!r}. Available: {available}')
+    return selector
+
+
+def run_selector(
+    selector: BaseFeatureSelector,
+    inputs: SelectorInputs,
+    feature_ids: Sequence[int] | None = None,
+) -> pd.DataFrame:
+    """Compatibility wrapper around selector.select_features."""
+    selected_features = selector.select_features(inputs, feature_ids=feature_ids)
+    required_cols = {'feature_id', 'importance_score', 'importance_rank', 'importance_metric'}
+    missing_cols = required_cols.difference(selected_features.columns)
+    if missing_cols:
+        missing = ', '.join(sorted(missing_cols))
+        raise ValueError(f'Selector result is missing required columns: {missing}')
+    return selected_features
+
+
+def run_selector_with_metrics(
+    selector: BaseFeatureSelector,
+    inputs: SelectorInputs,
+    feature_ids: Sequence[int] | None = None,
+) -> SelectorResult:
+    """Run selector and return selected features plus computed metric tables."""
+    result = selector.execute(inputs, feature_ids=feature_ids)
+    required_cols = {'feature_id', 'importance_score', 'importance_rank', 'importance_metric'}
+    missing_cols = required_cols.difference(result.selected_features.columns)
+    if missing_cols:
+        missing = ', '.join(sorted(missing_cols))
+        raise ValueError(f'Selector result is missing required columns: {missing}')
+    return result
+
+
+def selected_feature_ids(selected_features: pd.DataFrame) -> List[int]:
+    """Convenience helper to get ordered feature ids from selector output."""
+    return [int(x) for x in selected_features['feature_id'].tolist()]
+
+
+def _normalize_selector_configs(selector_configs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    if not selector_configs:
+        raise ValueError("selector_configs must contain at least one selector config")
+
+    normalized: List[Dict[str, Any]] = []
+    for selector_cfg in selector_configs:
+        if not isinstance(selector_cfg, dict):
+            raise ValueError("Each selector config must be a dict: {'name': str, 'params': dict, 'top_k': int}")
+
+        selector_name = str(selector_cfg.get('name', '')).strip()
+        if not selector_name:
+            raise ValueError("Each selector config must include non-empty 'name'")
+
+        if 'params' not in selector_cfg:
+            raise ValueError(f"Selector '{selector_name}': 'params' is required and must be a dict")
+        selector_params = selector_cfg['params']
+        if selector_params is None:
+            selector_params = {}
+        if not isinstance(selector_params, dict):
+            raise ValueError(f"Selector '{selector_name}': 'params' must be a dict")
+
+        if 'top_k' not in selector_cfg:
+            raise ValueError(f"Selector '{selector_name}': 'top_k' is required and must be positive")
+        selector_top_k = int(selector_cfg['top_k'])
+        if selector_top_k <= 0:
+            raise ValueError(f"Selector '{selector_name}': 'top_k' must be positive")
+
+        selector_feature_ids = selector_cfg.get('feature_ids')
+        if selector_feature_ids is not None:
+            if not isinstance(selector_feature_ids, (list, tuple)):
+                raise ValueError(f"Selector '{selector_name}': 'feature_ids' must be list/tuple or null")
+            selector_feature_ids = [int(fid) for fid in selector_feature_ids]
+
+        normalized.append({
+            'name': selector_name,
+            'params': selector_params,
+            'top_k': selector_top_k,
+            'feature_ids': selector_feature_ids,
+        })
+
+    return normalized
+
+
+def run_selectors_from_configs(
+    *,
+    selector_configs: List[Dict[str, Any]],
+    meta_df: Any,
+    features: FeatureMatrix,
+    pristine_meta_df: Any,
+    pristine_features: FeatureMatrix | None,
+    dataset: str,
+    sae_checkpoint_path: str,
+    cache_dir: str,
+    cache_paths: Any = None,
+    selector_registry: Optional[Dict[str, BaseFeatureSelector]] = None,
+    feature_weight_norms: Optional[Mapping[int, float]] = None,
+) -> Dict[str, SelectorResult]:
+    """Execute configured selectors and return full per-run selector outputs."""
+    from analysis.models import extract_decoder_weight_norms
+
+    normalized_configs = _normalize_selector_configs(selector_configs)
+
+    needs_pristine_inputs = any(
+        str(selector_cfg.get('name', '')).startswith('paired_delta')
+        for selector_cfg in normalized_configs
+    )
+    if not needs_pristine_inputs:
+        pristine_meta_df = None
+        pristine_features = None
+
+    registry = dict(selector_registry) if selector_registry is not None else get_selector_registry()
+    selector_entries = []
+    for selector_cfg in normalized_configs:
+        selector_name = selector_cfg['name']
+        if selector_name not in registry:
+            available = ', '.join(sorted(registry.keys()))
+            raise ValueError(f"Unknown selector '{selector_name}'. Available: {available}")
+        selector_entries.append((selector_name, registry[selector_name], selector_cfg))
+
+    needs_weighted_delta = False
+    for selector_cfg in normalized_configs:
+        selector_params = selector_cfg.get('params', {})
+        if selector_params and str(selector_params.get('ranking_mode', 'default')) == 'weighted_average':
+            needs_weighted_delta = True
+            break
+
+    resolved_weight_norms: Optional[Mapping[int, float]] = feature_weight_norms
+    if needs_weighted_delta and not resolved_weight_norms:
+        print('[run] Loading SAE decoder norms for weighted_average ranking mode...')
+        resolved_weight_norms = extract_decoder_weight_norms(
+            checkpoint_path=sae_checkpoint_path,
+            cache_dir=cache_dir,
+        )
+        print(f'[run] Loaded decoder norms for {len(resolved_weight_norms)} SAE features.')
+
+    selector_name_counts: Dict[str, int] = {}
+    for selector_name, _, _ in selector_entries:
+        selector_name_counts[selector_name] = selector_name_counts.get(selector_name, 0) + 1
+
+    selector_results: Dict[str, SelectorResult] = {}
+    for selector_idx, (selector_name, selector, selector_cfg) in enumerate(selector_entries, start=1):
+        selector_run_key = selector_name
+        if selector_name_counts[selector_name] > 1:
+            selector_run_key = f'{selector_name}__{selector_idx}'
+
+        selector_inputs = SelectorInputs(
+            meta_df=meta_df,
+            features=features,
+            dataset=dataset,
+            pristine_meta_df=pristine_meta_df,
+            pristine_features=pristine_features,
+            feature_weight_norms=resolved_weight_norms,
+            cache_paths=cache_paths,
+            selector_top_k=int(selector_cfg['top_k']),
+            selector_params=dict(selector_cfg.get('params', {})),
+        )
+        selector_feature_ids = selector_cfg.get('feature_ids')
+        selector_results[selector_run_key] = selector.execute(
+            selector_inputs,
+            feature_ids=selector_feature_ids,
+        )
+
+    return selector_results

analysis/features/feature_stats.py

@@ -0,0 +1,111 @@
+"""
+Статистики по SAE-признакам и их визуализация (bar charts).
+"""
+
+from typing import List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+def compute_feature_stats(features: FeatureMatrix) -> pd.DataFrame:
+    """
+    Вычисляет статистики для каждого признака SAE.
+
+    Параметры
+    ----------
+    features : CSR activations; ``column_feature_ids[j]`` = global SAE id
+
+    Возвращает
+    ----------
+    DataFrame с колонками: feature_id (global SAE id), mean, frequency, max, mean_acts
+        mean      — средняя активация по всем патчам
+        frequency — доля патчей с ненулевой активацией
+        max       — максимальная активация среди всех патчей
+        mean_acts — средняя активация по ненулевым патчам
+    """
+    codes = features.codes
+    mat = codes if codes.dtype == np.float32 else codes.astype(np.float32)
+
+    n_rows = mat.shape[0]
+    means = np.asarray(mat.mean(axis=0)).ravel()
+
+    freqs = np.asarray(mat.getnnz(axis=0), dtype=np.float32).ravel() / float(n_rows)
+
+    maxes = np.asarray(mat.tocsc().max(axis=0).toarray()).ravel()
+    mean_acts = np.where(freqs > 0, means / freqs, 0.0)
+
+    col_ids = [int(fid) for fid in features.column_feature_ids]
+    return pd.DataFrame({
+        'feature_id': col_ids,
+        'mean':       means,
+        'frequency':  freqs,
+        'max':        maxes,
+        'mean_acts':  mean_acts,
+    })
+
+
+def get_top_features(
+    stats: pd.DataFrame,
+    top_k: int = 20,
+    criterion: str = 'mean',
+    min_mean_acts: Optional[float] = None,
+) -> List[int]:
+    """
+    Возвращает индексы top-K признаков по выбранному критерию.
+
+    Параметры
+    ----------
+    stats         : DataFrame из compute_feature_stats
+    top_k         : число топ-признаков
+    criterion     : 'mean' | 'frequency' | 'max' | 'mean_acts'
+    min_mean_acts : предварительный фильтр по mean_acts
+
+    Возвращает
+    ----------
+    List[int] — feature_id в порядке убывания критерия
+    """
+    assert criterion in ('mean', 'frequency', 'max', 'mean_acts'), \
+        f"criterion must be one of 'mean', 'frequency', 'max', 'mean_acts', got {criterion!r}"
+    filtered = stats
+    if min_mean_acts is not None:
+        filtered = stats[stats['mean_acts'] >= min_mean_acts]
+    return filtered.nlargest(top_k, criterion)['feature_id'].tolist()
+
+
+def plot_top_features(
+    stats: pd.DataFrame,
+    top_k: int = 20,
+    criterion: str = 'mean',
+    figsize: Tuple[int, int] = (14, 4),
+    min_mean_acts: Optional[float] = None,
+) -> None:
+    """
+    Bar-chart топ-K признаков по выбранному критерию.
+    Над каждым баром подписывается frequency (доля ненулевых патчей).
+    """
+    top_ids = get_top_features(stats, top_k=top_k, criterion=criterion,
+                               min_mean_acts=min_mean_acts)
+    top_stats = stats.set_index('feature_id').loc[top_ids]
+
+    fig, ax = plt.subplots(figsize=figsize)
+    bars = ax.bar(range(len(top_ids)), top_stats[criterion].values, color='steelblue')
+
+    for bar, freq_val in zip(bars, top_stats['frequency'].values):
+        ax.text(
+            bar.get_x() + bar.get_width() / 2,
+            bar.get_height(),
+            f'{freq_val:.3f}',
+            ha='center', va='bottom', fontsize=6, color='black', rotation=90,
+        )
+
+    ax.set_xticks(range(len(top_ids)))
+    ax.set_xticklabels([str(i) for i in top_ids], rotation=90, fontsize=8)
+    ax.set_xlabel('Feature ID')
+    ax.set_ylabel(criterion)
+    ax.set_title(f'Top-{top_k} features by {criterion}')
+    plt.tight_layout()
+    plt.show()

analysis/metrics/correlations.py

@@ -0,0 +1,121 @@
+"""Correlation computations for SAE feature analysis."""
+
+from typing import List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+def _prepare_codes_and_categories(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str,
+    level: str,
+    binarize: bool,
+) -> Tuple[sp.csr_matrix | np.ndarray, np.ndarray, List[int], np.ndarray, np.ndarray]:
+    codes = features.codes
+    feature_names = [int(fid) for fid in features.column_feature_ids]
+
+    if binarize:
+        threshold = 0.2  # Такой был в PatchSAE
+        print(f'[binarize] Бинаризация с порогом {threshold}')
+        if sp.issparse(codes):
+            codes = codes.copy()
+            codes.data = (codes.data > threshold).astype(np.float32)
+            codes.eliminate_zeros()
+        else:
+            codes = (codes > threshold).astype(np.float32)
+
+    if level == 'image':
+        image_idx_arr = meta['image_idx'].values
+        rows, cats = [], []
+        groups = meta.groupby('image_idx').indices
+
+        for img_idx, _ in groups.items():
+            mask = np.where(image_idx_arr == img_idx)[0]
+            avg = np.asarray(codes[mask].mean(axis=0)).ravel().astype(np.float32)
+            rows.append(avg)
+            cats.append(meta.iloc[mask[0]][group_col])
+
+        work_codes = np.stack(rows)
+        work_categories = np.array(cats)
+    elif level == 'patch':
+        work_codes = codes
+        work_categories = meta[group_col].values
+    else:
+        raise ValueError(f"level must be 'image' or 'patch', got {level!r}")
+
+    unique_categories = np.unique(work_categories)
+    cat_to_idx = {c: i for i, c in enumerate(unique_categories)}
+    cat_idx = np.array([cat_to_idx[c] for c in work_categories], dtype=np.int32)
+
+    return work_codes, work_categories, feature_names, unique_categories, cat_idx
+
+
+def compute_distortion_correlations(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_group',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    binarize: bool = False,
+    level: str = 'image',
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    """Correlations between activations and distortion categories.
+
+    ``features.column_feature_ids[j]`` is the global SAE id for matrix column ``j``.
+    ``global_feature_ids`` restricts which globals to include (default: all columns).
+    """
+    cached = load_parquet_cache(cache_path, label='correlations')
+    if cached is not None:
+        return cached
+
+    work_features = features.subset(global_feature_ids)
+
+    work_codes, _, feature_names, unique_categories, cat_idx = _prepare_codes_and_categories(
+        meta=meta,
+        features=work_features,
+        group_col=group_col,
+        level=level,
+        binarize=binarize,
+    )
+
+    n_cats = len(unique_categories)
+
+    n_samples, _ = work_codes.shape
+
+    one_hot = sp.csr_matrix(
+        (np.ones(n_samples, dtype=np.float32), (cat_idx, np.arange(n_samples))),
+        shape=(n_cats, n_samples)
+    )
+    is_sparse = sp.issparse(work_codes)
+
+    feat_mean = np.asarray(work_codes.mean(axis=0)).ravel()
+    if is_sparse:
+        feat_sq_mean = np.asarray(work_codes.power(2).mean(axis=0)).ravel()
+    else:
+        feat_sq_mean = (work_codes ** 2).mean(axis=0)
+    feat_std = np.sqrt(feat_sq_mean - feat_mean ** 2).clip(min=1e-8)
+
+    product = one_hot @ work_codes
+    grouped_sum = product if is_sparse else np.asarray(product).astype(np.float32)
+
+    cat_counts = np.bincount(cat_idx, minlength=n_cats).astype(np.float32)
+    cat_mean = cat_counts / n_samples
+    cat_std = np.sqrt(cat_mean * (1 - cat_mean)).clip(min=1e-8)
+
+    cov = grouped_sum / n_samples - np.outer(cat_mean, feat_mean)
+    corr_mat = cov / np.outer(cat_std, feat_std)
+
+    result = pd.DataFrame(corr_mat.astype(np.float32),
+                          index=unique_categories,
+                          columns=feature_names)
+
+    save_parquet_cache(result, cache_path, label='correlations')
+
+    return result

analysis/metrics/iou_utils.py

@@ -0,0 +1,341 @@
+"""IoU (Intersection over Union) computation for feature selection based on spatial distortion masks."""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+from tqdm.auto import tqdm
+
+from analysis.datasets import distortion_types_mapping_qground, distortion_types_mapping_srground
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+_QGROUND_TYPE_TO_LABEL = {
+    dist_type: label_id
+    for label_id, dist_type in distortion_types_mapping_qground.items()
+    if label_id > 0
+}
+
+_SRGROUND_TYPE_TO_LABEL = {
+    dist_type: label_id
+    for label_id, dist_type in distortion_types_mapping_srground.items()
+    if label_id > 0
+}
+
+_SPATIAL_MASK_DATASETS = frozenset({'QGround', 'SRGround'})
+
+
+def _binarize_activations(
+    feature_activations: np.ndarray,
+) -> np.ndarray:
+    """
+    Binarize feature activations.
+    Uses threshold > 0 since ReLU already ensures non-negative values.
+    
+    Args:
+        feature_activations: Shape (n_patches,) or (n_samples, n_patches)
+        
+    Returns:
+        Binary array with same shape as input
+    """
+    return (feature_activations > 0).astype(np.uint8)
+
+
+def _mask_label_for_dist_type(dist_type: Optional[str]) -> Optional[int]:
+    if dist_type is None:
+        return None
+
+    dist_type = str(dist_type)
+    if dist_type == 'background':
+        return 0
+
+    if dist_type in _QGROUND_TYPE_TO_LABEL:
+        return _QGROUND_TYPE_TO_LABEL[dist_type]
+
+    if dist_type in _SRGROUND_TYPE_TO_LABEL:
+        return _SRGROUND_TYPE_TO_LABEL[dist_type]
+
+    # KADID / local KADID are binary masks: any non-background distortion is label 1.
+    return 1
+
+
+
+def _compute_patch_iou(
+    activation_binary: np.ndarray,
+    mask_binary: np.ndarray,
+) -> float:
+    """
+    Compute IoU (Jaccard index) between binary activation and binary mask.
+    
+    Args:
+        activation_binary: Shape (n_patches,) - binary feature activation map
+        mask_binary: Shape (n_patches,) - binary distortion mask
+        
+    Returns:
+        IoU score in [0, 1]
+    """
+    intersection = np.logical_and(activation_binary, mask_binary).sum()
+    union = np.logical_and(activation_binary, mask_binary).sum() + np.logical_xor(activation_binary, mask_binary).sum()
+    
+    if union == 0:
+        return 0.0
+    
+    return float(intersection / union)
+
+
+def _compute_iou_batch_torch(
+    activation_binary: torch.Tensor,
+    mask_binary: torch.Tensor,
+) -> torch.Tensor:
+    """Compute IoU for a batch of features against one binary mask.
+
+    activation_binary: bool tensor with shape (n_features, n_patches)
+    mask_binary: bool tensor with shape (n_patches,)
+    """
+    if activation_binary.ndim != 2:
+        raise ValueError(f'activation_binary must be 2D, got {activation_binary.shape}')
+    if mask_binary.ndim != 1:
+        raise ValueError(f'mask_binary must be 1D, got {mask_binary.shape}')
+
+    mask_row = mask_binary.unsqueeze(0)
+    intersection = torch.logical_and(activation_binary, mask_row).sum(dim=1)
+    union = torch.logical_or(activation_binary, mask_row).sum(dim=1)
+    iou = torch.zeros_like(union, dtype=torch.float32)
+    nonzero = union > 0
+    iou[nonzero] = intersection[nonzero].float() / union[nonzero].float()
+    return iou
+
+
+def _load_patch_mask_for_group(
+    group_df: pd.DataFrame,
+    target_dist_type: Optional[str] = None,
+    *,
+    dataset: str,
+) -> np.ndarray:
+    """Return a flattened binary patch mask for one image group.
+
+    Requires a cached patch-level label map.
+    """
+    if 'patch_mask_label' in group_df.columns:
+        patch_labels = group_df['patch_mask_label'].values.astype(np.uint8)
+    elif 'patch_is_distorted' in group_df.columns:
+        patch_labels = group_df['patch_is_distorted'].values.astype(np.uint8)
+    else:
+        raise ValueError(
+            "IoU requires either 'patch_mask_label' or 'patch_is_distorted' in metadata. "
+            "Rebuild the cache with patch-level labels or distorted-patch flags."
+        )
+
+    dataset_name = str(dataset).strip()
+    if dataset_name not in _SPATIAL_MASK_DATASETS:
+        return (patch_labels > 0).astype(np.uint8)
+
+    target_label = _mask_label_for_dist_type(target_dist_type)
+    if target_label is None:
+        return (patch_labels > 0).astype(np.uint8)
+    if target_label == 0:
+        return np.zeros_like(patch_labels, dtype=np.uint8)
+    return (patch_labels == target_label).astype(np.uint8)
+
+CREATE_IOU_CACHE_IF_MISSING: bool = True
+
+def compute_iou_per_feature(
+    features: FeatureMatrix,
+    meta_df: pd.DataFrame,
+    spatial_shape: Tuple[int, int],
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    feature_batch_size: int = 1024,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+    dataset: str,
+    create_cache_if_missing: Optional[bool] = None,
+) -> pd.DataFrame:
+    """
+    Compute per-feature IoU scores grouped by distortion type or group.
+    
+    This function:
+    1. Iterates through each feature
+    2. For each image with mask data, computes IoU between feature activation and mask
+    3. Groups IoU scores by distortion type/group
+    4. Computes median IoU per feature per distortion type/group
+    
+    Args:
+        features: Activations with global SAE id per column
+        meta_df: DataFrame with metadata including 'dist_type', 'dist_group', 'image_idx', 'patch_is_distorted'
+        spatial_shape: Tuple (patch_h, patch_w) for activation grid size
+        group_col: Column to group by ('dist_type' or 'dist_group')
+        global_feature_ids: optional subset of globals to compute (default: all columns)
+        feature_batch_size: Number of features to evaluate together on torch
+        
+    Returns:
+        DataFrame with shape (n_groups, n_features) where rows are distortion types/groups 
+        and columns are feature IDs, values are median IoU
+    """
+    cached = load_parquet_cache(cache_path, label='iou')
+    if cached is not None:
+        return cached
+
+    # Decide whether to create cache when missing. Priority:
+    # 1) explicit function arg, 2) module-level global flag, 3) default True
+    if create_cache_if_missing is None:
+        create_cache = bool(CREATE_IOU_CACHE_IF_MISSING)
+    else:
+        create_cache = bool(create_cache_if_missing)
+
+    if not create_cache:
+        # If we shouldn't build missing caches, return empty DataFrame so
+        # callers can detect absence of IoU data without triggering heavy compute.
+        return pd.DataFrame()
+
+    work = features.subset(global_feature_ids)
+    feature_codes = work.codes
+    global_labels = list(work.column_feature_ids)
+    column_indices = list(range(work.n_features))
+
+    n_samples, _n_features = feature_codes.shape
+    _patch_h, _patch_w = spatial_shape
+    n_patches = _patch_h * _patch_w
+    
+    if 'patch_mask_label' not in meta_df.columns and 'patch_is_distorted' not in meta_df.columns:
+        raise ValueError(
+            "Meta DataFrame must contain either 'patch_mask_label' or 'patch_is_distorted' column for IoU computation. "
+            "Rebuild the cache with patch-level labels or distorted-patch flags."
+        )
+    
+    # Group metadata by image for efficient processing
+    image_groups = meta_df.groupby('image_idx')
+    
+    # Track IoU scores: {group_value: {feature_id: [iou_scores]}}
+    iou_scores_by_group: Dict[str, Dict[int, List[float]]] = {}
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    feature_batch_size = max(1, int(feature_batch_size))
+
+    image_payloads = []
+    for image_idx, group_df in image_groups:
+        sample_indices = group_df.index.to_numpy()
+        if len(group_df) != n_patches:
+            raise ValueError(
+                f"Patch row count mismatch for image_idx={image_idx}: got {len(group_df)}, expected {n_patches}. "
+                "Rebuild the cache."
+            )
+        image_payloads.append((group_df, sample_indices))
+
+    dataset_name = str(dataset).strip()
+    spatial_mask_dataset = dataset_name in _SPATIAL_MASK_DATASETS
+
+    for feature_start in tqdm(
+        range(0, len(column_indices), feature_batch_size),
+        desc='Computing IoU batches',
+        unit='batch',
+        disable=not show_progress_bars,
+    ):
+        batch_col_indices = column_indices[feature_start:feature_start + feature_batch_size]
+        batch_global_labels = global_labels[feature_start:feature_start + feature_batch_size]
+        for group_df, sample_indices in image_payloads:
+            img_feature_block = feature_codes[sample_indices][:, batch_col_indices]
+            if sp.issparse(img_feature_block):
+                img_feature_block = img_feature_block.toarray()
+            else:
+                img_feature_block = np.asarray(img_feature_block)
+
+            act_binary = torch.as_tensor(
+                img_feature_block.T > 0,
+                device=device,
+                dtype=torch.bool,
+            )
+
+            group_values = group_df[group_col].unique()
+            for group_val in group_values:
+                group_key = str(group_val)
+                if spatial_mask_dataset:
+                    target_dist_type = group_key
+                elif 'dist_type' in group_df.columns and len(group_df) > 0:
+                    target_dist_type = str(group_df['dist_type'].iloc[0])
+                else:
+                    target_dist_type = None
+
+                patch_masks = _load_patch_mask_for_group(
+                    group_df,
+                    target_dist_type=target_dist_type,
+                    dataset=dataset,
+                )
+                mask_binary = torch.as_tensor(patch_masks > 0, device=device, dtype=torch.bool)
+                iou_batch = _compute_iou_batch_torch(act_binary, mask_binary).detach().cpu().numpy()
+
+                if group_key not in iou_scores_by_group:
+                    iou_scores_by_group[group_key] = {}
+
+                group_bucket = iou_scores_by_group[group_key]
+                for local_idx, global_id in enumerate(batch_global_labels):
+                    if global_id not in group_bucket:
+                        group_bucket[global_id] = []
+                    group_bucket[global_id].append(float(iou_batch[local_idx]))
+    
+    # Aggregate scores: compute median per feature per group
+    result_data = {}
+    for group_key in sorted(iou_scores_by_group.keys()):
+        row = {}
+        for global_id in global_labels:
+            scores = iou_scores_by_group[group_key].get(global_id, [])
+            if scores:
+                row[global_id] = float(np.median(scores))
+            else:
+                row[global_id] = np.nan
+        result_data[group_key] = row
+    
+    if not result_data:
+        return pd.DataFrame()
+    
+    # Create result DataFrame
+    iou_df = pd.DataFrame(result_data).T
+    iou_df.index.name = group_col
+    save_parquet_cache(iou_df, cache_path, label='iou')
+    return iou_df
+
+
+def compute_iou_per_distortion_type_and_feature(
+    features: FeatureMatrix,
+    meta_df: pd.DataFrame,
+    spatial_shape: Tuple[int, int],
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    feature_batch_size: int = 1024,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+    dataset: str,
+    create_cache_if_missing: Optional[bool] = None,
+) -> pd.DataFrame:
+    """
+    Compute per-feature median IoU scores, organized by distortion type or group.
+    
+    Args:
+        feature_codes: Sparse matrix of shape (n_samples, n_features)
+        meta_df: Metadata DataFrame with 'dist_type', 'dist_group', 'image_idx', 'patch_is_distorted' columns
+        spatial_shape: (patch_h, patch_w) grid dimensions
+        group_col: Column name for grouping ('dist_type' or 'dist_group')
+        global_feature_ids: optional subset of globals to compute
+
+    Returns:
+        DataFrame with shape (n_distortion_types, n_features) where values are median IoU
+    """
+    return compute_iou_per_feature(
+        features=features,
+        meta_df=meta_df,
+        spatial_shape=spatial_shape,
+        group_col=group_col,
+        global_feature_ids=global_feature_ids,
+        feature_batch_size=feature_batch_size,
+        show_progress_bars=show_progress_bars,
+        cache_path=cache_path,
+        dataset=dataset,
+        create_cache_if_missing=create_cache_if_missing,
+    )

analysis/metrics/mutual_information.py

@@ -0,0 +1,115 @@
+"""Mutual information computations for SAE feature analysis."""
+
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from tqdm import tqdm
+
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+from analysis.metrics.correlations import _prepare_codes_and_categories
+
+
+def _build_contingency_bincount(
+    x_bin: np.ndarray,
+    cat_idx: np.ndarray,
+    n_cats: int,
+    n_bins: int,
+) -> np.ndarray:
+    """Build contingency[c, b, f] counts using a fast bincount kernel."""
+    n_samples, n_features = x_bin.shape
+    if cat_idx.shape[0] != n_samples:
+        raise ValueError('cat_idx length must match number of samples in x_bin')
+
+    contingency = np.empty((n_cats, n_bins, n_features), dtype=np.float64)
+    packed_base = cat_idx.astype(np.int32, copy=False) * n_bins
+
+    for fid in range(n_features):
+        packed = packed_base + x_bin[:, fid]
+        counts = np.bincount(packed, minlength=n_cats * n_bins)
+        contingency[:, :, fid] = counts.reshape(n_cats, n_bins)
+
+    return contingency
+
+
+def compute_distortion_mutual_information(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_group',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    binarize: bool = False,
+    level: str = 'image',
+    n_bins: int = 16,
+    feature_chunk_size: int = 512,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    """
+    Compute mutual information between SAE feature activations and
+    a binary indicator for each distortion category.
+
+    Returns a DataFrame of shape (n_categories x n_features), values >= 0.
+
+    The computation is chunked over features to keep memory bounded.
+    """
+    cached = load_parquet_cache(cache_path, label='mutual information')
+    if cached is not None:
+        return cached
+
+    work_features = features.subset(global_feature_ids)
+
+    work_codes, _, feature_names, unique_categories, cat_idx = _prepare_codes_and_categories(
+        meta=meta,
+        features=work_features,
+        group_col=group_col,
+        level=level,
+        binarize=binarize,
+    )
+
+    _, n_features = work_codes.shape
+    n_cats = len(unique_categories)
+
+    if feature_chunk_size <= 0:
+        raise ValueError(f'feature_chunk_size must be > 0, got {feature_chunk_size}')
+
+    quantiles = np.linspace(0, 1, n_bins + 1)
+    mi_mat = np.empty((n_cats, n_features), dtype=np.float32)
+
+    for start in tqdm(range(0, n_features, feature_chunk_size), desc='Computing MI', unit='chunk', disable=not show_progress_bars):
+        end = min(start + feature_chunk_size, n_features)
+        if sp.issparse(work_codes):
+            x_values = work_codes[:, start:end].toarray().astype(np.float32, copy=False)
+        else:
+            x_values = np.asarray(work_codes[:, start:end], dtype=np.float32)
+
+        edges = np.quantile(x_values, quantiles, axis=0)
+        x_bin = np.zeros_like(x_values, dtype=np.int16)
+        for b in range(n_bins - 1):
+            x_bin += (x_values > edges[b + 1]).astype(np.int16)
+
+        contingency = _build_contingency_bincount(
+            x_bin=x_bin,
+            cat_idx=cat_idx,
+            n_cats=n_cats,
+            n_bins=n_bins,
+        )
+
+        total = contingency.sum(axis=1, keepdims=True)
+        pxy = contingency / total.clip(min=1e-12)
+        px = pxy.sum(axis=0, keepdims=True)
+        py = pxy.sum(axis=1, keepdims=True)
+
+        with np.errstate(divide='ignore', invalid='ignore'):
+            log_ratio = np.where(pxy > 0, np.log(pxy / (py * px).clip(min=1e-12)), 0.0)
+
+        mi_chunk = (pxy * log_ratio).sum(axis=1)
+        mi_mat[:, start:end] = mi_chunk.astype(np.float32)
+
+    result = pd.DataFrame(mi_mat, index=unique_categories, columns=feature_names)
+
+    save_parquet_cache(result, cache_path, label='mutual information')
+
+    return result

analysis/metrics/paired_deltas.py

@@ -0,0 +1,210 @@
+"""Helpers for paired original-vs-distorted SAE activation comparisons."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Dict, List, Mapping, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+def infer_original_image_path(distorted_img_path: str | Path) -> str | None:
+    """Infer the KADID-style original image path from a distorted image path."""
+    path = Path(distorted_img_path)
+    match = re.match(r'(I\d+)_\d+_\d+\.png$', path.name, re.IGNORECASE)
+    if match:
+        return str(path.with_name(f'{match.group(1)}.png'))
+
+    local_match = re.match(r'(I\d+)_\d+_dist\.png$', path.name, re.IGNORECASE)
+    if local_match:
+        return str(path.with_name(f'{local_match.group(1)}.png'))
+
+    return None
+
+
+def _extract_pair_path(dist_row: pd.Series, distorted_path_col: str) -> str | None:
+    if 'original_img_path' in dist_row and pd.notna(dist_row['original_img_path']):
+        return str(dist_row['original_img_path'])
+    if distorted_path_col in dist_row and pd.notna(dist_row[distorted_path_col]):
+        return infer_original_image_path(str(dist_row[distorted_path_col]))
+    return None
+
+
+def _delta_vector(
+    distorted_vec: np.ndarray,
+    original_vec: np.ndarray,
+    delta_mode: str,
+) -> np.ndarray:
+    if delta_mode == 'abs':
+        return np.abs(distorted_vec - original_vec)
+    if delta_mode == 'signed':
+        return distorted_vec - original_vec
+    if delta_mode == 'relative':
+        denom = np.abs(original_vec) + 1e-8
+        return (distorted_vec - original_vec) / denom
+    raise ValueError("delta_mode must be one of {'abs', 'signed', 'relative'}")
+
+
+def _aggregate_by_category(
+    category_to_rows: Mapping[str, List[np.ndarray]],
+    feature_names: Sequence[object],
+) -> pd.DataFrame:
+    rows = []
+    index = []
+    for category, category_rows in category_to_rows.items():
+        if not category_rows:
+            continue
+        stacked = np.stack(category_rows, axis=0).astype(np.float32, copy=False)
+        rows.append(stacked.mean(axis=0).astype(np.float32, copy=False))
+        index.append(category)
+
+    if not rows:
+        return pd.DataFrame(columns=feature_names)
+
+    return pd.DataFrame(np.stack(rows, axis=0), index=index, columns=feature_names)
+
+
+def build_paired_delta_tables(
+    distorted_meta: pd.DataFrame,
+    distorted: FeatureMatrix,
+    original_meta: pd.DataFrame,
+    original: FeatureMatrix,
+    *,
+    group_cols: Sequence[str] = ('dist_type', 'dist_group'),
+    global_feature_ids: Sequence[int] | None = None,
+    delta_mode: str = 'relative',
+    level: str = 'patch',
+    distorted_path_col: str = 'distorted_img_path',
+    original_path_col: str = 'original_img_path',
+    cache_prefix: Optional[str] = None,
+) -> Dict[str, pd.DataFrame]:
+    """Build category-by-feature paired delta tables from matched original/distorted samples.
+
+    Samples are matched by original image path. If a row does not expose the original path
+    explicitly, it is inferred from the distorted filename.
+    """
+    if level not in {'image', 'patch'}:
+        raise ValueError(f"level must be 'image' or 'patch', got {level!r}")
+
+    if distorted.column_feature_ids != original.column_feature_ids:
+        raise ValueError('distorted and original FeatureMatrix must share column_feature_ids')
+    work_distorted = distorted.subset(global_feature_ids)
+    work_original = original.subset(global_feature_ids)
+    distorted_codes = work_distorted.codes
+    original_codes = work_original.codes
+    feature_names = [int(fid) for fid in work_distorted.column_feature_ids]
+
+    if distorted_meta.shape[0] != distorted_codes.shape[0]:
+        raise ValueError('distorted_meta and distorted codes must have the same number of rows')
+    if original_meta.shape[0] != original_codes.shape[0]:
+        raise ValueError('original_meta and original codes must have the same number of rows')
+
+    if distorted_path_col not in distorted_meta.columns:
+        raise ValueError(f'distorted_meta must contain {distorted_path_col!r}')
+
+    if original_path_col not in original_meta.columns:
+        fallback_cols = [col for col in ('distorted_img_path', 'img_path', 'image_path')
+                         if col in original_meta.columns]
+        if not fallback_cols:
+            raise ValueError(
+                f'original_meta must contain {original_path_col!r} or a fallback path column; '
+                f'available columns: {sorted(original_meta.columns.tolist())}'
+            )
+        original_path_col = fallback_cols[0]
+
+    original_groups = original_meta.groupby(original_path_col).indices
+    distorted_groups = distorted_meta.groupby(distorted_path_col).indices
+
+    tables: Dict[str, pd.DataFrame] = {}
+    for group_col in group_cols:
+        table_key = f'paired_{delta_mode}_{group_col}_df'
+        table_cache_path = None
+        if cache_prefix is not None:
+            table_cache_path = f'{cache_prefix}_{delta_mode}_{group_col}_{level}.parquet'
+            cached = load_parquet_cache(table_cache_path, label=table_key)
+            if cached is not None:
+                tables[table_key] = cached
+                continue
+
+        if group_col not in distorted_meta.columns:
+            raise ValueError(f'distorted_meta must contain {group_col!r}')
+
+        category_to_rows: Dict[str, List[np.ndarray]] = {}
+        for distorted_path, distorted_indices in distorted_groups.items():
+            first_dist_idx = int(distorted_indices[0])
+            original_path = _extract_pair_path(distorted_meta.iloc[first_dist_idx], distorted_path_col)
+            if original_path is None:
+                continue
+
+            original_indices = original_groups.get(original_path)
+            if original_indices is None or len(original_indices) == 0:
+                continue
+
+            distorted_rows = np.asarray(distorted_indices, dtype=np.int64)
+            original_rows = np.asarray(original_indices, dtype=np.int64)
+
+            if level == 'patch' and 'patch_idx' in distorted_meta.columns and 'patch_idx' in original_meta.columns:
+                dist_patch_to_idx = {
+                    int(distorted_meta.iloc[row_idx]['patch_idx']): int(row_idx)
+                    for row_idx in distorted_rows
+                }
+                orig_patch_to_idx = {
+                    int(original_meta.iloc[row_idx]['patch_idx']): int(row_idx)
+                    for row_idx in original_rows
+                }
+                common_patches = sorted(set(dist_patch_to_idx).intersection(orig_patch_to_idx))
+                if not common_patches:
+                    continue
+
+                distorted_rows = np.asarray([dist_patch_to_idx[patch_idx]
+                                            for patch_idx in common_patches], dtype=np.int64)
+                original_rows = np.asarray([orig_patch_to_idx[patch_idx]
+                                           for patch_idx in common_patches], dtype=np.int64)
+            else:
+                n_rows = min(len(distorted_rows), len(original_rows))
+                if n_rows == 0:
+                    continue
+                distorted_rows = distorted_rows[:n_rows]
+                original_rows = original_rows[:n_rows]
+
+            distorted_block = distorted_codes[distorted_rows]
+            original_block = original_codes[original_rows]
+            if sp.issparse(distorted_block):
+                distorted_block = distorted_block.toarray()
+            if sp.issparse(original_block):
+                original_block = original_block.toarray()
+
+            if distorted_block.shape != original_block.shape:
+                n_rows = min(distorted_block.shape[0], original_block.shape[0])
+                n_cols = min(distorted_block.shape[1], original_block.shape[1])
+                distorted_block = distorted_block[:n_rows, :n_cols]
+                original_block = original_block[:n_rows, :n_cols]
+
+            if level == 'image':
+                distorted_vec = np.asarray(distorted_block.mean(axis=0)).ravel().astype(np.float32, copy=False)
+                original_vec = np.asarray(original_block.mean(axis=0)).ravel().astype(np.float32, copy=False)
+                delta_rows = [_delta_vector(distorted_vec, original_vec, delta_mode)]
+            else:
+                delta_rows = [
+                    _delta_vector(
+                        np.asarray(distorted_block[row_idx], dtype=np.float32),
+                        np.asarray(original_block[row_idx], dtype=np.float32),
+                        delta_mode,
+                    )
+                    for row_idx in range(distorted_block.shape[0])
+                ]
+
+            category_value = str(distorted_meta.iloc[first_dist_idx][group_col])
+            category_to_rows.setdefault(category_value, []).extend(delta_rows)
+
+        table_df = _aggregate_by_category(category_to_rows, feature_names)
+        save_parquet_cache(table_df, table_cache_path, label=table_key)
+        tables[table_key] = table_df
+
+    return tables

analysis/metrics/precision_recall.py

@@ -0,0 +1,154 @@
+"""Precision/Recall computations for SAE feature analysis."""
+
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from tqdm import tqdm
+
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+from analysis.metrics.correlations import _prepare_codes_and_categories
+
+
+def _compute_binary_metric_chunk(
+    scores: np.ndarray,
+    positives_mask: np.ndarray,
+    metric_name: str,
+    activation_threshold: float,
+) -> np.ndarray:
+    """Compute precision/recall for a block of features."""
+    preds = scores > float(activation_threshold)
+    positives_col = positives_mask[:, None]
+
+    tp = np.logical_and(preds, positives_col).sum(axis=0).astype(np.float32)
+    pred_pos = preds.sum(axis=0).astype(np.float32)
+    true_pos_total = float(positives_mask.sum())
+
+    if metric_name == 'precision':
+        out = np.divide(tp, pred_pos, out=np.zeros_like(tp), where=pred_pos > 0)
+    elif metric_name == 'recall':
+        if true_pos_total <= 0:
+            out = np.zeros_like(tp)
+        else:
+            out = tp / true_pos_total
+    else:
+        raise ValueError(f"Unsupported metric_name={metric_name!r}. Allowed: ('precision', 'recall')")
+
+    return np.clip(out, 0.0, 1.0).astype(np.float32)
+
+
+def _compute_distortion_binary_metric(
+    metric_name: str,
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    level: str = 'image',
+    feature_chunk_size: int = 512,
+    activation_threshold: float = 0.0,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    cache_label = metric_name
+    cached = load_parquet_cache(cache_path, label=cache_label)
+    if cached is not None:
+        return cached
+
+    work_features = features.subset(global_feature_ids)
+    work_codes, _, feature_names, unique_categories, cat_idx = _prepare_codes_and_categories(
+        meta=meta,
+        features=work_features,
+        group_col=group_col,
+        level=level,
+        binarize=False,
+    )
+
+    _n_samples, n_features = work_codes.shape
+    n_cats = len(unique_categories)
+
+    if feature_chunk_size <= 0:
+        raise ValueError(f'feature_chunk_size must be > 0, got {feature_chunk_size}')
+
+    metric_mat = np.empty((n_cats, n_features), dtype=np.float32)
+
+    for start in tqdm(
+        range(0, n_features, feature_chunk_size),
+        desc=f'Computing {metric_name}',
+        unit='chunk',
+        disable=not show_progress_bars,
+    ):
+        end = min(start + feature_chunk_size, n_features)
+        if sp.issparse(work_codes):
+            x_values = work_codes[:, start:end].toarray().astype(np.float32, copy=False)
+        else:
+            x_values = np.asarray(work_codes[:, start:end], dtype=np.float32)
+
+        for cat_i in range(n_cats):
+            positives_mask = cat_idx == cat_i
+            metric_mat[cat_i, start:end] = _compute_binary_metric_chunk(
+                scores=x_values,
+                positives_mask=positives_mask,
+                metric_name=metric_name,
+                activation_threshold=activation_threshold,
+            )
+
+    result = pd.DataFrame(metric_mat, index=unique_categories, columns=feature_names)
+    save_parquet_cache(result, cache_path, label=cache_label)
+    return result
+
+
+def compute_distortion_precision(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    level: str = 'image',
+    feature_chunk_size: int = 512,
+    activation_threshold: float = 0.0,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    """Compute one-vs-rest precision for distortion categories by feature."""
+    return _compute_distortion_binary_metric(
+        metric_name='precision',
+        meta=meta,
+        features=features,
+        group_col=group_col,
+        global_feature_ids=global_feature_ids,
+        level=level,
+        feature_chunk_size=feature_chunk_size,
+        activation_threshold=activation_threshold,
+        show_progress_bars=show_progress_bars,
+        cache_path=cache_path,
+    )
+
+
+def compute_distortion_recall(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    level: str = 'image',
+    feature_chunk_size: int = 512,
+    activation_threshold: float = 0.0,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    """Compute one-vs-rest recall for distortion categories by feature."""
+    return _compute_distortion_binary_metric(
+        metric_name='recall',
+        meta=meta,
+        features=features,
+        group_col=group_col,
+        global_feature_ids=global_feature_ids,
+        level=level,
+        feature_chunk_size=feature_chunk_size,
+        activation_threshold=activation_threshold,
+        show_progress_bars=show_progress_bars,
+        cache_path=cache_path,
+    )

analysis/metrics/roc_auc.py

@@ -0,0 +1,86 @@
+"""ROC-AUC computations for SAE feature analysis."""
+
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from scipy.stats import rankdata
+from tqdm import tqdm
+
+from analysis.cache_utils import load_parquet_cache, save_parquet_cache
+from analysis.features.feature_indexing import FeatureMatrix
+from analysis.metrics.correlations import _prepare_codes_and_categories
+
+
+def _auc_from_scores(scores: np.ndarray, positives_mask: np.ndarray) -> float:
+    """Compute binary ROC-AUC via rank statistics with neutral fallback for degenerate labels."""
+    n_pos = int(positives_mask.sum())
+    n_total = int(scores.shape[0])
+    n_neg = n_total - n_pos
+    if n_pos == 0 or n_neg == 0:
+        return 0.5
+
+    ranks = rankdata(scores, method='average')
+    sum_pos_ranks = float(ranks[positives_mask].sum())
+    auc = (sum_pos_ranks - (n_pos * (n_pos + 1) / 2.0)) / (n_pos * n_neg)
+    if not np.isfinite(auc):
+        return 0.5
+    return float(np.clip(auc, 0.0, 1.0))
+
+
+def compute_distortion_roc_auc(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    group_col: str = 'dist_type',
+    *,
+    global_feature_ids: Optional[List[int]] = None,
+    level: str = 'image',
+    feature_chunk_size: int = 512,
+    show_progress_bars: bool = True,
+    cache_path: Optional[str] = None,
+) -> pd.DataFrame:
+    """Compute one-vs-rest ROC-AUC between feature activations and distortion categories.
+
+    Returns a DataFrame with shape (n_categories, n_features), values in [0, 1].
+    """
+    cached = load_parquet_cache(cache_path, label='roc auc')
+    if cached is not None:
+        return cached
+
+    work_features = features.subset(global_feature_ids)
+    work_codes, _, feature_names, unique_categories, cat_idx = _prepare_codes_and_categories(
+        meta=meta,
+        features=work_features,
+        group_col=group_col,
+        level=level,
+        binarize=False,
+    )
+
+    n_samples, n_features = work_codes.shape
+    n_cats = len(unique_categories)
+
+    if feature_chunk_size <= 0:
+        raise ValueError(f'feature_chunk_size must be > 0, got {feature_chunk_size}')
+
+    auc_mat = np.empty((n_cats, n_features), dtype=np.float32)
+
+    for start in tqdm(range(0, n_features, feature_chunk_size), desc='Computing ROC-AUC', unit='chunk', disable=not show_progress_bars):
+        end = min(start + feature_chunk_size, n_features)
+        if sp.issparse(work_codes):
+            x_values = work_codes[:, start:end].toarray().astype(np.float32, copy=False)
+        else:
+            x_values = np.asarray(work_codes[:, start:end], dtype=np.float32)
+
+        for cat_i in range(n_cats):
+            positives_mask = (cat_idx == cat_i)
+            chunk_auc = np.empty((end - start,), dtype=np.float32)
+            for local_fid in range(end - start):
+                chunk_auc[local_fid] = _auc_from_scores(x_values[:, local_fid], positives_mask)
+            auc_mat[cat_i, start:end] = chunk_auc
+
+    result = pd.DataFrame(auc_mat, index=unique_categories, columns=feature_names)
+
+    save_parquet_cache(result, cache_path, label='roc auc')
+
+    return result

analysis/models.py

@@ -0,0 +1,298 @@
+import json
+import math
+import os
+from pathlib import Path
+from typing import Any, Dict, Mapping, Optional, Tuple
+
+import pandas as pd
+import torch
+
+from analysis.qalign_utils import QAlignVisionOnlyWrapper, flatten_blc_drop_cls
+
+_iqa_activations: Dict[str, torch.Tensor] = {}
+_iqa_activation_grids: Dict[str, Tuple[int, int]] = {}
+_hook_handle = None
+
+
+def _make_hook(name: str, sequence_layout: str = 'blc'):
+    """Forward hook для ARNIQA/MANIQA/LIQE.
+
+    Поддерживаемые форматы выхода слоя:
+    - (B, C, H, W) -> (B*H*W, C)
+    - (B, L, C)    -> (B*L, C)   (MANIQA)
+    - (L, B, C)    -> (B*L, C)   (LIQE/LIQE-MIX)
+    """
+    def hook(module, inp, out):
+        if hasattr(out, 'last_hidden_state'):
+            out_detached = out.last_hidden_state.detach() # Пока костыль
+        else:
+            out_detached = out.detach()
+        if out_detached.ndim == 4:
+            tensor_perm = out_detached.permute(0, 2, 3, 1)       # (B, H, W, C)
+            _iqa_activation_grids[name] = tuple(tensor_perm.shape[1:3])
+            _iqa_activations[name] = tensor_perm.flatten(0, -2)  # (B*H*W, C)
+            return
+
+        if out_detached.ndim == 3:
+            if sequence_layout == 'lbc':
+                # LIQE CLIP resblocks usually output (L, B, C).
+                tokens, batch, channels = out_detached.shape
+                # Drop CLS token when sequence length is 1 + square (e.g. 50 = 1 + 49).
+                spatial_wo_cls = int(math.isqrt(max(tokens - 1, 0)))
+                if spatial_wo_cls * spatial_wo_cls == (tokens - 1):
+                    out_detached = out_detached[1:, :, :]
+                    tokens = tokens - 1
+                flat_acts = out_detached.permute(1, 0, 2).reshape(batch * tokens, channels)
+            elif sequence_layout == 'blc':
+                # MANIQA Swin layers usually output (B, L, C).
+                batch, tokens, channels = out_detached.shape
+                flat_acts = out_detached.reshape(batch * tokens, channels)
+            elif sequence_layout == 'blc_drop_cls':
+                batch, tokens, channels = out_detached.shape
+                flat_acts = flatten_blc_drop_cls(out_detached)
+                tokens = flat_acts.shape[0] // batch
+            else:
+                raise ValueError(
+                    f'Unsupported sequence_layout={sequence_layout!r} for layer {name!r}; '
+                    "expected one of ('blc', 'lbc')"
+                )
+
+            spatial = int(math.isqrt(tokens))
+            if spatial * spatial == tokens:
+                _iqa_activation_grids[name] = (spatial, spatial)
+            else:
+                _iqa_activation_grids[name] = (tokens, 1)
+            _iqa_activations[name] = flat_acts
+            return
+
+        raise ValueError(
+            f'Unsupported hooked activation ndim={out_detached.ndim} for layer {name!r}; '
+            'expected 3D or 4D output'
+        )
+    return hook
+
+
+SAE_CONFIG_FILENAME = "sae_config.json"
+
+
+def read_sae_config(
+    checkpoint_path: str,
+    config_path: Optional[str] = None,
+    **overrides: Any,
+) -> Dict[str, Any]:
+    """Read SAE JSON config next to a checkpoint (with optional field overrides)."""
+    if config_path is None:
+        config_path = os.path.join(
+            os.path.dirname(os.path.abspath(checkpoint_path)),
+            SAE_CONFIG_FILENAME,
+        )
+
+    with open(config_path) as f:
+        cfg: Dict[str, Any] = json.load(f)
+    cfg.update(overrides)
+    return cfg
+
+
+def load_sae(
+    checkpoint_path: str,
+    config_path: Optional[str] = None,
+    device: str = 'cpu',
+    dtype: torch.dtype = torch.float32,
+    sae_config: Optional[Dict[str, Any]] = None,
+    **overrides: Any,
+) -> torch.nn.Module:
+    """Загружает SAE из чекпоинта, определяя архитектуру из JSON-конфига.
+
+    Parameters
+    ----------
+    checkpoint_path : str
+        Путь к директории чекпоинта (output_dir/checkpoint-<step>/) или файлу весов.
+    config_path : str, optional
+        Путь к sae_config.json. По умолчанию ищется в родительской директории
+        чекпоинта (т.е. output_dir/sae_config.json).
+    device : str
+        Устройство для загрузки.
+    dtype : torch.dtype
+        Тип данных для загрузки.
+    sae_config : dict, optional
+        Уже загруженный конфиг.
+    **overrides
+        Переопределяют отдельные поля конфига (например, mp_threshold=0.05).
+    """
+    import accelerate
+    from train_code.sae import SAE, MatchingPursuitSAE
+
+    if sae_config is None:
+        cfg = read_sae_config(checkpoint_path, config_path=config_path, **overrides)
+    else:
+        cfg = dict(sae_config)
+        cfg.update(overrides)
+
+    sae_type = cfg.get('sae_type', 'sae')
+    sae_input_dim: int = cfg['sae_input_dim']
+    inner_dim: int = cfg['inner_dim']
+
+    if sae_type == 'mp_sae':
+        model: torch.nn.Module = MatchingPursuitSAE(
+            in_dim=sae_input_dim,
+            inner_dim=inner_dim,
+            threshold=cfg.get('mp_threshold', 0.1),
+            normalize=cfg.get('mp_normalize', True),
+        )
+    else:
+        model = SAE(
+            in_dim=sae_input_dim,
+            inner_dim=inner_dim,
+            weight_norm_init=cfg.get('weight_norm_init', 0.3),
+        )
+
+    accelerate.load_checkpoint_in_model(model, checkpoint_path)
+    model.to(device, dtype)
+    model.eval()
+    return model
+
+
+def load_iqa_model(
+    layer_num: int,
+    device: str = 'cuda',
+    iqa_metric: str = 'arniqa-kadid',
+    swin_num: int = 2,
+):
+    """
+    Загружает IQA-модель и регистрирует hook на выбранный слой.
+
+    Поддерживаемые метрики:
+    - arniqa-kadid: hook на iqa.net.encoder[layer_num]
+    - maniqa: hook на iqa.net.swintransformer{swin_num}.layers[layer_num]
+    - liqe / liqe_mix: hook на iqa.net.clip_model.visual.transformer.resblocks[layer_num]
+
+    Возвращает
+    ----------
+    iqa        : загруженная IQA-модель
+    layer_name : строковый ключ, под которым активации хранятся в _iqa_activations
+    """
+    import pyiqa
+
+    global _hook_handle
+    metric_key = iqa_metric.lower()
+    create_kwargs: Dict[str, Any] = {
+        'as_loss': False,
+        'device': device,
+        'loss_reduction': 'none',
+    }
+    if metric_key == 'maniqa':
+        create_kwargs['test_sample'] = 1
+
+    iqa = pyiqa.create_metric(metric_key, **create_kwargs)
+    if iqa_metric == 'qalign':
+        iqa = QAlignVisionOnlyWrapper(iqa)
+    iqa.eval()
+
+    if _hook_handle is not None:
+        _hook_handle.remove()
+
+    hook_layout = 'blc'
+    if metric_key == 'arniqa-kadid':
+        layer_name = f'arniqa_enc_{layer_num}'
+        target_layer = iqa.net.encoder[layer_num]
+    elif metric_key == 'maniqa':
+        layer_name = f'maniqa_swintransformer{swin_num}_layers_{layer_num}'
+        target_layer = getattr(iqa.net, f'swintransformer{swin_num}').layers[layer_num]
+    elif metric_key == 'qalign':
+        layer_name = 'visual_abstractor'
+        target_layer = iqa.visual_abstractor
+    elif metric_key in {'liqe', 'liqe_mix'}:
+        # Keep behavior close to train_script_liqe: skip layers after target block.
+        num_visual_layers = len(iqa.net.clip_model.visual.transformer.resblocks)
+        for i in range(layer_num + 1, num_visual_layers):
+            iqa.net.clip_model.visual.transformer.resblocks[i] = torch.nn.Identity()
+        layer_name = f'liqe_resblock_{layer_num}'
+        target_layer = iqa.net.clip_model.visual.transformer.resblocks[layer_num]
+        hook_layout = 'lbc'
+    else:
+        raise ValueError(
+            f'Unsupported iqa_metric={iqa_metric!r}; expected one of '
+            "('arniqa-kadid', 'maniqa', 'qalign', 'liqe', 'liqe_mix')"
+        )
+    if metric_key == 'qalign':
+        hook_layout = 'blc_drop_cls'
+
+    _hook_handle = target_layer.register_forward_hook(
+        _make_hook(layer_name, sequence_layout=hook_layout)
+    )
+    return iqa, layer_name
+
+
+def _save_decoder_weight_norms_cache(cache_dir: str, norms: Mapping[int, float]) -> None:
+    cache_path = Path(cache_dir) / 'sae_decoder_weight_norms.parquet'
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    rows = [{'feature_id': int(k), 'norm': float(v)} for k, v in norms.items()]
+    pd.DataFrame(rows, columns=['feature_id', 'norm']).to_parquet(cache_path, index=False)
+
+
+def _load_decoder_weight_norms_cache(cache_dir: str) -> Optional[Dict[int, float]]:
+    cache_path = Path(cache_dir) / 'sae_decoder_weight_norms.parquet'
+    if not cache_path.exists():
+        return None
+    df = pd.read_parquet(cache_path)
+    if df.empty:
+        return None
+    return {int(row.feature_id): float(row.norm) for row in df.itertuples(index=False)}
+
+
+def extract_decoder_weight_norms(checkpoint_path: str, cache_dir: str) -> Dict[int, float]:
+    cached_norms = _load_decoder_weight_norms_cache(cache_dir)
+    if cached_norms is not None:
+        return cached_norms
+
+    sae_model = load_sae(checkpoint_path=checkpoint_path, device='cpu')
+    try:
+        decoder_weight = sae_model.decoder.weight
+    except AttributeError:
+        decoder_weight = sae_model.W.mT
+
+    decoder_weight = decoder_weight.detach().cpu()
+    norms = decoder_weight.norm(dim=0).numpy()
+    norms_dict = {int(i): float(norms[i]) for i in range(int(norms.shape[0]))}
+    _save_decoder_weight_norms_cache(cache_dir, norms_dict)
+    return norms_dict
+
+
+def extract_model_hyperparameters(
+    sae_config: Optional[Dict[str, Any]],
+    checkpoint_path: str,
+) -> Dict[str, Any]:
+    """Extract model hyperparameters from runtime/config and SAE config json."""
+    hyperparams: Dict[str, Any] = {
+        'iqa_layer': 3,
+        'iqa_metric': 'arniqa-kadid',
+    }
+
+    if sae_config is None:
+        checkpoint = Path(checkpoint_path).resolve()
+        config_dir = checkpoint.parent if checkpoint.is_file() else checkpoint
+        config_path = config_dir / SAE_CONFIG_FILENAME
+        if not config_path.exists():
+            return hyperparams
+        try:
+            with config_path.open('r', encoding='utf-8') as f:
+                sae_config = json.load(f)
+        except Exception as exc:
+            print(f'[warn] Failed to load SAE config from {config_path}: {exc}')
+            return hyperparams
+
+    if sae_config and isinstance(sae_config, dict):
+        hyperparams['iqa_layer'] = sae_config.get('layer_num', hyperparams['iqa_layer'])
+        hyperparams['iqa_metric'] = sae_config.get('iqa_metric', hyperparams['iqa_metric'])
+        hyperparams['sae_type'] = sae_config.get('sae_type', 'sae')
+        hyperparams['lambda_param'] = sae_config.get('lambda_param', 'unknown')
+        hyperparams['sae_inner_dim'] = sae_config.get('inner_dim', 'unknown')
+        hyperparams['sae_input_dim'] = sae_config.get('sae_input_dim', 'unknown')
+
+        if hyperparams['sae_type'] == 'mp_sae':
+            hyperparams['mp_threshold'] = sae_config.get('mp_threshold', 0.1)
+            hyperparams['mp_normalize'] = sae_config.get('mp_normalize', True)
+        else:
+            hyperparams['weight_norm_init'] = sae_config.get('weight_norm_init', 0.3)
+
+    return hyperparams

analysis/qalign_utils.py

@@ -0,0 +1,86 @@
+from typing import Callable, Union
+
+import pyiqa
+import torch
+from torchvision.transforms import functional as TF
+
+
+class QAlignVisionOnlyWrapper(torch.nn.Module):
+    """Wrapper that only runs vision encoder up to visual_abstractor, skipping LLM.
+
+    Architecture path:
+    - base_model: InferenceModel
+    - base_model.net: QAlign
+    - base_model.net.model: MPLUGOwl2LlamaForCausalLM
+    - base_model.net.model.model: MPLUGOwl2LlamaModel
+    - vision_model: MplugOwlVisionModel
+    - visual_abstractor: MplugOwlVisualAbstractorModel
+    """
+
+    def __init__(self, base_model: torch.nn.Module):
+        super().__init__()
+        self.base_model = base_model
+        self.vision_model = base_model.net.model.model.vision_model
+        self.visual_abstractor = base_model.net.model.model.visual_abstractor
+
+        self.base_model.eval()
+        for parameter in self.base_model.parameters():
+            parameter.requires_grad_(False)
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        self.base_model.eval()
+        return self
+
+    def eval(self):
+        super().eval()
+        self.base_model.eval()
+        return self
+
+    def forward(self, images: torch.Tensor):
+        device = next(self.vision_model.parameters()).device
+        if not isinstance(images, torch.Tensor):
+            images = torch.stack(list(images))
+        pixel_values = images.to(device=device, dtype=next(self.vision_model.parameters()).dtype)
+
+        with torch.no_grad():
+            hidden_states = self.vision_model(pixel_values).last_hidden_state
+            abstract_output = self.visual_abstractor(encoder_hidden_states=hidden_states)
+
+        return abstract_output
+
+
+def flatten_blc_drop_cls(hidden_states: torch.Tensor) -> torch.Tensor:
+    """Drop the last CLS token from (B, L, C) and flatten to (B*L, C)."""
+    if hidden_states.dim() == 3 and hidden_states.shape[1] > 1:
+        hidden_states = hidden_states[:, :-1, :]
+    return hidden_states.reshape(-1, hidden_states.shape[-1])
+
+
+def _center_crop_with_padding(image: torch.Tensor, crop_size: int) -> torch.Tensor:
+    _, height, width = image.shape
+    pad_height = max(0, crop_size - height)
+    pad_width = max(0, crop_size - width)
+    if pad_height > 0 or pad_width > 0:
+        padding = [pad_width // 2, pad_height // 2, pad_width - pad_width // 2, pad_height - pad_height // 2]
+        image = TF.pad(image, padding, fill=0, padding_mode="constant")
+    return TF.center_crop(image, [crop_size, crop_size])
+
+
+def qalign_image_transform(crop_size: int) -> Callable[[torch.Tensor], torch.Tensor]:
+    def _transform(image: torch.Tensor) -> torch.Tensor:
+        return _center_crop_with_padding(image, crop_size)
+
+    return _transform
+
+
+def create_qalign_model(device: Union[str, torch.device]) -> QAlignVisionOnlyWrapper:
+    """Create Q-Align metric wrapped for vision-only activation extraction."""
+    base_model = pyiqa.api_helpers.create_metric("qalign")
+    base_model.to(device)
+    return QAlignVisionOnlyWrapper(base_model)
+
+
+def get_visual_abstractor(model: torch.nn.Module):
+    """Get visual_abstractor from a QAlignVisionOnlyWrapper."""
+    return model.visual_abstractor

analysis/utils.py

@@ -0,0 +1,120 @@
+"""
+Общие вспомогательные утилиты пакета analysis.
+"""
+
+from typing import List
+
+import numpy as np
+import pandas as pd
+
+from analysis.features.feature_indexing import FeatureMatrix
+
+
+def get_top_images_for_feature(
+    features: FeatureMatrix,
+    meta: pd.DataFrame,
+    feature_id: int,
+    top_n: int = 10,
+    aggregation: str = 'mean_acts',
+) -> List[int]:
+    """
+    Возвращает индексы изображений, на которых признак feature_id активируется сильнее всего.
+
+    Активации патчей внутри одного изображения агрегируются в одно скалярное значение,
+    после чего изображения сортируются по убыванию.
+
+    Параметры
+    ----------
+    features    : CSR activations with global id per column
+    meta        : DataFrame с колонкой 'image_idx' (один патч — одна строка)
+    feature_id  : global SAE feature id
+    top_n       : число возвращаемых изображений
+    aggregation : 'mean_acts' | 'max' | 'sum'
+                  mean_acts — среднее по патчам с активацией > 0
+                  max       — максимальная активация среди патчей
+                  sum       — сумма всех активаций
+
+    Возвращает
+    ----------
+    List[int] — image_idx в порядке убывания агрегированной активации
+    """
+    assert aggregation in ('mean_acts', 'max', 'sum'), (
+        f"aggregation must be 'mean_acts', 'max' or 'sum', got {aggregation!r}"
+    )
+
+    col = features.column_for(feature_id)
+    feature_acts = np.asarray(features.codes[:, col].todense()).ravel()  # (n_patches,)
+    image_idx_arr = meta['image_idx'].values
+
+    unique_images = np.unique(image_idx_arr)
+    scores = np.empty(len(unique_images), dtype=np.float32)
+
+    for i, img_idx in enumerate(unique_images):
+        mask = image_idx_arr == img_idx
+        vals = feature_acts[mask]
+        if aggregation == 'mean_acts':
+            active = vals[vals > 0]
+            scores[i] = active.mean() if len(active) > 0 else 0.0
+        elif aggregation == 'max':
+            scores[i] = vals.max()
+        else:
+            scores[i] = vals.sum()
+
+    order = np.argsort(scores)[::-1]
+    return unique_images[order[:top_n]].tolist()
+
+
+def get_top_images_for_feature_by_iou(
+    features: FeatureMatrix,
+    meta: pd.DataFrame,
+    feature_id: int,
+    top_n: int = 10,
+    dataset: str | None = None,
+) -> List[int]:
+    """
+    Возвращает индексы изображений с наибольшим IoU между бинарной картой
+    активаций признака и маской искажений для каждого изображения.
+
+    Требует, чтобы в `meta` были колонки `image_idx` и либо `patch_mask_label`,
+    либо `patch_is_distorted` (см. iou_utils._load_patch_mask_for_group).
+    """
+    from analysis.metrics import iou_utils
+
+    col = features.column_for(feature_id)
+    feature_acts = np.asarray(features.codes[:, col].todense()).ravel()
+    image_groups = meta.groupby('image_idx')
+
+    scores = []  # list of (image_idx, iou)
+    for image_idx, group_df in image_groups:
+        sample_indices = group_df.index.to_numpy()
+        vals = feature_acts[sample_indices]
+        # binary activation map for this image (per-patch)
+        act_binary = (vals > 0).astype(np.uint8)
+
+        try:
+            patch_masks = iou_utils._load_patch_mask_for_group(
+                group_df,
+                target_dist_type=None,
+                dataset=(dataset or ''),
+            )
+        except Exception:
+            # If no patch masks available, IoU is undefined — treat as 0
+            scores.append((int(image_idx), 0.0))
+            continue
+
+        if patch_masks.shape[0] != act_binary.shape[0]:
+            # mismatch: skip or treat as 0
+            scores.append((int(image_idx), 0.0))
+            continue
+
+        try:
+            iou = float(iou_utils._compute_patch_iou(act_binary, patch_masks))
+        except Exception:
+            iou = 0.0
+        scores.append((int(image_idx), iou))
+
+    if not scores:
+        return []
+
+    scores.sort(key=lambda x: x[1], reverse=True)
+    return [img for img, _ in scores[:top_n]]

analysis/viz/umap_plot.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import io
+import base64
+from typing import Any, List, Tuple
+
+import numpy as np
+from PIL import Image
+import plotly.graph_objects as go
+
+
+CATEGORY_PALETTE = [
+    '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#17becf', '#bcbd22',
+    '#8c564b', '#e377c2', '#7f7f7f', '#9467bd', '#aec7e8', '#ffbb78', '#98df8a',
+]
+
+
+def _img_to_base64(img: Image.Image, fmt: str = 'JPEG', quality: int = 80) -> str:
+    buf = io.BytesIO()
+    img.save(buf, format=fmt, quality=quality)
+    return base64.b64encode(buf.getvalue()).decode('utf-8')
+
+
+def make_thumb_b64(img_idx: int, kadid_ds: Any, size: int = 140) -> str | None:
+    """Return base64 thumbnail for image index using dataset object providing `images`.
+
+    `kadid_ds` is expected to expose sequence-like `.images` where each entry is a path.
+    """
+    try:
+        p = str(kadid_ds.images[int(img_idx)])
+        img = Image.open(p).convert('RGB').resize((size, size), Image.LANCZOS)
+        return _img_to_base64(img)
+    except Exception as e:
+        print(f"Error occurred while processing image {img_idx}: {e}")
+        return None
+
+
+def _align_umap_series_length(
+    emb: np.ndarray,
+    labels: np.ndarray,
+    customdata: np.ndarray | list | None,
+    hovertext: list[str] | None,
+    marker_sizes: np.ndarray | None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | list | None, list[str] | None, np.ndarray | None]:
+    n = int(emb.shape[0])
+    labels_arr = np.asarray(labels)
+    if labels_arr.shape[0] == n:
+        return emb, labels_arr, customdata, hovertext, marker_sizes
+
+    n = min(n, int(labels_arr.shape[0]))
+    emb = emb[:n]
+    labels_arr = labels_arr[:n]
+    if customdata is not None:
+        customdata = list(customdata)[:n]
+    if hovertext is not None:
+        hovertext = list(hovertext)[:n]
+    if marker_sizes is not None:
+        marker_sizes = np.asarray(marker_sizes)[:n]
+    return emb, labels_arr, customdata, hovertext, marker_sizes
+
+
+def build_categorical_traces(
+    emb: np.ndarray,
+    labels: np.ndarray,
+    selected_category: str | None = None,
+    palette: List[str] | None = None,
+    customdata: np.ndarray | list | None = None,
+    hovertext: list[str] | None = None,
+    marker_sizes: np.ndarray | None = None,
+) -> Tuple[list[go.Scatter], list[str]]:
+    emb, labels, customdata, hovertext, marker_sizes = _align_umap_series_length(
+        emb, labels, customdata, hovertext, marker_sizes
+    )
+    palette = CATEGORY_PALETTE if palette is None else palette
+    unique_cats = sorted(str(x) for x in np.unique(labels))
+    selected_category = None if selected_category in {None, ''} else str(selected_category)
+    if selected_category not in unique_cats:
+        selected_category = None
+
+    traces: list[go.Scatter] = []
+    for i, cat in enumerate(unique_cats):
+        if selected_category is not None and cat != selected_category:
+            continue
+
+        mask = np.asarray([str(v) == cat for v in labels], dtype=bool)
+        if not np.any(mask):
+            continue
+        img_indices = np.flatnonzero(mask).astype(np.int32)
+
+        # Prepare customdata as list of [img_idx, label]
+        if customdata is None:
+            trace_customdata = [[int(idx), str(labels[int(idx)])] for idx in img_indices]
+        else:
+            trace_customdata = [customdata[int(idx)] for idx in img_indices]
+
+        trace_hovertext = None
+        if hovertext is not None:
+            trace_hovertext = [hovertext[int(idx)] for idx in img_indices]
+
+        trace_sizes = None
+        if marker_sizes is not None:
+            trace_sizes = [float(marker_sizes[int(idx)]) for idx in img_indices]
+
+        color = palette[i % len(palette)]
+        traces.append(
+            go.Scatter(
+                x=emb[mask, 0],
+                y=emb[mask, 1],
+                mode='markers',
+                name=cat,
+                legendgroup=cat,
+                marker=dict(size=trace_sizes if trace_sizes is not None else 7, color=color, opacity=0.85, line=dict(width=0)),
+                customdata=trace_customdata,
+                hoverinfo='none',
+                hovertext=trace_hovertext if trace_hovertext is not None else [f"image_idx={int(idx)} | category={cat}" for idx in img_indices],
+                showlegend=True,
+            )
+        )
+
+    return traces, unique_cats
+
+
+def build_umap_figure(
+    embedding_2d: np.ndarray,
+    labels: np.ndarray,
+    *,
+    title: str,
+    selected_category: str | None = None,
+    palette: List[str] | None = None,
+    customdata: np.ndarray | list | None = None,
+    hovertext: list[str] | None = None,
+    marker_sizes: np.ndarray | None = None,
+) -> Tuple[go.Figure, list[str]]:
+    traces, unique_cats = build_categorical_traces(
+        embedding_2d,
+        labels,
+        selected_category=selected_category,
+        palette=palette,
+        customdata=customdata,
+        hovertext=hovertext,
+        marker_sizes=marker_sizes,
+    )
+    fig = go.Figure(data=traces)
+    fig.update_layout(
+        title=dict(text=title, x=0.01, xanchor='left', y=0.97, yanchor='top', font=dict(size=16)),
+        xaxis_title='UMAP 1',
+        yaxis_title='UMAP 2',
+        legend=dict(
+            orientation='h',
+            yanchor='top',
+            y=-0.16,
+            xanchor='left',
+            x=0.0,
+            title_text='',
+            font=dict(size=11),
+            tracegroupgap=8,
+        ),
+        margin=dict(l=40, r=24, t=96, b=92),
+        hovermode='closest',
+        autosize=True,
+    )
+    return fig, unique_cats
+
+
+def plot_categorical_scatter_matplotlib(
+    emb: np.ndarray,
+    labels: np.ndarray,
+    title: str,
+    legend_title: str = 'category',
+) -> None:
+    import matplotlib.pyplot as plt
+
+    unique_cats = sorted(dict.fromkeys(str(x) for x in labels))
+    cmap = plt.get_cmap('tab10')
+    n = max(1, len(unique_cats))
+    colors = [cmap(i % cmap.N) for i in range(n)]
+    plt.figure(figsize=(10, 7))
+    for i, cat in enumerate(unique_cats):
+        mask = np.array([str(v) == cat for v in labels])
+        if not np.any(mask):
+            continue
+        plt.scatter(
+            emb[mask, 0],
+            emb[mask, 1],
+            c=[colors[i]],
+            label=cat,
+            s=28,
+            alpha=0.85,
+            edgecolors='none',
+        )
+    plt.title(title)
+    plt.xlabel('UMAP 1')
+    plt.ylabel('UMAP 2')
+    plt.grid(True, alpha=0.3)
+    plt.legend(title=legend_title, bbox_to_anchor=(1.05, 1), loc='upper left')
+    plt.tight_layout()
+    plt.show()

analysis/viz/umap_utils.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import pickle
+import re
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, Hashable, Mapping, Optional
+
+import numpy as np
+import umap
+
+_UMAP_CACHE: Dict[Hashable, Dict[str, Any]] = {}
+
+
+def _json_default(value: object) -> object:
+    if isinstance(value, (np.integer, np.floating)):
+        return value.item()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return str(value)
+
+
+def _sanitize_cache_id(cache_id: str) -> str:
+    safe = re.sub(r'[^\w.\-]+', '_', str(cache_id).strip())
+    return safe or 'umap'
+
+
+def build_umap_cache_key(
+    *,
+    cache_id: str,
+    umap_params: Mapping[str, object],
+    signature: Mapping[str, object] | None = None,
+    x_shape: tuple[int, ...] | None = None,
+) -> str:
+    """Build a stable short hash for a UMAP embedding cache entry."""
+    payload: dict[str, object] = {
+        'cache_id': str(cache_id),
+        'umap_params': dict(umap_params),
+    }
+    if signature is not None:
+        payload['signature'] = dict(signature)
+    if x_shape is not None:
+        payload['x_shape'] = tuple(int(v) for v in x_shape)
+    raw = json.dumps(payload, sort_keys=True, default=_json_default).encode('utf-8')
+    return hashlib.sha256(raw).hexdigest()[:8]
+
+
+def resolve_umap_cache_path(cache_dir: str | Path, cache_id: str, cache_key: str) -> Path:
+    """Resolve on-disk cache path for one UMAP embedding."""
+    safe_id = _sanitize_cache_id(cache_id)
+    return Path(cache_dir) / f'{safe_id}_{cache_key}.npz'
+
+
+def _load_umap_from_disk(cache_path: Path) -> Dict[str, Any] | None:
+    if not cache_path.exists():
+        return None
+    try:
+        with np.load(cache_path, allow_pickle=False) as data:
+            embedding_2d = np.asarray(data['embedding_2d'], dtype=np.float32)
+            seconds = float(np.asarray(data['seconds']).item())
+            x_shape = tuple(int(v) for v in np.asarray(data['x_shape']).tolist())
+        return {
+            'embedding_2d': embedding_2d,
+            'seconds': seconds,
+            'x_shape': x_shape,
+        }
+    except Exception as exc:
+        print(f'[cache][umap] invalid cache {cache_path}: {exc}. Recomputing...')
+        return None
+
+
+def _save_umap_to_disk(cache_path: Path, result: Mapping[str, Any]) -> None:
+    cache_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        cache_path,
+        embedding_2d=np.asarray(result['embedding_2d'], dtype=np.float32),
+        seconds=np.float64(result.get('seconds', 0.0)),
+        x_shape=np.asarray(result['x_shape'], dtype=np.int64),
+    )
+    print(f'[cache][umap] saved: {cache_path}')
+
+
+def _resolve_disk_cache_path(
+    *,
+    cache_dir: str | Path,
+    cache_id: str,
+    umap_params: Mapping[str, object],
+    signature: Mapping[str, object] | None,
+    x_shape: tuple[int, ...] | None,
+) -> Path:
+    cache_key = build_umap_cache_key(
+        cache_id=cache_id,
+        umap_params=umap_params,
+        signature=signature,
+        x_shape=x_shape,
+    )
+    return resolve_umap_cache_path(cache_dir, cache_id, cache_key)
+
+
+def compute_umap_from_features(features: np.ndarray, umap_params: Dict) -> Dict[str, Any]:
+    """Compute a 2D UMAP embedding for the given feature matrix.
+
+    Returns a dict with keys: 'embedding_2d', 'seconds'.
+    """
+    if umap is None:
+        raise RuntimeError("umap package is not available")
+
+    t0 = time.perf_counter()
+    reducer = umap.UMAP(**umap_params, verbose=False)
+    embedding_2d = reducer.fit_transform(features)
+    dt = time.perf_counter() - t0
+
+    return {"embedding_2d": embedding_2d.astype(np.float32), "seconds": float(dt)}
+
+
+def get_or_compute_umap_with_builder(
+    key: Hashable,
+    build_features_fn: Callable[[], np.ndarray],
+    umap_params: Dict,
+    cache: Optional[Dict[Hashable, Dict[str, Any]]] = None,
+    *,
+    cache_dir: str | Path | None = None,
+    cache_id: str | None = None,
+    cache_signature: Mapping[str, object] | None = None,
+) -> Dict[str, Any]:
+    """Get or compute UMAP embedding using a builder function to produce features.
+
+    - ``key``: a hashable cache key (e.g. (feature_id, params_tuple)).
+    - ``build_features_fn``: zero-arg callable that returns a np.ndarray of shape (n_samples, n_features).
+    - ``umap_params``: parameters passed to ``umap.UMAP``.
+    - ``cache``: optional dict to use for caching; if None, module-level cache is used.
+    - ``cache_dir``: optional directory for on-disk NPZ cache (like feature filters).
+    - ``cache_id``: stable logical id used in cache filenames (defaults to ``str(key)``).
+    - ``cache_signature``: optional data fingerprint for disk cache lookup before feature build.
+
+    Returns a dict with embedding and timing and other meta fields. Adds 'cache_hit' flag.
+    """
+    cache = _UMAP_CACHE if cache is None else cache
+    if key in cache:
+        cached = dict(cache[key])
+        cached["cache_hit"] = True
+        return cached
+
+    resolved_cache_id = _sanitize_cache_id(cache_id if cache_id is not None else str(key))
+
+    if cache_dir is not None and cache_signature is not None:
+        disk_cache_path = _resolve_disk_cache_path(
+            cache_dir=cache_dir,
+            cache_id=resolved_cache_id,
+            umap_params=umap_params,
+            signature=cache_signature,
+            x_shape=None,
+        )
+        loaded = _load_umap_from_disk(disk_cache_path)
+        if loaded is not None:
+            print(f'[cache][umap] hit: {disk_cache_path}')
+            result = {**loaded, "cache_hit": True, "disk_cache_hit": True}
+            cache[key] = dict(result)
+            return result
+
+    features = build_features_fn()
+    x_shape = tuple(features.shape)
+
+    if cache_dir is not None:
+        disk_cache_path = _resolve_disk_cache_path(
+            cache_dir=cache_dir,
+            cache_id=resolved_cache_id,
+            umap_params=umap_params,
+            signature=cache_signature,
+            x_shape=None if cache_signature is not None else x_shape,
+        )
+        loaded = _load_umap_from_disk(disk_cache_path)
+        if loaded is not None:
+            print(f'[cache][umap] hit: {disk_cache_path}')
+            result = {**loaded, "x_shape": x_shape, "cache_hit": True, "disk_cache_hit": True}
+            cache[key] = dict(result)
+            return result
+    else:
+        disk_cache_path = None
+
+    info = compute_umap_from_features(features, umap_params)
+    result = {**info, "x_shape": x_shape, "cache_hit": False, "disk_cache_hit": False}
+    cache[key] = dict(result)
+
+    if cache_dir is not None and disk_cache_path is not None:
+        try:
+            _save_umap_to_disk(disk_cache_path, result)
+        except Exception as exc:
+            print(f'[cache][umap] failed to save {disk_cache_path}: {exc}')
+
+    return result
+
+
+def save_cache(path: str, cache: Optional[Dict[Hashable, Dict[str, Any]]] = None) -> None:
+    """Save in-memory cache to disk using pickle."""
+    cache = _UMAP_CACHE if cache is None else cache
+    with open(path, "wb") as f:
+        pickle.dump(cache, f)
+
+
+def load_cache(path: str) -> Dict[Hashable, Dict[str, Any]]:
+    """Load in-memory cache from disk (returns the cache dict)."""
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+    if not isinstance(data, dict):
+        raise ValueError("Cache file does not contain a dict")
+    return data
+
+
+def clear_cache(cache: Optional[Dict[Hashable, Dict[str, Any]]] = None) -> None:
+    cache = _UMAP_CACHE if cache is None else cache
+    cache.clear()

analysis/viz/vis_correlations.py

@@ -0,0 +1,98 @@
+"""Visualization helpers for correlation and metric heatmaps."""
+
+from pathlib import Path
+from typing import Optional, Tuple
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+
+
+def plot_correlation_heatmap(
+    corr_df: pd.DataFrame,
+    top_k: int = 30,
+    figsize: Optional[Tuple[int, int]] = None,
+    title: str = 'Feature-Distortion Correlation',
+    save_path: Optional[str] = None,
+    show: bool = True,
+) -> Optional[str]:
+    """
+    Heatmap correlations (n_categories x top_k_features).
+
+    Top-k features are selected by maximum absolute correlation
+    with at least one category.
+    """
+    return plot_metric_heatmap(
+        metric_df=corr_df,
+        top_k=top_k,
+        figsize=figsize,
+        title=title,
+        metric_label='Pearson r',
+        use_abs_ranking=True,
+        center=0.0,
+        vmin=-1.0,
+        vmax=1.0,
+        save_path=save_path,
+        show=show,
+    )
+
+
+def plot_metric_heatmap(
+    metric_df: pd.DataFrame,
+    top_k: int = 30,
+    figsize: Optional[Tuple[int, int]] = None,
+    title: str = 'Feature-Distortion Metric',
+    metric_label: str = 'Metric value',
+    use_abs_ranking: bool = True,
+    center: Optional[float] = None,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+    save_path: Optional[str] = None,
+    show: bool = True,
+    cmap: Optional[str] = None,
+) -> Optional[str]:
+    """Generic heatmap for correlations, MI, and other metrics."""
+    if use_abs_ranking:
+        rank_series = metric_df.abs().max(axis=0)
+    else:
+        rank_series = metric_df.max(axis=0)
+
+    top_features = rank_series.nlargest(top_k).index.tolist()
+    plot_data = metric_df.loc[:, top_features]
+
+    if figsize is None:
+        figsize = (max(12, top_k // 2), max(4, len(metric_df.index)))
+
+    fig, ax = plt.subplots(figsize=figsize)
+    annot = len(top_features) <= 20
+    kwargs = dict(
+        ax=ax,
+        cmap= cmap if cmap is not None else 'RdBu_r',
+        center=center,
+        vmin=vmin,
+        vmax=vmax,
+        linewidths=0.3,
+        annot=annot,
+        cbar_kws={'label': metric_label},
+    )
+    if annot:
+        kwargs['fmt'] = '.2f'
+
+    sns.heatmap(plot_data, **kwargs)
+    ax.set_xlabel('Feature ID')
+    ax.set_ylabel('Distortion category')
+    ax.set_title(title)
+    plt.tight_layout()
+
+    saved: Optional[str] = None
+    if save_path is not None:
+        save_path_obj = Path(save_path)
+        save_path_obj.parent.mkdir(parents=True, exist_ok=True)
+        fig.savefig(save_path_obj, dpi=200, bbox_inches='tight')
+        saved = str(save_path_obj)
+
+    if show:
+        plt.show()
+
+    plt.close(fig)
+    return saved

analysis/viz/vis_heatmaps.py

@@ -0,0 +1,770 @@
+"""
+Визуализация тепловых карт SAE-признаков поверх изображений KADID-10k.
+"""
+
+import io
+import math
+import re as _re
+from pathlib import Path
+
+from analysis.features.feature_indexing import FeatureMatrix
+from typing import Dict, List, Optional, Union
+from IPython.display import display
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+import matplotlib.pyplot as plt
+from matplotlib.patches import Patch, Rectangle
+from PIL import Image
+from torchvision import transforms
+
+from overcomplete.visualization.plot_utils import show, interpolate_cv2, get_image_dimensions
+from overcomplete.visualization.cmaps import VIRIDIS_ALPHA, TAB10_ALPHA
+
+from analysis.datasets import (
+    QGROUND_DISTORTION_TYPES,
+    SRGROUND_DISTORTION_TYPES,
+    SRGROUND_LEGEND_LABELS,
+    available_distortions,
+    distortion_types_mapping,
+    _image_rel_from_meta_row,
+    srground_mask_rgb_for_meta_row,
+    srground_prominences_by_image_paths,
+)
+
+# Show per-image filenames (and distortion subtitles where enabled) under heatmap panels.
+SHOW_HEATMAP_IMAGE_TITLES = False
+
+MASK_LEGEND_FONTSIZE = 20
+MASK_LEGEND_HANDLE_HEIGHT = 1.15
+MASK_LEGEND_HANDLE_LENGTH = 1.55
+
+
+def _get_distortion_info(img_path: str):
+    """
+    По имени файла KADID (I04_07_05.png) возвращает
+    (distortion_name, distortion_group) или (None, None) для оригиналов.
+    """
+    name = Path(img_path).name
+    m = _re.match(r'I\d+_(\d+)_(\d+)\.png$', name, _re.IGNORECASE)
+    if m:
+        dist_id = int(m.group(1))
+        dist_name = distortion_types_mapping.get(dist_id, f'dist_{dist_id}')
+        dist_group = available_distortions.get(dist_name, '?')
+        return dist_name, dist_group
+    return None, None
+
+
+def _get_original_path(distorted_path: str) -> Optional[str]:
+    """
+    Из пути к искажённому изображению KADID (I04_07_05.png)
+    возвращает путь к оригиналу (I04.png).
+    """
+    p = Path(distorted_path)
+    match = _re.match(r'(I\d+)_\d+_\d+\.png$', p.name, _re.IGNORECASE)
+    if match:
+        return str(p.parent / f'{match.group(1)}.png')
+
+    # Local-KADID presaved naming: I04_001_dist.png -> I04.png
+    local_match = _re.match(r'(I\d+)_\d+_dist\.png$', p.name, _re.IGNORECASE)
+    if local_match:
+        return str(p.parent / f'{local_match.group(1)}.png')
+
+    return None
+
+
+def _meta_path_value(meta_row: Optional[pd.Series], column_name: str) -> Optional[str]:
+    if meta_row is None or column_name not in meta_row:
+        return None
+    value = meta_row.get(column_name)
+    if value is None or pd.isna(value):
+        return None
+
+    value_str = str(value).strip()
+    return value_str or None
+
+
+def _build_image_title(
+    img_path: str,
+    meta_row: Optional[pd.Series] = None,
+    *,
+    include_filename: bool | None = None,
+    include_distortion_subtitle: bool = True,
+) -> str:
+    show_filename = SHOW_HEATMAP_IMAGE_TITLES if include_filename is None else include_filename
+    title_lines: list[str] = []
+    if show_filename:
+        title_lines.append(Path(img_path).name)
+
+    if include_distortion_subtitle:
+        dist_name = _meta_path_value(meta_row, 'dist_type')
+        dist_group = _meta_path_value(meta_row, 'dist_group')
+        if dist_name and dist_name.lower() != 'background':
+            if dist_group and dist_group != dist_name:
+                title_lines.append(f'{dist_name} [{dist_group}]')
+            else:
+                title_lines.append(dist_name)
+        elif meta_row is not None:
+            ann_id = _meta_path_value(meta_row, 'qground_ann_id')
+            if ann_id:
+                title_lines.append(f'ann {ann_id}')
+
+    return '\n'.join(title_lines)
+
+
+def _set_image_subplot_title(
+    ax,
+    img_path: str,
+    meta_row: Optional[pd.Series] = None,
+    *,
+    include_filename: bool | None = None,
+    include_distortion_subtitle: bool = True,
+) -> None:
+    title = _build_image_title(
+        img_path,
+        meta_row,
+        include_filename=include_filename,
+        include_distortion_subtitle=include_distortion_subtitle,
+    )
+    if title:
+        ax.set_title(title, fontsize=10)
+
+
+def _build_meta_lookup(meta: pd.DataFrame) -> dict[int, pd.Series]:
+    if 'image_idx' not in meta.columns:
+        return {}
+
+    first_rows = meta.groupby('image_idx', sort=False).first()
+    lookup: dict[int, pd.Series] = {}
+    for image_idx, row in first_rows.iterrows():
+        try:
+            lookup[int(image_idx)] = row
+        except (TypeError, ValueError):
+            continue
+    return lookup
+
+
+FigureSaveResult = Union[str, bytes, None]
+
+
+def _finalize_figure(
+    fig,
+    *,
+    save_path: Optional[str],
+    show_img: bool,
+    dpi: int = 200,
+    **savefig_kw,
+) -> FigureSaveResult:
+    """Save figure to disk, return PNG bytes, or only display (notebook)."""
+    if save_path is not None:
+        save_path_obj = Path(save_path)
+        save_path_obj.parent.mkdir(parents=True, exist_ok=True)
+        fig.savefig(save_path_obj, dpi=dpi, bbox_inches='tight', **savefig_kw)
+        result: FigureSaveResult = str(save_path_obj)
+    elif show_img:
+        result = None
+    else:
+        buffer = io.BytesIO()
+        fig.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight', **savefig_kw)
+        result = buffer.getvalue()
+
+    if show_img:
+        display(fig)
+
+    plt.close(fig)
+    return result
+
+
+def _append_heatmap_artifact(
+    artifacts: List[Dict[str, object]],
+    *,
+    kind: str,
+    feature_id: int,
+    saved: FigureSaveResult,
+) -> None:
+    if saved is None:
+        return
+    entry: Dict[str, object] = {'kind': kind, 'feature_id': str(feature_id)}
+    if isinstance(saved, bytes):
+        entry['bytes'] = saved
+    else:
+        entry['path'] = saved
+    artifacts.append(entry)
+
+
+def _plot_heatmap_grid(
+    imgs_list: List[torch.Tensor],
+    codes_tensor: torch.Tensor,
+    feature_col: int,
+    image_paths: List[str],
+    img_size_inches: float = 4.0,
+    meta_rows: Optional[List[Optional[pd.Series]]] = None,
+    save_path: Optional[str] = None,
+    show_img: bool = True,
+    *,
+    title_feature_id: int | None = None,
+) -> FigureSaveResult:
+    """
+    Рисует сетку overlay-хитмапов для одного feature_id.
+
+    Параметры
+    ----------
+    imgs_list      : список тензоров изображений (C, H, W)
+    codes_tensor   : тензор активаций (B, spatial, spatial, inner_dim)
+    feature_id     : индекс признака SAE
+    image_paths    : пути к PNG-файлам (для заголовков)
+    img_size_inches: размер одной ячейки в дюймах
+    """
+    B = len(imgs_list)
+    inner_concepts = codes_tensor.shape[-1]
+    # cols = min(B, 5)
+    cols = B
+    # rows = math.ceil(B / cols)
+    rows = 1
+
+    display_id = int(title_feature_id) if title_feature_id is not None else int(feature_col)
+    if inner_concepts < 10:
+        cmap = TAB10_ALPHA[display_id % len(TAB10_ALPHA)]
+    else:
+        cmap = VIRIDIS_ALPHA
+
+    if meta_rows is None:
+        meta_rows = [None] * B
+
+    fig, axes = plt.subplots(rows, cols,
+                             figsize=(img_size_inches * cols, img_size_inches * rows * 1.25))
+    axes_flat = np.array(axes).flatten() if B > 1 else [axes]
+
+    for ax, img_t, code_row, img_path, meta_row in zip(axes_flat, imgs_list, codes_tensor, image_paths, meta_rows):
+        plt.sca(ax)
+        width, height = get_image_dimensions(img_t)
+        heatmap_patch = code_row[:, :, feature_col]
+        if isinstance(heatmap_patch, torch.Tensor):
+            heatmap_patch = heatmap_patch.numpy()
+        heatmap = interpolate_cv2(heatmap_patch, (width, height))
+        show(img_t)
+        show(heatmap, cmap=cmap, alpha=1.0)
+        _set_image_subplot_title(ax, img_path, meta_row)
+        ax.axis('off')
+
+    for ax in axes_flat[B:]:
+        ax.axis('off')
+
+    fig.suptitle(f'Feature {display_id}', fontsize=13, x=0.5, ha='center')
+    plt.tight_layout()
+    return _finalize_figure(fig, save_path=save_path, show_img=show_img)
+
+
+def _plot_diff_grid(
+    imgs_list: List[torch.Tensor],
+    orig_tensors: List[Optional[torch.Tensor]],
+    feature_id: int,
+    image_paths: List[str],
+    img_size_inches: float = 4.0,
+    meta_rows: Optional[List[Optional[pd.Series]]] = None,
+    save_path: Optional[str] = None,
+    show_img: bool = True,
+) -> FigureSaveResult:
+    """
+    Рисует сетку |distorted − original| для одного feature_id.
+
+    Параметры
+    ----------
+    imgs_list      : список тензоров искажённых изображений (C, H, W)
+    orig_tensors   : список тензоров оригиналов (или None, если файл не найден)
+    feature_id     : индекс признака SAE (для заголовка)
+    image_paths    : пути к PNG-файлам (для заголовков)
+    img_size_inches: размер одной ячейки в дюймах
+    """
+    B = len(imgs_list)
+    fig_diff, axes = plt.subplots(1, B, figsize=(img_size_inches * B, img_size_inches * 1.25))
+    if B == 1:
+        axes = [axes]
+    if meta_rows is None:
+        meta_rows = [None] * B
+
+    for ax, dist_t, orig_t, img_path, meta_row in zip(axes, imgs_list, orig_tensors, image_paths, meta_rows):
+        if orig_t is not None:
+            diff = (dist_t.float() - orig_t.float()).abs()
+            diff_np = diff.permute(1, 2, 0).numpy().mean(axis=-1)
+            im = ax.imshow(diff_np, vmin=0, vmax=1)
+            fig_diff.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+            _set_image_subplot_title(ax, img_path, meta_row)
+        else:
+            ax.text(0.5, 0.5, 'original not found', ha='center', va='center',
+                    transform=ax.transAxes, fontsize=9)
+        ax.axis('off')
+    fig_diff.suptitle(f'|distorted − original|  (feature {feature_id})',
+                      fontsize=12, x=0.5, ha='center')
+    plt.tight_layout()
+    return _finalize_figure(fig_diff, save_path=save_path, show_img=show_img)
+
+
+def _plot_qground_mask_grid(
+    mask_paths: List[Optional[str]],
+    feature_id: int,
+    image_paths: List[str],
+    crop_size: int,
+    img_size_inches: float = 4.0,
+    meta_rows: Optional[List[Optional[pd.Series]]] = None,
+    save_path: Optional[str] = None,
+    show_img: bool = True,
+) -> FigureSaveResult:
+    B = len(mask_paths)
+    fig, axes = plt.subplots(1, B, figsize=(img_size_inches * B, img_size_inches * 1.35))
+    if B == 1:
+        axes = [axes]
+    if meta_rows is None:
+        meta_rows = [None] * B
+
+    for ax, mask_path, img_path, meta_row in zip(axes, mask_paths, image_paths, meta_rows):
+        if mask_path is not None and Path(mask_path).exists():
+            mask_rgb = np.asarray(Image.open(mask_path).convert('RGB'), dtype=np.uint8)
+            ax.imshow(mask_rgb)
+            height, width = mask_rgb.shape[:2]
+            rect_width = float(crop_size)
+            rect_height = float(crop_size)
+            rect_x = (width - rect_width) / 2.0
+            rect_y = (height - rect_height) / 2.0
+            ax.add_patch(
+                Rectangle(
+                    (rect_x, rect_y),
+                    rect_width,
+                    rect_height,
+                    fill=False,
+                    edgecolor='black',
+                    linewidth=2.0,
+                    linestyle='--',
+                )
+            )
+            _set_image_subplot_title(ax, img_path, meta_row)
+        else:
+            ax.text(0.5, 0.5, 'QGround mask not found', ha='center', va='center',
+                    transform=ax.transAxes, fontsize=9)
+        ax.axis('off')
+
+    handles = [Patch(facecolor='black', edgecolor='none', label='background')]
+    for label_name, rgb in QGROUND_DISTORTION_TYPES.items():
+        color = tuple((np.asarray(rgb, dtype=np.float32) / 255.0).tolist())
+        handles.append(Patch(facecolor=color, edgecolor='none', label=label_name))
+
+    _add_mask_legend(fig, handles)
+    fig.suptitle(f'QGround annotation mask (feature {feature_id})', fontsize=12, x=0.5, ha='center')
+    fig.tight_layout(rect=(0, 0.14, 1, 1))
+    return _finalize_figure(fig, save_path=save_path, show_img=show_img)
+
+
+def _add_mask_legend(fig, handles: list[Patch], *, ncol: int = 3) -> None:
+    fig.legend(
+        handles=handles,
+        loc='lower center',
+        ncol=ncol,
+        frameon=False,
+        bbox_to_anchor=(0.5, -0.02),
+        fontsize=MASK_LEGEND_FONTSIZE,
+        handleheight=MASK_LEGEND_HANDLE_HEIGHT,
+        handlelength=MASK_LEGEND_HANDLE_LENGTH,
+        labelspacing=0.55,
+        borderpad=0.4,
+    )
+
+
+def srground_legend_handles(*, include_sr_artifact: bool = True) -> list[Patch]:
+    handles = [Patch(facecolor='black', edgecolor='none', label='background')]
+    for dist_name, rgb in SRGROUND_DISTORTION_TYPES.items():
+        if dist_name == 'sr_artifact' and not include_sr_artifact:
+            continue
+        color = tuple((np.asarray(rgb, dtype=np.float32) / 255.0).tolist())
+        label = SRGROUND_LEGEND_LABELS.get(dist_name, dist_name)
+        handles.append(Patch(facecolor=color, edgecolor='none', label=label))
+    return handles
+
+
+def _plot_overlay_and_srground_mask_rows(
+    imgs_list: List[torch.Tensor],
+    codes_tensor: torch.Tensor,
+    feature_col: int,
+    image_paths: List[str],
+    mask_rgb_list: List[Optional[np.ndarray]],
+    img_size_inches: float,
+    meta_rows: Optional[List[Optional[pd.Series]]],
+    *,
+    title_feature_id: int,
+    include_sr_artifact: bool,
+    crop_size: int,
+    save_path: Optional[str] = None,
+    show_img: bool = False,
+) -> FigureSaveResult:
+    B = len(imgs_list)
+    inner_concepts = codes_tensor.shape[-1]
+    display_id = int(title_feature_id)
+
+    if inner_concepts < 10:
+        cmap = TAB10_ALPHA[display_id % len(TAB10_ALPHA)]
+    else:
+        cmap = VIRIDIS_ALPHA
+
+    if meta_rows is None:
+        meta_rows = [None] * B
+
+    fig, axes = plt.subplots(2, B, figsize=(img_size_inches * B, img_size_inches * 2.55))
+    if B == 1:
+        overlay_axes = [axes[0]]
+        mask_axes = [axes[1]]
+    else:
+        overlay_axes = list(axes[0])
+        mask_axes = list(axes[1])
+
+    for ax, img_t, code_row, img_path, meta_row in zip(
+        overlay_axes, imgs_list, codes_tensor, image_paths, meta_rows
+    ):
+        plt.sca(ax)
+        width, height = get_image_dimensions(img_t)
+        heatmap_patch = code_row[:, :, feature_col]
+        if isinstance(heatmap_patch, torch.Tensor):
+            heatmap_patch = heatmap_patch.numpy()
+        heatmap = interpolate_cv2(heatmap_patch, (width, height))
+        show(img_t)
+        show(heatmap, cmap=cmap, alpha=1.0)
+        _set_image_subplot_title(
+            ax,
+            img_path,
+            meta_row,
+            include_distortion_subtitle=False,
+        )
+        ax.axis('off')
+
+    for ax, mask_rgb in zip(mask_axes, mask_rgb_list):
+        if mask_rgb is not None:
+            ax.imshow(mask_rgb)
+            height, width = mask_rgb.shape[:2]
+            rect_size = float(crop_size)
+            rect_x = (width - rect_size) / 2.0
+            rect_y = (height - rect_size) / 2.0
+            ax.add_patch(
+                Rectangle(
+                    (rect_x, rect_y),
+                    rect_size,
+                    rect_size,
+                    fill=False,
+                    edgecolor='black',
+                    linewidth=2.0,
+                    linestyle='--',
+                )
+            )
+        else:
+            ax.text(
+                0.5,
+                0.5,
+                'annotation mask not found',
+                ha='center',
+                va='center',
+                transform=ax.transAxes,
+                fontsize=9,
+            )
+        ax.axis('off')
+
+    _add_mask_legend(
+        fig,
+        srground_legend_handles(include_sr_artifact=include_sr_artifact),
+    )
+    fig.suptitle(f'Feature {display_id}', fontsize=13, x=0.5, ha='center')
+    fig.tight_layout(rect=(0, 0.16, 1, 0.96))
+    return _finalize_figure(fig, save_path=save_path, show_img=show_img)
+
+
+def render_top_feature_panel_srground(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    image_indices: List[int],
+    image_paths: List[str],
+    feature_id: int,
+    *,
+    patches_per_image: Optional[int] = None,
+    crop_size: int = 224,
+    img_size_inches: float = 3.5,
+    include_sr_artifact: bool = True,
+    datasets_root: str | None = None,
+) -> bytes | None:
+    """Two-row panel: SAE heatmap overlays and SRGround annotation masks with legend."""
+    pil_images = [Image.open(p).convert('RGB') for p in image_paths]
+    preprocess = transforms.Compose(
+        [
+            transforms.CenterCrop(int(crop_size)),
+            transforms.ToTensor(),
+        ]
+    )
+    imgs_list = [preprocess(img) for img in pil_images]
+
+    if patches_per_image is None:
+        patches_per_image = int(meta['patch_idx'].max()) + 1
+    spatial = int(math.isqrt(patches_per_image))
+    inner_dim = features.codes.shape[1]
+    image_idx_arr = meta['image_idx'].values
+    codes_list = []
+    for img_idx in image_indices:
+        row_mask = image_idx_arr == img_idx
+        codes_list.append(features.codes[row_mask].toarray().astype(np.float32))
+    codes_tensor = torch.from_numpy(np.stack(codes_list)).view(len(image_indices), spatial, spatial, inner_dim)
+
+    meta_lookup = _build_meta_lookup(meta)
+    meta_rows = [meta_lookup.get(int(img_idx)) for img_idx in image_indices]
+    image_rels = [_image_rel_from_meta_row(row) for row in meta_rows]
+    prom_map = srground_prominences_by_image_paths(
+        [rel for rel in image_rels if rel],
+        datasets_root=datasets_root,
+    )
+    mask_rgb_list = [
+        srground_mask_rgb_for_meta_row(
+            row,
+            datasets_root=datasets_root,
+            include_sr_artifact=include_sr_artifact,
+            crop_size=crop_size,
+            prominences=prom_map.get(image_rel) if image_rel else None,
+        )
+        for row, image_rel in zip(meta_rows, image_rels)
+    ]
+
+    matrix_col = features.column_for(int(feature_id))
+    saved = _plot_overlay_and_srground_mask_rows(
+        imgs_list,
+        codes_tensor,
+        matrix_col,
+        image_paths,
+        mask_rgb_list,
+        img_size_inches,
+        meta_rows,
+        title_feature_id=int(feature_id),
+        include_sr_artifact=include_sr_artifact,
+        crop_size=crop_size,
+        save_path=None,
+        show_img=False,
+    )
+    return saved if isinstance(saved, bytes) else None
+
+
+def visualize_feature_heatmaps(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    image_indices: List[int],
+    image_paths: List[str],
+    feature_ids: List[int],
+    patches_per_image: Optional[int] = None,
+    crop_size: int = 224,
+    img_size_inches: float = 4.0,
+    show_diff: bool = True,
+    save_dir: Optional[str] = None,
+    file_prefix: str = '',
+    show_img: bool = True,
+) -> List[Dict[str, object]]:
+    """
+    Визуализирует тепловые карты SAE-признаков из предвычисленных sparse активаций.
+
+    Для каждого feature_id отображается:
+    - overlay-хитмап поверх искажённого изображения
+    - (опционально) |distorted − original| попиксельно
+
+    Параметры
+    ----------
+    meta             : DataFrame с метаданными (image_idx, patch_idx, ...)
+    features         : CSR activations with global id per column
+    image_indices    : индексы изображений (image_idx) для визуализации
+    image_paths      : пути к соответствующим PNG-файлам
+    feature_ids      : global SAE feature ids to visualize
+    patches_per_image: число патчей на изображ��ние; если None — из patch_idx
+    crop_size        : размер кропа при кэшировании
+    img_size_inches  : размер одного изображения на фигуре (в дюймах)
+    show_diff        : показывать ли |distorted − original|
+    save_dir         : директория для сохранения PNG; если None — PNG в памяти (ключ ``bytes``)
+    file_prefix      : префикс имени файла (например, stage/agg)
+    show_img         : показывать ли фигуры inline
+
+    Возвращает
+    ----------
+    List[Dict[str, object]] — артефакты с ключом ``path`` (диск) или ``bytes`` (память)
+    """
+    assert len(image_indices) == len(image_paths), \
+        "image_indices и image_paths должны иметь одинаковую длину"
+
+    codes = features.codes
+    inner_dim = codes.shape[1]
+
+    if patches_per_image is None:
+        patches_per_image = int(meta['patch_idx'].max()) + 1
+
+    spatial = int(math.isqrt(patches_per_image))
+    assert spatial * spatial == patches_per_image, (
+        f"Число патчей {patches_per_image} не является точным квадратом — "
+        f"проверьте crop_size или слой модели"
+    )
+
+    preprocess = transforms.Compose(
+        [
+            transforms.CenterCrop(crop_size),
+            transforms.ToTensor(),
+        ]
+    )
+
+    pil_images = [Image.open(p).convert("RGB") for p in image_paths]
+    return visualize_feature_heatmaps_from_images(
+        meta=meta,
+        features=features,
+        image_indices=image_indices,
+        images=pil_images,
+        image_names=image_paths,
+        feature_ids=feature_ids,
+        patches_per_image=patches_per_image,
+        crop_size=crop_size,
+        img_size_inches=img_size_inches,
+        show_diff=show_diff,
+        save_dir=save_dir,
+        file_prefix=file_prefix,
+        show_img=show_img,
+        preprocess=preprocess,
+    )
+
+
+def visualize_feature_heatmaps_from_images(
+    meta: pd.DataFrame,
+    features: FeatureMatrix,
+    image_indices: List[int],
+    images: List[Image.Image],
+    image_names: List[str],
+    feature_ids: List[int],
+    patches_per_image: Optional[int] = None,
+    crop_size: int = 224,
+    img_size_inches: float = 4.0,
+    show_diff: bool = True,
+    save_dir: Optional[str] = None,
+    file_prefix: str = '',
+    show_img: bool = True,
+    *,
+    preprocess: Optional[object] = None,
+) -> List[Dict[str, str]]:
+    """Same as `visualize_feature_heatmaps`, but accepts already opened PIL images.
+
+    This is useful for Dash uploads to avoid writing temporary image files.
+    """
+    assert len(image_indices) == len(images) == len(image_names), (
+        "image_indices, images, and image_names must have the same length"
+    )
+
+    codes = features.codes
+    inner_dim = codes.shape[1]
+
+    if patches_per_image is None:
+        patches_per_image = int(meta['patch_idx'].max()) + 1
+
+    spatial = int(math.isqrt(patches_per_image))
+    assert spatial * spatial == patches_per_image, (
+        f"Число патчей {patches_per_image} не является точным квадратом — "
+        f"проверьте crop_size или слой модели"
+    )
+
+    if preprocess is None:
+        preprocess = transforms.Compose(
+            [
+                transforms.CenterCrop(crop_size),
+                transforms.ToTensor(),
+            ]
+        )
+
+    imgs_list = [preprocess(img.convert("RGB")) for img in images]
+
+    meta_lookup = _build_meta_lookup(meta)
+    meta_rows = [meta_lookup.get(int(img_idx)) for img_idx in image_indices]
+
+    orig_tensors: List[Optional[torch.Tensor]] = []
+    mask_paths: List[Optional[str]] = []
+    has_qground_masks = False
+    if show_diff:
+        for img_name, meta_row in zip(image_names, meta_rows):
+            orig_path = _get_original_path(img_name) or _meta_path_value(meta_row, 'original_img_path')
+            if orig_path and Path(orig_path).exists():
+                orig_tensors.append(preprocess(Image.open(orig_path).convert('RGB')))
+            else:
+                orig_tensors.append(None)
+
+            mask_path = _meta_path_value(meta_row, 'mask_path')
+            if mask_path and Path(mask_path).exists():
+                has_qground_masks = True
+            mask_paths.append(mask_path)
+
+    image_idx_arr = meta['image_idx'].values
+    codes_list = []
+    for img_idx in image_indices:
+        mask = image_idx_arr == img_idx
+        codes_list.append(codes[mask].toarray().astype(np.float32))
+
+    codes_np = np.stack(codes_list)  # (B, P, inner_dim)
+    codes_tensor = torch.from_numpy(codes_np).view(
+        len(image_indices), spatial, spatial, inner_dim
+    )
+
+    save_root = Path(save_dir) if save_dir is not None else None
+    artifacts: List[Dict[str, object]] = []
+
+    for global_feature_id in feature_ids:
+        matrix_col = features.column_for(int(global_feature_id))
+        safe_prefix = f'{file_prefix}_' if file_prefix else ''
+        overlay_name = f'{safe_prefix}feature_{global_feature_id}_overlay.png'
+        overlay_path = str(save_root / overlay_name) if save_root is not None else None
+        saved_overlay = _plot_heatmap_grid(
+            imgs_list,
+            codes_tensor,
+            matrix_col,
+            image_names,
+            img_size_inches,
+            meta_rows=meta_rows,
+            save_path=overlay_path,
+            show_img=show_img,
+            title_feature_id=int(global_feature_id),
+        )
+        _append_heatmap_artifact(
+            artifacts,
+            kind='overlay',
+            feature_id=int(global_feature_id),
+            saved=saved_overlay,
+        )
+        if show_diff and any(t is not None for t in orig_tensors):
+            diff_name = f'{safe_prefix}feature_{global_feature_id}_diff.png'
+            diff_path = str(save_root / diff_name) if save_root is not None else None
+            saved_diff = _plot_diff_grid(
+                imgs_list,
+                orig_tensors,
+                global_feature_id,
+                image_names,
+                img_size_inches,
+                meta_rows=meta_rows,
+                save_path=diff_path,
+                show_img=show_img,
+            )
+            _append_heatmap_artifact(
+                artifacts,
+                kind='diff',
+                feature_id=int(global_feature_id),
+                saved=saved_diff,
+            )
+        elif show_diff and has_qground_masks:
+            mask_name = f'{safe_prefix}feature_{global_feature_id}_mask.png'
+            mask_path = str(save_root / mask_name) if save_root is not None else None
+            saved_mask = _plot_qground_mask_grid(
+                mask_paths,
+                global_feature_id,
+                image_names,
+                meta_rows=meta_rows,
+                img_size_inches=img_size_inches,
+                crop_size=crop_size,
+                save_path=mask_path,
+                show_img=show_img,
+            )
+            _append_heatmap_artifact(
+                artifacts,
+                kind='mask',
+                feature_id=int(global_feature_id),
+                saved=saved_mask,
+            )
+
+    return artifacts

analysis/viz/vis_metrics.py

@@ -0,0 +1,147 @@
+"""Generalized visualization for relationship metrics (correlation, MI, etc.)."""
+
+from pathlib import Path
+from typing import Optional, Tuple
+
+import pandas as pd
+
+from analysis.viz.vis_correlations import plot_metric_heatmap
+
+
+def plot_similarity_heatmap(
+    metric_df: pd.DataFrame,
+    metric_type: str = 'custom',
+    top_k: int = 30,
+    figsize: Optional[Tuple[int, int]] = None,
+    title: str = 'Feature-Relationship Heatmap',
+    metric_label: Optional[str] = None,
+    use_abs_ranking: Optional[bool] = None,
+    center: Optional[float] = None,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+    cmap: Optional[str] = None,
+    save_path: Optional[str] = None,
+    show: bool = True,
+) -> Optional[str]:
+    """
+    Generic heatmap visualization for feature-relationship metrics.
+
+    Adapts visualization parameters (colormap, center, vmin/vmax, ranking)
+    based on metric type, with ability to override for custom types.
+
+    Parameters
+    ----------
+    metric_df : pd.DataFrame
+        DataFrame with shape (n_categories, n_features), values are metric scores.
+    metric_type : {'correlation', 'mutual_information', 'custom'}, default 'custom'
+        Type of metric, determines automatic parameter selection:
+        - 'correlation': uses RdBu_r colormap, center=0.0, vmin=-1, vmax=1, abs ranking
+        - 'mutual_information': uses viridis colormap, center=None, vmin=0, vmax=max, regular ranking
+        - 'custom': uses explicit parameters (use_abs_ranking, center, vmin, vmax, cmap)
+    top_k : int, default 30
+        Number of top features to display (selected by importance rank).
+    figsize : tuple, optional
+        Figure size (width, height). If None, computed automatically.
+    title : str, default 'Feature-Relationship Heatmap'
+        Plot title.
+    metric_label : str, optional
+        Colorbar label. If None, defaults based on metric_type.
+    use_abs_ranking : bool, optional
+        If True, ranks features by absolute max value. If None, defaults based on metric_type.
+    center : float, optional
+        Value to center colormap symmetrically (for diverging maps like RdBu_r).
+        If None, colormap is not centered.
+    vmin : float, optional
+        Minimum value for colormap scale. If None, defaults based on metric_type.
+    vmax : float, optional
+        Maximum value for colormap scale. If None, defaults based on metric_type.
+    cmap : str, optional
+        Matplotlib colormap name. If None, defaults based on metric_type.
+    save_path : str, optional
+        If provided, save figure to this path.
+    show : bool, default True
+        If True, display the figure in the notebook/interface.
+
+    Returns
+    -------
+    str or None
+        Path to saved figure if save_path was provided, otherwise None.
+
+    Examples
+    --------
+    >>> # Correlation heatmap (auto-configured)
+    >>> plot_similarity_heatmap(corr_df, metric_type='correlation', top_k=30, title='Correlations')
+    >>>
+    >>> # Mutual Information heatmap (auto-configured)
+    >>> plot_similarity_heatmap(mi_df, metric_type='mutual_information', top_k=30, title='MI')
+    >>>
+    >>> # Custom metric with explicit colormap parameters
+    >>> plot_similarity_heatmap(
+    ...     custom_metric_df,
+    ...     metric_type='custom',
+    ...     cmap='coolwarm',
+    ...     center=0.5,
+    ...     vmin=0,
+    ...     vmax=1,
+    ...     use_abs_ranking=False
+    ... )
+    """
+
+    # Auto-configure parameters based on metric_type
+    if metric_type == 'correlation':
+        if metric_label is None:
+            metric_label = 'Pearson r'
+        if use_abs_ranking is None:
+            use_abs_ranking = True
+        if center is None:
+            center = 0.0
+        if vmin is None:
+            vmin = -1.0
+        if vmax is None:
+            vmax = 1.0
+        if cmap is None:
+            cmap = 'RdBu_r'
+
+    elif metric_type == 'mutual_information':
+        if metric_label is None:
+            metric_label = 'Mutual Information (nats)'
+        if use_abs_ranking is None:
+            use_abs_ranking = False
+        if center is None:
+            center = None
+        if vmin is None:
+            vmin = 0.0
+        if vmax is None:
+            vmax = metric_df.max().max()
+        if cmap is None:
+            cmap = 'viridis'
+
+    elif metric_type == 'custom':
+        # For custom type, use explicit parameters or raise if required ones missing
+        if metric_label is None:
+            metric_label = 'Metric value'
+        if use_abs_ranking is None:
+            use_abs_ranking = True
+        if cmap is None:
+            cmap = 'viridis'
+        # center, vmin, vmax can remain None for auto-scaling
+
+    else:
+        available = "'correlation', 'mutual_information', 'custom'"
+        raise ValueError(f"Unknown metric_type: {metric_type!r}. Available: {available}")
+
+    # Delegate to the core plotting function
+    return plot_metric_heatmap(
+        metric_df=metric_df,
+        top_k=top_k,
+        figsize=figsize,
+        title=title,
+        metric_label=metric_label,
+        use_abs_ranking=use_abs_ranking,
+        center=center,
+        vmin=vmin,
+        vmax=vmax,
+        cmap=cmap,
+        save_path=save_path,
+        show=show,
+    )

analysis/viz/vis_scatter.py

@@ -0,0 +1,484 @@
+"""
+Scatter-plot активаций SAE-признаков, часть функций тоже из PatchSAE:
+log(sparsity) × log(mean_acts), окрашенные по энтропии меток искажений.
+"""
+
+from pathlib import Path
+
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+import plotly.express as px
+
+from analysis.features.feature_indexing import FeatureMatrix
+from analysis.metrics.iou_utils import compute_iou_per_feature
+from analysis.metrics.precision_recall import compute_distortion_precision, compute_distortion_recall
+from analysis.metrics.roc_auc import compute_distortion_roc_auc
+
+def calculate_entropy(
+    top_val: torch.Tensor,
+    top_label: torch.Tensor,
+    ignore_label_idx: Optional[int] = None,
+    eps: float = 1e-9,
+) -> torch.Tensor:
+    dict_size = top_label.shape[0]
+    entropy = torch.zeros(dict_size)
+
+    for i in range(dict_size):
+        unique_labels, counts = top_label[i].unique(return_counts=True)
+        if ignore_label_idx is not None:
+            mask = unique_labels != ignore_label_idx
+            counts = counts[mask]
+            unique_labels = unique_labels[mask]
+
+        if len(unique_labels) == 0 or counts.sum().item() < 10:
+            entropy[i] = -1
+            continue
+
+        summed_probs = torch.zeros(len(unique_labels), dtype=top_val.dtype)
+        for j, label in enumerate(unique_labels):
+            summed_probs[j] = top_val[i][top_label[i] == label].sum().item()
+
+        summed_probs = summed_probs / summed_probs.sum()
+        entropy[i] = -torch.sum(summed_probs * torch.log(summed_probs + eps))
+
+    return entropy
+
+
+def get_top_k_patches(
+    codes: sp.csr_matrix,
+    meta: pd.DataFrame,
+    top_k: Optional[int] = 50,
+    label_col: str = 'dist_group',
+) -> Tuple[torch.Tensor, torch.Tensor, List[str]]:
+    """
+    Для каждого признака SAE выбирает top-K патчей с наибольшими активациями.
+
+    Параметры
+    ----------
+    codes     : CSR матрица (n_patches, n_features)
+    meta      : DataFrame с колонкой label_col
+    top_k     : число патчей; None — все ненулевые (с паддингом до max nnz)
+    label_col : 'dist_group' или 'dist_type'
+
+    Возвращает
+    ----------
+    top_val     : FloatTensor (n_features, effective_k)
+    top_label   : LongTensor  (n_features, effective_k)
+    label_names : List[str] — расшифровка числовых меток
+    """
+    n_patches, n_features = codes.shape
+
+    raw_labels = meta[label_col].values
+    label_names, label_int = np.unique(raw_labels, return_inverse=True)
+    label_int = torch.from_numpy(label_int.astype(np.int64))
+
+    codes_csc = codes.tocsc()
+
+    if top_k is None:
+        nnz_per_feat = np.diff(codes_csc.indptr)
+        effective_k = int(nnz_per_feat.max()) if codes_csc.nnz > 0 else 1
+    else:
+        effective_k = min(top_k, n_patches)
+
+    top_val = torch.zeros(n_features, effective_k, dtype=torch.float32)
+    top_label = torch.zeros(n_features, effective_k, dtype=torch.long)
+
+    for feat_idx in range(n_features):
+        col = codes_csc.getcol(feat_idx)
+        col_data = col.data
+        col_rows = col.indices
+
+        if len(col_data) == 0:
+            continue
+
+        if top_k is None:
+            k = len(col_data)
+            order = np.argsort(col_data)[::-1]
+        else:
+            k = min(effective_k, len(col_data))
+            order = np.argpartition(col_data, -k)[-k:]
+            order = order[np.argsort(col_data[order])[::-1]]
+
+        top_val[feat_idx, :k] = torch.from_numpy(col_data[order[:k]].astype(np.float32))
+        top_label[feat_idx, :k] = label_int[col_rows[order[:k]]]
+
+    return top_val, top_label, list(label_names)
+
+
+def prepare_scatter_stats(
+    codes: sp.csr_matrix,
+    meta: pd.DataFrame,
+    top_k: Optional[int] = 50,
+    label_col: str = 'dist_group',
+    group_filter: Optional[str] = None,
+) -> Tuple[Dict[str, torch.Tensor], List[str]]:
+    """
+    Собирает словарь статистик, совместимый с get_stats_scatter_plot.
+
+    Параметры
+    ----------
+    codes        : CSR матрица (n_patches, n_features)
+    meta         : DataFrame с метаданными
+    top_k        : число top-патчей для расчёта энтропии
+    label_col    : 'dist_group' или 'dist_type'
+    group_filter : если задана — фильтрует строки по dist_group == group_filter
+
+    Возвращает
+    ----------
+    stats       : dict с ключами 'sparsity', 'mean_acts', 'top_entropy' — FloatTensors
+    label_names : List[str]
+    """
+    if group_filter is not None:
+        row_mask = (meta['dist_group'] == group_filter).values
+        codes = codes[row_mask]
+        meta = meta[row_mask].reset_index(drop=True)
+
+    codes_f32 = codes.astype(np.float32)
+    sparsity = torch.from_numpy(np.asarray((codes_f32 > 0).mean(axis=0)).ravel())
+    mean_acts = torch.from_numpy(np.asarray(codes_f32.mean(axis=0)).ravel())
+
+    top_val, top_label, label_names = get_top_k_patches(codes, meta, top_k=top_k, label_col=label_col)
+    top_entropy = calculate_entropy(top_val, top_label)
+
+    stats = {
+        'sparsity': sparsity.float(),
+        'mean_acts': mean_acts.float(),
+        'top_entropy': top_entropy.float(),
+    }
+    return stats, label_names
+
+
+def _metric_df_to_feature_vector(
+    metric_df: pd.DataFrame,
+    features: FeatureMatrix,
+    category_aggregation: str = 'max',
+) -> np.ndarray:
+    n_features = features.n_features
+    metric_values = np.full((n_features,), np.nan, dtype=np.float32)
+    if metric_df.empty:
+        return metric_values
+    category_aggregation = str(category_aggregation).strip().lower()
+    if category_aggregation not in {'max', 'mean'}:
+        raise ValueError(f"Unsupported category aggregation: {category_aggregation!r}. Allowed: ('max', 'mean')")
+    global_to_col = {int(gid): col for col, gid in enumerate(features.column_feature_ids)}
+    numeric_df = metric_df.apply(pd.to_numeric, errors='coerce')
+    for column in numeric_df.columns:
+        try:
+            global_id = int(column)
+        except (TypeError, ValueError):
+            continue
+        col = global_to_col.get(global_id)
+        if col is None:
+            continue
+        column_values = numeric_df[column].to_numpy(dtype=np.float32)
+        finite_values = column_values[np.isfinite(column_values)]
+        if finite_values.size == 0:
+            continue
+        if category_aggregation == 'max':
+            aggregated_value = float(np.max(finite_values))
+        else:
+            aggregated_value = float(np.mean(finite_values))
+        metric_values[col] = aggregated_value
+    return metric_values
+
+
+def _compute_color_metric_values(color_metric: str,
+                                 features: FeatureMatrix,
+                                 meta: pd.DataFrame,
+                                 label_col: str,
+                                 stats: Dict[str, torch.Tensor],
+                                 category_aggregation: str = 'max',
+                                 activation_threshold: float = 0.0,
+                                 cache_path: Optional[str] = None,
+                                 *,
+                                 dataset: str) -> torch.Tensor:
+    '''
+    Вычисляет вектор значений метрики для окрашивания точек на scatter plot.
+    Поддерживает 'entropy', 'roc_auc', 'iou', 'precision' и 'recall'.
+    Для табличных метрик возвращает вектор длины n_features,
+    агрегируя значения по категориям из label_col.
+    '''
+    color_metric = str(color_metric).strip().lower()
+    if color_metric == 'entropy':
+        return stats['top_entropy']
+    if color_metric == 'roc_auc':
+        auc_df = compute_distortion_roc_auc(
+            meta=meta,
+            features=features,
+            group_col=label_col,
+            level='patch',
+            cache_path=cache_path,
+        )
+        metric_values = _metric_df_to_feature_vector(
+            metric_df=auc_df,
+            features=features,
+            category_aggregation=category_aggregation,
+        )
+        return torch.from_numpy(metric_values).float()
+    if color_metric == 'iou':
+        spatial_shape = 1, meta.groupby('image_idx').size().max()
+        iou_df = compute_iou_per_feature(
+            features=features,
+            meta_df=meta,
+            spatial_shape=spatial_shape,
+            group_col=label_col,
+            cache_path=cache_path,
+            dataset=dataset,
+        )
+        metric_values = _metric_df_to_feature_vector(
+            metric_df=iou_df,
+            features=features,
+            category_aggregation=category_aggregation,
+        )
+        return torch.from_numpy(metric_values).float()
+    if color_metric == 'precision':
+        precision_df = compute_distortion_precision(
+            meta=meta,
+            features=features,
+            group_col=label_col,
+            level='patch',
+            activation_threshold=activation_threshold,
+            cache_path=cache_path,
+        )
+        metric_values = _metric_df_to_feature_vector(
+            metric_df=precision_df,
+            features=features,
+            category_aggregation=category_aggregation,
+        )
+        return torch.from_numpy(metric_values).float()
+    if color_metric == 'recall':
+        recall_df = compute_distortion_recall(
+            meta=meta,
+            features=features,
+            group_col=label_col,
+            level='patch',
+            activation_threshold=activation_threshold,
+            cache_path=cache_path,
+        )
+        metric_values = _metric_df_to_feature_vector(
+            metric_df=recall_df,
+            features=features,
+            category_aggregation=category_aggregation,
+        )
+        return torch.from_numpy(metric_values).float()
+    raise ValueError(
+        f"Unsupported scatter color metric: {color_metric!r}. Allowed: ('entropy', 'iou', 'roc_auc', 'precision', 'recall')"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Scatter plot
+# ---------------------------------------------------------------------------
+
+
+def get_stats_scatter_plot(stats: Dict[str, torch.Tensor], color_metric_values: Optional[torch.Tensor] = None, color_metric_label: str = 'entropy', mask: Optional[torch.Tensor] = None, backend: str = 'plotly', save_directory: Optional[str] = None, file_stem: str = 'scatter_plot', save_html: bool = False, show: bool = True, eps: float = 1e-9, title: Optional[str] = None) -> Dict[str, Optional[str]]:
+    if color_metric_values is None:
+        color_metric_values = stats['top_entropy']
+    if color_metric_values.shape != stats['sparsity'].shape:
+        raise ValueError('color_metric_values must match stats feature dimension: {} vs {}'.format(tuple(color_metric_values.shape), tuple(stats['sparsity'].shape)))
+    if mask is None:
+        mask = torch.ones_like(stats['sparsity'], dtype=torch.bool)
+    indices = torch.where(mask)[0]
+    plotting_data = torch.stack([torch.log10(stats['sparsity'][mask] + eps), torch.log10(stats['mean_acts'][mask] + eps), color_metric_values[mask], indices.float()], dim=0).T
+    x_label = 'log10(sparsity)'
+    y_label = 'log10(mean_acts)'
+    color_label = str(color_metric_label)
+    hover_label = 'index'
+    df = pd.DataFrame(plotting_data.numpy(), columns=[x_label, y_label, color_label, hover_label])
+    backend = str(backend).lower().strip()
+    if backend not in {'plotly', 'matplotlib'}:
+        raise ValueError(f"backend must be 'plotly' or 'matplotlib', got {backend!r}")
+    saved_png: Optional[str] = None
+    saved_html: Optional[str] = None
+    if backend == 'plotly':
+        fig = px.scatter(df, x=x_label, y=y_label, color=color_label, marginal_x='histogram', marginal_y='histogram', opacity=0.5, hover_data=[hover_label])
+        if title is not None:
+            fig.update_layout(title=title)
+        if save_directory is not None:
+            save_dir = Path(save_directory)
+            save_dir.mkdir(parents=True, exist_ok=True)
+            png_path = save_dir / f'{file_stem}.png'
+            try:
+                fig.write_image(png_path)
+                saved_png = str(png_path)
+            except Exception as exc:
+                print(f'[warn] Failed to save Plotly PNG to {png_path}: {exc}')
+                print('[warn] Install kaleido to enable Plotly PNG export.')
+            if save_html:
+                html_path = save_dir / f'{file_stem}.html'
+                fig.write_html(html_path, include_plotlyjs='cdn')
+                saved_html = str(html_path)
+        if show:
+            fig.show()
+    else:
+        fig, ax = plt.subplots(figsize=(10, 7))
+        scatter = ax.scatter(df[x_label], df[y_label], c=df[color_label], cmap='viridis', alpha=0.6, s=14, edgecolors='none')
+        ax.set_xlabel(x_label)
+        ax.set_ylabel(y_label)
+        ax.set_title(title or 'Scatter of SAE feature stats')
+        cbar = fig.colorbar(scatter, ax=ax)
+        cbar.set_label(color_label)
+        fig.tight_layout()
+        if save_directory is not None:
+            save_dir = Path(save_directory)
+            save_dir.mkdir(parents=True, exist_ok=True)
+            png_path = save_dir / f'{file_stem}.png'
+            fig.savefig(png_path, dpi=200, bbox_inches='tight')
+            saved_png = str(png_path)
+            if save_html:
+                print('[warn] save_html=True ignored for matplotlib backend.')
+        if show:
+            plt.show()
+        else:
+            plt.close(fig)
+    return {'png_path': saved_png, 'html_path': saved_html}
+
+
+def scatter_plot_by_distortion(features: FeatureMatrix,
+                               meta: pd.DataFrame,
+                               mode: str = 'group', 
+                               group_name: Optional[str] = None, 
+                               top_k: Optional[int] = 50, 
+                               color_metric: str = 'entropy', 
+                               per_category: bool = True, 
+                               activation_threshold: float = 0.0,
+                               backend: str = 'plotly', 
+                               save_directory: Optional[str] = None, 
+                               file_stem: Optional[str] = None, 
+                               save_html: bool = False, 
+                               show: bool = True,
+                               cache_path: Optional[str] = None,
+                               *,
+                               dataset: str,
+                               title: Optional[str] = None) -> Dict[str, Optional[str]]:
+    if mode == 'group':
+        label_col, group_filter = 'dist_group', None
+        title_suffix = 'by distortion group'
+    elif mode == 'type_in_group':
+        if group_name is None:
+            raise ValueError("mode='type_in_group' requires group_name")
+        label_col, group_filter = 'dist_type', group_name
+        title_suffix = f'by distortion type within "{group_name}"'
+    else:
+        raise ValueError(f"mode must be 'group' or 'type_in_group', got {mode!r}")
+
+    print(f'Preparing stats ({title_suffix})...')
+    codes = features.codes
+    stats, label_names = prepare_scatter_stats(codes, meta, top_k=top_k, label_col=label_col, group_filter=group_filter)
+
+    if group_filter is None:
+        features_slice = features
+        meta_slice = meta
+    else:
+        row_mask = (meta['dist_group'] == group_filter).values
+        features_slice = features.with_row_mask(row_mask)
+        meta_slice = meta[row_mask].reset_index(drop=True)
+
+    print(f'  Labels ({len(label_names)}): {label_names}')
+    print(f'  Features: {stats["sparsity"].shape[0]}')
+    print(f'  Color metric: {str(color_metric).lower()} (per_category={per_category})')
+
+    if file_stem is None:
+        if mode == 'type_in_group' and group_name is not None:
+            normalized_group = group_name.lower().replace(' ', '_')
+            file_stem = f'scatter_{mode}_{normalized_group}'
+        else:
+            file_stem = f'scatter_{mode}'
+
+    plot_title = title
+    if plot_title is None:
+        plot_title = f'Scatter of SAE feature stats ({title_suffix})'
+
+    metric_name = str(color_metric).strip().lower()
+
+    # Per-category plotting for table metrics
+    global_to_col = {
+        int(gid): col for col, gid in enumerate(features.column_feature_ids)
+    }
+
+    if per_category and metric_name in {'roc_auc', 'iou', 'precision', 'recall'}:
+        if metric_name == 'roc_auc':
+            metric_df = compute_distortion_roc_auc(
+                meta=meta_slice,
+                features=features_slice,
+                group_col=label_col,
+                level='patch',
+                cache_path=cache_path,
+            )
+        elif metric_name == 'iou':
+            spatial_shape = 1, meta_slice.groupby('image_idx').size().max()
+            metric_df = compute_iou_per_feature(
+                features=features_slice,
+                meta_df=meta_slice,
+                spatial_shape=spatial_shape,
+                group_col=label_col,
+                cache_path=cache_path,
+                dataset=dataset,
+            )
+        elif metric_name == 'precision':
+            metric_df = compute_distortion_precision(
+                meta=meta_slice,
+                features=features_slice,
+                group_col=label_col,
+                level='patch',
+                activation_threshold=activation_threshold,
+                cache_path=cache_path,
+            )
+        else:
+            metric_df = compute_distortion_recall(
+                meta=meta_slice,
+                features=features_slice,
+                group_col=label_col,
+                level='patch',
+                activation_threshold=activation_threshold,
+                cache_path=cache_path,
+            )
+
+        results = {'png_path': None, 'html_path': None}
+        n_cols = features.n_features
+        for category in metric_df.index.tolist():
+            row = metric_df.loc[category]
+            metric_values = np.full((n_cols,), np.nan, dtype=np.float32)
+            for col_name in metric_df.columns:
+                try:
+                    global_id = int(col_name)
+                except Exception:
+                    continue
+                matrix_col = global_to_col.get(global_id)
+                if matrix_col is None:
+                    continue
+                val = row[col_name]
+                metric_values[matrix_col] = np.nan if pd.isna(val) else float(val)
+
+            color_tensor = torch.from_numpy(metric_values).float()
+            category_label = str(category)
+            cat_file_stem = f"{file_stem}_{category_label}" if file_stem is not None else f"scatter_{metric_name}_{category_label}"
+
+            out = get_stats_scatter_plot(stats, color_metric_values=color_tensor, color_metric_label=f"{metric_name}:{category_label}", mask=None, backend=backend, save_directory=save_directory, file_stem=cat_file_stem, save_html=save_html, show=show, title=f'{plot_title} | {category_label}')
+            if out.get('png_path'):
+                results['png_path'] = out['png_path']
+            if out.get('html_path'):
+                results['html_path'] = out['html_path']
+
+        return results
+
+    # Default single-color-metric path
+    color_metric_values = _compute_color_metric_values(
+        color_metric=color_metric,
+        features=features_slice,
+        meta=meta_slice,
+        label_col=label_col,
+        stats=stats,
+        activation_threshold=activation_threshold,
+        cache_path=cache_path,
+        dataset=dataset,
+    )
+
+    category_label = 'all' if metric_name == 'entropy' else label_col
+    return get_stats_scatter_plot(stats, color_metric_values=color_metric_values, color_metric_label=str(color_metric).lower(), mask=None, backend=backend, save_directory=save_directory, file_stem=file_stem, save_html=save_html, show=show, title=f'{plot_title} | {category_label}')
+

assets/examples

@@ -0,0 +1 @@
+/data/examples

assets/style.css

@@ -0,0 +1,1137 @@
+:root {
+    --bg: #f6f8fb;
+    --bg-soft: rgba(245, 247, 250, 0.9);
+    --panel: #ffffff;
+    --panel-strong: #f8fafc;
+    --border: rgba(15, 23, 42, 0.06);
+    --text: #0b1220;
+    --text-soft: #4b5563;
+    --accent: #2563eb;
+    --accent-strong: #1e40af;
+    --shadow: 0 18px 40px rgba(15, 23, 42, 0.06);
+}
+
+html,
+body {
+    margin: 0;
+    min-height: 100%;
+    background: linear-gradient(180deg, var(--bg) 0%, #ffffff 100%);
+    color: var(--text);
+    font-family: "Inter", "DM Sans", "Segoe UI", sans-serif;
+}
+
+.app-shell {
+    position: relative;
+    max-width: 1480px;
+    margin: 0 auto;
+    padding: 40px 24px 56px;
+}
+
+.page-back-nav {
+    position: absolute;
+    top: 12px;
+    left: 24px;
+    z-index: 30;
+    margin: 0;
+}
+
+.back-link {
+    display: inline-flex;
+    align-items: center;
+    padding: 8px 14px;
+    font-size: 0.9rem;
+    font-weight: 500;
+    color: var(--text-soft);
+    text-decoration: none;
+    background: var(--panel);
+    border: 1px solid var(--border);
+    border-radius: 999px;
+    box-shadow: 0 4px 12px rgba(15, 23, 42, 0.06);
+    transition: color 0.15s ease, border-color 0.15s ease, box-shadow 0.15s ease;
+}
+
+.back-link:hover {
+    color: var(--accent);
+    border-color: rgba(37, 99, 235, 0.28);
+    box-shadow: 0 6px 16px rgba(37, 99, 235, 0.1);
+}
+
+.hero-card,
+.form-card,
+.mock-card,
+.placeholder-card {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    box-shadow: var(--shadow);
+    border-radius: 14px;
+    white-space: pre-wrap;
+}
+
+.hero-card {
+    padding: 28px 30px 24px;
+    margin-bottom: 22px;
+}
+
+.page-back-nav + .hero-card {
+    padding-top: 48px;
+}
+
+.app-kicker {
+    font-size: 0.78rem;
+    text-transform: uppercase;
+    letter-spacing: 0.14em;
+    color: var(--accent);
+    margin-bottom: 10px;
+}
+
+.app-title {
+    margin: 0;
+    font-size: clamp(2rem, 4vw, 3.45rem);
+    line-height: 0.98;
+    letter-spacing: -0.04em;
+}
+
+.home-grid,
+.mock-grid {
+    display: grid;
+    gap: 18px;
+}
+
+.home-grid {
+    grid-template-columns: minmax(0, 1.7fr) minmax(280px, 0.9fr);
+    align-items: start;
+}
+
+.mock-grid {
+    grid-template-columns: minmax(260px, 0.7fr) minmax(0, 1.3fr);
+    align-items: start;
+}
+
+.form-card,
+.info-card,
+.mock-card,
+.placeholder-card {
+    padding: 24px;
+}
+
+.form-block {
+    margin-bottom: 18px;
+}
+
+.section-label,
+.panel-label,
+.preview-label {
+    display: block;
+    margin-bottom: 10px;
+    color: var(--text-soft);
+    font-size: 0.78rem;
+    text-transform: uppercase;
+    letter-spacing: 0.16em;
+}
+
+.metric-switch {
+    display: flex;
+    gap: 10px;
+    flex-wrap: wrap;
+}
+
+.metric-pill {
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    min-width: 110px;
+    padding: 10px 14px;
+    border-radius: 999px;
+    border: 1px solid rgba(37, 99, 235, 0.12);
+    background: rgba(37, 99, 235, 0.06);
+    color: var(--text);
+    cursor: pointer;
+    transition: transform 120ms ease, border-color 120ms ease, color 120ms ease, background 120ms ease;
+}
+
+.metric-pill:hover {
+    transform: translateY(-1px);
+    border-color: rgba(125, 211, 252, 0.36);
+}
+
+.metric-descriptions {
+    margin-top: 14px;
+    display: flex;
+    flex-direction: column;
+    gap: 10px;
+}
+
+.metric-desc-block {
+    margin-top: 2px;
+}
+
+.metric-desc-name {
+    font-weight: 600;
+    color: var(--text);
+    letter-spacing: 0.02em;
+    font-size: 0.9rem;
+    margin-bottom: 6px;
+}
+
+.metric-desc-text {
+    margin: 0 0 8px;
+    font-size: 0.88rem;
+    line-height: 1.45;
+    color: var(--text-soft);
+}
+
+.metric-desc-text:last-child {
+    margin-bottom: 0;
+}
+
+.dash-dropdown .Select-control,
+.dash-dropdown .Select-menu-outer,
+.dash-dropdown .VirtualizedSelectOption,
+.dash-dropdown .Select-value-label,
+.dash-dropdown .Select-placeholder {
+    background-color: transparent !important;
+    color: var(--text) !important;
+    border-color: var(--border) !important;
+}
+
+.dash-dropdown .Select-control {
+    min-height: 58px;
+    border-radius: 16px !important;
+}
+
+.dash-dropdown .Select-input>input {
+    color: var(--text) !important;
+}
+
+.model-dropdown .Select-menu-outer {
+    padding: 8px 0;
+}
+
+/* Model option label classes (moved from inline styles in Python) */
+.dash-dropdown .model-option {
+    padding: 0.15rem 0;
+}
+
+.dash-dropdown .model-option-exp {
+    font-size: 1.05rem;
+    font-weight: 300;
+    line-height: 1.1;
+    color: var(--text) !important;
+    letter-spacing: 0.01em;
+}
+
+.dash-dropdown .model-option-meta {
+    font-size: 0.78rem;
+    color: var(--text-soft) !important;
+    margin-top: 0.2rem;
+    line-height: 1.2;
+}
+
+.selection-preview {
+    margin: 16px 0 18px;
+    padding: 14px 16px;
+    border-radius: 12px;
+    background: var(--panel-strong);
+    border: 1px solid rgba(15, 23, 42, 0.04);
+}
+
+.preview-value {
+    white-space: pre-wrap;
+}
+
+.preview-value,
+.home-status,
+.mock-summary-meta,
+.mock-placeholder {
+    color: var(--text-soft);
+}
+
+.preview-value {
+    font-size: 0.95rem;
+    line-height: 1.45;
+}
+
+.action-row {
+    display: flex;
+    align-items: center;
+    gap: 14px;
+}
+
+.confirm-button {
+    border: 0;
+    border-radius: 999px;
+    padding: 12px 20px;
+    background: linear-gradient(135deg, var(--accent), var(--accent-strong));
+    color: #ffffff;
+    font-weight: 700;
+    cursor: pointer;
+    letter-spacing: 0.02em;
+    box-shadow: 0 8px 20px rgba(37, 99, 235, 0.12);
+    text-decoration: none;
+}
+
+.confirm-button:hover,
+.confirm-button:visited,
+.confirm-button:focus {
+    color: #ffffff;
+    text-decoration: none;
+}
+
+.confirm-button:hover {
+    transform: translateY(-1px);
+}
+
+.confirm-button:disabled {
+    opacity: 0.65;
+    cursor: wait;
+    transform: none;
+}
+
+/* Upload area styled like a button */
+.upload-card {
+    position: relative;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    gap: 10px;
+    padding: 18px 18px 14px;
+    border-radius: 12px;
+    border: 1px dashed var(--border);
+    background: var(--panel-strong);
+    cursor: pointer;
+    transition: transform 120ms ease, border-color 120ms ease, background 120ms ease;
+}
+
+.upload-card:hover {
+    transform: translateY(-2px);
+    border-color: rgba(37, 99, 235, 0.2);
+    background: rgba(37, 99, 235, 0.03);
+}
+
+.upload-card a {
+    display: inline-block;
+    padding: 8px 14px;
+    border-radius: 999px;
+    background: linear-gradient(135deg, var(--accent), var(--accent-strong));
+    color: #ffffff;
+    font-weight: 700;
+    text-decoration: none;
+}
+
+.upload-card a:hover {
+    filter: brightness(0.95);
+}
+
+.upload-help {
+    position: absolute;
+    top: 10px;
+    right: 10px;
+    padding: 3px 8px;
+    border-radius: 999px;
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    background: rgba(37, 99, 235, 0.12);
+    color: var(--accent-strong);
+    font-size: 0.68rem;
+    font-weight: 800;
+    line-height: 1;
+    border: 1px solid rgba(37, 99, 235, 0.2);
+    cursor: help;
+    user-select: none;
+    white-space: nowrap;
+}
+
+.upload-help:hover::after {
+    content: attr(data-tooltip);
+    position: absolute;
+    top: calc(100% + 8px);
+    right: 0;
+    z-index: 2;
+    min-width: 220px;
+    max-width: 280px;
+    padding: 10px 12px;
+    border-radius: 12px;
+    background: rgba(15, 23, 42, 0.95);
+    color: #ffffff;
+    font-size: 0.82rem;
+    font-weight: 600;
+    line-height: 1.35;
+    box-shadow: 0 12px 28px rgba(15, 23, 42, 0.2);
+    white-space: normal;
+    text-align: left;
+    pointer-events: none;
+}
+
+.upload-help:hover::before {
+    content: "";
+    position: absolute;
+    top: calc(100% + 2px);
+    right: 14px;
+    border: 6px solid transparent;
+    border-bottom-color: rgba(15, 23, 42, 0.95);
+    pointer-events: none;
+}
+
+/* Ensure the feature overlay image is left-aligned rather than centered */
+#fm-overlay-img {
+    display: block;
+    margin: 0 auto;
+    /* center horizontally */
+    width: auto;
+    max-width: 720px;
+}
+
+.home-nav-loading {
+    min-height: 52px;
+}
+
+.home-status.is-loading::before {
+    content: "Loading";
+    display: inline-block;
+    margin-right: 10px;
+    padding: 4px 10px;
+    border-radius: 999px;
+    background: rgba(37, 99, 235, 0.12);
+    color: var(--accent-strong);
+    font-size: 0.76rem;
+    font-weight: 700;
+    letter-spacing: 0.04em;
+    text-transform: uppercase;
+}
+
+.page-transition-loading {
+    position: relative;
+    min-height: 55vh;
+}
+
+.page-transition-loading>._dash-loading-callback {
+    display: flex !important;
+    align-items: center;
+    justify-content: center;
+    min-height: 55vh;
+}
+
+.skeleton-line {
+    height: 0.9rem;
+    margin-bottom: 8px;
+    border-radius: 6px;
+    background: rgba(15, 23, 42, 0.07);
+}
+
+.skeleton-line-short {
+    width: 55%;
+}
+
+.viz-model-info-loading {
+    min-height: 88px;
+}
+
+.umap-graph-loading {
+    flex: 1 1 auto;
+    min-height: 430px;
+    overflow: hidden;
+}
+
+.feature-detail-body-loading {
+    min-height: 140px;
+}
+
+.info-list {
+    margin: 0;
+    padding-left: 18px;
+    color: var(--text-soft);
+    line-height: 1.6;
+}
+
+.mock-summary-lines {
+    display: grid;
+    gap: 8px;
+    margin: 10px 0 14px;
+    color: var(--text);
+    line-height: 1.5;
+}
+
+.mock-summary-hyperparams {
+    font-size: 0.96rem;
+    line-height: 1.55;
+    color: var(--text-soft);
+    white-space: pre-wrap;
+}
+
+.viz-summary-divider {
+    margin: 14px 0 12px;
+    border: 0;
+    border-top: 1px solid var(--border);
+}
+
+.viz-sae-card,
+.viz-umap-card {
+    position: relative;
+}
+
+.viz-sae-help-wrap,
+.viz-umap-help-wrap {
+    position: absolute;
+    top: 22px;
+    right: 22px;
+    z-index: 3;
+}
+
+.viz-sae-help-wrap .viz-sae-help-trigger {
+    position: relative;
+    top: auto;
+    right: auto;
+}
+
+.viz-sae-help-popover {
+    display: none;
+    position: absolute;
+    top: calc(100% + 8px);
+    right: 0;
+    z-index: 4;
+    min-width: 260px;
+    max-width: 340px;
+    padding: 12px 14px;
+    border-radius: 12px;
+    background: rgba(15, 23, 42, 0.95);
+    color: #ffffff;
+    font-size: 0.86rem;
+    line-height: 1.45;
+    box-shadow: 0 12px 28px rgba(15, 23, 42, 0.2);
+}
+
+.viz-sae-help-wrap:hover .viz-sae-help-popover,
+.viz-sae-help-wrap:focus-within .viz-sae-help-popover {
+    display: block;
+}
+
+.viz-sae-help-popover::before {
+    content: "";
+    position: absolute;
+    top: -6px;
+    right: 14px;
+    border: 6px solid transparent;
+    border-bottom-color: rgba(15, 23, 42, 0.95);
+}
+
+.viz-help-popover-title {
+    margin-bottom: 8px;
+    font-size: 0.8rem;
+    font-weight: 700;
+    letter-spacing: 0.02em;
+}
+
+.viz-help-popover-lead {
+    margin-bottom: 8px;
+    font-size: 0.8rem;
+    font-weight: 500;
+    color: rgba(255, 255, 255, 0.9);
+}
+
+.viz-help-popover-subtitle {
+    margin: 10px 0 6px;
+    font-size: 0.76rem;
+    font-weight: 700;
+    letter-spacing: 0.04em;
+    text-transform: uppercase;
+    color: rgba(255, 255, 255, 0.72);
+}
+
+.viz-umap-help-popover {
+    min-width: 280px;
+    max-width: 400px;
+}
+
+.viz-sae-help-popover .help-desc-line {
+    margin-bottom: 6px;
+}
+
+.viz-sae-help-popover .help-desc-line:last-child {
+    margin-bottom: 0;
+}
+
+.viz-sae-help-popover strong,
+.viz-sae-help-popover .help-desc-line strong {
+    font-weight: 700;
+    color: #ffffff;
+}
+
+.viz-sae-help-popover .help-desc {
+    font-weight: 400;
+    color: rgba(255, 255, 255, 0.82);
+}
+
+.help-desc {
+    font-weight: 400;
+    color: var(--text-soft);
+}
+
+.viz-filter-intro strong,
+.viz-filter-intro-list strong,
+.viz-filter-line strong {
+    font-weight: 600;
+    color: var(--text);
+}
+
+.viz-filter-intro {
+    margin-bottom: 12px;
+    font-size: 0.98rem;
+    line-height: 1.55;
+    color: var(--text-soft);
+}
+
+.viz-filter-intro-lead {
+    margin: 0 0 8px;
+}
+
+.viz-filter-intro-list {
+    margin: 0 0 10px;
+    padding-left: 1.15rem;
+}
+
+.viz-filter-intro-list li {
+    margin-bottom: 6px;
+}
+
+.viz-filter-summary {
+    display: grid;
+    gap: 8px;
+    font-size: 0.98rem;
+    line-height: 1.55;
+}
+
+.viz-filter-line {
+    margin: 0;
+    color: var(--text);
+}
+
+.mock-placeholder {
+    margin-top: 14px;
+    padding: 14px;
+    border-radius: 12px;
+    border: 1px dashed rgba(15, 23, 42, 0.06);
+    background: var(--panel-strong);
+    min-height: 500px;
+    display: flex;
+    flex-direction: column;
+    gap: 8px;
+}
+
+.feature-panel {
+    display: grid;
+    gap: 14px;
+    margin-top: 18px;
+    padding: 18px;
+    border-radius: 18px;
+    background: linear-gradient(180deg, rgba(255, 255, 255, 0.94), rgba(248, 250, 252, 0.98));
+    border: 1px solid rgba(37, 99, 235, 0.08);
+    box-shadow: inset 0 1px 0 rgba(255, 255, 255, 0.7);
+}
+
+.feature-workbench {
+    display: grid;
+    gap: 14px;
+    margin-top: 18px;
+    width: 100%;
+}
+
+.feature-workbench-title-row {
+    position: relative;
+    padding-right: 40px;
+}
+
+.feature-workbench-title-row .panel-label {
+    margin-bottom: 0;
+}
+
+.viz-feature-workbench-help-wrap {
+    top: 0;
+    right: 0;
+}
+
+.viz-feature-workbench-help-popover {
+    min-width: 280px;
+    max-width: 400px;
+}
+
+.feature-workbench-header {
+    display: grid;
+    gap: 10px;
+}
+
+.feature-workbench-toolbar {
+    display: grid;
+    grid-template-columns: minmax(220px, 1.35fr) minmax(280px, 1.15fr) minmax(170px, 0.75fr);
+    gap: 14px;
+    align-items: end;
+}
+
+.feature-workbench-controls {
+    display: grid;
+    gap: 10px;
+    padding: 14px 16px;
+    border-radius: 18px;
+    background: linear-gradient(180deg, rgba(255, 255, 255, 0.96), rgba(248, 250, 252, 0.95));
+    border: 1px solid rgba(37, 99, 235, 0.08);
+    box-shadow: var(--shadow);
+    min-width: 0;
+}
+
+.feature-workbench-controls .section-label {
+    margin-bottom: 8px;
+}
+
+.feature-workbench-controls-sort .dash-dropdown,
+.feature-workbench-controls-ranking .dash-dropdown {
+    width: 100%;
+}
+
+.feature-nav-row {
+    display: grid;
+    grid-template-columns: auto minmax(180px, 1fr) auto;
+    gap: 10px;
+    align-items: center;
+}
+
+.nav-button {
+    border: 1px solid rgba(37, 99, 235, 0.16);
+    border-radius: 999px;
+    padding: 10px 16px;
+    background: #ffffff;
+    color: var(--accent-strong);
+    font-weight: 700;
+    cursor: pointer;
+    transition: transform 120ms ease, border-color 120ms ease, background 120ms ease;
+}
+
+.nav-button:hover {
+    transform: translateY(-1px);
+    border-color: rgba(37, 99, 235, 0.34);
+    background: rgba(37, 99, 235, 0.04);
+}
+
+.feature-jump-input .Select-control,
+.feature-jump-input input {
+    min-height: 42px;
+}
+
+/* Hide +/- steppers in feature id number input (browser-native spinners). */
+.feature-jump-input input[type="number"] {
+    -moz-appearance: textfield;
+    appearance: textfield;
+}
+
+.feature-jump-input input[type="number"]::-webkit-outer-spin-button,
+.feature-jump-input input[type="number"]::-webkit-inner-spin-button {
+    -webkit-appearance: none;
+    margin: 0;
+}
+
+.baseline-grid {
+    display: grid;
+    /* allow more columns so the baseline box appears more horizontal */
+    grid-template-columns: repeat(auto-fill, minmax(120px, 1fr));
+    gap: 12px;
+    margin-top: 6px;
+}
+
+/* Page-specific helper to vertically stack the input card above overlay controls */
+.mock-grid-vertical {
+    grid-template-columns: 1fr;
+}
+
+.baseline-thumb {
+    border: 1px solid rgba(15, 23, 42, 0.08);
+    background: rgba(248, 250, 252, 0.8);
+    border-radius: 14px;
+    padding: 10px;
+    cursor: pointer;
+    text-align: left;
+    transition: transform 120ms ease, border-color 120ms ease, background 120ms ease;
+}
+
+.baseline-thumb:hover {
+    transform: translateY(-1px);
+    border-color: rgba(37, 99, 235, 0.25);
+    background: rgba(37, 99, 235, 0.06);
+}
+
+.baseline-thumb-img {
+    width: 100%;
+    height: 120px;
+    object-fit: cover;
+    border-radius: 10px;
+    border: 1px solid rgba(15, 23, 42, 0.08);
+    background: #ffffff;
+    display: block;
+}
+
+.baseline-thumb-meta {
+    margin-top: 8px;
+    display: grid;
+    gap: 4px;
+}
+
+.baseline-thumb-label {
+    font-size: 0.82rem;
+    color: var(--text);
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+
+.baseline-thumb-badge {
+    font-size: 0.72rem;
+    color: var(--text-soft);
+}
+
+.feature-status {
+    min-height: 20px;
+    color: var(--text-soft);
+    font-size: 0.92rem;
+}
+
+.feature-detail-card {
+    display: grid;
+    gap: 14px;
+    min-height: 420px;
+    padding: 20px;
+    border-radius: 18px;
+    background: linear-gradient(135deg, rgba(37, 99, 235, 0.06), rgba(255, 255, 255, 0.98) 42%);
+    border: 1px solid rgba(37, 99, 235, 0.1);
+}
+
+.feature-detail-hero {
+    display: grid;
+    gap: 6px;
+}
+
+.feature-detail-hero-header {
+    display: flex;
+    align-items: flex-start;
+    justify-content: space-between;
+    gap: 16px;
+}
+
+.feature-image-ranking-dropdown {
+    width: 100%;
+}
+
+.feature-detail-kicker,
+.feature-placeholder-title {
+    color: var(--text-soft);
+    text-transform: uppercase;
+    letter-spacing: 0.14em;
+    font-size: 0.72rem;
+}
+
+.feature-detail-number {
+    font-size: clamp(2.5rem, 6vw, 4.8rem);
+    line-height: 0.95;
+    letter-spacing: -0.05em;
+    font-weight: 800;
+}
+
+.feature-detail-subtitle {
+    color: var(--text-soft);
+    font-size: 0.98rem;
+}
+
+.feature-placeholder {
+    display: grid;
+    gap: 8px;
+    min-height: 124px;
+    padding: 16px;
+    border-radius: 16px;
+    border: 1px dashed rgba(37, 99, 235, 0.18);
+    background: rgba(255, 255, 255, 0.7);
+}
+
+.feature-placeholder-copy,
+.feature-empty {
+    color: var(--text-soft);
+    line-height: 1.5;
+}
+
+.feature-comparison-block {
+    display: grid;
+    gap: 12px;
+}
+
+.feature-comparison-title {
+    color: var(--text-soft);
+    text-transform: uppercase;
+    letter-spacing: 0.14em;
+    font-size: 0.72rem;
+}
+
+.feature-comparison-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(170px, 1fr));
+    gap: 12px;
+}
+
+/* Top images grid */
+.feature-top-images-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(160px, 1fr));
+    gap: 12px;
+    align-items: start;
+}
+
+.feature-image-card {
+    position: relative;
+    border-radius: 12px;
+    overflow: hidden;
+    background: #fff;
+    border: 1px solid rgba(15, 23, 42, 0.06);
+}
+
+.feature-image-card img {
+    width: 100%;
+    height: auto;
+    display: block;
+}
+
+.feature-image-caption {
+    padding: 6px 8px;
+    font-size: 0.82rem;
+    color: var(--text-soft);
+}
+
+.feature-image-caption:empty {
+    display: none;
+}
+
+.comparison-card {
+    display: grid;
+    gap: 8px;
+    min-height: 110px;
+    padding: 14px 16px;
+    border-radius: 16px;
+    border: 1px solid rgba(15, 23, 42, 0.07);
+    background: rgba(255, 255, 255, 0.88);
+}
+
+.comparison-card-active {
+    border-color: rgba(37, 99, 235, 0.3);
+    background: linear-gradient(135deg, rgba(37, 99, 235, 0.09), rgba(255, 255, 255, 0.96));
+}
+
+.comparison-card-header {
+    display: flex;
+    align-items: center;
+    justify-content: space-between;
+    gap: 10px;
+}
+
+.comparison-selector-name {
+    font-size: 0.92rem;
+    font-weight: 700;
+    color: var(--text);
+    overflow: hidden;
+    text-overflow: ellipsis;
+}
+
+.comparison-score {
+    font-size: 1.25rem;
+    font-weight: 800;
+    letter-spacing: -0.03em;
+}
+
+.comparison-rank {
+    color: var(--text-soft);
+    font-size: 0.88rem;
+}
+
+.umap-controls {
+    display: grid;
+    gap: 14px;
+    margin-bottom: 12px;
+}
+
+.umap-mode-primary {
+    display: grid;
+    gap: 10px;
+    padding: 14px 16px;
+    border: 1px solid var(--border);
+    border-radius: 12px;
+    background: var(--panel-strong);
+}
+
+.umap-mode-label {
+    font-size: 0.82rem;
+    text-transform: uppercase;
+    letter-spacing: 0.14em;
+    color: var(--accent-strong);
+    font-weight: 700;
+}
+
+.umap-mode-switch {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 10px;
+}
+
+.umap-mode-pill {
+    display: inline-flex;
+    align-items: center;
+    gap: 8px;
+    padding: 10px 16px;
+    border: 1px solid var(--border);
+    border-radius: 999px;
+    background: #fff;
+    cursor: pointer;
+    font-size: 0.95rem;
+    font-weight: 600;
+    transition: border-color 0.15s ease, background 0.15s ease, color 0.15s ease;
+}
+
+.umap-mode-input {
+    accent-color: var(--accent);
+}
+
+.umap-mode-pill:has(.umap-mode-input:checked) {
+    border-color: rgba(37, 99, 235, 0.35);
+    background: rgba(37, 99, 235, 0.08);
+    color: var(--accent-strong);
+}
+
+.umap-controls-secondary {
+    padding-left: 2px;
+}
+
+.umap-controls-row {
+    display: flex;
+    flex-wrap: wrap;
+    align-items: flex-end;
+    gap: 16px 20px;
+}
+
+.umap-control-block {
+    display: grid;
+    gap: 6px;
+    min-width: 180px;
+    flex: 1 1 180px;
+}
+
+.umap-features-color-fixed {
+    padding: 10px 12px;
+    border: 1px solid var(--border);
+    border-radius: 10px;
+    background: var(--panel-strong);
+    font-size: 0.92rem;
+    font-weight: 600;
+    color: var(--accent-strong);
+}
+
+.umap-status {
+    margin-bottom: 10px;
+    font-size: 0.82rem;
+    color: var(--text-soft);
+    line-height: 1.4;
+}
+
+.umap-status.is-loading::before {
+    content: "Updating UMAP…";
+    display: inline-block;
+    margin-right: 10px;
+    padding: 4px 10px;
+    border-radius: 999px;
+    background: rgba(37, 99, 235, 0.12);
+    color: var(--accent-strong);
+    font-size: 0.76rem;
+    font-weight: 700;
+    letter-spacing: 0.04em;
+    text-transform: uppercase;
+}
+
+.umap-hover-shell {
+    position: relative;
+    display: flex;
+    flex-direction: column;
+    flex: 1 1 auto;
+    min-height: 0;
+    overflow: hidden;
+}
+
+.loading-block {
+    width: 100%;
+}
+
+.feature-detail-loading {
+    min-height: 220px;
+}
+
+.feature-images-loading {
+    min-height: 180px;
+}
+
+#feature-top-images[data-dash-is-loading="true"]::before {
+    content: "Loading...";
+    display: inline-block;
+    margin-bottom: 10px;
+    padding: 4px 10px;
+    border-radius: 999px;
+    background: rgba(37, 99, 235, 0.12);
+    color: var(--accent-strong);
+    font-size: 0.76rem;
+    font-weight: 700;
+    letter-spacing: 0.04em;
+    text-transform: uppercase;
+}
+
+#feature-top-images[data-dash-is-loading="true"]::before {
+    content: "Loading top images...";
+}
+
+.umap-graph {
+    flex: 1 1 auto;
+    min-height: 430px;
+    overflow: hidden;
+}
+
+.umap-graph .js-plotly-plot,
+.umap-graph .plot-container,
+.umap-graph .svg-container {
+    height: 100% !important;
+}
+
+@media (max-width: 980px) {
+
+    .home-grid,
+    .mock-grid {
+        grid-template-columns: 1fr;
+    }
+
+    .app-shell {
+        padding: 24px 14px 42px;
+    }
+
+    .page-back-nav {
+        top: 8px;
+        left: 14px;
+    }
+
+    .feature-stat-grid {
+        grid-template-columns: 1fr 1fr;
+    }
+
+    .feature-nav-row {
+        grid-template-columns: 1fr;
+    }
+
+    .feature-workbench-toolbar {
+        grid-template-columns: 1fr;
+        align-items: stretch;
+    }
+
+    .mock-placeholder {
+        min-height: 420px;
+    }
+
+    .umap-graph {
+        min-height: 360px;
+    }
+}