How do you optimize this code (without vectorizing, as this leads up to using the semantics of the calculation, which is quite often far from being non-trivial):
slow_lib.py:
import numpy as np
def foo():
size = 200
np.random.seed(1000031212)
bar = np.random.rand(size, size)
moo = np.zeros((size,size), dtype = np.float)
for i in range(0,size):
for j in range(0,size):
val = bar[j]
moo += np.outer(val, val)
The point is that such kind loops correspond quite often to operations where you have double sums over some vector operation.
This is quite slow:
>>t = timeit.timeit('foo()', 'from slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 41.165681839
Ok, so then let's cynothize it and add type annotations likes there is no tomorrow:
c_slow_lib.pyx:
import numpy as np
cimport numpy as np
import cython
#cython.boundscheck(False)
#cython.wraparound(False)
def foo():
cdef int size = 200
cdef int i,j
np.random.seed(1000031212)
cdef np.ndarray[np.double_t, ndim=2] bar = np.random.rand(size, size)
cdef np.ndarray[np.double_t, ndim=2] moo = np.zeros((size,size), dtype = np.float)
cdef np.ndarray[np.double_t, ndim=1] val
for i in xrange(0,size):
for j in xrange(0,size):
val = bar[j]
moo += np.outer(val, val)
>>t = timeit.timeit('foo()', 'from c_slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 42.3104710579
... ehr... what? Numba to the rescue!
numba_slow_lib.py:
import numpy as np
from numba import jit
size = 200
np.random.seed(1000031212)
bar = np.random.rand(size, size)
#jit
def foo():
bar = np.random.rand(size, size)
moo = np.zeros((size,size), dtype = np.float)
for i in range(0,size):
for j in range(0,size):
val = bar[j]
moo += np.outer(val, val)
>>t = timeit.timeit('foo()', 'from numba_slow_lib import foo', number = 10)
>>print("took: "+str(t))
took: 40.7327859402
So is there really no way to speed this up? The point is:
if I convert the inner loop into a vectorized version (building a larger matrix representing the inner loop and then calling np.outer on the larger matrix) I get much faster code.
if I implement something similar in Matlab (R2016a) this performs quite well due to JIT.
Here's the code for outer:
def outer(a, b, out=None):
a = asarray(a)
b = asarray(b)
return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis,:], out)
So each call to outer involves a number of python calls. Those eventually call compiled code to perform the multiplication. But each incurs an overhead that has nothing to do with the size of your arrays.
So 200 (200**2?) calls to outer will have all that overhead, whereas one call to outer with all 200 rows has one overhead set, followed by one fast compiled operation.
cython and numba don't compile or otherwise bypass the Python code in outer. All they can do is streamline the iteration code that you wrote - and that isn't consuming much time.
Without getting into details, the MATLAB jit must be able to replace the 'outer' with faster code - it rewrites the iteration. But my experience with MATLAB dates from a time before its jit.
For real speed improvements with cython and numba you need to use primitive numpy/python code all the way down. Or better yet focus your effort on slow inner pieces.
Replacing your outer with a streamlined version cuts run time about in half:
def foo1(N):
size = N
np.random.seed(1000031212)
bar = np.random.rand(size, size)
moo = np.zeros((size,size), dtype = np.float)
for i in range(0,size):
for j in range(0,size):
val = bar[j]
moo += val[:,None]*val
return moo
With the full N=200 your function took 17s per loop. If I replace the inner two lines with pass (no calculation), time drops to 3ms per loop. In other words, the outer loop mechanism is not a big time consumer, at least not compared to many calls to outer().
Memory permitting, you can use np.einsum to perform those heavy calculations in a vectorized manner, like so -
moo = size*np.einsum('ij,ik->jk',bar,bar)
One can also use np.tensordot -
moo = size*np.tensordot(bar,bar,axes=(0,0))
Or simply np.dot -
moo = size*bar.T.dot(bar)
Many tutorials and demonstrations of Cython, Numba, etc. make it seem as if these tools can speed up your code automagically, but in practice, this is often not the case: You'll need to modify your code a little to extract the best performance. If you had already implemented some degree of vectorization, it usually means writing out ALL the loops. Reasons Numpy array operations are non-optimal include:
Lots of temporary arrays are created and looped over;
Significant per-call overhead if the arrays are small;
Short-circuiting logic can't be implemented, because arrays are processed as a whole;
Sometimes the optimal algorithm can't be expressed using array expressions and you settle for an algorithm with a worse time complexity.
Using Numba or Cython wont optimize these problems away! Instead, these tools allow you to write loopy code that is much faster than plain Python.
Also, for Numba specifically, you should be aware of the difference between "object mode" and "nopython mode". The tight loops from your example have to run in nopython mode to provide any significant speedup. However, numpy.outer is not yet supported by Numba, resulting in the function to be compiled in object mode. Decorate with jit(nopython=True) to let such cases throw an exception.
Example to demonstrate a speedup is indeed possible:
import numpy as np
from numba import jit
#jit
def foo_nb(bar):
size = bar.shape[0]
moo = np.zeros((size, size))
for i in range(0,size):
for j in range(0,size):
val = bar[j]
moo += np.outer(val, val)
return moo
#jit
def foo_nb2(bar):
size = bar.shape[0]
moo = np.zeros((size, size))
for i in range(size):
for j in range(size):
for k in range(0,size):
for l in range(0,size):
moo[k,l] += bar[j,k] * bar[j,l]
return moo
size = 100
bar = np.random.rand(size, size)
np.allclose(foo_nb(bar), foo_nb2(bar))
# True
%timeit foo_nb(bar)
# 1 loop, best of 3: 816 ms per loop
%timeit foo_nb2(bar)
# 10 loops, best of 3: 176 ms per loop
The example you show us is kind of inefficient algorithm, since you calculate the same outer product multiple times. The resulting time complexity is O(n^4). It can be reduced to n^3.
for i in range(0,size):
val = bar[i]
moo += size * np.outer(val, val)
Related
I am trying to loop over 2 2d arrays in Cython. The arrays have the following shape:
ranges_1 is a 6000x3 array of int64, while ranges_2 is a 2000x2 of int64. This iteration needs to be perfomed around 10000 times. This would mean that the total number of calculations inside the nested for loop would be around 2000x6000x10000 = 120 billions times.
This is the code I am using to generate the "dummy" data:
import numpy as np
ranges_1 = np.stack([np.random.randint(0, 10_000, 6_000), np.random.randint(0, 10_000, 6_000), np.arange(0, 6_000)], axis=1)
ranges_2 = np.stack([np.random.randint(0, 10_000, 2_000), np.random.randint(0, 10_000, 2_000)], axis=1)
Which gives 2 arrays like these:
array([[6131, 1478, 0],
[9317, 7263, 1],
[7938, 6249, 2],
...,
[5153, 426, 5997],
[9164, 9211, 5998],
[1695, 1792, 5999]])
and:
array([[ 433, 558],
[3420, 2494],
[6367, 7916],
...,
[8693, 1692],
[1256, 9013],
[4096, 1860]])
The first implementation I tried is a "naive" version and it's the following (the function inside is just a test function which uses all the data in the array):
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func(np.ndarray[DTYPE_t, ndim = 2] ranges_1, np.ndarray[DTYPE_t, ndim=2] ranges_2, int n ):
k = 0
for i in range(n):
for j in range(len(ranges_1)):
r1 = ranges_1[j]
a = r1[0]
b = r1[1]
c = r1[2]
for f in range(len(ranges_2)):
r2 = ranges_2[f]
d = r2[0]
e = r2[1]
k = (a + b + c + d + e)/(d+e)
return k
This takes about 5 seconds per each of the 10_000 outside loops.
So I then tried flatteing out the arrays and, since I know the dimension on the other axis, accessing the items like this:
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func_flattened(np.ndarray[DTYPE_t, ndim = 1] ranges_1_, np.ndarray[DTYPE_t, ndim=1] ranges_2_, int n ):
k = 0
for i in range(n):
for j in range(0, len(ranges_1_), 3):
a = ranges_1_[j]
b = ranges_1_[j+1]
c = ranges_1_[j+2]
for f in range(0, len(ranges_2_), 2):
d = ranges_2_[f]
e = ranges_2_[f+1]
k = (a + b + c + d + e)/(d+e)
return k
But that provided no speed up at all. The time to perform one single iteration of the 10_000 seems too high, considered that for one single iteration it's just 12_000_000 operations inside the loop. I also tried implementing a much simpler example both in Cython and in python which was then compiled with numba:
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_1(int n ):
cdef k = 0
cdef a = 0
for i in range(n):
a = i +1
return a
This took 15s to run with n = 1_000_000_000.
While with numba:
def test_1_python(n ):
k = 0
a = 0
for i in range(n):
if i % 2 == 0:
a = a + 1
else:
a = a - 1
return a
test_1_numba= numba.jit(test_1_python)
%%time
test_1_numba(120_000_000_000)
The full run with n = 120bln took about 6s, (albeit the function inside is simpler) this means that it would be 500 times faster than Cython, could this be possible?
I am new to Cython so I probably am missing something obvious, but since the numba version (without the array accessing) is tha much faster I think that the difference in speed might come from the overhead associated with accessing the items in the array.
Is this a wrong assumption?
If not, what would be the best of going about looping over a 2D list of integers in Cython?
What you measure in your benchmarks is mainly compilation artefacts and overheads.
First of all, Cython use the default compiler installed on your machine preferred by the Python stack. On Linux, it should be GCC. On Windows, it is certainly MSVC if installed, otherwise MinGW (if any). Meanwhile, Numba is based on LLVM-Lite which is based on the LLVM stack like Clang. Thus, in your case, it is very likely that different compilers are used resulting in different binary with different performance. If you want to make a fair benchmark, you need to use Clang to build you Cython program.
Additionally, the default optimization for Cython is -O2 while it is -O3 for Numba. The former should not enable auto-vectorization while the later does (this is dependent of the target compilers -- newer version of GCC changes this behaviour). Furthermore, Cython does not enable machine-specific non-portable optimizations by default (since binaries may be packaged for other machines like with pip). This means Cython can only use the old SSE2 SIMD instruction set by default on x86-64 processors. Meanwhile, LLVM-JIT can use of the much faster AVX2/AVX-512 SIMD instruction set. You need to enable such optimization manually with Cython so for the benchmark to be fair (ie. -march=native on GCC/Clang).
In fact, on my x86-64 mainstream Intel machine, Numba does use the AVX2 instruction set and Cython does not on your last benchmark. Here is for example the main loop generated by the Numba JIT:
.LBB0_7:
vptestnmq %ymm5, %ymm4, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm16, %ymm4, %ymm18
vptestnmq %ymm5, %ymm18, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm17, %ymm4, %ymm4
addq $-2, %rdx
jne .LBB0_7
For the benchmark doing a = i + 1 in a loop, it is flawed as a good compiler can just optimize out the whole loop (ie. remove it) and replace it with only one assignment since only the last iteration matter. In fact, the same thing applies with k = (a + b + c + d + e)/(d+e): only the last iteration matters. The i variable of for i in range(n) is not even used. Clang and GCC often do such kind of optimizations.
Finally, the speed of your initial first code will be memory-bound if it is modified to compute something meaningful in a real-world use-case and multiple threads are used.
Note that division are very expensive and you can precompute the reciprocal so to perform multiplications instead in your main loop.
I have been playing with jax lately, and it is very impressive, but then the following set of experiments confused me greatly:
First, we set up the timer utility:
import time
def timefunc(foo, *args):
tic = time.perf_counter()
tmp = foo(*args)
toc = time.perf_counter()
print(toc - tic)
return tmp
Now, let’s see what happens when we compute the eigenvalues of a random symmetric matrix matrix, thus (jnp is jax.numpy, so the eigh is done on the GPU)
def jfunc(n):
tmp = np.random.randn(n, n)
return jnp.linalg.eigh(tmp + tmp.T)
def nfunc(n):
tmp = np.random.randn(n, n)
return np.linalg.eigh(tmp + tmp.T)
Now for the timings (the machine is an nVidia DGX box, so the GPU is an A100, while the CPUs are some AMD EPYC2 parts.
>>> e1 = timefunc(nfunc, 10)
0.0002442029945086688
>>> e2 = timefunc(jfunc, 10)
0.013523647998226807
>>> e1 = timefunc(nfunc, 100)
0.11742364699603058
>>> e2 = timefunc(jfunc, 100)
0.11005625998950563
>>> e1 = timefunc(nfunc, 1000)
0.6572738009999739
>>> e2 = timefunc(jfunc, 1000)
0.5530761769914534
>>> e1 = timefunc(nfunc, 10000)
36.22587636699609
>>> e2 = timefunc(jfunc, 10000)
8.867857075005304
You will notice that the crossover is somewhere around 1000. Initially, I thought this was because of the overhead of moving stuff to/from the GPU, but if you define yet another function:
def jjfunc(n):
key=jax.random.PRNGKey(0)
tmp = jax.random.normal(key, [n, n])
return jnp.linalg.eigh(tmp + tmp.T)
>>> e1=timefunc(jjfunc, 10)
0.01886096798989456
>>> e1=timefunc(jjfunc, 100)
0.2756766739912564
>>> e1=timefunc(jjfunc, 1000)
0.7205733209993923
>>> e1=timefunc(jjfunc, 10000)
6.8624101399909705
Note that the small examples are actually (much) slower than moving the numpy array to the GPU and back.
So, my question is: what is going on, and is there a silver bullet? Is this a jax implementation bug?
I don't think your timings are reflective of actual JAX vs. numpy performance, for a few reasons:
JAX's computation model uses Asynchronous Dispatch, which means that JAX operations return before the computation is finished. As mentioned at that link, you can use the block_until_ready() method to ensure you are timing the computation rather than the dispatch.
Because operations like eigh are JIT-compiled by default, the first time you run them for a given size will incur the one-time compilation cost. Subsequent runs will be faster as JAX caches previous compilations.
Your computations are indeed being foiled by device transfer costs. It's easiest to see if you measure it directly:
def transfer(n):
tmp = np.random.randn(n, n)
return jnp.array(tmp).block_until_ready()
timefunc(transfer, 10000);
# 4.600406924000026
Your jjfunc combines the eigh call with the jax.random.normal call. The latter is slower than numpy's random number generation, and I believe is dominating the difference for small n.
Unrelated to JAX, but in general using time.time for profiling Python code can give you misleading results. Modules like timeit are much better for this kind of thing, particularly when you're dealing with microbenchmarks that complete in fractions of a second.
If you're interested in accurate benchmarks of JAX vs. Numpy versions of algorithms, I'd suggest isolating exactly the operations you're interested in benchmarking (i.e. generate the data & do any device transfer outside the benchmarks). Read up on the advice in Asynchronous Dispatch in JAX as it relates to benchmarking, and check out Python's timeit Docs for tips on getting accurate timings of small code snippets (though I find the %timeit magic more convenient if working IPython or Jupyter notebook).
Any suggestions on how to speed up this function?
def smooth_surface(z,c):
hph_arr_list = []
for x in xrange(c,len(z)-(c+1)):
new_arr = np.hstack(z[x-c:x+c])
hph_arr_list.append(np.percentile(new_arr[((new_arr >= np.percentile(new_arr,15)) & (new_arr <= np.percentile(new_arr,85)))],99))
return np.array(map(float,hph_arr_list))
The variable z has a length of ~15 million and c is value of window size + and -. The function is basically a sliding window that calculates a percentile value per iteration. Any help would be appreciated! z is an array of arrays (hence the np.hstack). Maybe any idea if numba would help with this. If so, how to implement?
The slow part of the computation appear to be the line np.percentile(new_arr[((new_arr >= np.percentile(new_arr,15)) & (new_arr <= np.percentile(new_arr,85)))],99). This is due to the unexpectedly slow np.percentile on small arrays as well as the creation of several intermediate arrays.
Since new_arr is actually quite small, it is much faster to just sort it and make the interpolation yourself. Moreover, numba can also help to speed the computation up.
#njit #Use #njit instead of #jit to increase speed
def filter(arr):
arr = arr.copy() # This line can be removed to modify arr in-place
arr.sort()
lo = int(math.ceil(len(arr)*0.15))
hi = int(len(arr)*0.85)
interp = 0.99 * (hi - 1 - lo)
interp = interp - int(interp)
assert lo <= hi-2
return arr[hi-2]* (1.0 - interp) + arr[hi-1] * interp
This code is 160 times faster with arrays of size 20 on my machine and should produce the same result.
Finally, you can speed up smooth_surface too by using automatic parallelization in numba (see here for more information). Here is an untested prototype:
#jit(parallel=True)
def smooth_surface(z,c):
hph_arr = np.zeros(len(z)-(c+1)-c)
for x in prange(c,len(z)-(c+1)):
hph_arr[x-c] = filter(np.hstack(z[x-c:x+c]))
return hph_arr
I am doing a statistical calculation of the following form:
where
d (data) is a matrix of size [i=30, 100x100]
m (model) is a matrix of size [i=30, 100x100]
C-1 (covariance) is a 1002 by 1002 symmetric matrix
d and C-1 are constants and do not have any structure apart from C-1 being symmetric. m changes, but is always sparse. The output of the calculation is just a floating-point number.
I need to perform this calculation many times in a Monte Carlo simulation, so speed is of the essence. Using sparse array multiplication techniques with m speeds things up considerably over a naive matrix dot product. However, each iteration of the slow function below still takes about 0.1 seconds to run. The vast majority of time (>98%) is spent in the matrix multiplications, and not the model generation function (generate_model). I would like to speed this up by an order of magnitude if possible.
The code and output are pasted below.
Things that do not work include:
Upgrading to the Intel MKL linear algebra routines (few percent speedup, surprisingly small)
Using numpy.linalg.multi_dot
Taking advantage of the fact that C-1 is symmetric (this does not work even in principle, see this mathoverflow question)
Things that sort of work include:
Precalculating C-1d, gives a ~40% speedup
How can I speed up this code? Solutions relying on packages like cython, numba, etc are most welcome as well as "standard" scipy/numpy solutions. Thanks in advance!
from __future__ import division
import numpy as np
import scipy.sparse
import sys
import timeit
def generate_model(n, size, hw = 8):
#model for the data--squares at random locations
output = np.zeros((n, size, size))
for i in range(n):
randx = np.random.randint(hw, size-hw)
randy = np.random.randint(hw, size-hw)
output[i,(randx-hw):(randx+hw), (randy-hw):(randy+hw)]=np.random.random((hw*2, hw*2))
return output
def slow_function(datacube, invcovmatrix, size):
model = generate_model(30, size)
output = 0
for i in range(model.shape[0]):
data = datacube[i,:,:].flatten()
mu = model[i,:,:].flatten()
sparsemu = scipy.sparse.csr_matrix(mu)
output += -0.5* (
np.float(-2.0*sparsemu.dot(invcovmatrix).dot(data)) +
np.float(sparsemu.dot(sparsemu.dot(invcovmatrix).T))
)
return output
def wrapper(func, *args, **kwargs):
def wrapped():
return func(*args, **kwargs)
return wrapped
if __name__ == "__main__":
size = 100
invcovmat = np.random.random((size**2, size**2))
#make symmetric for consistency
invcovmat = (invcovmat+invcovmat.T)/2
datacube = np.random.random((30, size, size))
#TIMING
wrapped = wrapper(slow_function, datacube, invcovmat, size)
times = []
for i in range(20):
print i
times.append(timeit.timeit(wrapped, number = 1))
times.sort()
print '\n', np.mean(times[0:3]), ' s/iteration; best of 3'
Output:
0.10408163070678711 s/iteration; best of 3
I am trying to figure out if Python/Numpy is a viable alternative to develop my numerical software which is already available in C++. In order to get performance in Python/Numpy, one need to "vectorize" the code. But it turns out that as soon as I move away from very simple examples, I struggle to vectorize the code (I am not talking about SIMD instructions but "efficient Numpy code" without loops). Here is an algorithm that I want to get efficiently in Python/Numpy.
Create an numpy array containing: 1.0, 1.0 + 1/n, 1.0 + 2/n, ..., 2.0
For every u in the array, compute the root of x^2 - u, using a Newton method, stopping when |dx| <= 1.0e-7. Store the result in an array result.
Sum all the elements of the result array
Here is the algorithm in Python I want to speed up
import numpy as np
n = 1000000
data = np.arange(1.0, 2.0, 1.0 / n)
def newton(u):
x = 2.0
while True:
f = x**2 - u
df_dx = 2 * x
dx = f / df_dx
if (abs(dx) <= 1.0e-7):
break
x -= dx
return x
result = map(newton, data)
print result[n - 1]
Here is a version of the algorithm in C++11
#include <iostream>
#include <vector>
#include <cmath>
int main (int argc, char const *argv[]) {
auto n = std::size_t{100000000};
auto v = std::vector<double>(n + 1);
for(size_t k = 0; k < v.size(); ++k) {
v[k] = 1.0 + static_cast<double>(k) / n;
}
auto result = std::vector<double>(n + 1);
for(size_t k = 0; k < v.size(); ++k) {
auto x = double{2.0};
while(true) {
auto f = double{x * x - v[k]};
auto df_dx = double{2 * x};
auto dx = double{f / df_dx};
if (std::abs(dx) <= 1.0e-7) {
break;
}
x -= dx;
}
result[k] = x;
}
auto somme = double{0.0};
for(size_t k = 0; k < result.size(); ++k) {
somme += result[k];
}
std::cout << somme << std::endl;
return 0;
}
It takes 2.9 seconds to run on my machine. Is there a way to make a fast Python/Numpy algorithm that does the same thing (I am willing to get something that is less than 5 times slower).
Thanks.
You can do step 1. with numpy efficiently:
1.0 + np.arange(n + 1) / n
however I think you would need the np.vectorize() method to feed back x into your calculated values and it's not an efficient function (basically a wrapper for a python loop). If you can use scipy then there are built in methods that might do what you want http://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.optimize.newton.html
EDIT: Having thought a bit more about this I followed up on #ev-br's point and tried some alternatives. The masking uses too much processing but the abs().max() is pretty fast so a compromise might be to "divide the problem into blocks" both in the 1st dimension of the array and in iteration direction. The following doesn't do too badly (< 20s) on my pretty low power laptop - certainly much faster than np.vectorize() or any of the scipy solving systems I could find. (If I set m too big it runs out of something (memory?) and grinds to a complete halt!)
n = 100000000
m = 5000000
block = 3
u = 1.0 + np.arange(n + 1) / n
x = np.full(u.shape, 2.0)
dx = np.ones(u.shape)
for i in range(0, n, m):
while np.abs(dx[i:i+m]).max() > 1.0e-7:
for j in range(block):
dx[i:i+m] = (x[i:i+m] ** 2 - u[i:i+m]) / (2 * x[i:i+m])
x[i:i+m] -= dx[i:i+m]
Here's a toy example. Notice that often vectorization means writing your code as if you're manipulating numbers, and letting numpy do its magic:
>>> import numpy as np
>>> a = np.array([1., 2., 3.])
>>> def f(x):
... return x**2 - a, 2.*x # function and derivative
>>>
>>> def newt(f, x0):
... x = np.asarray(x0)
... for _ in range(5): # hardcode the number of iterations (I know)
... v, dv = f(x)
... x -= v / dv
... return x
>>>
>>> newt(f, [1., 1., 1.])
array([ 1. , 1.41421356, 1.73205081])
If this is a performance bottleneck, this is unlikely to be competetive with hand-written C++ code: First of all, you're manipulating python objects with all the overhead; then numpy is likely doing a bunch of array allocations under the hood.
An often viable strategy is to start by writing things in python/numpy, and then move bottlenecks into a compiled code --- eg Cython or C++ wrapped by Cython. In this particular case since you already have the C++ code, just wrapping it with Cython is likely easiest but YMMV.
I'm not looking to wave small snippets of code as a solution, but here's something to get you started. I have a strong suspicion that you're having troubles just declaring such an array in python without spending too much time on it, so I'll mostly help you out there.
As far as the square roots come in, please add your example python code and I'll see what I can help optimize from that point on. In my example roots and sums are found with the default numpy functions/methods.
def summing():
n = 1000000
ar = np.arange(0, n)
ar = ar/float(n)
ar = ar + np.ones(n)
sqrt = np.sqrt(ar)
return np.sum(ar)
In short, to get the starting array it's best to use a "workaround".
initialize an array ar with values `[1,2,3,....n]
divide ar with n. This gets us the 1/n, 2/n ... members
add to that an array of same dimensions that contain just the number 1.0
This gets us the full array [ 1., 1.000001, 1.000002, ..., 1.999998, 1.999999]) we're after. If I understood you right.
find square roots, sum it
Average of 10 sequential execution times is 0.018786 seconds.
Obviously I'm 6 years late to this party, but this question is a common stumbling block for people in effectively using numpy for real scientific work. The basic idea is covered in #ev-br's answer. The OP points out that the solution offered there (even modified to stop iterating when a convergence criterion is met rather than after a fixed number of iterations) takes the same number of passes for each element of u. I want to show how you can avoid that objection using pure numpy code, making explicit the mask suggestion in #ev-br's comment.
However, I also want to point out that in many real world situations, the number of passes for Newton-like iteration to converge varies so little that this general technique I illustrate here will actually slow numpy code down significantly. If the average number of iterations will be within a factor of two or three of the maximum number of iterations, you should stick with something closer to #ev-br's answer (including his first comment).
The numpy performance numbers you need to understand are these: Loops over array indices will run 200 to 500 times slower in pure numpy code than in compiled code. On the other hand, if you manage to use numpy's array syntax to avoid all index loops, you can get within about a factor of 5 of compiled speed. (The factor of 5 is partly because of memory management as #ev-br mentions, but also because optimized compiled code overlaps many different arithmetical operations inside each index loop, while numpy just performs a single arithmetic operation, storing everything back to memory after each operation.) The point is that factor of 100 difference means that it often pays to do substantial amounts of "extra" work in numpy code: Even if you do 10 times the number of floating point operations in vectorized numpy code, it will still run 10 times faster than the index-loop code that avoids the "extra" work. (Incidentally, the python map function is implemented as an interpreted index loop - it has nothing to do with numpy array operations.)
from numpy import asfarray, broadcast_arrays, arange
# Begin by defining the function to be inverted by Newton's method.
def f_dfdx(x):
x = asfarray(x) # always avoid repeated type conversions
return x**2, 2.*x
# First, the simplest algorithm to find x such that f(x)=y.
# We must supply a starting guess x0 for x.
def f_inverse0(f_dfdx, y, x0, tol=1.e-7):
y, x = broadcast_arrays(asfarray(y), asfarray(x0))
x = x.copy() # without this may clobber input x0
for npass in range(20):
f, dfdx = f_dfdx(x)
dx = (f - y) / dfdx
if (abs(dx) <= tol).all():
break # iterate all x until all have converged
x -= dx
else:
raise RuntimeError("failed to converge")
return x
# A frequently slower algorithm that avoids extra iterations.
def f_inverse1(f_dfdx, y, x0, tol=1.e-7):
y, x = broadcast_arrays(asfarray(y), asfarray(x0))
shape = x.shape
y, x = y.ravel(), x.flatten() # avoid clobbering x0
unconverged = arange(y.size)
for npass in range(20):
f, dfdx = f_dfdx(x[unconverged])
dx = (f - y[unconverged]) / dfdx
unc = abs(dx) > tol
unconverged = unconverged[unc]
if not unconverged.size:
break # iterate all x until all have converged
x[unconverged] -= dx[unc]
else:
raise RuntimeError("failed to converge")
return x.reshape(shape)
On my machine, the OP's C++ program runs in 2.03 s (1.64+0.38 user+sys). For n=100 million as for the C++ program, f_inverse0 runs in 20.4 s (4.7+15.6 user+sys). As expected, f_inverse1 is slower, 51.3 s (11.5+39.8 user+sys). Again, don't automatically try to minimize total operation count when you are writing numpy code. The high system overhead is probably due to heavy memory management - every vector temporary is 0.8 GB and the memory manager is struggling.
Cutting the array size to n = 1 million elements (8 MB), then multiplying the runtime by 100 brings the system time down by a large factor, f_inverse0 now takes 16.1 s (12.5+3.6), while f_inverse1 takes 22.3 s (16.2+5.1). This factor of 8 to 10 slower than compiled code is not unreasonable to expect for numpy performance.