Trying to understand how to optimize rescaling with weave

weave_rescale.py

#vim: fileencoding=utf-8
import time
import timeit
import numpy as np
from scipy import weave


dtype = np.int8


def rescale_np(data, scale, offset):
    return ((data - offset) * scale).astype(dtype)


def rescale_np_savealloc(data, scale, offset):
    aux = data - offset
    aux *= scale
    return aux.astype(dtype, copy=False)


def rescale_weave(data, scale, offset):
    size = data.size
    input = np.ascontiguousarray(data)
    output = np.empty((size,), dtype=dtype)
    code = ("for (int i = 0; i < size; ++i) "
            "{ output[i] = (((double)input[i] - offset) * scale); }")
    weave.inline(code, ["size", "input", "output", "offset", "scale"],
                 compiler="gcc", extra_compile_args=["-march=native -mtune=native -O3"])
    return output.reshape(data.shape)


def main():
    generate_code = "np.empty((512, 512))"
    code = "rescale_{{0}}({}, 1.23, 0.123)".format(generate_code)
    setup = "import numpy as np; from __main__ import rescale_{0}; " + code
    number = 100
    timer = time.clock
    def print_timeit(*args, **kwargs):
        times = timeit.repeat(*args, **kwargs)
        print "{:.2f} µs".format(min(times) / number * 1e6)
    print_timeit(generate_code, setup="import numpy as np",
                 number=number, timer=timer) # time used for data generation
    print_timeit(code.format("np"), setup=setup.format("np"),
                 number=number, timer=timer)
    print_timeit(code.format("np_savealloc"), setup=setup.format("np_savealloc"),
                 number=number, timer=timer)
    print_timeit(code.format("weave"), setup=setup.format("weave"),
                 number=number, timer=timer)


def check():
    t = np.random.random((8,))
    print(t)
    print(rescale_np(t, 10, -1))
    print(rescale_weave(t, 10, -1))


if __name__ == "__main__":
    main()
    check()