I was in a very similar situation and I found a way to use a Numba-JITed function inside of a class.

The trick is to use a static method, since this kind of methods are not called prepending the object instance to the argument list. The downside of not having access to self is that you cannot use variables defined outside of the method. So you have to pass them to the static method from a calling method that has access to self. In my case I did not need to define a wrapper method. I just had to split the method I wanted to JIT compile into two methods.

In the case of your example, the solution would be:

from numba import jit

class MyClass:
    def __init__(self):
        self.k = 1

    def calculation(self):
        k = self.k
        return self.complicated([1,2,3],k)

    @staticmethod
    @jit(nopython=True)                             
    def complicated(x,k):                                  
        for a in x:
            b = a**2 .+ a**3 .+ k

I modified your code a bit and was able to get the function to run. I removed the Particle class and changed all the list instances to numpy arrays.

This is the output for the Zavg, Zrelm and counter:

Zavg: [0.07047501 0.06735052 0.06728123 ... 0.39516435 0.3947497  0.39433495] 
Zrelm: [-0.04179043 -0.04461464 -0.0394889  ... -0.11080628 -0.11087257
     -0.11093883] 
Counter: 538

Here is the modified code:

import numpy as np
import matplotlib.pyplot as plt
import random
from math import *
from random import *
from numba import jit


"Dynamics"

@jit(nopython=True)
def f(SP, Zavgi, Zreli, alf, dt, n):
    "Time"
    counter = 0;
    sum1, sum2 = 0, 0;
    Zavg = np.array([Zavgi]);
    Zrelm = np.array([Zreli]);
    T_plot = [0];
    for i in range(1, iter + 1):

        t = i * dt;
        T_plot.append(t)
        Z = [];
        Up = [];
        Down = [];
        c, s = cos(t), sin(t);
        c1, s1 = cos(t - dt), sin(t - dt);
        for j in range(n - 1):
            collchk = ((c * (SP[j][0]) + s * (SP[j][1])) - (c * (SP[j + 1][0]) + s * (SP[j + 1][1]))) * (
                    c1 * (SP[j][0]) + s1 * (SP[j][1]) - (c1 * (SP[j + 1][0]) + s1 * (SP[j + 1][1])));

            prel = ((c * (SP[j][1]) - s * (SP[j][0])) - (c * (SP[j + 1][1]) - s * (SP[j + 1][0]))) / 2;

            rcoeff = 1 / (1 + (prel * alf) ** 2);
            rand_value = random();

            if collchk < 0:

                SP[j], SP[j + 1] = SP[j + 1], SP[j];

                if rcoeff > rand_value:
                    counter = counter + 1
                    SP[j][2], SP[j + 1][2] = SP[j + 1][2], SP[j][2];
            if SP[j][2] == 1:
                Up.append(c * (SP[j][0]) + s * (SP[j][1]))
            else:
                Down.append(c * (SP[j][0]) + s * (SP[j][1]))
            Z.append(c * (SP[j][0]) + s * (SP[j][1]))

        Zrel = np.sum(np.array(Up)) / len(Up) - np.sum(np.array(Down)) / len(Down);
        Zrelm = np.append(Zrelm, Zrel)

        Zm = np.sum(np.array(Z)) / len(Z)
        Zavg = np.append(Zavg, Zm)


    return Zavg, Zrelm, counter, T_plot



if __name__ == '__main__':

    n = 1000
    mu = np.random.uniform(0, 1, n)
    r = [sqrt(-2 * log(1 - i)) for i in mu]
    eta = np.random.uniform(0, 1, n)
    theta = 2 * pi * eta;
    cuz = [cos(i) for i in theta]
    suz = [sin(i) for i in theta]
    Zinitial = [a * b for a, b in zip(r, cuz)];
    Pinitial = [a * b for a, b in zip(r, suz)];

    iter = 10 ** (5);
    dt = 1 / (2 * n);
    alf = sqrt(n);


    SP = np.array(sorted(np.array([  np.array([i,j,choice([0,1])]) for i, j in zip(Zinitial, Pinitial)]),
                key=lambda x: x[0]))
    Upi = [];
    Downi = [];
    count_plot = [];
    for j in range(len(SP)):
        if SP[j][2] == 1:
            Upi.append(SP[j][0])
        else:
            Downi.append(SP[j][0])
    Zavgi = sum(Zinitial) / len(Zinitial)
    Zreli = sum(Upi) / len(Upi) - sum(Downi) / len(Downi)


    Zavg, Zrelm, counter, T_plot = f(SP, Zavgi, Zreli, alf, dt, n)

    print(Zavg, Zrelm, counter)
    plt.plot(T_plot, Zrelm)
    plt.show()

This is how the plot looks like:

There are a couple factors here affecting your results. The main two are

• Your calculation of run time includes the numba compilation, which is a one-time overhead cost
• Your functions are not optimized to work well with numba or numpy resulting in functions that don't benefit from numpy efficiency and compile poorly

I was able to partially reproduce your results, and I offer some suggestions for changes to make the code operate better after numba compilation as well as restructuring the testing code so that the differences between the compilation step and the regular execution can be seen a little easier.

To test the impact of numba compilation, I time and call the function once and then again in a loop to get an average run time after compilation. The run time after compilation is many times faster in some cases because that compilation represents a large but one-time cost.

To show the impact of better function design, I rewrote a couple of the functions to use numpy.ndarrays instead of lists both internally and as inputs, and it shows significant improvements in performance, in some cases >100x faster.

The last note, numpy is already highly optimized, where you can use numpy functions it may be difficult to attain better performance otherwise, even with numba compilation. When we look at the last function, which uses numpy.square, I would not expect to see much improvement (once the input is adjusted to be an expected numpy type), and the testing confirms that expectation.

Here is the testing code:

Copyimport time

import numpy as np
import numba

import matplotlib.pyplot as plt

def makeSquareLoopFor(a):
    result = []
    for i, v in enumerate(a):
        result.append(a[i] ** 2)
    return result

@numba.njit
def makeSquareLoopForJIT(a):
    result = []
    for i, v in enumerate(a):
        result.append(a[i] ** 2)
    return result

def makeSquareLoopForOpt(a):
    result = np.empty(shape=(len(a),), dtype=np.int64)
    for i, v in enumerate(a):
        result[i] = a[i] ** 2
    return result

@numba.njit
def makeSquareLoopForOptJIT(a):
    result = np.empty(shape=(len(a),), dtype=np.int64)
    for i, v in enumerate(a):
        result[i] = a[i] ** 2
    return result

def makeSquareLoopWhile(a):
    result = []
    i = 0
    while i < len(a):
        result.append(a[i] ** 2)
        i += 1
    return result

@numba.njit
def makeSquareLoopWhileJIT(a):
    result = []
    i = 0
    while i < len(a):
        result.append(a[i] ** 2)
        i += 1
    return result

def makeSquareLoopWhileOpt(a):
    result = np.empty(shape=(len(a),), dtype=np.int64)
    i = 0
    while i < len(a):
        result[i] = a[i] ** 2
        i += 1
    return result

@numba.njit
def makeSquareLoopWhileOptJIT(a):
    result = np.empty(shape=(len(a),), dtype=np.int64)
    i = 0
    while i < len(a):
        result[i] = a[i] ** 2
        i += 1
    return result
def makeSquareListComprehension(a):
    return [v ** 2 for v in a]

@numba.njit
def makeSquareListComprehensionJIT(a):
    return [v ** 2 for v in a]

def makeSquareListMap(a):
    return list(map(lambda v: v ** 2, a))

@numba.njit
def makeSquareListMapJIT(a):
    return list(map(lambda v: v ** 2, a))

def makeSquareListNumpy(a):
    return np.square(a)

@numba.njit
def makeSquareListNumpyJIT(a):
    return np.square(a)

def _main():

    trials = 10

    array_len = 5_000_000
    a_list = list(range(1, array_len, 1))
    a_array = np.array(a_list)

    print(f'{"function":<35}{"input type":<12}{"first run":>12}{"trial runs":>12}')
    function_names = []
    results = {}
    for pythonfun, numbafun, humanname in (
            (makeSquareLoopFor, makeSquareLoopForJIT, 'for'),
            (makeSquareLoopForOpt, makeSquareLoopForOptJIT, 'for opt'),
            (makeSquareLoopWhile, makeSquareLoopWhileJIT, 'while'),
            (makeSquareLoopWhileOpt, makeSquareLoopWhileOptJIT, 'while opt'),
            (makeSquareListComprehension, makeSquareListComprehensionJIT, 'comp'),
            (makeSquareListMap, makeSquareListMapJIT, 'map'),
            (makeSquareListNumpy, makeSquareListNumpyJIT, 'numpy fun')
    ):
        function_names.append(humanname)
        for fun, fun_type in (
                (pythonfun, 'python'),
                (numbafun, 'numba')
        ):


            for a_data, a_type in (
                (a_list, 'list'),
                (a_array, 'array')
            ):
                try:
                    start_first = time.perf_counter_ns()
                    fun(a_data)
                    end_first = time.perf_counter_ns()
                    start_trials = time.perf_counter_ns()
                    for i in range(trials):
                        fun(a_data)
                    end_trials = time.perf_counter_ns()

                    run_time_first = (end_first - start_first) / 10 ** 6
                    run_time_trials = ((end_trials - start_trials) / trials) / 10 ** 6

                    if fun_type not in results:
                        results[fun_type] = {}
                    if a_type not in results[fun_type]:
                        results[fun_type][a_type] = {
                            'first': [],
                            'trials': []
                        }
                    results[fun_type][a_type]['first'].append(run_time_first)
                    results[fun_type][a_type]['trials'].append(run_time_trials)

                    print(f'{fun.__name__:<35}{a_type:<12}{run_time_first:12.3f}{run_time_trials:12.3f}')
                except numba.core.errors.TypingError:
                    print(f'{fun.__name__:<35}{a_type:<12}{"failed":>12}')
        print()

    fig, axs = plt.subplots(2)
    for funtype, fundata in results.items():
        for inputtype, inputdata in fundata.items():
            for i, (runtype, rundata) in enumerate(inputdata.items()):
                axs[i].plot(rundata, '-x', label=f'{funtype} - {inputtype}')

    axs[0].legend()
    axs[1].legend()
    axs[0].set_title('With compilation step')
    axs[1].set_title('Without compilation step')
    axs[0].set_yscale('log')
    axs[1].set_yscale('log')
    axs[0].get_xaxis().set_visible(False)
    axs[1].set_xticks(range(len(rundata)), labels=function_names)
    plt.show()


if __name__ == '__main__':
    _main()

Running that code produces the following console print:

function                           input type     first run  trial runs
makeSquareLoopFor                  list             485.755     481.474
makeSquareLoopFor                  array            893.265     857.773
makeSquareLoopForJIT               list            4436.258    3915.449
makeSquareLoopForJIT               array            245.782     164.247

makeSquareLoopForOpt               list             738.646     672.020
makeSquareLoopForOpt               array           1123.080    1127.280
makeSquareLoopForOptJIT            list            4013.135    3812.577
makeSquareLoopForOptJIT            array             85.737       8.752

makeSquareLoopWhile                list             588.821     586.823
makeSquareLoopWhile                array            931.613     899.923
makeSquareLoopWhileJIT             list            4055.000    3963.027
makeSquareLoopWhileJIT             array            253.426     167.106

makeSquareLoopWhileOpt             list             791.360     790.028
makeSquareLoopWhileOpt             array           1185.349    1133.906
makeSquareLoopWhileOptJIT          list            3889.629    3778.378
makeSquareLoopWhileOptJIT          array             92.073       8.482

makeSquareListComprehension        list             314.250     312.969
makeSquareListComprehension        array            447.048     443.911
makeSquareListComprehensionJIT     list            3981.916    3908.984
makeSquareListComprehensionJIT     array            238.009     164.230

makeSquareListMap                  list             471.389     455.676
makeSquareListMap                  array            600.735     615.717
makeSquareListMapJIT               list            4212.756    3950.546
makeSquareListMapJIT               array            457.831     242.098

makeSquareListNumpy                list             182.675     182.450
makeSquareListNumpy                array              9.080       9.116
makeSquareListNumpyJIT             list              failed
makeSquareListNumpyJIT             array            127.094       9.239

And the following graph:

They key takeaways here are:

• When written in a way that does not work well with numba, numba is consistently the slowest
• When written to target numba compilation, numba calculation is consistently the fastest, excepting only already highly optimized numpy functions (see functions with opt in name which when numba compiled hugely outperform their unoptimized counterpart at about 20x faster)
• When discounting the compilation step, numba compilation of appropriate functions offers extremely large performance improvements over pure python implementations.
• Even without numba compilation, simply performing all operations on numpy objects can offer major performance improvements (see makeSquareListNumpy operating on a list vs a numpy.ndarray)

Whenever working with numpy and numba I highly recommend profiling, experimenting and comparing to see what solution works the best. Often, just using numpy functions will offer the best performance as was the case here, but sometimes numba compilation will really make a difference, and some problems are just better solved with pure python. Comparing is the surest way to get the best performance.