Cannot make sense of timing of numba-compiled functions

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()