<reponame>twj2417/srf<gh_stars>0
import math

import numpy as np
from numba import jit, cuda


@cuda.jit(device = True)
def norm(x, y):
    return (x ** 2 + y ** 2) ** 0.5


@cuda.jit(device = True)
def get_block_theta(event):
    before_theta = event[0] / norm(event[0], event[1])
    return math.acos(before_theta) / math.pi * 180


@cuda.jit(device = True)
def block_id(event, nb_crystals_per_ring, block_grid_y, num_block):
    theta_event = get_block_theta(event)
    if event[1] < 0:
        theta_event = 360 - theta_event
    fixed_theta = (theta_event + 180 / nb_crystals_per_ring * block_grid_y) % 360
    return int(fixed_theta / 360 * num_block)


@cuda.jit(device = True)
def crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius, grid_block,
               size_block, event):
    id_block = block_id(event, nb_crystals_per_ring, grid_block[1], nb_blocks_per_ring)
    theta_block = 2 * math.pi * id_block / nb_blocks_per_ring
    block_center_x = (inner_radius + outer_radius) / 2 * math.cos(theta_block)
    block_center_y = (inner_radius + outer_radius) / 2 * math.sin(theta_block)
    dist = (block_center_x - event[0]) * math.cos(theta_block) + (
            block_center_y - event[1]) * math.sin(theta_block)
    pro_point_x = dist * math.cos(theta_block) + event[0]
    pro_point_y = dist * math.sin(theta_block) + event[1]
    edge_block_x = block_center_x - math.cos(theta_block + math.pi / 2) * size_block[1] / 2
    edge_block_y = block_center_y - math.sin(theta_block + math.pi / 2) * size_block[1] / 2
    id_crystal = int(
        norm(pro_point_x - edge_block_x, pro_point_y - edge_block_y) / size_block[1] * grid_block[
            1])
    return id_block * grid_block[1] + id_crystal


@cuda.jit(device = True)
def ring_id(nb_rings, axial_length, grid_block, size_block, event_z):
    return int((event_z + nb_rings * size_block[2] / 2) / axial_length * nb_rings * grid_block[2])


@cuda.jit(device = True)
def position2crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings, inner_radius,
                       outer_radius, axial_length, grid_block,
                       size_block, event):
    ring_id1 = ring_id(nb_rings, axial_length, grid_block, size_block, event[2])
    crystal_id1 = crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius,
                             grid_block, size_block, event[:2])
    ring_id2 = ring_id(nb_rings, axial_length, grid_block, size_block, event[5])
    crystal_id2 = crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius,
                             grid_block, size_block, event[3:5])
    return (ring_id1 * nb_crystals_per_ring + crystal_id1), (
            ring_id2 * nb_crystals_per_ring + crystal_id2)


@cuda.jit(device = True)
def change_id(id1, id2):
    if id1 < id2:
        return id2, id1
    else:
        return id1, id2


@cuda.jit
def crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings, inner_radius, outer_radius,
              axial_length, grid_block, size_block, events, final_id):
    c, j = cuda.grid(2)
    i = 512 * 16 * c + j
    if i < events.shape[0]:
        id1, id2 = position2crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings,
                                      inner_radius, outer_radius, axial_length,
                                      grid_block, size_block, events[i, :])
        fixed_id1, fixed_id2 = change_id(id1, id2)
        final_id[i, 0] = fixed_id1
        final_id[i, 1] = fixed_id2


@jit(nopython = True)
def cal_sinogram(list_mode_data, nb_rings, grid_block_z, nb_crystals_per_ring):
    total_num_crystal = nb_rings * grid_block_z * nb_crystals_per_ring
    weight = np.zeros((int(total_num_crystal * (total_num_crystal - 1) / 2), 1))
    for i in range(int(list_mode_data.shape[0])):
        pos = int((list_mode_data[i, 0] - 1) * list_mode_data[i, 0] / 2 + list_mode_data[i, 1])
        weight[pos] = weight[pos] + list_mode_data[i, 2]
    return weight


@jit(nopython = True)
def get_all_lors_id(total_num_crystal):
    lors = np.zeros((int(total_num_crystal * (total_num_crystal - 1) / 2), 2), dtype = np.int32)
    num = 0
    for i in range(total_num_crystal):
        for j in range(i):
            lors[num, 0] = i
            lors[num, 1] = j
            num = num + 1
    return lors


def get_crystal_xy(grid_block, num_block, crystal_id, block_id, r_inner, size_block):
    size_crystal = size_block / grid_block
    angle_block = math.pi * 2 / num_block * block_id
    center = np.hstack((((r_inner) * np.cos(angle_block)).reshape(block_id.size, 1),
                        ((r_inner) * np.sin(angle_block)).reshape(block_id.size, 1)))
    vector_tang = np.hstack(
        (np.cos(angle_block + math.pi / 2 * np.ones_like(angle_block)).reshape(block_id.size, 1),
         np.sin(angle_block + math.pi / 2 * np.ones_like(angle_block)).reshape(block_id.size, 1)))
    dist_from_center = size_crystal[1] * (
            crystal_id - (grid_block[1] - 1) / 2 * np.ones_like(crystal_id))
    return center + (np.hstack((dist_from_center.reshape(dist_from_center.size, 1),
                                dist_from_center.reshape(dist_from_center.size, 1)))
                     * vector_tang)


def get_crystal_z(ring_id, grid_block, size_block, nb_rings):
    size_crystal = size_block / grid_block
    dist_from_center = (ring_id - (grid_block[2] * nb_rings - 1) / 2 * np.ones_like(ring_id)) * \
                       size_crystal[2]
    return dist_from_center


def get_center(scanner, crystal_id_whole_scanner):
    ring_id = np.floor(
        crystal_id_whole_scanner / scanner.nb_blocks_per_ring / scanner.blocks.shape[1])
    crystal_per_ring_id = crystal_id_whole_scanner - ring_id * scanner.nb_blocks_per_ring * \
                          scanner.blocks.shape[1]
    block_id = crystal_per_ring_id // scanner.blocks.shape[1]
    crystal_id = crystal_per_ring_id % scanner.blocks.shape[1]
    center_xy = get_crystal_xy(np.array(scanner.blocks.shape), scanner.nb_blocks_per_ring,
                               crystal_id, block_id, scanner.inner_radius,
                               np.array(scanner.blocks.size))
    center_z = get_crystal_z(ring_id, np.array(scanner.blocks.shape), np.array(scanner.blocks.size),
                             scanner.nb_rings)
    return np.hstack((center_xy, center_z.reshape(center_z.size, 1)))


# @jit
def lm2sino(listmode, scanner):
    listmode_pos = listmode.lors.data
    listmode_id = np.zeros((listmode_pos.shape[0], 2), dtype = np.int32)
    crystalid[(512, 512), (16, 16)](scanner.nb_blocks_per_ring * scanner.blocks.shape[1],
                                    scanner.nb_blocks_per_ring, scanner.nb_rings,
                                    scanner.inner_radius,
                                    scanner.outer_radius, scanner.axial_length,
                                    np.array(scanner.blocks.shape, dtype = np.int32),
                                    np.array(scanner.blocks.size, dtype = np.float32),
                                    listmode_pos, listmode_id)
    listmode_data = np.hstack((listmode_id, listmode.data.reshape(-1, 1)))
    return cal_sinogram(listmode_data, scanner.nb_rings, scanner.blocks.shape[2],
                        scanner.nb_blocks_per_ring * scanner.blocks.shape[1])


def sino2lm(scanner, sinogram, lors):
    index = np.where(sinogram > 0)[0]
    all_position = np.zeros((len(index), 6))
    all_position[:, :3] = get_center(scanner, lors[index, 0])
    all_position[:, 3:6] = get_center(scanner, lors[index, 1])
    return np.hstack((all_position, sinogram[index].reshape(-1, 1)))


__all__ = []

DEFAULT_GROUP_NAME = 'listmode_data'

DEFAULT_COLUMNS = ['fst', 'snd', 'weight', 'tof']

import math

import numpy as np
from numba import jit, cuda


@cuda.jit(device = True)
def norm(x, y):
    return (x ** 2 + y ** 2) ** 0.5


@cuda.jit(device = True)
def get_block_theta(event):
    before_theta = event[0] / norm(event[0], event[1])
    return math.acos(before_theta) / math.pi * 180


@cuda.jit(device = True)
def block_id(event, nb_crystals_per_ring, block_grid_y, num_block):
    theta_event = get_block_theta(event)
    if event[1] < 0:
        theta_event = 360 - theta_event
    fixed_theta = (theta_event + 180 / nb_crystals_per_ring * block_grid_y) % 360
    return int(fixed_theta / 360 * num_block)


@cuda.jit(device = True)
def crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius, grid_block,
               size_block, event):
    id_block = block_id(event, nb_crystals_per_ring, grid_block[1], nb_blocks_per_ring)
    theta_block = 2 * math.pi * id_block / nb_blocks_per_ring
    block_center_x = (inner_radius + outer_radius) / 2 * math.cos(theta_block)
    block_center_y = (inner_radius + outer_radius) / 2 * math.sin(theta_block)
    dist = (block_center_x - event[0]) * math.cos(theta_block) + (
            block_center_y - event[1]) * math.sin(theta_block)
    pro_point_x = dist * math.cos(theta_block) + event[0]
    pro_point_y = dist * math.sin(theta_block) + event[1]
    edge_block_x = block_center_x - math.cos(theta_block + math.pi / 2) * size_block[1] / 2
    edge_block_y = block_center_y - math.sin(theta_block + math.pi / 2) * size_block[1] / 2
    id_crystal = int(
        norm(pro_point_x - edge_block_x, pro_point_y - edge_block_y) / size_block[1] * grid_block[
            1])
    return id_block * grid_block[1] + id_crystal


@cuda.jit(device = True)
def ring_id(nb_rings, axial_length, grid_block, size_block, event_z):
    return int((event_z + nb_rings * size_block[2] / 2) / axial_length * nb_rings * grid_block[2])


@cuda.jit(device = True)
def position2crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings, inner_radius,
                       outer_radius, axial_length, grid_block,
                       size_block, event):
    ring_id1 = ring_id(nb_rings, axial_length, grid_block, size_block, event[2])
    crystal_id1 = crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius,
                             grid_block, size_block, event[:2])
    ring_id2 = ring_id(nb_rings, axial_length, grid_block, size_block, event[5])
    crystal_id2 = crystal_id(nb_crystals_per_ring, nb_blocks_per_ring, inner_radius, outer_radius,
                             grid_block, size_block, event[3:5])
    return (ring_id1 * nb_crystals_per_ring + crystal_id1), (
            ring_id2 * nb_crystals_per_ring + crystal_id2)


@cuda.jit(device = True)
def change_id(id1, id2):
    if id1 < id2:
        return id2, id1
    else:
        return id1, id2


@cuda.jit
def crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings, inner_radius, outer_radius,
              axial_length, grid_block, size_block, events, final_id):
    c, j = cuda.grid(2)
    i = 512 * 16 * c + j
    if i < events.shape[0]:
        id1, id2 = position2crystalid(nb_crystals_per_ring, nb_blocks_per_ring, nb_rings,
                                      inner_radius, outer_radius, axial_length,
                                      grid_block, size_block, events[i, :])
        fixed_id1, fixed_id2 = change_id(id1, id2)
        final_id[i, 0] = fixed_id1
        final_id[i, 1] = fixed_id2


@jit(nopython = True)
def cal_sinogram(list_mode_data, nb_rings, grid_block_z, nb_crystals_per_ring):
    total_num_crystal = nb_rings * grid_block_z * nb_crystals_per_ring
    weight = np.zeros((int(total_num_crystal * (total_num_crystal - 1) / 2), 1))
    for i in range(int(list_mode_data.shape[0])):
        pos = int((list_mode_data[i, 0] - 1) * list_mode_data[i, 0] / 2 + list_mode_data[i, 1])
        weight[pos] = weight[pos] + list_mode_data[i, 2]
    return weight


@jit(nopython = True)
def get_all_lors_id(total_num_crystal):
    lors = np.zeros((int(total_num_crystal * (total_num_crystal - 1) / 2), 2), dtype = np.int32)
    num = 0
    for i in range(total_num_crystal):
        for j in range(i):
            lors[num, 0] = i
            lors[num, 1] = j
            num = num + 1
    return lors


def get_crystal_xy(grid_block, num_block, crystal_id, block_id, r_inner, size_block):
    size_crystal = size_block / grid_block
    angle_block = math.pi * 2 / num_block * block_id
    center = np.hstack((((r_inner) * np.cos(angle_block)).reshape(block_id.size, 1),
                        ((r_inner) * np.sin(angle_block)).reshape(block_id.size, 1)))
    vector_tang = np.hstack(
        (np.cos(angle_block + math.pi / 2 * np.ones_like(angle_block)).reshape(block_id.size, 1),
         np.sin(angle_block + math.pi / 2 * np.ones_like(angle_block)).reshape(block_id.size, 1)))
    dist_from_center = size_crystal[1] * (
            crystal_id - (grid_block[1] - 1) / 2 * np.ones_like(crystal_id))
    return center + (np.hstack((dist_from_center.reshape(dist_from_center.size, 1),
                                dist_from_center.reshape(dist_from_center.size, 1)))
                     * vector_tang)


def get_crystal_z(ring_id, grid_block, size_block, nb_rings):
    size_crystal = size_block / grid_block
    dist_from_center = (ring_id - (grid_block[2] * nb_rings - 1) / 2 * np.ones_like(ring_id)) * \
                       size_crystal[2]
    return dist_from_center


def get_center(scanner, crystal_id_whole_scanner):
    ring_id = np.floor(
        crystal_id_whole_scanner / scanner.nb_blocks_per_ring / scanner.blocks.shape[1])
    crystal_per_ring_id = crystal_id_whole_scanner - ring_id * scanner.nb_blocks_per_ring * \
                          scanner.blocks.shape[1]
    block_id = crystal_per_ring_id // scanner.blocks.shape[1]
    crystal_id = crystal_per_ring_id % scanner.blocks.shape[1]
    center_xy = get_crystal_xy(np.array(scanner.blocks.shape), scanner.nb_blocks_per_ring,
                               crystal_id, block_id, scanner.inner_radius,
                               np.array(scanner.blocks.size))
    center_z = get_crystal_z(ring_id, np.array(scanner.blocks.shape), np.array(scanner.blocks.size),
                             scanner.nb_rings)
    return np.hstack((center_xy, center_z.reshape(center_z.size, 1)))


# @jit
def lm2sino(listmode, scanner):
    listmode_pos = listmode.lors.data
    listmode_id = np.zeros((listmode_pos.shape[0], 2), dtype = np.int32)
    crystalid[(512, 512), (16, 16)](scanner.nb_blocks_per_ring * scanner.blocks.shape[1],
                                    scanner.nb_blocks_per_ring, scanner.nb_rings,
                                    scanner.inner_radius,
                                    scanner.outer_radius, scanner.axial_length,
                                    np.array(scanner.blocks.shape, dtype = np.int32),
                                    np.array(scanner.blocks.size, dtype = np.float32),
                                    listmode_pos, listmode_id)
    listmode_data = np.hstack((listmode_id, listmode.data.reshape(-1, 1)))
    return cal_sinogram(listmode_data, scanner.nb_rings, scanner.blocks.shape[2],
                        scanner.nb_blocks_per_ring * scanner.blocks.shape[1])


def sino2lm(scanner, sinogram, lors):
    index = np.where(sinogram > 0)[0]
    all_position = np.zeros((len(index), 6))
    all_position[:, :3] = get_center(scanner, lors[index, 0])
    all_position[:, 3:6] = get_center(scanner, lors[index, 1])
    return np.hstack((all_position, sinogram[index].reshape(-1, 1)))


__all__ = []

DEFAULT_GROUP_NAME = 'listmode_data'

DEFAULT_COLUMNS = ['fst', 'snd', 'weight', 'tof']