jit-test.py

import random
from numba import jit, cuda
import numpy as np


@jit
def calculate_probability():
    valid_count = 0
    iterations = 10000000  # 10M
    for _ in range(0, iterations):
        person_chosen = list(np.random.choice(100, 3, replace=False))
        class_chosen = np.array([])
        class_chosen = [element // 20 for element in person_chosen]
        if np.sum(class_chosen == 0) == np.sum(class_chosen == 1) == np.sum(class_chosen == 2) == np.sum(class_chosen == 3) == np.sum(class_chosen == 4) == 2:
            valid_count += 1
    print(valid_count / iterations)


def main():
    calculate_probability()
    cuda.synchronize()


if __name__ == "__main__":
    main()

'''

from numba import cuda


def cpu_print(N):
    for i in range(0, N):
        print(i)


@cuda.jit
def gpu_print(N):
    for i in range(0, N):
        print(i)


def main():
    print("gpu print:")
    gpu_print[1, 1](8)
    cuda.synchronize()
    print("cpu print:")
    cpu_print(8)


if __name__ == "__main__":
    main()
'''