analysis/features/feature_matrix.py

from __future__ import annotations

from typing import List, Sequence

import numpy as np
import scipy.sparse as sp

from analysis.features.feature_indexing import FeatureMatrix


def _axis0_to_dense(values) -> np.ndarray:
    if hasattr(values, 'toarray'):
        return np.asarray(values.toarray()).ravel().astype(np.float32)
    return np.asarray(values).ravel().astype(np.float32)


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 only).
    """
    return _zscore_columns(
        build_image_feature_matrix_raw(
            features,
            image_row_indices,
            n_images_used,
            aggregation_mode,
        )
    )


def _to_dense_image_features(values) -> np.ndarray:
    if sp.issparse(values):
        return np.asarray(values.toarray(), dtype=np.float32)
    return np.asarray(values, dtype=np.float32)


def _image_group_matrix(image_idx_arr: np.ndarray, n_images_used: int) -> sp.csr_matrix:
    """Map each patch row to one image column (n_patches, n_images)."""
    n_patches = int(len(image_idx_arr))
    rows = np.arange(n_patches, dtype=np.int32)
    cols = np.asarray(image_idx_arr, dtype=np.int32)
    return sp.csr_matrix(
        (np.ones(n_patches, dtype=np.float32), (rows, cols)),
        shape=(n_patches, int(n_images_used)),
    )


def _grouped_sum(codes, group: sp.csr_matrix) -> np.ndarray:
    return _to_dense_image_features(group.T @ codes)


def _grouped_mean_acts(codes, group: sp.csr_matrix) -> np.ndarray:
    if sp.issparse(codes):
        positive = codes.multiply(codes > 0)
        pos_count = group.T @ (codes > 0).astype(np.float32)
    else:
        codes_arr = np.asarray(codes, dtype=np.float32)
        positive = codes_arr * (codes_arr > 0)
        pos_count = group.T @ (codes_arr > 0).astype(np.float32)

    pos_sum = _to_dense_image_features(group.T @ positive)
    pos_count = _to_dense_image_features(pos_count)
    out = np.zeros_like(pos_sum, dtype=np.float32)
    nonzero = pos_count > 0
    out[nonzero] = pos_sum[nonzero] / pos_count[nonzero]
    return out


def _grouped_max(codes, image_idx_arr: np.ndarray, n_images_used: int, n_features_total: int) -> np.ndarray:
    x = np.zeros((int(n_images_used), int(n_features_total)), dtype=np.float32)
    for img_i in range(int(n_images_used)):
        patch_mask = image_idx_arr == img_i
        if not np.any(patch_mask):
            continue
        chunk = codes[patch_mask, :]
        x[img_i] = _axis0_to_dense(chunk.max(axis=0))
    return x


def _aggregate_codes_by_image_idx(
    codes,
    image_idx_arr: np.ndarray,
    n_images_used: int,
    aggregation_mode: str,
) -> np.ndarray:
    """Aggregate patch-level codes per image using dense image_idx labels."""
    mode = str(aggregation_mode)
    if mode not in {'max', 'sum', 'mean_acts'}:
        raise ValueError(f'Unknown aggregation mode: {mode!r}')

    image_idx_arr = np.asarray(image_idx_arr, dtype=np.int32)
    n_images_used = int(n_images_used)
    n_features_total = int(codes.shape[1])

    if mode == 'max':
        return _grouped_max(codes, image_idx_arr, n_images_used, n_features_total)

    group = _image_group_matrix(image_idx_arr, n_images_used)
    if mode == 'sum':
        return _grouped_sum(codes, group)
    return _grouped_mean_acts(codes, group)


def build_image_feature_matrix_from_image_idx(
    codes,
    image_idx_arr: np.ndarray,
    n_images_used: int,
    aggregation_mode: str = 'max',
) -> np.ndarray:
    """Build z-scored image-level matrix using per-patch ``image_idx`` labels."""
    return _zscore_columns(
        _aggregate_codes_by_image_idx(
            codes,
            image_idx_arr,
            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'}:
        raise ValueError(f'Unknown aggregation mode: {mode!r}')

    codes_subset = features.codes
    image_idx_arr = np.empty(int(codes_subset.shape[0]), dtype=np.int32)
    for img_i, row_ids in enumerate(image_row_indices):
        if len(row_ids) == 0:
            continue
        image_idx_arr[list(row_ids)] = int(img_i)

    return _aggregate_codes_by_image_idx(
        codes_subset,
        image_idx_arr,
        n_images_used,
        mode,
    )