modified setup and short dot

setup.py

@@ -1,6 +1,16 @@
+#!/usr/bin/env python
+
 from distutils.core import setup
+from distutils.extension import Extension
 from Cython.Build import cythonize
+from Cython.Distutils import build_ext
+
+import numpy
+
 
 setup(
-    ext_modules = cythonize("shortdot.pyx")
+    cmdclass = {'build_ext': build_ext},
+    ext_modules = [Extension("shortdot",
+                         sources=["shortdot.pyx"],
+                         include_dirs=[numpy.get_include()])],
 )

shortdot.pyx

@@ -4,9 +4,10 @@ cimport numpy as np
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef np.ndarray[double, ndim=1] shortdot(
+def shortdot(
     np.ndarray[float, ndim=2] A,
     np.ndarray[float, ndim=1] B,
+    np.ndarray[double, ndim=1] C,
     float threshold = -3.0):
     '''
     Matrix-vector multiplication.
@@ -20,14 +21,11 @@ cdef np.ndarray[double, ndim=1] shortdot(
         int A_m = A.shape[1]
         int B_n = B.shape[0]
         double tot = 0.0
-        np.ndarray[double, ndim=1] C
 
     # Initialize the results matrix.
-    C = np.zeros(A_n)
     for i in xrange(A_n):
         for j in xrange(A_m):
             tot += A[i, j] * B[j]
             if j > 10 and tot < threshold:
                 break
         C[i] = tot
-    return C