My numba code is slower than my original code

Numba
1

I am new at using Numba and just trying to learn about it. I had a code based on scipy and python dictionaries, which I have written as follows to speed it up using numba but it is way slower than the original code. I believe it is running into the nopython mode and I have tried some variations as well. I don’t know what am I doing wrong. Can someone help?
@export
class ModuleCrystalsPMTPulses:

# Define structured array for parameters
parameters_array = np.array(
    [
        (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
        (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
        (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
        (0.5, 0.2), (0.5, 0.2)
    ], dtype=[('mu', 'f8'), ('sigma', 'f8')]
)

@staticmethod
@njit(parallel=True, nogil=True)
def compute_pulse_area(pulses, photoelectrons, params):
    
    for i in prange(photoelectrons.shape[0]):
        mu = params[i]['mu']
        sigma = params[i]['sigma']
        pulses[i] = np.random.normal(
            loc=(photoelectrons[i] * mu),
            scale=(np.sqrt(photoelectrons[i]) * sigma)
        )


     
def action(self, photoelectrons, pmt_indices):
   
    # Initialize array to store results
    pulse_areas = np.zeros_like(photoelectrons, dtype=np.float64)

    # Call the static method for computation
    self.compute_pulse_area(pulse_areas, photoelectrons, self.parameters_array[pmt_indices])

    return pulse_areas
2

I have updated your code to this

import numpy as np
import time
from numba import njit, prange


class ModuleCrystalsPMTPulses:
    # Define structured array for parameters
    parameters_array = np.array(
        [
            (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
            (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
            (0.5, 0.2), (0.5, 0.2), (0.5, 0.2), (0.5, 0.2),
            (0.5, 0.2), (0.5, 0.2)
        ], dtype=[('mu', 'f8'), ('sigma', 'f8')]
    )

    @staticmethod
    @njit(cache=True, parallel=True, nogil=True)
    def compute_pulse_area(pulses, photoelectrons, params):
        for i in prange(photoelectrons.shape[0]):
            mu = params[i]['mu']
            sigma = params[i]['sigma']
            pulses[i] = np.random.normal(
                loc=(photoelectrons[i] * mu),
                scale=(np.sqrt(photoelectrons[i]) * sigma)
            )

    def action(self, photoelectrons, pmt_indices):
        # Initialize array to store results
        pulse_areas = np.zeros_like(photoelectrons, dtype=np.float64)

        start_time = time.time_ns()
        self.compute_pulse_area(pulse_areas, photoelectrons, self.parameters_array[pmt_indices])
        end_time = time.time_ns()

        execution_time = end_time - start_time
        return pulse_areas, execution_time


def test_hello():
    # Create an instance of the class
    module = ModuleCrystalsPMTPulses()

    # Define dummy values
    num_values = 10 * 1024 * 1024
    iterations = 128
    execution_times = []

    for _ in range(iterations):
        photoelectrons = np.random.randint(1, 1024 * 1024, size=num_values)
        pmt_indices = np.random.randint(0, 14, size=num_values)  # 14 PMTs available

        # Run the action method
        _, execution_time = module.action(photoelectrons, pmt_indices)
        execution_times.append(execution_time / 1e9)

    # print("Execution times (ns):", execution_times1)

    # Compute percentiles
    percentiles = np.percentile(execution_times, [25, 50, 75, 90, 95, 99])
    print("Percentiles (ns):")
    print("25th:", percentiles[0])
    print("50th (median):", percentiles[1])
    print("75th:", percentiles[2])
    print("90th:", percentiles[3])
    print("95th:", percentiles[4])
    print("99th:", percentiles[5])


test_hello()

It doesn’t look that bad to me?

Percentiles (ns):
25th: 0.07273679599999999
50th (median): 0.07727650850000001
75th: 0.08527318725
90th: 0.0989272276
95th: 0.10974934099999999
99th: 0.12633157909000003

Process finished with exit code 0

Can you try enabling cache and try again? It could simply the time it takes to compile the function.