numpy 转 torch , GPU 并行计算

代码小白, 写代码的时候直接使用 numpy , 没考虑到之后要用 GPU

目前是打算用 Colab 跑这些代码

求各位大神帮帮忙, 怎么改动才能把这些需要多次调用的函数使用 GPU 计算 (在后面的代码中也会用到 PSO 之类的东西)

import numpy as np
import pandas as pd

import warnings
warnings.filterwarnings('ignore')


# @title 阴影确定

# ### 1]. 判断点是否在阴影内
# Displaying all the modified functions
def is_left(P0, P1, P2):
    """Determine if point P2 is to the left of the line segment from P0 to P1."""
    return (P1[0] - P0[0]) * (P2[1] - P0[1]) - (P2[0] - P0[0]) * (P1[1] - P0[1])

# 绕射法判断点是否在多边形内部
def winding_number_revised(P, polygon_corners):
    """Revised winding number function."""
    wn = 0
    num_vertices = len(polygon_corners)

    # Iterate through the edges of the polygon
    for i in range(num_vertices - 1):
        if polygon_corners[i][1] <= P[1] < polygon_corners[i+1][1]:  # Upward crossing
            if is_left(polygon_corners[i], polygon_corners[i+1], P) > 0:
                wn += 1
        elif polygon_corners[i+1][1] <= P[1] < polygon_corners[i][1]:  # Downward crossing
            if is_left(polygon_corners[i], polygon_corners[i+1], P) < 0:
                wn -= 1

    # Convert wn to binary (inside/outside) value
    return 1 if wn != 0 else 0



# ### 2]. 角点计算
# 角点排序
def sort_corners(unsorted_corners, center):
    """Sort corners in counter-clockwise order."""
    vectors = [np.array(corner) - np.array(center) for corner in unsorted_corners]
    angles = [np.arctan2(vector[1], vector[0]) for vector in vectors]
    sorted_corners = [corner for _, corner in sorted(zip(angles, unsorted_corners))]
    return sorted_corners

# 确保凸四边形的情况下计算角点
def debug_compute_corners_corrected(center, normal, K, L):
    """Compute corners ensuring they form a convex parallelogram."""
    v = np.array([0, 0, 1])  # Assuming z-axis as a reference vector
    if np.allclose(normal, v) or np.allclose(normal, -v):  # If normal is aligned with z-axis, take another reference
        v = np.array([1, 0, 0])

    dir1 = np.cross(normal, v)
    dir1 = dir1 / np.linalg.norm(dir1)
    dir2 = np.cross(normal, dir1)
    dir2 = dir2 / np.linalg.norm(dir2)

    half_K = K / 2
    half_L = L / 2
    corner1 = center + half_K * dir1 + half_L * dir2
    corner2 = center + half_K * dir1 - half_L * dir2
    corner3 = center - half_K * dir1 - half_L * dir2  # Swapped the ordering to ensure convex shape
    corner4 = center - half_K * dir1 + half_L * dir2  # Swapped the ordering to ensure convex shape

    return corner1, corner2, corner3, corner4

# 存储角点
def compute_matrix_corners_sorted_v4(data, K, L):
    """Computes the matrix corners with the corrected function and includes z-coordinates."""
    corners = data.apply(lambda row: debug_compute_corners_corrected([row["x(m)"], row["y(m)"], row["z(m)"]],
                                                                     [row["x_normal"], row["y_normal"], row["z_normal"]],
                                                                     K, L), axis=1)

    # Sort the corners
    corners = corners.apply(lambda x: sort_corners(x, [0, 0, 0]))

    data['Corner1_x'] = corners.apply(lambda x: x[0][0])
    data['Corner1_y'] = corners.apply(lambda x: x[0][1])
    data['Corner1_z'] = corners.apply(lambda x: x[0][2])
    data['Corner2_x'] = corners.apply(lambda x: x[1][0])
    data['Corner2_y'] = corners.apply(lambda x: x[1][1])
    data['Corner2_z'] = corners.apply(lambda x: x[1][2])
    data['Corner3_x'] = corners.apply(lambda x: x[2][0])
    data['Corner3_y'] = corners.apply(lambda x: x[2][1])
    data['Corner3_z'] = corners.apply(lambda x: x[2][2])
    data['Corner4_x'] = corners.apply(lambda x: x[3][0])
    data['Corner4_y'] = corners.apply(lambda x: x[3][1])
    data['Corner4_z'] = corners.apply(lambda x: x[3][2])

    return data



# ### 3]. 反射光线计算
# 每枚镜子反射光线计算
def compute_reflected_light_direction(V, B, data_with_plane_equation):
    """Compute the direction of the reflected light."""
    # 通过镜子座标检索 data_with_plane_equation 中对应的镜子法向量
    matching_row = data_with_plane_equation[(data_with_plane_equation["x(m)"] == B[0]) &
                                            (data_with_plane_equation["y(m)"] == B[1]) &
                                            (data_with_plane_equation["z(m)"] == B[2])]

    if matching_row.empty:
        return None

    # 提取镜子的法向量
    normal = np.array([matching_row["x_normal"].values[0],
                       matching_row["y_normal"].values[0],
                       matching_row["z_normal"].values[0]])

    # 计算反射光线的方向向量
    V_reflected = V - 2 * np.dot(V, normal) * normal

    return V_reflected


# ### 4]. 交点转化判断
# 计算平面方程
def compute_plane_equation(data):
    """Compute the equation of the plane for each mirror."""
    A_values, B_values, C_values, D_values = [], [], [], []
    for index, row in data.iterrows():
        P1 = np.array([row["Corner1_x"], row["Corner1_y"], row["Corner1_z"]])
        P2 = np.array([row["Corner2_x"], row["Corner2_y"], row["Corner2_z"]])
        P3 = np.array([row["Corner3_x"], row["Corner3_y"], row["Corner3_z"]])
        v1 = P2 - P1
        v2 = P3 - P1
        n = np.cross(v1, v2)
        D = -np.dot(n, P1)
        A_values.append(n[0])
        B_values.append(n[1])
        C_values.append(n[2])
        D_values.append(D)
    updated_data = data.copy()
    updated_data['A'] = A_values
    updated_data['B'] = B_values
    updated_data['C'] = C_values
    updated_data['D'] = D_values
    return updated_data

# 计算线面交点
def compute_intersection(P0, d, plane_equation):
    """Compute the intersection point of line and plane."""
    A, B, C, D = plane_equation
    n = np.array([A, B, C])
    P0 = np.array(P0)
    t = -(np.dot(n, P0) + D) / np.dot(n, d)
    intersection = P0 + t * d
    return intersection

# 计算反射光线
def reflective_light_function_v2(x, y, z, d, plane_equation_df):
    """Determine if light is reflected from a given point to the viewer."""
    for _, row in plane_equation_df.iterrows():
        A, B, C, D = row['A'], row['B'], row['C'], row['D']
        corners = [
            (row["Corner1_x"], row["Corner1_y"]),
            (row["Corner2_x"], row["Corner2_y"]),
            (row["Corner3_x"], row["Corner3_y"]),
            (row["Corner4_x"], row["Corner4_y"]),
            (row["Corner1_x"], row["Corner1_y"])
        ]
        P0 = [x, y, z]
        intersection_point = compute_intersection(P0, d, (A, B, C, D))
        if winding_number_revised(intersection_point[:2], corners):
            return 1
    return 0


# ### 5]. 功能函数
# Displaying the remaining modified functions
def find_points_in_hemisphere(B, df, R):
    """Find mirrors that are within a hemisphere facing point B."""
    x_B, y_B, z_B = B.flatten()  # Flattening to ensure it's a 1D array
    vec_OB = np.array([x_B, y_B, z_B])
    rows_list = []

    for index, row in df.iterrows():
        x, y, z = row['x(m)'], row['y(m)'], row['z(m)']
        vec_P = np.array([x - x_B, y - y_B, z - z_B])
        distance_to_B = np.linalg.norm(vec_P)
        dot_product = np.dot(vec_OB, vec_P)

        if distance_to_B <= R and dot_product < 0:
            new_row = {
                'x(m)': row['x(m)'],
                'y(m)': row['y(m)'],
                'z(m)': row['z(m)'],
                'x_normal': row['x_normal'],
                'y_normal': row['y_normal'],
                'z_normal': row['z_normal']
            }
            rows_list.append(new_row)

    df_inside_hemisphere = pd.DataFrame(rows_list)
    return df_inside_hemisphere

# 分块扫描并记录
def block_scanning(data_all, data, plane_equation_data, V, num_blocks, K, L):
    """Block scanning function with corrected corner computations."""
    blocks = np.zeros((num_blocks, num_blocks))

    matching_row = data_all[(data_all["x(m)"] == data[0]) &
                            (data_all["y(m)"] == data[1]) &
                            (data_all["z(m)"] == data[2])]

    if matching_row.empty:
        return None

    normal = np.array([matching_row["x_normal"].values[0], matching_row["y_normal"].values[0], matching_row["z_normal"].values[0]])
    center = np.array([matching_row["x(m)"].values[0], matching_row["y(m)"].values[0], matching_row["z(m)"].values[0]])

    # Compute corners using the corrected function
    corner1, corner2, corner3, corner4 = debug_compute_corners_corrected(center, normal, K, L)

    corners = [corner1, corner2, corner3, corner4]

    dir1 = np.array(corner2) - np.array(corner1)
    dir2 = np.array(corner4) - np.array(corner1)

    step1 = dir1 / num_blocks
    step2 = dir2 / num_blocks

    d = compute_reflected_light_direction(V, data, data_all)

    for i in range(num_blocks):
        for j in range(num_blocks):
            block_center = np.array(corner1) + (i + 0.5) * step1 + (j + 0.5) * step2
            x, y, z = block_center

            A, B, C, D = matching_row["A"].values[0], matching_row["B"].values[0], matching_row["C"].values[0], matching_row["D"].values[0]
            z = (-D - A*x - B*y) / C

            # The reflective light function is omitted for now, but can be added back.
            result = reflective_light_function_v2(x, y, z, d, plane_equation_data)
            blocks[i, j] = result

    ratio = np.sum(blocks) / np.size(blocks)
    return ratio