Looking at the documentation, it's a pretty simple function, which is pretty easy to write in Python. I renamed the function to add the missing ‘e’ because it was annoying me. Anyway:
def quantize(signal, partitions, codebook):
indices = []
quanta = []
for datum in signal:
index = 0
while index < len(partitions) and datum > partitions[index]:
index += 1
indices.append(index)
quanta.append(codebook[index])
return indices, quanta
Trying it with the example in the documentation:
>>> index, quants = quantize([3, 34, 84, 40, 23], range(10, 90, 10), range(10, 100, 10))
>>> index
[0, 3, 8, 3, 2]
>>> quants
[10, 40, 90, 40, 30]
For a slightly more efficient but less flexible version, we can bypass the ranges and just use mathematics:
from __future__ import division
import math
def opt_quantize(signal, num_quanta, partition_start, partition_step,
codebook_start, codebook_step):
indices = []
quanta = []
for datum in signal:
index = int(math.floor((datum - partition_start) / partition_step + 1))
if index < 0:
index = 0
if index >= num_quanta:
index = num_quanta - 1
indices.append(index)
quanta.append(codebook_start + codebook_step * index)
return indices, quanta
Trying it with the example in the documentation:
>>> index, quants = opt_quantize([3, 34, 84, 40, 23], 9, 10, 10, 10, 10)
>>> index
[0, 3, 8, 4, 2]
>>> quants
[10, 40, 90, 50, 30]
So the results are a tiny bit difference in the case where a datum is exactly on a partition due to floating-point error, but it works if nothing's on a partition.
So that reduces the running time, where n is the length of the signal and m is the number of partitions from O(mn) to O(n). That should give you a significant performance boost. Can we do better?
Yes. With our new mathematics-based approach, the code is easily vectorized, and we can make Numpy do the hard work:
import numpy as np
def np_quantize(signal, num_quanta, partition_start, partition_step,
codebook_start, codebook_step):
signal = np.asarray(signal, dtype=float)
indices = np.empty_like(signal, dtype=int)
np.floor_divide((signal - partition_start + partition_step), \
partition_step, indices)
np.clip(indices, 0, num_quanta - 1, indices)
quanta = np.asarray(indices, dtype=float) * codebook_step + codebook_start
return indices, quanta
I did, by chance, benchmark it, and it appears that each of my optimizations has made it slower, so either I'm doing something horribly wrong or I'm not testing on data large enough to amortize the constant.
~$ python -m timeit -s 'from quantize import orig_quantize' 'orig_quantize([-3, -2, -1, 0, 1, 2, 3], [-0.5, 0.5], [-1, 0, 1])'
100000 loops, best of 3: 8.58 usec per loop
~$ python -m timeit -s 'from quantize import opt_quantize' 'opt_quantize([-3, -2, -1, 0, 1, 2, 3], 3, -0.5, 1, -1, 1)'
100000 loops, best of 3: 10.8 usec per loop
~$ python -m timeit -s 'from quantize import np_quantize' 'np_quantize([-3, -2, -1, 0, 1, 2, 3], 3, -0.5, 1, -1, 1)'
10000 loops, best of 3: 57.4 usec per loop
For kicks, I tried using Cython as well as Numpy:
cimport cython
cimport numpy as np
cdef extern from "math.h":
float floorf(float)
@cython.boundscheck(False)
def cynp_quantize(np.ndarray[float, ndim=1] signal, int num_quanta,
float partition_start, float partition_step,
float codebook_start, float codebook_step):
cdef int i
cdef int index
cdef np.ndarray[np.int_t, ndim=1] indices = np.empty_like(signal, dtype=int)
cdef np.ndarray[float, ndim=1] quanta = np.empty_like(signal)
for i in range(signal.shape[0]):
index = <int>floorf((signal[i] - partition_start)
/ partition_step + 1.0)
if index < 0:
index = 0
if index >= num_quanta:
index = num_quanta - 1
indices[i] = index
quanta[i] = codebook_start + index * codebook_step
return indices, quanta
From what I gather, Cython also experimentally supports OpenMP, which would let it do everything with multiple threads. I was unable to test the performance of this Cython solution, though, with or without threads (I'm missing a header file needed to compile the result).