# How can I make this arithmetic run faster?

**URL:** <https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106>\
**Category:** Support: How do I do ...?\
**Created:** [September 3, 2023, 3:57pm UTC](https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106 "2023-09-03T15:57:31Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![ofk123](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/ofk123/32/892_2.png) [@ofk123](https://numba.discourse.group/u/ofk123)\
**Post date:** [September 3, 2023, 3:57pm UTC](https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106/1 "2023-09-03T15:57:31Z")

</div>

> **MRE**
>
> ```auto
> import numpy as np
> from numba import jit, f8
> 
> def function( A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = ((A1-B1)/const1)**2 + ((A2-B2)/const2)**2 + ((A3-B3)/const3)**2
> return x
> 
> @jit(f8[:,:]( f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8, f8, f8 ), nopython=True, parallel=False)
> def function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = ((A1-B1)/const1)**2 + ((A2-B2)/const2)**2 + ((A3-B3)/const3)**2
> return x
> 
> @jit(f8[:,:]( f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8, f8, f8 ), nopython=True, fastmath=True, parallel=False)
> def function_jit_loop_fastmath( A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = np.empty_like(A1)
> for i in nb.prange(A1.shape[0]):
> for j in range(A1.shape[1]):
> x[i,j] = ((A1[i,j]-B1[i,j])/const1)**2 + ((A2[i,j]-B2[i,j])/const2)**2 + ((A3[i,j]-B3[i,j])/const3)**2
> return x
> 
> n = 200
> 
> A1 = np.arange(n**2, dtype=np.float64).reshape(n, n)
> A2 = A1 * np.random.uniform()
> A3 = A1 * np.random.uniform()
> B1 = A1.T
> B2 = A2.T
> B3 = A3.T
> 
> const1 = np.random.uniform()*10
> const2 = np.random.uniform()*10
> const3 = np.random.uniform()*100
> 
> %timeit function( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> %timeit function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> %timeit function_jit_loop_fastmath( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> print(np.allclose( function( A1, A2, A3, B1, B2, B3, const1, const2, const3 ), function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 )))
> print(np.allclose( function( A1, A2, A3, B1, B2, B3, const1, const2, const3 ), function_jit_loop_fastmath( A1, A2, A3, B1, B2, B3, const1, const2, const3 )))
> 
> 416 µs ± 2.29 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 624 µs ± 3.11 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 261 µs ± 4.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> True
> True
> 
> ```

After reading [this answer](https://numba.discourse.group/t/how-can-i-rewrite-the-indexing-so-the-function-compiles/2104/2) I created a loop. Also `fastmath=True` seems to help. Together this cuts 38% off the original runtime. Is it possible to cut more from the original runtime?

---

<div class="post-metadata">

**Author:** ![sschaer](https://avatars.discourse-cdn.com/v4/letter/s/5daacb/32.png) [@sschaer](https://numba.discourse.group/u/sschaer)\
**Post date:** [September 6, 2023, 8:08pm UTC](https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106/2 "2023-09-06T20:08:37Z")

</div>

Hi @ofk123

There are a few things to note here:

- In most cases, it is better not to give the `jit` decorator a signature and let Numba infer the types. For example, in your case, you tell Numba that the arrays can be any 2-dimensional array of type `float64`. However, Numba can do better if you allow it to compile a more specific function for your particular array layout (the array layout is part of the type). This is the reason why your vectorized jitted function runs slower than the pure Python version.
- There is no real need to traverse your data in two dimensions. It may be better to pretend the array is one dimensional. An example of this will be shown later.
- Make sure your tests are actually representative. The `A` arrays have C layout, while the `B` arrays have F layout. Also, the `B` arrays share data with the `A` arrays. Such things have a big impact on performance since the memory access patterns are different. How big the effect is depends on your hardware and the size of the array. You can run the code below to see what I mean.

> **see code here**
>
> ```auto
> import numpy as np
> import numba as nb
> 
> opts = dict(
> fastmath=True,
> parallel=False,
> error_model="numpy",
> )
> 
> def function(A1, A2, A3, B1, B2, B3, const1, const2, const3):
> return (
> ((A1 - B1) / const1) ** 2  
> + ((A2 - B2) / const2) ** 2  
> + ((A3 - B3) / const3) ** 2
> )
> 
> @nb.njit(**opts)
> def function_jit1(A1, A2, A3, B1, B2, B3, const1, const2, const3):
> return (
> ((A1 - B1) / const1) ** 2  
> + ((A2 - B2) / const2) ** 2  
> + ((A3 - B3) / const3) ** 2
> )
> 
> @nb.njit(**opts)
> def function_jit2(A1, A2, A3, B1, B2, B3, const1, const2, const3):
> x = np.empty_like(A1)
> for i in nb.prange(x.shape[0]):
> for j in range(x.shape[1]):
> x[i, j] = (
> ((A1[i, j] - B1[i, j]) / const1) ** 2  
> + ((A2[i, j] - B2[i, j]) / const2) ** 2  
> + ((A3[i, j] - B3[i, j]) / const3) ** 2
> )
> return x
> 
> @nb.njit(**opts)
> def function_jit3(A1, A2, A3, B1, B2, B3, const1, const2, const3):
> x = np.empty_like(A1)
> for i in nb.prange(x.size):
> x.flat[i] = (
> ((A1.flat[i] - B1.flat[i]) / const1) ** 2  
> + ((A2.flat[i] - B2.flat[i]) / const2) ** 2  
> + ((A3.flat[i] - B3.flat[i]) / const3) ** 2
> )
> return x
> 
> for n in (50, 200, 600):
> print(f"\nArray size: {n} by {n}")
> A1 = np.random.rand(n, n)
> A2 = np.random.rand(n, n)
> A3 = np.random.rand(n, n)
> 
> c1 = np.random.rand()
> c2 = np.random.rand()
> c3 = np.random.rand()
> 
> for B1, B2, B3, case in (
> (A1.copy(), A2.copy(), A3.copy(), "C order, data not shared"),
> (A1.copy().T, A2.copy().T, A3.copy().T, "\nF order, data not shared"),
> (A1, A2, A3, "\nC order, shared data"),
> (A1.T, A2.T, A3.T, "\nF order, shared data"),
> ):
> print(case)
>         
> assert np.allclose(function(A1, A2, A3, B1, B2, B3, c1, c2, c3), function_jit1(A1, A2, A3, B1, B2, B3, c1, c2, c3))
> assert np.allclose(function(A1, A2, A3, B1, B2, B3, c1, c2, c3), function_jit2(A1, A2, A3, B1, B2, B3, c1, c2, c3))
> assert np.allclose(function(A1, A2, A3, B1, B2, B3, c1, c2, c3), function_jit3(A1, A2, A3, B1, B2, B3, c1, c2, c3))
> 
> %timeit function(A1, A2, A3, B1, B2, B3, c1, c2, c3)
> %timeit function_jit1(A1, A2, A3, B1, B2, B3, c1, c2, c3)
> %timeit function_jit2(A1, A2, A3, B1, B2, B3, c1, c2, c3)
> %timeit function_jit3(A1, A2, A3, B1, B2, B3, c1, c2, c3)
> 
> ```

---

<div class="post-metadata">

**Author:** ![ofk123](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/ofk123/32/892_2.png) [@ofk123](https://numba.discourse.group/u/ofk123)\
**Post date:** [September 6, 2023, 11:14pm UTC](https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106/3 "2023-09-06T23:14:49Z")

</div>

Hi @sschaer, I ran your code and it looks very interesting!  
`function_jit3` with C-order is definitely the fastest combination, however the input is always `A1, A2, A3, A1.T, A2.T, A3.T, const1, const2, const3` so I will not be able to get the runtime-reduction seen in `functon_jit3`.

The timeit-results tell me I should use `function_jit2`. This does traverse the array in 2 dimensions, and is equal to the function in my original post.

> **timeit-results**
>
> ```auto
> Array size: 50 by 50
> C order, data not shared
> 24.1 µs ± 123 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 24.4 µs ± 176 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 14.6 µs ± 57.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 4.06 µs ± 43.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 
> F order, data not shared
> 31.7 µs ± 235 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 23.9 µs ± 155 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 17.5 µs ± 55.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 52.3 µs ± 102 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 
> C order, shared data
> 23 µs ± 77.6 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 24.2 µs ± 183 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 14.7 µs ± 75.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 3.63 µs ± 24.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 
> F order, shared data
> 31.5 µs ± 169 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 23.7 µs ± 71.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 17.6 µs ± 97.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
> 52.7 µs ± 250 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 
> Array size: 200 by 200
> C order, data not shared
> 291 µs ± 3.73 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 360 µs ± 2.54 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 206 µs ± 731 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 56.7 µs ± 826 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 
> F order, data not shared
> 381 µs ± 2.36 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 484 µs ± 2.41 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 319 µs ± 2.03 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 902 µs ± 6.75 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 
> C order, shared data
> 256 µs ± 1.53 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 359 µs ± 1.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 206 µs ± 1.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 37.3 µs ± 448 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> 
> F order, shared data
> 361 µs ± 2.19 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 449 µs ± 6.94 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 292 µs ± 1.18 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 846 µs ± 2.54 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 
> Array size: 600 by 600
> C order, data not shared
> 7.21 ms ± 91 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 3.3 ms ± 6.72 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 1.91 ms ± 19 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 1.07 ms ± 41.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 
> F order, data not shared
> 9.73 ms ± 171 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 13.9 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 12 ms ± 98.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 18.2 ms ± 154 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 
> C order, shared data
> 6.28 ms ± 54.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 3.28 ms ± 5.03 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 1.85 ms ± 15.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 410 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 
> F order, shared data
> 9.12 ms ± 213 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 13.6 ms ± 237 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 11.6 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> 17.8 ms ± 220 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
> CPU times: user 4min 33s, sys: 8.92 s, total: 4min 42s
> Wall time: 4min 42s
> 
> ```

Thanks alot for the feedback though! The bulletpoints and code was very helpful. Are there any other methods that can help with this arithmetic?

---

<div class="post-metadata">

**Author:** ![ofk123](https://yyz2.discourse-cdn.com/free1/user_avatar/numba.discourse.group/ofk123/32/892_2.png) [@ofk123](https://numba.discourse.group/u/ofk123)\
**Post date:** [February 29, 2024, 12:40am UTC](https://numba.discourse.group/t/how-can-i-make-this-arithmetic-run-faster/2106/4 "2024-02-29T00:40:50Z")

</div>

Since the input-2D-arrays actually have some of its data repeated: Instead of passing the 2D-arrays as arguments, I ended up reducing them to 1D arrays including only the unique values, and loop through them as if they were 2D-arrays. It reduces time to a half, which is great, but I would love to improve it further.

In each array’s case, I need to calculate the difference between each element, and all other elements.  
I edited the example a bit, and provided the function I use today, named `function_1darrays()` .

> **Summary**
>
> ```auto
> import numpy as np
> import numpy.matlib as nm
> from numba import njit
> 
> def function( A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = ((A1-B1)/const1)**2 + ((A2-B2)/const2)**2 + ((A3-B3)/const3)**2
> return x
> 
> @njit(parallel=False)
> def function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = ((A1-B1)/const1)**2 + ((A2-B2)/const2)**2 + ((A3-B3)/const3)**2
> return x
> 
> @njit(fastmath=True, parallel=False)
> def function_jit_loop_fastmath(n, A1, A2, A3, B1, B2, B3, const1, const2, const3 ):
> x = np.zeros_like(A1)
> for i in range(n):
> for j in range(n):
> x[i,j] = ((A1[i,j]-B1[i,j])/const1)**2 + ((A2[i,j]-B2[i,j])/const2)**2 + ((A3[i,j]-B3[i,j])/const3)**2
> return x
> 
> @njit( fastmath = True, parallel=False )
> def function_1darrays( n, a1, a2, a3, const1, const2, const3 ) -> np.ndarray:
> x = np.empty((n,n))
> for i in range(n):
> a1i = a1[i]
> a2i = a2[i]
> a3i = a3[i]
> for j in range(n):
> x[i,j] = ((a1i - a1[j])/const1)**2 + ((a2i - a2[j])/const2)**2 + ((a3i - a3[j])/const3)**2
> return x
> 
> n = 200
> 
> a1 = np.arange(1, n+1, dtype=np.float64)
> a2 = np.arange(n+1, 2*n+1, dtype=np.float64)
> a3 = np.arange(2*n+1,3*n+1, dtype=np.float64)
> A1 = nm.repmat(a1, n, 1)
> A2 = nm.repmat(a2, n, 1)
> A3 = nm.repmat(a3, n, 1)
> B1 = A1.T
> B2 = A2.T
> B3 = A3.T
> 
> const1 = 1.0
> const2 = 1.0
> const3 = 1.0
> 
> K1 = function( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> K2 = function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> K3 = function_jit_loop_fastmath( n, A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> K4 = function_1darrays( n, a1, a2, a3, const1, const2, const3 )
> 
> %timeit function( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> %timeit function_jit( A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> %timeit function_jit_loop_fastmath( n, A1, A2, A3, B1, B2, B3, const1, const2, const3 )
> %timeit function_1darrays( n, a1, a2, a3, const1, const2, const3 )
> 
> print(np.array_equal( K1, K2))
> print(np.array_equal( K1, K3))
> print(np.array_equal( K1, K4))
> print(np.array_equal( K2, K3))
> print(np.array_equal( K2, K4))
> print(np.array_equal( K3, K4))
> 
> ```
> 
> Out:
> 
> ```auto
> 376 µs ± 2.25 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 488 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 270 µs ± 1.47 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
> 125 µs ± 709 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
> True
> True
> True
> True
> True
> True
> 
> ```