Cannot make sense of timing of numba-compiled functions - numpy

Question

I'm running some simulations, where I use numba to compile my python code to speed up the simulations. I wrote a function that will overwrite one of the input arrays, and therefore I would like to pass in a copy of that array instead. However, this makes the code much slower, and far slower than the time it takes to make the copy.
Here are the timing results:
> population_ = population.copy()
> %timeit _ = run_simulation(population_, Tmax, dt, Nskip = Nskip)
64.6 ms ± 215 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
> %timeit _ = run_simulation(population.copy(), Tmax, dt, Nskip = Nskip)
87.4 ms ± 778 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
> %timeit _ = population.copy()
442 ns ± 10.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
So calling run_simulation directly with the result of .copy() as an argument is about 23 milliseconds slower, despite the fact that making the copy only takes about 0.0004 milliseconds. I don't understand why this is the case.
For background, here is the full code:
import numpy as np
from numba import jit, int32, int64, float64
#jit('int32[:,:,:](int32[:,:,:], float64)', nopython=True)
def one_step(population, dt):
# Hard-coding model parameters here
beta = 0.55
tau = 10
# This probabilty doesn't depend on the other states
pIR = 1 - np.exp(-dt/tau)
# Double for loop over towns and towns
for i in range(population.shape[0]):
I = np.sum(population[i,1,:])
N = np.sum(population[i,:,:])
# Transition probability from susceptible to infected
pSI = 1 - np.exp(-dt*beta*I/N)
for j in range(population.shape[1]):
# Unpack variables for convenience
S, I, R = population[i,j,:]
S2I = np.random.binomial(S, pSI)
I2R = np.random.binomial(I, pIR)
# Calculate new values
S = S - S2I
I = I + S2I - I2R
R = R + I2R
population[i,j,:] = (S, I, R)
return population
#jit('int32[:,:,:](int32[:,:,:], float64, float64, int64)', nopython=True)
def run_simulation(population, Tmax, dt, Nskip = 10):
Nt = int(Tmax/dt)
history = np.zeros((population.shape[0], 3, int((Tmax/dt)/Nskip) + 1), dtype = np.int32)
history[:,:,0] = np.sum(population, axis = 1)
t = 0
for i in range(1, Nt+1):
population = one_step(population, dt)
t += dt
if i % Nskip == 0:
history[:,:,int(i/Nskip)] = np.sum(population, axis = 1)
return history
# Initial state
population = np.random.randint(low = 0, high = 1000, size = (10,10,3), dtype = np.int32)
# Run simulation for 100 days
Tmax = 100
dt = 0.01
# Only store once per day
Nskip = int(1/dt)
# Call one timestep to compile numba-decorated functions
# prior to measuring timing
_ = run_simulation(population, 1.0, 1.0, Nskip = 1)
# Run timing
population_ = population.copy()
%timeit _ = run_simulation(population_, Tmax, dt, Nskip = Nskip)
# Run timing
%timeit _ = run_simulation(population.copy(), Tmax, dt, Nskip = Nskip)
# Run timing
%timeit _ = population.copy()

Answer 1

What you are referring to isn't really related to numba.
Consider the following simple example:
import numpy as np
def run_simulation_2(population):
return population.sum(axis=0)
# Initial state
population = np.random.randint(low = 0, high = 1000, size = (10,10,3), dtype = np.int32)
# Run timing
population_ = population.copy()
%timeit _ = run_simulation_2(population_)
# Run timing
%timeit _ = run_simulation_2(population.copy())
# Run timing
%timeit _ = population.copy()
Timing results are:
3.45 µs ± 193 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
4.34 µs ± 91.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
680 ns ± 23.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
So you then is an overhead of about 25% even without numba which is about the same overhead you saw yourself.
Therefore I think it is not related to numba, but to different "behind the scenes" stuff that happens when you send aa variable as an argument vs result of a function.
Unfortunately I can't offer you a good explanation why it happens, but I hope that the fact that it's not related to numba is a well enough explanation for your needs.