Skip to content

nemes-inc/thrml-act

Repository files navigation

THRML Chirplet Transform

Implementation combining THRML discrete sampling with gradient refinement for adaptive chirplet decomposition.

Getting Started

Note: This implementation has only been tested on Mac M series laptops. Compatibility with other platforms is not guaranteed. By default, the code runs on CPU for maximum compatibility.

# Create the virtual environment
uv venv .venv

# Sync dependencies
uv sync

# Activate the virtual environment
source .venv/bin/activate

# Run the visualization
python run_visualization.py

# Run the test suite
python act-challenging-test.py
python act-profile-test.py
python act-synth-test.py

Usage

Basic Decomposition

from act import decompose_signal

components, parameters, residual = decompose_signal(
    signal,
    sampling_rate=256.0,
    n_components=10,
    residual_threshold=0.01,
    use_thrml=True
)

Configuration

from act import THRMLChirpletOptimizer

optimizer = THRMLChirpletOptimizer(signal, sampling_rate=256.0)
results = optimizer.optimize(
    n_components=5,
    use_coarse_sampling=True,
    n_refinement_candidates=10
)

Parameters

  • n_levels: Discretization resolution (default: 32)
  • beta: Sampling temperature (default: 0.5-1.0)
  • n_samples: Number of MCMC samples (default: 5000)
  • n_warmup: Burn-in iterations (default: 1000)
  • n_refinement_candidates: Number of candidates for BFGS refinement (default: 10)

API

decompose_signal()

Decompose a signal into chirplet components.

Parameters

  • signal: Input signal array
  • sampling_rate: Sampling frequency (Hz)
  • n_components: Maximum number of components
  • residual_threshold: Stop when residual < threshold * signal energy
  • use_thrml: Whether to use THRML sampling (vs pure BFGS)

Returns: (components, parameters, residual)

ChirpletParams

Dataclass containing chirplet parameters

  • tc: Time center (samples)
  • fc: Frequency center (Hz)
  • logDt: Log duration scale
  • c: Chirp rate (Hz/s)

Mathematical

The Adaptive Chirplet Transform uses four continuous parameters to define Gaussian chirplet atoms.

$$g_{t_c,f_c,\log\Delta t,c}(t) = \frac{1}{\sqrt{2\pi}\Delta t} \exp\left(-\frac{1}{2}\left(\frac{t-t_c}{\Delta t}\right)^2\right) \exp\left(j2\pi\left[c(t-t_c)^2 + f_c(t-t_c)\right]\right)$$

Note: Some references use $\exp(-\frac{1}{2}(t/\Delta t)^2)$ for the Gaussian envelope, but this would center the atom at $t=0$ regardless of $t_c$.

Components

  • Gaussian envelope: $\frac{1}{\sqrt{2\pi}\Delta t} \exp\left(-\frac{1}{2}\left(\frac{t-t_c}{\Delta t}\right)^2\right)$ provides time localization centered at $t_c$
  • Chirp modulation: $\exp\left(j2\pi\left[c(t-t_c)^2 + f_c(t-t_c)\right]\right)$ encodes frequency sweep
    • Linear term $f_c(t-t_c)$: base frequency
    • Quadratic term $c(t-t_c)^2$: frequency chirp rate

How It Works

The implementation uses a hybrid coarse-to-fine optimization strategy.

  1. THRML Coarse Sampling: Explores discrete parameter space (32 levels per parameter) using probabilistic graphical model with 20 binary spin nodes. Samples 5000 configurations to find promising regions.
  2. BFGS Refinement: Gradient-based optimization from top THRML candidates (typically 10 candidates refined with L-BFGS).
  3. Matching Pursuit: Extracts components iteratively from residual signal until convergence threshold is met.

When to Use THRML

THRML provides maximum benefit when

  • Highly multimodal signals (overlapping chirps with similar characteristics)
  • Noisy measurements (SNR < 10 dB) where gradient estimates are unreliable

Skip THRML when,

  • Real-time processing (< 10ms latency) - use pure BFGS
  • High-SNR signals (> 30 dB) with well-separated components

Implementation Notes

  • Discretization: 32 levels (5 bits) per parameter provides decent balance between resolution and tractability
  • Sampling: Requires 500-1000 warmup iterations for convergence
  • Memory: On-the-fly chirplet generation recommended for large signals to avoid 4GB+ GPU memory usage

Dictionary Search Strategy Comparison

C++ ACT (Exhaustive Search)

The reference C++ implementation uses exhaustive dictionary search

// GEMV: scores = A^T * signal  (600k atoms × signal length)
act::blas::gemv_colmajor_trans(m, n, alpha, A, x, ...);
int best_idx = cblas_isamax(n, scores, 1);  // Find max correlation

Python THRML (Probabilistic Sampling)

This implementation uses probabilistic sampling via THRML

# Sample parameter space probabilistically
sampled_params = sample_chirplet_parameters(...)  # 3000 samples via MCMC
correlations = [compute_correlation(residual, generate_chirplet(p)) 
                for p in sampled_params]
top_candidates = top_k(correlations)  # Select best 10-12

Performance Comparison

Implementation Dictionary Size Search Method Speedup
C++ CPU (BLAS) 600k atoms GEMV (all atoms) 2-3.5× real-time (Mac M1)
C++ GPU (MLX) 600k atoms Matrix multiply 31-54× real-time (M1/RTX4090)
Python THRML ~3k samples Probabilistic Not measured

However, the search strategy is fundamentally different

  • C++ ACT: Exhaustive dictionary search (optimized via BLAS/GPU)
  • Python THRML: Probabilistic sampling (optimized via THRML+BFGS hybrid)

The Python implementation prioritizes sample efficiency (fewer evaluations) over raw throughput (GPU-accelerated exhaustive search), making it suitable for exploration and research rather than real-time processing.

References

About

Implementation of the core Adaptive Chirplet Transform (ACT) algorithm using THRML (Thermodynamic sampling) to efficiently find the best matching atom in the continuous space

Topics

Resources

Stars

Watchers

Forks

Contributors

Languages