Python中大规模球体无重叠随机移动模拟的性能优化实践

import numpy as np
from scipy.spatial import cKDTree
import numba as nb
import math

# 假设Rmax, Zmin, Zmax是全局变量或通过参数传入
# 为了演示，这里定义一些示例值
Rmax = 100.0
Zmin = -50.0
Zmax = 50.0

@nb.njit()
def in_cylinder(point, Rmax_sq, Zmin, Zmax):
    """
    检查一个点是否在圆柱体内。
    使用Rmax_sq (Rmax的平方) 避免不必要的平方根计算。
    point: 单个点的坐标数组 [x, y, z]
    """
    # 径向距离的平方
    radial_distance_sq = point[0]**2 + point[1]**2
    return (radial_distance_sq <= Rmax_sq) & \
           (Zmin <= point[2]) & (point[2] <= Zmax)

@nb.njit()
def generate_random_vector(max_magnitude):
    """
    生成一个随机方向和随机大小的3D向量。
    """
    # 生成一个随机方向向量
    direction = np.random.randn(3)
    direction_norm = np.linalg.norm(direction)
    # 避免除以零
    if direction_norm == 0:
        direction = np.array([1.0, 0.0, 0.0]) # 默认方向
    else:
        direction /= direction_norm

    # 生成一个随机大小
    magnitude = np.random.uniform(0, max_magnitude)
    return direction * magnitude

@nb.njit()
def euclidean_distance(vec_a, vec_b):
    """
    计算两个3D向量之间的欧几里得距离。
    """
    acc = 0.0
    for i in range(vec_a.shape[0]):
        acc += (vec_a[i] - vec_b[i]) ** 2
    return math.sqrt(acc)

@nb.njit()
def any_neighbor_in_range(new_center, all_centers, neighbors_indices, threshold, ignore_idx):
    """
    检查新球体中心是否与任何潜在邻居重叠。
    all_centers: 所有球体的中心点数组
    neighbors_indices: 潜在邻居的索引列表
    threshold: 距离阈值 (2 * r_spheres)
    ignore_idx: 当前移动球体的索引，避免与自身比较
    """
    for neighbor_idx in neighbors_indices:
        if neighbor_idx == ignore_idx:
            # 忽略自身
            continue
        distance = euclidean_distance(new_center, all_centers[neighbor_idx])
        if distance < threshold:
            return True # 发现重叠
    return False # 没有重叠

def move_spheres_optimized(centers, r_spheres, motion_coef, N_motions):
    """
    优化后的球体随机移动函数。
    centers: 初始球体中心点数组 (N, 3)
    r_spheres: 球体半径
    motion_coef: 运动系数，用于计算最大移动距离
    N_motions: 模拟的总步数
    """
    n_spheres = len(centers)
    updated_centers = np.copy(centers)
    motion_magnitude = motion_coef * r_spheres
    overlap_threshold = 2 * r_spheres # 两个球体不重叠的最小距离
    Rmax_sq = Rmax ** 2 # 预计算Rmax的平方

    for motion_step in range(N_motions):
        # 每步重新构建KDTree，因为球体位置可能发生变化
        # 使用updated_centers构建KDTree
        tree = cKDTree(updated_centers)

        # 批处理查询所有球体的潜在邻居，利用多核并行
        # 查询半径为 2*r_spheres + 2*motion_magnitude，这是最大可能重叠的范围
        potential_neighbors_batch = tree.query_ball_point(
            updated_centers, 
            overlap_threshold + 2 * motion_magnitude, # 考虑最大移动距离后的潜在邻居范围
            workers=-1 # 利用所有可用CPU核心
        )

        updated_count = 0
        for i in range(n_spheres):
            # 生成随机移动向量
            vector = generate_random_vector(motion_magnitude)

            # 预测新中心位置
            new_center = updated_centers[i] + vector

            # 检查空间边界
            if in_cylinder(new_center, Rmax_sq, Zmin, Zmax):
                # 获取当前球体的潜在邻居索引
                neighbors_indices = np.array(potential_neighbors_batch[i], dtype=np.int64)

                # 检查是否与任何邻居重叠
                overlap = any_neighbor_in_range(
                    new_center, 
                    updated_centers, 
                    neighbors_indices, 
                    overlap_threshold, 
                    i
                )

                # 如果没有重叠，则更新球体位置
                if not overlap:
                    updated_centers[i] = new_center
                    updated_count += 1
            # else:
                # print('out of cylinder') # 调试信息，在生产代码中通常移除

        print(f"Motion Step {motion_step + 1}/{N_motions}: Updated {updated_count} spheres ({updated_count/n_spheres:.2%})")
    return updated_centers

# 示例用法 (需要先定义初始球体数据)
if __name__ == "__main__":
    # 示例数据
    num_spheres = 10000 # 减少数量以便快速测试
    sphere_radius = 1.0
    initial_centers = np.random.rand(num_spheres, 3) * 200 - 100 # 随机分布在 [-100, 100] 范围内

    # 确保初始球体不重叠 (此处简化，实际应用中需要更复杂的初始化过程)
    # 假设initial_centers已经是非重叠的

    motion_coefficient = 0.1 # 每次移动最大半径的10%
    num_motions = 5

    print(f"Starting simulation for {num_spheres} spheres...")
    final_centers = move_spheres_optimized(initial_centers, sphere_radius, motion_coefficient, num_motions)
    print("Simulation finished.")
    # print("Final sphere centers:\n", final_centers)