fbscatnet 1.0.0__py3-none-any.whl

fbscatnet/__init__.py

@@ -0,0 +1,14 @@
+import logging
+
+from .generate_bank import FourierBesselWaveletBank
+from .scatnet import FourierBesselScatNet
+
+
+def setup_logger(name: str) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.addHandler(logging.NullHandler())
+
+    return logger
+
+
+__all__ = ["FourierBesselWaveletBank", "FourierBesselScatNet"]

fbscatnet/generate_bank.py

@@ -0,0 +1,231 @@
+from collections.abc import KeysView, ValuesView
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from .logger_config import setup_logger
+from .math_utils import (
+    _find_neumann_root_muller,
+    _generate_fourier_bessel_wavelet,
+    _generate_fourier_low_pass_filter,
+)
+
+logger = setup_logger(__name__)
+
+
+class FourierBesselWaveletBank:
+    """A bank of Fourier-Bessel wavelets indexed by parameters m and k.
+
+    Attributes:
+        size (int): Image size.
+        m (int): Maximum order.
+        k (int): Maximum angular index.
+        sigma (float): Scale parameter for the wavelets.
+        verbose (bool): Display wavelet diagnostic prints.
+        wavelet_bank (dict[str, np.ndarray]): Stored dictionary of wavelets.
+        mk_to_key (dict[tuple, str]): Mapping from (m, k) tuples to string keys.
+    """
+
+    def __init__(
+        self, size: int, m: int, k: int, sigma: float = 0.3, norm: str = "l1", verbose: bool = False
+    ) -> None:
+        """Initialise the FourierBesselWaveletBank.
+
+        Args:
+            size (int): Image size.
+            m (int): Maximum order.
+            k (int): Maximum angular index.
+            sigma (float, optional): Scale parameter for the wavelets.
+            norm (str, optional): Wavelet normalisation
+            verbose (bool, optional): Display wavelet diagnostic prints. Defaults to False.
+
+        Raises:
+            ValueError: If angular order k is greater than m.
+            ValueError: If input parameters are negative
+            ValueError: If norm is not 'l1' or 'l2'
+        """
+        self.size = size
+        self.m = m
+        self.k = k
+        self.sigma = sigma
+        self.sigma2 = sigma**2
+        self.m_values = np.arange(0, m)
+        self.k_values = np.arange(0, k)
+        self.verbose = verbose
+        self.norm = norm
+
+        if self.m < self.k:
+            raise ValueError("m <= k condition is not respected")
+
+        if self.m < 0 or self.k < 0 or self.sigma < 0 or self.size <= 0:
+            raise ValueError("Cannot accept negative parameters")
+
+        if self.norm not in ("l1", "l2"):
+            raise ValueError(f"Invalid norm: Must be 'l1' or 'l2' not {self.norm}")
+
+        self.lambda_max = _find_neumann_root_muller(
+            0, int(self.k_values.max()) if len(self.k_values) > 0 else 0
+        )
+
+        wavelet_bank: dict[str, np.ndarray] = {}
+        mk_to_key: dict[tuple, str] = {}
+
+        self.freq_limit = int(self.lambda_max + 2 / self.sigma)
+
+        for k_val in self.k_values:
+            for m_val in self.m_values:
+                if np.abs(m_val) > k_val:
+                    continue
+
+                if k_val == 0:
+                    _, Z = _generate_fourier_low_pass_filter(
+                        size=self.size,
+                        sigma=self.sigma,
+                        norm=self.norm,
+                        freq_limit=self.freq_limit,
+                        verbose=self.verbose,
+                    )
+                else:
+                    _, _, Z = _generate_fourier_bessel_wavelet(
+                        m_val,
+                        k_val,
+                        size=self.size,
+                        sigma=self.sigma,
+                        norm=self.norm,
+                        freq_limit=self.freq_limit,
+                        verbose=self.verbose,
+                    )
+
+                key_name = f"m_{m_val}_k_{k_val}_s{self.sigma}"
+                mk_to_key[m_val, k_val] = key_name
+                wavelet_bank[key_name] = Z
+
+        self.mk_to_key = mk_to_key
+        self.wavelet_bank = wavelet_bank
+
+    def __getitem__(self, key_or_indices: str | tuple) -> np.ndarray:
+        """Retrieve a specific wavelet by its string key or (m_index, k_index) tuple.
+
+        Args:
+            key_or_indices (str | tuple): Index to retrieve either as string "m_k_sigma"
+                or by an (m, k) tuple.
+
+        Returns:
+            np.ndarray: The requested wavelet array.
+
+        Raises:
+            KeyError: If the key or (m, k) combination does not exist.
+            TypeError: If an invalid index type is provided.
+        """
+        if isinstance(key_or_indices, str):
+            if key_or_indices not in self.wavelet_bank:
+                raise KeyError(f"Wavelet key '{key_or_indices}' not found.")
+            return self.wavelet_bank[key_or_indices]
+
+        elif isinstance(key_or_indices, tuple) and len(key_or_indices) == 2:
+            m_index, k_index = key_or_indices
+
+            key = self.mk_to_key.get((m_index, k_index))
+            if key is None:
+                raise KeyError(f"Wavelet with parameters m={m_index}, k={k_index} not found.")
+
+            return self.wavelet_bank[key]
+
+        raise TypeError("Invalid index type. Use a string key or an (m, k) tuple.")
+
+    def __len__(self) -> int:
+        """Retrieve the number of wavelets in the bank.
+
+        Returns:
+            int: Total count of wavelets.
+        """
+        return len(self.wavelet_bank)
+
+    def get_keys(self) -> KeysView[str]:
+        """Return the dictionary keys representing individual wavelets.
+
+        Returns:
+            KeysView[str]: View of the string keys.
+        """
+        return self.wavelet_bank.keys()
+
+    def get_values(self) -> ValuesView[np.ndarray]:
+        """Return the collection of wavelet arrays.
+
+        Returns:
+            ValuesView[np.ndarray]: View of the wavelet numpy arrays.
+        """
+        return self.wavelet_bank.values()
+
+    def summary(self, verbose: bool = True) -> tuple[int, int, float]:
+        """Print an optional summary of the wavelet bank parameters and size.
+
+        Args:
+            verbose (bool, optional): Whether to print summary details. Defaults to True.
+
+        Returns:
+            tuple[int, int, float]: A tuple containing (m, k, sigma).
+        """
+        if verbose:
+            print()
+            logger.info("Fourier-Bessel Wavelet bank summary:")
+            logger.info(
+                "Parameters: m = %d, k = %d, sigma = %.2f, norm = %s",
+                self.m,
+                self.k,
+                self.sigma,
+                self.norm,
+            )
+            logger.info("Total wavelets: %d", len(self.wavelet_bank))
+            logger.info("Frequency limit: %.2f", self.freq_limit)
+            logger.info("Key naming structure: {m_{m_val}_k_{k_val}_s{sigma}}")
+
+        return self.m, self.k, self.sigma
+
+    def plot_bank(self) -> None:
+        """Plot the grid of wavelets using matplotlib."""
+        _, axes = plt.subplots(
+            len(self.k_values), len(self.m_values), figsize=(8, 6), sharex=True, sharey=True
+        )
+
+        axes = np.atleast_2d(axes)
+        for row_idx, k_val in enumerate(self.k_values):
+            for col_idx, m_val in enumerate(self.m_values):
+                ax = axes[row_idx, col_idx]
+                ax.set_xlim(-self.freq_limit, self.freq_limit)
+                ax.set_ylim(-self.freq_limit, self.freq_limit)
+                ax.axis("off")
+
+                if row_idx == 0:
+                    ax.set_title(f"m = {m_val}", fontsize=12, fontweight="bold")
+                if col_idx == 0:
+                    ax.text(
+                        -self.freq_limit * 1.2,
+                        0,
+                        f"k = {k_val}",
+                        fontsize=12,
+                        fontweight="bold",
+                        ha="right",
+                        va="center",
+                    )
+
+                if np.abs(m_val) > k_val:
+                    continue
+
+                if f"m_{m_val}_k_{k_val}_s{self.sigma}" not in self.wavelet_bank:
+                    continue
+
+                Z = self.wavelet_bank[f"m_{m_val}_k_{k_val}_s{self.sigma}"]
+                z_max: float = float(np.max(np.abs(Z)))
+                ax.imshow(
+                    np.real(Z),
+                    extent=[-self.freq_limit, self.freq_limit, -self.freq_limit, self.freq_limit],
+                    cmap="inferno",
+                    origin="lower",
+                    vmin=-z_max,
+                    vmax=z_max,
+                )
+
+        plt.tight_layout()
+        plt.subplots_adjust(wspace=0.05, hspace=0.05)
+        plt.show()

fbscatnet/logger_config.py

@@ -0,0 +1,47 @@
+import logging
+import sys
+
+
+class LogColors:
+    RESET = "\033[0m"
+    DEBUG = "\033[36m"  # Cyan
+    INFO = "\033[32m"  # Green
+    WARNING = "\033[33m"  # Yellow
+    ERROR = "\033[31m"  # Red
+    CRITICAL = "\033[35m"  # Magenta
+
+
+class ColoredFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        color = getattr(LogColors, record.levelname, LogColors.RESET)
+        colored_level = f"{color}%(levelname)s{LogColors.RESET}"
+        log_fmt = f"%(asctime)s [{colored_level}] %(message)s"
+        formatter = logging.Formatter(log_fmt, datefmt="%Y-%m-%d %H:%M:%S")
+        return formatter.format(record)
+
+
+def setup_logger(name: str) -> logging.Logger:
+    """Internal function to set up the library logger with a NullHandler."""
+    logger = logging.getLogger(name)
+    if not logger.handlers:
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def enable_colored_logs(level: int = logging.INFO) -> None:
+    """
+    Helper function for users to easily enable colored console logging for fbscatnet.
+
+    Usage:
+        import fbscatnet
+        fbscatnet.enable_colored_logs()
+    """
+    # Get the top-level logger for the package
+    logger = logging.getLogger("fbscatnet")
+    logger.setLevel(level)
+
+    # Check if we already added a StreamHandler to avoid duplicates
+    if not any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(ColoredFormatter())
+        logger.addHandler(handler)

fbscatnet/math_utils.py

@@ -0,0 +1,236 @@
+from typing import cast
+
+import numpy as np
+import numpy.typing as npt
+import scipy.special
+
+from .logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+### BESSEL WRAPPERS VIA SCIPY ###
+
+
+def _first_kind_bessel(X: npt.ArrayLike, order: int | np.integer) -> float | np.ndarray:
+    """Vectorised Bessel function supporting both scalar points and arrays."""
+    X_arr = np.asarray(X)
+    abs_order = int(abs(int(order)))
+
+    if order < 0:
+        val = (-1) ** abs_order * scipy.special.jv(abs_order, X_arr)
+    else:
+        val = scipy.special.jv(int(order), X_arr)
+
+    # If the input was a scalar, return a Python float; otherwise the array
+    if np.isscalar(X):
+        return float(np.asarray(val).item())
+    return cast(np.ndarray, np.asarray(val))
+
+
+def _first_kind_bessel_deriv(X: npt.ArrayLike, order: int | np.integer) -> float | np.ndarray:
+    """Derivative supporting both scalar points and arrays using recurrence relations."""
+    val = 0.5 * (_first_kind_bessel(X, order - 1) - _first_kind_bessel(X, order + 1))
+    return val
+
+
+def _first_modified_bessel(X: npt.ArrayLike, order: int | np.integer) -> float | np.ndarray:
+    """Modified Bessel function supporting both scalar points and arrays."""
+
+    X_arr = np.asarray(X)
+    val = scipy.special.iv(order, X_arr)
+
+    # If the input was a scalar, return a Python float; otherwise the array
+    if np.isscalar(X):
+        return float(np.asarray(val).item())
+    return cast(np.ndarray, np.asarray(val))
+
+
+### MULLER'S METHOD TO FIND EIGENVALUE ###
+
+
+def _mcmahon_seed(m: int | np.integer, k: int | np.integer) -> float:
+    """Approximate the k-th positive root using McMahon's asymptotic expansion."""
+    m = int(m)
+    k = int(k)
+    beta: float
+    if m == 0:
+        s = k
+        nu = 1
+        beta = (s + nu / 2.0 - 0.25) * np.pi
+        return beta - (4.0 * nu**2 - 1.0) / (8.0 * beta)
+    else:
+        s = k - m + 1
+        beta = (s + m / 2.0 - 0.75) * np.pi
+        if beta <= 0:
+            return float(m) + 0.5
+        return beta - (4.0 * m**2 + 3.0) / (8.0 * beta)
+
+
+def _find_neumann_root_muller(
+    m: int | np.integer, k: int | np.integer, thresh: float = 1e-12, max_iter: int = 200
+) -> float:
+    """Find the k-th positive root of using Muller's method."""
+
+    x3 = 0.0
+
+    if k < m or (m == 0 and k == 0):
+        return x3
+
+    x_seed = _mcmahon_seed(m, k)
+    x0 = max(0.01, x_seed - 0.15)
+    x1 = x_seed
+    x2 = x_seed + 0.15
+
+    def f(x: npt.ArrayLike) -> float | np.ndarray:
+        return _first_kind_bessel_deriv(x, m)
+
+    d0: float = float(f(x0))
+    d1: float = float(f(x1))
+    d2: float = float(f(x2))
+
+    converged = False
+
+    for _ in range(max_iter):
+        h1 = x1 - x0
+        h2 = x2 - x1
+        if h1 == 0 or h2 == 0:
+            break
+        delta1 = (d1 - d0) / h1
+        delta2 = (d2 - d1) / h2
+        if (h2 + h1) == 0:
+            break
+        d_coef = (delta2 - delta1) / (h2 + h1)
+        a = d_coef
+        b = delta2 + h2 * d_coef
+        c = d2
+        disc = np.lib.scimath.sqrt(b**2 - 4 * a * c)
+
+        dx = -2 * c / (b + disc) if np.real(b) >= 0 else -2 * c / (b - disc)
+
+        x3 = np.real(x2 + dx).item()
+        d3 = float(f(x3))
+        if abs(d3) < thresh or abs(dx) < thresh:
+            converged = True
+            break
+        x0, x1, x2 = x1, x2, x3
+        d0, d1, d2 = d1, d2, d3
+
+    if not converged:
+        logger.warning(
+            "Root finder did not converge for m=%d, k=%d after %d iterations "
+            "(residual=%.2e). Consider increasing max_iter.",
+            int(m),
+            int(k),
+            max_iter,
+            abs(d3) if "d3" in locals() else float("nan"),
+        )
+    return float(x3)
+
+
+### FOURIER WAVELET ###
+
+
+def _generate_fourier_bessel_wavelet(
+    m: int | np.integer,
+    k: int | np.integer,
+    size: int = 50,
+    sigma: float = 0.1,
+    norm: str = "l1",
+    freq_limit: int = 20,
+    verbose: bool = False,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+
+    abs_m = np.abs(m)
+    freq = np.linspace(-freq_limit, freq_limit, size)
+    Kx: np.ndarray
+    Ky: np.ndarray
+    Kx, Ky = np.meshgrid(freq, freq, indexing="ij")
+
+    Q = np.sqrt(Kx**2 + Ky**2)
+    Psi = np.arctan2(Ky, Kx)
+
+    eigenvalue = _find_neumann_root_muller(abs_m, k)
+    sigma2 = sigma**2
+    eig2 = eigenvalue**2
+    angular_profile = np.exp(1j * m * Psi)
+    K = np.exp(-(eig2 * sigma2) / 2)
+
+    mod_bessel = _first_modified_bessel((eig2 * sigma2) / 2, m)
+    mod_bessel_freq = _first_modified_bessel(eigenvalue * sigma2 * Q, m)
+
+    start = (1j**m) * angular_profile
+    left_hand = sigma2 * np.exp(-(sigma2 * (eig2 + Q**2)) / 2) * mod_bessel_freq
+
+    if m == 0:
+        bracket = K * mod_bessel - 2 * np.exp(-(3 * sigma2 * eig2) / 4) + np.exp(-sigma2 * eig2)
+        norm_term = 1 / (np.sqrt((np.pi * sigma2) * bracket))
+
+    else:
+        norm_term = 1 / (np.sqrt(np.pi * sigma2 * K * mod_bessel))
+
+    if m == 0:
+        Z = start * norm_term * (left_hand - K * sigma2 * np.exp(-(sigma2 * Q**2) / 2))
+    else:
+        Z = start * norm_term * left_hand
+
+    if norm == "l1":
+        z_max = np.max(np.abs(Z))
+        Z /= z_max
+
+    # --- DIAGNOSTIC PRINTS ---
+    if verbose:
+        dx = freq[1] - freq[0]
+        mean_val = np.abs(np.mean(Z))
+        if norm == "l2":
+            post_norm_energy = np.sum(np.abs(Z) ** 2) * (dx * dx)
+        else:
+            post_norm_energy = np.max(np.abs(Z))
+
+        logger.info(
+            f"Wavelet (m={m}, k={k}) | "
+            f"Mean (~0): {mean_val:.2e} | "
+            f"{norm.upper()}: {post_norm_energy:.4f} | "
+            f"Eigenvalue: {eigenvalue:.4f}"
+        )
+    # -------------------------
+
+    return Kx, Ky, Z
+
+
+def _generate_fourier_low_pass_filter(
+    size: int = 50,
+    sigma: float = 0.1,
+    norm: str = "l1",
+    freq_limit: int = 20,
+    verbose: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    freq = np.linspace(-freq_limit, freq_limit, size)
+    Kx: np.ndarray
+    Ky: np.ndarray
+    Kx, Ky = np.meshgrid(freq, freq, indexing="ij")
+    Q = np.sqrt(Kx**2 + Ky**2)
+
+    sigma2 = sigma**2
+    Z = np.exp(-(sigma2) * (Q**2) / 2)
+
+    if norm == "l1":
+        z_max = np.max(np.abs(Z))
+        Z /= z_max
+
+    # --- DIAGNOSTIC PRINTS ---
+    if verbose:
+        dx = freq[1] - freq[0]
+        mean_val = np.abs(np.mean(Z))
+        if norm == "l2":
+            post_norm_energy = np.sum(np.abs(Z) ** 2) * (dx * dx)
+        else:
+            post_norm_energy = np.max(np.abs(Z))
+
+        logger.info(
+            f"Low pass  (m=0, k=0) | "
+            f"Mean (>0): {mean_val:.2e} | "
+            f"{norm.upper()}: {post_norm_energy:.4f}"
+        )
+    # -------------------------
+
+    return Q, Z

fbscatnet/scatnet.py

@@ -0,0 +1,468 @@
+import os
+import re
+import warnings
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.fft as spfft
+from joblib import Parallel, delayed
+from tqdm import tqdm
+
+from .generate_bank import FourierBesselWaveletBank
+from .logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+try:
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        import cupy as cp
+
+    try:
+        HAS_CUPY = cp.cuda.runtime.getDeviceCount() > 0
+    except Exception as exc:
+        logger.warning("CuPy is installed but no usable GPU was detected: %s", exc)
+        HAS_CUPY = False
+except ImportError:
+    cp = None
+    HAS_CUPY = False
+
+
+def _extract_k(key_str: str) -> int:
+    """Extract the integer 'k' index encoded in a filter bank key string."""
+    match = re.search(r"k[=_ ]*(\d+)", key_str)
+    if match:
+        return int(match.group(1))
+    else:
+        logger.warning("Valid k key not found in '%s', defaulting to 0", key_str)
+        return 0
+
+
+class FourierBesselScatNet:
+    """A scattering network based on Fourier-Bessel wavelets for generating image embeddings.
+
+    Attributes:
+        size (int): Image spatial size.
+        bank (FourierBesselWaveletBank): The bank of Fourier-Bessel wavelets.
+        num_filters (int): Total number of filters in the wavelet bank.
+        bank_keys (list[str]): List of string keys representing the filters.
+        low_pass (np.ndarray): Low-pass filter array retrieved from the bank.
+        final_features (np.ndarray): Generated feature embeddings from the last run.
+    """
+
+    def __init__(self, bank: FourierBesselWaveletBank, backend: str = "cpu") -> None:
+        """Initialise the FourierBesselScatNet.
+
+        Args:
+            bank (FourierBesselWaveletBank): A configured bank of Fourier-Bessel wavelets.
+            backend (str): Device to operate on
+
+        Raises:
+            ImportError: If CuPy requested but no GPU.
+            ValueError: If the bank contains no band-pass wavelets (i.e. only
+                the (0, 0) low-pass filter), since the scattering cascade has
+                nothing to filter with in that case.
+        """
+        self.backend = backend.lower()
+        if self.backend == "gpu" and not HAS_CUPY:
+            raise ImportError(
+                "GPU backend was requested, but 'cupy' is not"
+                ""
+                "installed or no compatible GPU was found."
+            )
+        if len(bank) <= 1:
+            raise ValueError(
+                "The wavelet bank contains no band-pass wavelets (only the "
+                "low-pass filter). FourierBesselScatNet needs at least one "
+                "wavelet - try constructing the bank with m > 1 and/or k > 1."
+            )
+        self.xp: Any = cp if self.backend == "gpu" else np
+        self.size = bank[0, 0].shape[0]
+        self.bank = bank
+        self.num_filters = len(bank)
+        self.bank_keys = list(bank.get_keys())
+        self.low_pass = self.xp.asarray(bank[0, 0], dtype=self.xp.float32)
+        self.final_features: Any | None = None
+
+        # Precompute everything that is constant across batches
+        xp = self.xp
+        self.order_1_keys = self.bank_keys[1:]
+        self.num_order_1_maps = len(self.order_1_keys)
+
+        self._k_map = {key: _extract_k(key) for key in self.order_1_keys}
+
+        self._order2_children: dict[str, list[str]] = {
+            key1: [key2 for key2 in self.order_1_keys if self._k_map[key2] < self._k_map[key1]]
+            for key1 in self.order_1_keys
+        }
+        self.num_order_2_maps = sum(len(v) for v in self._order2_children.values())
+
+        self._filters_1 = {
+            key: xp.asarray(bank[key], dtype=xp.complex64) for key in self.order_1_keys
+        }
+        self._low_pass_c = xp.asarray(bank[0, 0], dtype=xp.complex64)
+
+        # Stacked filter tensors for vectorised filtering.
+        # Shape: (num_order_1_maps, H, W)
+        self._filters_1_stack = xp.stack([self._filters_1[k] for k in self.order_1_keys])
+
+        # Per-key1 stacked children filters for vectorised order-2 filtering.
+        # Shape per entry: (num_children, H, W), or None if no valid children.
+        self._filters_2_stack_by_key1: dict[str, Any] = {
+            key1: (xp.stack([self._filters_1[k2] for k2 in children]) if children else None)
+            for key1, children in self._order2_children.items()
+        }
+
+    def _fft2(self, x: Any, xp: Any) -> Any:
+        """Backend-aware 2D FFT (uses scipy on CPU for multi-threaded speed, cupy on GPU)."""
+        if self.backend == "gpu":
+            return xp.fft.fft2(x, axes=(-2, -1))
+        return spfft.fft2(x, axes=(-2, -1), workers=-1)
+
+    def _ifft2(self, x: Any, xp: Any) -> Any:
+        """Backend-aware inverse 2D FFT, mirroring `_fft2`."""
+        if self.backend == "gpu":
+            return xp.fft.ifft2(x, axes=(-2, -1))
+        return spfft.ifft2(x, axes=(-2, -1), workers=-1)
+
+    def _filter_and_modulus(self, xp: Any, freq_signal: Any, filters: Any) -> Any:
+        """Apply wavelet filter(s) in the frequency domain and take the spatial modulus.
+
+        Args:
+            xp: The array module to use (numpy or cupy) for this batch.
+            freq_signal: Input already in fftshifted frequency domain,
+                broadcastable against `filters`.
+            filters: One or a stack of frequency-domain filters to apply.
+
+        Returns:
+            The fftshifted frequency-domain modulus.
+        """
+        filtered_fft = freq_signal * filters
+        shifted_freq = xp.fft.ifftshift(filtered_fft, axes=(-2, -1))
+        spatial_complex = self._ifft2(shifted_freq, xp)
+        modulus_spatial = xp.abs(spatial_complex).astype(xp.float32)
+        return xp.fft.fftshift(self._fft2(modulus_spatial, xp), axes=(-2, -1)).astype(xp.complex64)
+
+    def _smooth_and_pool(
+        self,
+        xp: Any,
+        modulus_fft: Any,
+        low_pass_c: Any,
+        batch_size: int,
+        n_maps: int,
+        d_size: int,
+        downsize: int,
+    ) -> Any:
+        """Low-pass smooth a modulus map and block-mean-pool it down to `d_size`.
+
+        Args:
+            xp: The array module to use (numpy or cupy) for this batch.
+            modulus_fft: Output of `_filter_and_modulus`, shape (batch, n_maps, H, W).
+            low_pass_c: Broadcastable complex low-pass filter.
+            batch_size: Number of samples in this batch.
+            n_maps: Number of filter channels being pooled.
+            d_size: Output spatial size after downsampling (H // downsize).
+            downsize: Block-mean pooling factor.
+
+        Returns:
+            Real-valued pooled features of shape (batch_size, d_size, d_size, n_maps).
+        """
+        filtered_low_pass = modulus_fft * low_pass_c
+        smoothed_shifted = xp.fft.ifftshift(filtered_low_pass, axes=(-2, -1))
+        smoothed_spatial = self._ifft2(smoothed_shifted, xp)
+        return (
+            xp.real(smoothed_spatial)
+            .reshape(batch_size, n_maps, d_size, downsize, d_size, downsize)
+            .mean(axis=(3, 5))
+            .transpose(0, 2, 3, 1)
+        )
+
+    def generate_embeddings(
+        self,
+        data: np.ndarray,
+        downsize: int,
+        batch_size: int = 32,
+        use_multiprocessing: bool = False,
+    ) -> np.ndarray:
+        """Generate scattering network feature embeddings for a given dataset.
+
+        Args:
+            data (np.ndarray): Input image dataset of shape (num_samples, height, width).
+            downsize (int): Spatial downsampling factor via block mean pooling.
+            batch_size (int, optional): Number of samples per batch. Defaults to 32.
+            use_multiprocessing (bool, optional): Whether to use multiple CPU cores
+                (Ignored if backend='gpu'). Defaults to False.
+
+        Returns:
+            np.ndarray: Flattened feature embeddings of shape (num_samples, feature_dim).
+        """
+        xp = self.xp
+        data = np.asarray(data, dtype=np.float32)
+
+        num_samples = data.shape[0]
+        d_size = int(self.size / downsize)
+        if (self.size % downsize) != 0:
+            raise ValueError(
+                f"Image size ({self.size}) must be exactly divisible by downsize ({downsize})."
+            )
+
+        order_1_keys = self.order_1_keys
+        num_order_1_maps = self.num_order_1_maps
+        num_order_2_maps = self.num_order_2_maps
+
+        def _process_batch(batch_data: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+            b_xp: Any = cp if self.backend == "gpu" else np
+            b_data = (
+                b_xp.asarray(batch_data, dtype=b_xp.float32)
+                if self.backend == "gpu"
+                else batch_data
+            )
+            curr_batch_size = b_data.shape[0]
+
+            batch_fft = b_xp.fft.fftshift(self._fft2(b_data, b_xp), axes=(-2, -1)).astype(
+                b_xp.complex64
+            )
+
+            low_pass_c = self._low_pass_c
+
+            # ORDER 0: plain low-pass response, no modulus non-linearity.
+            low_pass_spatial = self._ifft2(
+                b_xp.fft.ifftshift(batch_fft * low_pass_c, axes=(-2, -1)), b_xp
+            )
+            low_pass_down = (
+                b_xp.real(low_pass_spatial)
+                .reshape(-1, d_size, downsize, d_size, downsize)
+                .mean(axis=(2, 4))
+            )
+            order_0_res = low_pass_down.reshape(curr_batch_size, -1)
+
+            # ORDER 1: filter with every wavelet, take modulus, smooth, pool.
+            bank_1_stack = self._filters_1_stack  # (num_filters, H, W)
+            modulus_fft_1 = self._filter_and_modulus(
+                b_xp, batch_fft[:, None, :, :], bank_1_stack[None, :, :, :]
+            )
+            batch_pooled_order_1 = self._smooth_and_pool(
+                b_xp,
+                modulus_fft_1,
+                low_pass_c[None, None, :, :],
+                curr_batch_size,
+                num_order_1_maps,
+                d_size,
+                downsize,
+            )
+
+            # ORDER 2: for each order-1 channel, filter again with only the
+            # wavelets of strictly lower angular index k
+            batch_pooled_order_2 = b_xp.zeros(
+                (curr_batch_size, d_size, d_size, num_order_2_maps), dtype=b_xp.float32
+            )
+            order_2_idx = 0
+
+            for i, key1 in enumerate(order_1_keys):
+                children = self._order2_children[key1]
+                if not children:
+                    continue
+
+                filters_2_stack = self._filters_2_stack_by_key1[key1]  # (n_children, H, W)
+                mod_fft_1_single = modulus_fft_1[:, i]  # (batch, H, W)
+
+                modulus_fft_2 = self._filter_and_modulus(
+                    b_xp, mod_fft_1_single[:, None, :, :], filters_2_stack[None, :, :, :]
+                )
+                n_children = len(children)
+                down = self._smooth_and_pool(
+                    b_xp,
+                    modulus_fft_2,
+                    low_pass_c[None, None, :, :],
+                    curr_batch_size,
+                    n_children,
+                    d_size,
+                    downsize,
+                )
+                batch_pooled_order_2[..., order_2_idx : order_2_idx + n_children] = down
+                order_2_idx += n_children
+
+            return (
+                order_0_res,
+                batch_pooled_order_1.reshape(curr_batch_size, -1),
+                batch_pooled_order_2.reshape(curr_batch_size, -1),
+            )
+
+        # Build batches index ranges
+        batches = [
+            data[start : min(start + batch_size, num_samples)]
+            for start in range(0, num_samples, batch_size)
+        ]
+
+        # Execution Selection
+        if self.backend == "cpu" and use_multiprocessing:
+            nb_cpu: int = os.cpu_count() or 1
+            n_jobs = max(1, nb_cpu - 1) if os.cpu_count() and nb_cpu > 1 else 1
+            results = self._run_multiprocess(batches, _process_batch, n_jobs)
+        else:
+            results = []
+            for batch in tqdm(
+                batches, desc=f"Generating Embeddings on {self.backend.upper()}", unit="batch"
+            ):
+                results.append(_process_batch(batch))
+
+        # Re-assemble outputs from batch results
+        order_0_list, order_1_list, order_2_list = zip(*results, strict=True)
+
+        order_0_features = xp.concatenate(order_0_list, axis=0)
+        first_order_features = xp.concatenate(order_1_list, axis=0)
+        second_order_features = xp.concatenate(order_2_list, axis=0)
+
+        final_features = xp.concatenate(
+            (order_0_features, first_order_features, second_order_features), axis=1
+        )
+        self.final_features = final_features
+
+        if self.backend == "gpu":
+            return np.asarray(final_features.get())
+
+        return np.asarray(final_features)
+
+    def _run_multiprocess(
+        self, batches: list[np.ndarray], _process_batch: Any, n_jobs: int
+    ) -> list[Any]:
+        """Run batches across multiple processes with a progress bar"""
+        try:
+            # Joblib generator evaluation wrapped inside tqdm
+            with tqdm(
+                total=len(batches), desc="Generating Embeddings (Multiprocessing)", unit="batch"
+            ) as pbar:
+                job_results = Parallel(n_jobs=n_jobs, return_as="generator")(
+                    delayed(_process_batch)(b) for b in batches
+                )
+                results = []
+                for res in job_results:
+                    results.append(res)
+                    pbar.update(1)
+        except TypeError:
+            # Fallback for older joblib versions without generator support
+            results = []
+            for res in tqdm(
+                Parallel(n_jobs=n_jobs)(delayed(_process_batch)(b) for b in batches),
+                desc="Generating Embeddings (Multiprocessing)",
+                unit="batch",
+            ):
+                results.append(res)
+
+        return results
+
+    def save_embeddings(self, path: str) -> None:
+        """Save the generated feature embeddings to a compressed .npz file."""
+        if self.final_features is None:
+            raise ValueError("No embeddings found. Run generate_embeddings() first.")
+
+        m, k, sigma = self.bank.summary(verbose=False)
+        os.makedirs(path, exist_ok=True)
+        save_path = rf"{path}/embedding_m{m}_k{k}_sigma{sigma}.npz"
+
+        features = self.final_features.get() if self.backend == "gpu" else self.final_features
+        np.savez_compressed(save_path, embedding=features)
+
+        logger.info("Embedding successfully saved to '%s'", save_path)
+
+    def visualise_maps(self, image: np.ndarray, downsize: int) -> None:
+        """
+        Visualises Order 0 and Order 1 scattering maps.
+
+        Args:
+            image (np.ndarray): A single 2D image array (height, width).
+            downsize (int): Spatial downsampling factor. Set to 1 for no downsampling.
+        """
+
+        if image.ndim != 2:
+            raise ValueError("Please provide a single 2D image array of shape (height, width).")
+
+        d_size = int(self.size / downsize)
+
+        # Add fake batch dimension to match FFT logic
+        batch = image[None, ...]
+        batch_fft = np.fft.fftshift(np.fft.fft2(batch, axes=(-2, -1)), axes=(-2, -1))
+
+        # Dictionaries to hold maps for plotting
+        maps = {}
+
+        # ORDER 0
+        low_pass_coeffs = batch_fft * self.low_pass
+        low_pass_spatial = np.real(
+            np.fft.ifft2(np.fft.ifftshift(low_pass_coeffs, axes=(-2, -1)), axes=(-2, -1))
+        )[0]
+
+        low_pass_down = low_pass_spatial.reshape(d_size, downsize, d_size, downsize).mean(
+            axis=(1, 3)
+        )
+
+        maps["Order 0 (Low Pass)"] = low_pass_down
+
+        # ORDER 1
+        order_1_keys = self.bank_keys[1:]
+
+        for key1 in order_1_keys:
+            wavelet_fft = self.bank[key1]
+
+            # Standard scattering cascade
+            filtered_fft = batch_fft * wavelet_fft
+            spatial_complex = np.fft.ifft2(
+                np.fft.ifftshift(filtered_fft, axes=(-2, -1)), axes=(-2, -1)
+            )
+            modulus_spatial = np.abs(spatial_complex)
+
+            modulus_fft = np.fft.fftshift(
+                np.fft.fft2(modulus_spatial, axes=(-2, -1)), axes=(-2, -1)
+            )
+            filtered_low_pass = modulus_fft * self.low_pass
+
+            smoothed_spatial = np.real(
+                np.fft.ifft2(np.fft.ifftshift(filtered_low_pass, axes=(-2, -1)), axes=(-2, -1))
+            )[0]
+
+            # Downsample
+            smoothed_down = smoothed_spatial.reshape(d_size, downsize, d_size, downsize).mean(
+                axis=(1, 3)
+            )
+
+            maps[f"Order 1 ({key1})"] = smoothed_down
+
+        # DYNAMIC GRID CALCULATION
+        num_maps = len(maps)
+
+        # Find the best exact integer factors
+        best_factor = 1
+        for i in range(1, int(np.sqrt(num_maps)) + 1):
+            if num_maps % i == 0:
+                best_factor = i
+
+        rows = best_factor
+        cols = num_maps // best_factor
+
+        # Fallback for prime numbers or extremely stretched grids
+        # If the aspect ratio is wider than 3:1, use a square-ish grid instead.
+        if cols / rows > 3:
+            cols = int(np.ceil(np.sqrt(num_maps)))
+            rows = int(np.ceil(num_maps / cols))
+
+        # PLOTTING
+        _, axes = plt.subplots(rows, cols, figsize=(cols * 3, rows * 3))
+
+        # Flatten the axes array so we can iterate through it easily
+        if isinstance(axes, np.ndarray):
+            axes = axes.flatten()
+        else:
+            axes = [axes]  # Catches the edge case where num_maps == 1
+
+        for idx, (title, map_data) in enumerate(maps.items()):
+            axes[idx].imshow(map_data, cmap="inferno")
+            axes[idx].set_title(title, fontsize=10)
+            axes[idx].axis("off")
+
+        # Turn off the axes for any leftover empty subplots from the prime-number fallback
+        for empty_idx in range(num_maps, len(axes)):
+            axes[empty_idx].axis("off")
+
+        plt.tight_layout()
+        plt.show()

fbscatnet-1.0.0.dist-info/METADATA

@@ -0,0 +1,120 @@
+Metadata-Version: 2.4
+Name: fbscatnet
+Version: 1.0.0
+Summary: Fourier-Bessel wavelet scattering transforms
+Project-URL: Homepage, https://github.com/Smee18/FourierBesselWavelets
+Project-URL: Repository, https://github.com/Smee18/FourierBesselWavelets.git
+Project-URL: Issues, https://github.com/Smee18/FourierBesselWavelets/issues
+Author-email: Marcel Venturotti <mv514@bath.ac.uk>
+License: MIT
+License-File: LICENSE
+Keywords: fourier-bessel,image-processing,machine-learning,scattering-transform,wavelets
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.10
+Requires-Dist: joblib>=1.3
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: tqdm>=4.65
+Provides-Extra: dev
+Requires-Dist: build; extra == 'dev'
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: scipy-stubs; extra == 'dev'
+Requires-Dist: types-tqdm; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == 'docs'
+Requires-Dist: sphinx-autobuild; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme; extra == 'docs'
+Provides-Extra: gpu
+Requires-Dist: cupy-cuda12x>=11.0; extra == 'gpu'
+Description-Content-Type: text/markdown
+
+# fbscatnet
+
+[![PyPI version](https://img.shields.io/pypi/v/fbscatnet.svg)](https://pypi.org/project/fbscatnet/)
+[![Python Versions](https://img.shields.io/pypi/pyversions/fbscatnet.svg)](https://pypi.org/project/fbscatnet/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![CI](https://github.com/Smee18/FourierBesselWavelets/actions/workflows/ci.yml/badge.svg)](https://github.com/Smee18/FourierBesselWavelets/actions/workflows/ci.yml)
+[![Code style: ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+`fbscatnet` is a high-performance Python library for computing **Fourier-Bessel wavelet scattering transforms**. It generates robust, feature embeddings from 2D images, making it an ideal feature extractor for computer vision, biomedical imaging, and physics-based machine learning. This is a brand new project of mine that I have been working on for some time. I would love any criticism, improvements, corrections to both the code and the maths.
+
+## Features
+
+- **Novel Wavelet:** New wavelet relying on Bessel basis functions.
+- **Hardware Accelerated:** Seamlessly switch between multi-core CPU execution (`scipy`/`joblib`) and GPU acceleration (`cupy`).
+- **Highly Optimized:** Vectorised filtering in the frequency domain for maximum throughput across large image batches.
+- **ML-Ready:** Outputs flattened feature arrays directly compatible with `scikit-learn`, `xgboost`, or `pytorch`.
+
+## Installation
+
+Install the base package (CPU-only) via pip:
+
+```bash
+pip install fbscatnet
+```
+
+To enable **GPU acceleration**, install with the `gpu` extra (requires a CUDA-compatible GPU):
+
+```bash
+pip install fbscatnet[gpu]
+```
+
+## Quickstart
+
+Extracting features from a dataset takes just a few lines of code:
+
+```python
+import numpy as np
+import logging
+from fbscatnet import FourierBesselWaveletBank, FourierBesselScatNet, logger_config
+logger_config.enable_colored_logs(logging.DEBUG) # IMPORTANT: SET TO SEE LOGS
+
+# 1. Create some dummy image data (e.g., 10 grayscale images of size 64x64)
+images = np.random.rand(10, 64, 64)
+
+# 2. Instantiate a Fourier-Bessel Wavelet Bank
+# size: spatial dimension (64x64), m: angular order, k: radial roots
+bank = FourierBesselWaveletBank(size=64, m=2, k=2, sigma=0.1)
+
+# 3. Initialize the Scattering Network
+# Use backend="gpu" if you installed with CuPy
+net = FourierBesselScatNet(bank=bank, backend="cpu")
+
+# 4. Generate feature embeddings
+# downsize: spatial pooling factor (must evenly divide the image size)
+features = net.generate_embeddings(images, downsize=4, use_multiprocessing=True)
+
+print(f"Generated embeddings shape: {features.shape}")
+# Output: (10, feature_dimension)
+```
+
+## Visualising Wavelet Maps
+
+If you want to inspect how the wavelet filters are interacting with your image at the first scattering order, `fbscatnet` includes a built-in plotting tool:
+
+```python
+# Pass a single 2D image to visualize
+single_image = images[0]
+net.visualise_maps(single_image, downsize=4)
+```
+
+## API Overview
+
+Sphinx documentation at: https://smee18.github.io/FourierBesselWavelets/
+
+The `example` folder also contains a full pipeline, classifying MNIST using the library
+The pdf of my notes taken along this project contains an in-depth explanation of the mathematics behind these wavelets. Here you can find all derivations, proofs and useful information. Some stuff might seem trivial but my goal is to really expose all aspects so anyone new to wavelet theory can understand the mathematics behind these Fourier Bessel wavelets.
+
+## License
+This project is licensed under the MIT License.

fbscatnet-1.0.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+fbscatnet/__init__.py,sha256=WPRlnUw-NS_Xo0l-ihj0KrhlJF6Niqydc3nNN73Sjz0,339
+fbscatnet/generate_bank.py,sha256=Qo_ymKoxqT82yjr0BGGdM2Iy9bc22309kEzidKXCwF8,8394
+fbscatnet/logger_config.py,sha256=SPLTnoE05nkkN7fGTyIlcqrut7bNlfyD-KOBJT9OV24,1594
+fbscatnet/math_utils.py,sha256=rzB_zoPlpdNgOx5O-pvzFGN4JeF_qg_OzoykkHiXIs4,6906
+fbscatnet/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fbscatnet/scatnet.py,sha256=nRZHlMRdk5pBvSJUZyJKfn5i3hFqW4NCYNBYwnth4ts,18560
+fbscatnet-1.0.0.dist-info/METADATA,sha256=N4Yx345ZkKPQMx2dXg73orYKOfRb-tXqdVLl57OqBI8,5373
+fbscatnet-1.0.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+fbscatnet-1.0.0.dist-info/licenses/LICENSE,sha256=ppcDcZqSfk-QMGl6mrPhtOb-5fVTK45M0rEN-WqOhEg,1074
+fbscatnet-1.0.0.dist-info/RECORD,,

fbscatnet-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

fbscatnet-1.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Marcel Venturotti
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.