veloxquant_mlx.math package¶
Pure-numpy math helpers shared by the codebook, quantizer, and preconditioner layers.
Groups three independent pieces: closed-form/numerically-integrated PDFs
(distributions) used as Lloyd-Max inputs, the
Lloyd-Max scalar-quantizer solver itself
(lloyd_max), and random orthogonal/JL matrix
generators (rotation) used to build
preconditioners. Re-exports beta_pdf, gaussian_pdf,
polar_angle_pdf, lloyd_max, make_jl_matrix, and
make_rotation_matrix.
- veloxquant_mlx.math.beta_pdf(x, d)[source]
PDF of a single coordinate of a d-dimensional unit-sphere uniform sample.
After a random orthogonal rotation, each coordinate of a d-dimensional vector drawn uniformly from the unit sphere follows:
f_X(x) = Γ(d/2) / (√π · Γ((d-1)/2)) · (1 - x²)^((d-3)/2)
on [-1, 1]. In high dimensions (d ≥ 64), this converges to N(0, 1/d).
- veloxquant_mlx.math.gaussian_pdf(x, sigma=1.0)[source]
PDF of a zero-mean Gaussian with standard deviation sigma.
f(x) = 1 / (σ√(2π)) · exp(-x²/(2σ²))
- veloxquant_mlx.math.polar_angle_pdf(psi, level)[source]
PDF of polar angle at a given recursion level after random preconditioning.
From PolarQuant (Theorem 1 in 2502.02617):
Level 1: uniform on [0, 2π). f(ψ) = 1/(2π)
Level ℓ ≥ 2: f(ψ) ∝ sin^(2^(ℓ-1) - 1)(2ψ) on [0, π/2].
- Parameters:
- Return type:
- Returns:
Normalised PDF values at each psi. Points outside the support return 0.
- Raises:
ValueError – If level < 1.
- veloxquant_mlx.math.lloyd_max(pdf_fn, support, n_levels, n_iter=500, tol=1e-09, n_quad_points=10000)[source]
Solve the Lloyd-Max 1-D scalar quantisation problem.
Iterates the Lloyd-Max optimality conditions until convergence:
Initialise centroids uniformly over support.
Loop: a. Boundaries = midpoints between adjacent centroids, extended to ±∞. b. Update centroids via numerical integration:
c_i = ∫_{b_{i-1}}^{b_i} x · f(x) dx / ∫_{b_{i-1}}^{b_i} f(x) dx
Stop when max|c_new - c_old| < tol.
Compute and cache the final MSE cost.
The resulting
last_mse_costattribute stores:C = Σ_i ∫_{b_{i-1}}^{b_i} (x - c_i)² · f(x) dx- Parameters:
pdf_fn (
Callable[[ndarray],ndarray]) – Vectorised PDF callable: pdf_fn(x) -> non-negative values.support (
tuple[float,float]) – (lo, hi) defining the distribution’s effective support.n_levels (
int) – Number of quantisation levels k = 2^b.n_iter (
int) – Maximum number of Lloyd-Max iterations.tol (
float) – Convergence tolerance on centroid shift.n_quad_points (
int) – Number of quadrature points for numerical integration.
- Return type:
- Returns:
Tuple (centroids, boundaries) both sorted ascending. - centroids: shape (n_levels,). - boundaries: shape (n_levels + 1,) including ±∞ at the edges.
- Raises:
ValueError – If n_levels < 1 or support is invalid.
- veloxquant_mlx.math.make_jl_matrix(d, m, seed=42)[source]
Generate an m×d Gaussian JL projection matrix.
Each row is drawn i.i.d. from N(0, I_d). This is the correct construction for the QJL sign-based inner product estimator to be unbiased:
E[sqrt(pi/2)/m * ||k|| * sum_i sign(s_i·k)(s_i·q)] = <q, k>
Unlike orthogonal JL, Gaussian JL allows m > d.
- Parameters:
- Return type:
- Returns:
Float64 array of shape (m, d) with i.i.d. N(0,1) entries.
- Raises:
ValueError – If d < 1 or m < 1.
- veloxquant_mlx.math.make_rotation_matrix(d, seed=42)[source]
Generate a d×d random orthogonal rotation matrix via QR decomposition.
Draws G ~ N(0, 1)^{d×d} and returns Q from G = QR (economy QR). The resulting Q is a Haar-distributed orthogonal matrix.
- Parameters:
- Return type:
- Returns:
Float64 array of shape (d, d) with orthonormal rows (Q @ Q.T ≈ I).
- Raises:
ValueError – If d < 1.