File size: 5,437 Bytes
45eaa9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """Spectral feature maps for Hilbert-space learning.
The paper emphasizes Fourier, wavelet, scattering, and learnable spectral
representations. This module implements compact CPU-friendly versions for the
8x8 Digits benchmark.
"""
from __future__ import annotations
import numpy as np
import torch
def _as_images(X: np.ndarray) -> np.ndarray:
"""Return images as (n, 8, 8) float32 in [0, 1]."""
X = np.asarray(X, dtype=np.float32)
if X.ndim == 2 and X.shape[1] == 64:
X = X.reshape(-1, 8, 8)
elif X.ndim != 3:
raise ValueError(f"Expected (n,64) or (n,8,8), got {X.shape}")
# sklearn digits are in [0,16]. If already normalized this is harmless.
if X.min() >= 0.0 and X.max() > 1.0:
X = X / 16.0
return X.astype(np.float32)
def pixel_features(X: np.ndarray) -> np.ndarray:
"""Identity Hilbert embedding over pixels."""
imgs = _as_images(X)
return imgs.reshape(imgs.shape[0], -1).astype(np.float32)
def fourier_features(X: np.ndarray) -> np.ndarray:
"""Real-valued 2D Fourier feature map.
We concatenate real and imaginary FFT coefficients. This is a spectral
embedding in a global orthonormal basis.
"""
imgs = _as_images(X)
fft = np.fft.fft2(imgs, axes=(-2, -1), norm="ortho")
feat = np.concatenate([fft.real.reshape(len(imgs), -1), fft.imag.reshape(len(imgs), -1)], axis=1)
return feat.astype(np.float32)
def haar2d_single(img: np.ndarray, levels: int = 3) -> np.ndarray:
"""Orthonormal 2D Haar wavelet transform for square power-of-two images.
Coefficients are packed in-place using the standard quadrant layout.
For 8x8 images, levels=3 yields a full multiresolution transform.
"""
x = np.asarray(img, dtype=np.float32).copy()
h, w = x.shape
if h != w or (h & (h - 1)) != 0:
raise ValueError("haar2d_single expects square power-of-two images")
size = h
for _ in range(levels):
if size < 2:
break
block = x[:size, :size]
a = block[0::2, 0::2]
b = block[0::2, 1::2]
c = block[1::2, 0::2]
d = block[1::2, 1::2]
ll = (a + b + c + d) / 2.0
lh = (a - b + c - d) / 2.0
hl = (a + b - c - d) / 2.0
hh = (a - b - c + d) / 2.0
half = size // 2
x[:half, :half] = ll
x[:half, half:size] = lh
x[half:size, :half] = hl
x[half:size, half:size] = hh
size = half
return x
def inverse_haar2d_single(coeff: np.ndarray, levels: int = 3) -> np.ndarray:
"""Inverse of haar2d_single."""
x = np.asarray(coeff, dtype=np.float32).copy()
h, w = x.shape
max_levels = int(np.log2(h))
levels = min(levels, max_levels)
sizes = [2 ** i for i in range(1, levels + 1)]
for size in sizes:
half = size // 2
ll = x[:half, :half].copy()
lh = x[:half, half:size].copy()
hl = x[half:size, :half].copy()
hh = x[half:size, half:size].copy()
a = (ll + lh + hl + hh) / 2.0
b = (ll - lh + hl - hh) / 2.0
c = (ll + lh - hl - hh) / 2.0
d = (ll - lh - hl + hh) / 2.0
block = np.zeros((size, size), dtype=np.float32)
block[0::2, 0::2] = a
block[0::2, 1::2] = b
block[1::2, 0::2] = c
block[1::2, 1::2] = d
x[:size, :size] = block
return x
def haar2d_batch(X: np.ndarray, levels: int = 3) -> np.ndarray:
"""Batch Haar wavelet features."""
imgs = _as_images(X)
coeffs = np.stack([haar2d_single(img, levels=levels).reshape(-1) for img in imgs], axis=0)
return coeffs.astype(np.float32)
def scattering_like_features(X: np.ndarray, levels: int = 3) -> np.ndarray:
"""A compact deterministic scattering-like feature map.
True scattering networks require a bank of complex wavelets. For a small
offline repo, we approximate the paper's wavelet/scattering idea by using
Haar coefficients, coefficient magnitudes, and energy coefficients. This
gives deformation-stable, modulus-style spectral features while remaining
dependency-free.
"""
c = haar2d_batch(X, levels=levels)
return np.concatenate([c, np.abs(c), c * c], axis=1).astype(np.float32)
def torch_haar2d(x: torch.Tensor, levels: int = 3) -> torch.Tensor:
"""Differentiable 2D Haar transform for tensors shaped (B, 1, 8, 8).
This version avoids in-place writes so autograd remains fast and stable.
It returns coefficients flattened in multiresolution order rather than
quadrant-packed image order.
"""
if x.ndim != 4 or x.shape[1] != 1:
raise ValueError("Expected tensor of shape (B,1,H,W)")
current = x
details = []
for _ in range(levels):
if current.shape[-1] < 2 or current.shape[-2] < 2:
break
a = current[:, :, 0::2, 0::2]
b = current[:, :, 0::2, 1::2]
c = current[:, :, 1::2, 0::2]
d = current[:, :, 1::2, 1::2]
ll = (a + b + c + d) / 2.0
lh = (a - b + c - d) / 2.0
hl = (a + b - c - d) / 2.0
hh = (a - b - c + d) / 2.0
details.extend([lh.flatten(1), hl.flatten(1), hh.flatten(1)])
current = ll
details.append(current.flatten(1))
return torch.cat(details, dim=1)
def torch_wavelet_features(x: torch.Tensor, levels: int = 3) -> torch.Tensor:
"""Flattened differentiable Haar features for learnable spectral operator."""
return torch_haar2d(x, levels=levels)
|