KDExpress 1.0.1

Fast Kernel Density Estimation with FFT and JAX

Kernel Density Estimation accelerated with Fast Fourier Transform and JAX (JIT+GPU support+AD).

Features

• FFT-accelerated KDE (inspired by KernelDensity.jl):

  • fft_kde1d: 1D KDE with FFT convolution (requires regular grid)
  • fft_kde2d: 2D KDE with FFT convolution (requires regular grid)
  • fft_kde3d: 3D KDE with FFT convolution (requires regular grid)

  The KDE can be evaluated on non-regular grid by interpolation, such as jax.numpy.interp or jax.scipy.interpolate.RegularGridInterpolator

• Binned-accelerated implementation:

  • binned_kde1d: 1D KDE with data binning

• Bandwidth estimators:

  • Scott's rule (scott_bw1d/2d/3d)
  • Silverman's rule (silverman_bw1d/2d/3d)

Installation

1. Install JAX following the instructions at: https://github.com/jax-ml/jax, e.g

   pip install "jax[cpu]" # For CPU-only version
   

   or

   pip install "jax[cuda12]" # For GPU support (CUDA)
   

2. Install with pip

   pip install KDExpress
   

3. Alternative way:

   • Clone the KDExpress repo

     git clone https://github.com/mtagliazucchi/KDExpress
     

   • Install the code

     a. Editable install:

     cd KDExpress
     pip install -e .
     

     b. Or use it ad-hoc:

     import sys; sys.path.append("/path/to/KDExpress")
     from KDExpress import fft_kde1d
     

Usage and benchmarks

See the examples folder.

License

MIT © Matteo Tagliazucchi