Can not use numba in the class

Question

I recently want to use @njit(parallel=True) in package numba to speed up my nbodysimulation code, but when I separate the original function out of the class, my code can not work anymore. How to fix this problem?

The following block is the original code to calculate acceleration.

def _calculate_acceleration(self, mass, pos, rsoft):
        """
        Calculate the acceleration.
        """
        # TODO:
        N = self.particles
        rsoft = self.rsoft
        posx = pos[:,0]
        posy = pos[:,1]
        posz = pos[:,2]
        G    = self.G
        npts = self.nparticles
        acc  = np.zeros((npts, 3))

        for i in prange(npts):
            for j in prange(npts):
                if (j>i): 
                    x   = (posx[i]-posx[j])
                    y   = (posy[i]-posy[j])
                    z   = (posz[i]-posz[j])
                    rsq = x**2 + y**2 + z**2
                    req = np.sqrt(x**2 + y**2)

                    f   = -G*mass[i,0]*mass[j,0]/rsq
                    
                    theta = np.arctan2(y, x)
                    phi = np.arctan2(z, req)
                    fx  = f*np.cos(theta)*np.cos(phi)
                    fy  = f*np.sin(theta)*np.cos(phi)
                    fz  = f*np.sin(phi)

                    acc[i,0] += fx/mass[i]
                    acc[i,1] += fy/mass[i]
                    acc[i,2] += fz/mass[i]
                    acc[j,0] -= fx/mass[j]
                    acc[j,1] -= fy/mass[j]
                    acc[j,2] -= fz/mass[j]
        return acc

def initialRandomParticles(N = 100, total_mass = 10):
        """
        Initial particles

        """
        particles  = Particles(N)
        masses     = particles.masses
        mass       = total_mass/particles.nparticles
        particles.masses = (masses*mass)
        positions  = np.random.randn(N,3)
        velocities = np.random.randn(N,3)
        accelerations = np.random.randn(N,3)
        particles.positions = positions
        particles.velocities = velocities
        particles.accelerations = accelerations

        return particles

particles = initialRandomParticles(N = 10**5, total_mass = 20)
sim = NbodySimulation(particles)
sim.setup(G=G,method="RK4",io_freq=200,io_title=problem_name,io_screen=True,visualized=False, rsoft=0.01)
sim.evolve(dt=0.01,tmax=10)

# Particles and NbodySimulation are defined class.

I spearate the original function out of the class, and define another function to call it in the class, but it's still can not work.

The following is new code which is out of the class.

@njit(nopython=True, parallel=True)
def _calculate_acceleration(n, npts, G, mass, pos, rsoft):
    """
    Calculate the acceleration. This function is out of the class.
    """
    # TODO:
    posx = pos[:,0]
    posy = pos[:,1]
    posz = pos[:,2]
    acc  = np.zeros((n, 3))
    sqrt = np.sqrt

    for i in prange(npts):
        for j in prange(npts):
            if (j>i): 
                x   = (posx[i]-posx[j])
                y   = (posy[i]-posy[j])
                z   = (posz[i]-posz[j])
                rsq = x**2 + y**2 + z**2
                req = sqrt(x**2 + y**2 + z**2)

                f   = -G*mass[i,0]*mass[j,0]/(req + rsoft)**2
                fx  = f*x**2/rsq
                fy  = f*y**2/rsq
                fz  = f*z**2/rsq

                

                acc[i,0] = fx/mass[i] + acc[i,0]
                acc[i,1] = fy/mass[i] + acc[i,1]
                acc[i,2] = fz/mass[i] + acc[i,2]
                acc[j,0] = fx/mass[j] - acc[j,0]
                acc[j,1] = fy/mass[j] - acc[j,1]
                acc[j,2] = fz/mass[j] - acc[j,2]
    return acc

Here is the error:

TypingError                               Traceback (most recent call last)

   :428, in NbodySimulation._update_particles_rk4(self, dt, particles)
    426 position = particles.positions   # y0[0]
    427 velocity = particles.velocities # y0[1], k1[0]
--> 428 acceleration  = self._calculate_acceleration_inclass() # k1[1]
    430 position2     = position + 0.5*velocity * dt      # y1[0]
    431 velocity2     = velocity + 0.5*acceleration * dt  # y1[1], k2[0]

   :381, in NbodySimulation._calculate_acceleration_inclass(self)
    377 def _calculate_acceleration_inclass(self):
    378     """
    379     Calculate the acceleration.
...
    <source elided>

                acc[i,0] = fx/mass[i] + acc[i,0]
                ^