pdft

Learning parametric quantum Fourier transforms via manifold optimization, in JAX.

pdft approximates the discrete Fourier transform with a trainable, parameterized quantum circuit and optimizes it on the unitary manifold so an image becomes sparse in the learned basis. It is a faithful Python port of ParametricDFT.jl, with results checked against committed Julia goldens, and the reference implementation for the paper Fast Trainable Multilinear Bases for Image Compression.

API reference

Bases, Riemannian optimizers, losses, training loops, I/O and coherence tools, generated from the source docstrings.

API reference

Example gallery

Short runnable scripts that train bases and compare optimizers, with their loss curves rendered at build time.

Example gallery

Paper & citation

The arXiv paper this package accompanies, and the BibTeX entry to cite it.

Paper & citation

Installation

From PyPI (Python 3.11+):

pip install "pdft>=0.2.3"

Note: the older pdft==0.2.2 wheel predates DCT4Basis and the parametrization="u4" option of TEBDBasis / MERABasis, so it cannot run the paper’s DCT-IV, TEBD-U4, or MERA-U4 configurations. If pdft.__version__ reports 0.2.2, upgrade with pip install -U pdft.

From source:

git clone https://github.com/zazabap/pdft.git
cd pdft
pip install -e ".[dev]"

Optional extras: pdft[plot] adds matplotlib for the plotting helpers in pdft.viz, and pdft[gpu] installs a CUDA 12 build of JAX.

Note

Importing pdft turns on JAX’s 64-bit mode for the whole process. Julia parity depends on complex128 arithmetic, so import pdft before you create any JAX arrays.

Quick start

Train a parametric QFT basis on a target image with Riemannian gradient descent:

import jax
import jax.numpy as jnp
import pdft

target = jax.random.normal(jax.random.PRNGKey(7), (4, 4)).astype(jnp.complex128)
basis = pdft.QFTBasis(m=2, n=2)

result = pdft.train_basis(
    basis,
    target=target,
    loss=pdft.L1Norm(),
    optimizer=pdft.RiemannianGD(lr=0.01),
    steps=50,
    seed=0,
)
print(result.loss_history[0], "->", result.loss_history[-1])

What’s in the package

Area	Contents
Circuit bases	QFTBasis, EntangledQFTBasis, TEBDBasis, MERABasis, RichBasis, RealRichBasis, DCT4Basis
Block bases	BlockedBasis, freeze_as_blocked()
Optimizers	RiemannianGD (Armijo line search), RiemannianAdam
Losses	L1Norm, MSELoss with top-k truncation
Training	train_basis() (single target), train_basis_batched() (multi-image, cosine schedule, early stopping)
Coherence	coherence(), certify_flat_modulus()
I/O	pdft.io: JSON serialization compatible with Julia, and top-k compression