entrypoint_cuda.py

import math
import itertools
from collections import defaultdict

pi = math.pi

import numpy as np
from numba import cuda
from numba.cuda.cudadrv.devicearray import DeviceNDArray
import matplotlib.pyplot as plt

from functions.potential import potential_function_cuda
from kernels.potential_evaluator import PotentialEvaluator
from kernels.potential_minimum_searcher import PotentialMinimumSearcher
from utils.array_stacker import ArrayStacker

# Parameters for simulation ###################################################
NUMBER_OF_PHI_POINTS = 101
NUMBER_OF_FIELD_POINTS = 21
NUMBER_OF_FIELD_POINTS_PER_RUN = 21
NUMBER_OF_FIELD_RUNS = (
    NUMBER_OF_FIELD_POINTS - 1
) // NUMBER_OF_FIELD_POINTS_PER_RUN + 1
ALPHA = 1
LOWER = -0.5
UPPER = 1.5

lr_array = np.linspace(LOWER * 2 * pi, UPPER * 2 * pi, NUMBER_OF_FIELD_POINTS)
phixx_array = np.linspace(-pi, pi, NUMBER_OF_PHI_POINTS)

# Kernels #####################################################################
BLOCKS_PER_GRID = (NUMBER_OF_FIELD_POINTS_PER_RUN, NUMBER_OF_FIELD_POINTS_PER_RUN)

potential_evaluator = PotentialEvaluator(NUMBER_OF_PHI_POINTS, potential_function_cuda)
THREADS_PER_BLOCK_potential_evaluation = potential_evaluator.allocate_max_threads(8)

potential_minimum_searcher = PotentialMinimumSearcher(NUMBER_OF_PHI_POINTS)
THREADS_PER_BLOCK_potential_search = potential_minimum_searcher.allocate_max_threads()

# Execution ###################################################################
DEVICE_lr_array = cuda.to_device(lr_array)
DEVICE_phixx_array = cuda.to_device(phixx_array)
DEVICE_potential_array = cuda.device_array(
    shape=(
        NUMBER_OF_FIELD_POINTS_PER_RUN,
        NUMBER_OF_FIELD_POINTS_PER_RUN,
        NUMBER_OF_PHI_POINTS,
        NUMBER_OF_PHI_POINTS,
        NUMBER_OF_PHI_POINTS,
    ),
    dtype=np.float32,
)
DEVICE_grid_search_result_array = cuda.device_array(
    shape=(NUMBER_OF_FIELD_POINTS_PER_RUN, NUMBER_OF_FIELD_POINTS_PER_RUN, 4),
    dtype=np.float32,
)

# Go through teach of the field section and evaluate ##########################
quadrants = defaultdict(
    lambda: [[None] * NUMBER_OF_FIELD_RUNS for i in range(0, NUMBER_OF_FIELD_RUNS)]
)

for (L_RUN, R_RUN) in itertools.product(
    range(0, NUMBER_OF_FIELD_RUNS), range(0, NUMBER_OF_FIELD_RUNS)
):
    print(
        f"🦑 Running (L={L_RUN}/{NUMBER_OF_FIELD_RUNS - 1}), (R={R_RUN}/{NUMBER_OF_FIELD_RUNS - 1})"
    )
    L_OFFSET = int(L_RUN * NUMBER_OF_FIELD_POINTS_PER_RUN)
    R_OFFSET = int(R_RUN * NUMBER_OF_FIELD_POINTS_PER_RUN)
    potential_evaluator.kernel[BLOCKS_PER_GRID, THREADS_PER_BLOCK_potential_evaluation](
        DEVICE_phixx_array,
        DEVICE_lr_array,
        L_OFFSET,
        R_OFFSET,
        ALPHA,
        DEVICE_potential_array,
    )
    potential_minimum_searcher.kernel[
        BLOCKS_PER_GRID, THREADS_PER_BLOCK_potential_search
    ](DEVICE_potential_array, DEVICE_grid_search_result_array)

    grid_search_result_array = DEVICE_grid_search_result_array.copy_to_host()
    quadrants["potential"][L_RUN][R_RUN] = grid_search_result_array[:, :, 0]
    quadrants["phi01"][L_RUN][R_RUN] = phixx_array[
        grid_search_result_array[:, :, 1].astype(int)
    ]
    quadrants["phi02"][L_RUN][R_RUN] = phixx_array[
        grid_search_result_array[:, :, 2].astype(int)
    ]
    quadrants["phi03"][L_RUN][R_RUN] = phixx_array[
        grid_search_result_array[:, :, 3].astype(int)
    ]

result = {}
for key, value in quadrants.items():
    print(f"🦑--------------------{key}--------------------")
    result[key] = ArrayStacker.stack_into_square(value)
    print(result[key])


###############################################################################
#   sudo yum install python36-tkinter and do ssh -X to print on own computer  #
###############################################################################
fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(5, 5))

im = ax.imshow(
    # result["potential"],
    result["phi01"],
    extent=[LOWER, UPPER, LOWER, UPPER],
    origin="lower",
    cmap="cividis",
    # cmap='YlGnBu'
    # interpolation='spline36'
)
cbar = fig.colorbar(im)
plt.show()