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