优化Python中大量球体无重叠随机运动模拟的策略

tree = cKDTree(centers)
potential_neighbors = [tree.query_ball_point(center, search_radius) for center in updated_centers]

tree = cKDTree(centers)
potential_neighbors_batch = tree.query_ball_point(updated_centers, 2*r_spheres + 2*motion_magnitude, workers=-1)

potential_neighbors_batch = tree.query_ball_point(updated_centers, 2*r_spheres + 2*motion_magnitude, workers=-1)

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

# 定义空间边界参数 (示例值，实际应用中需根据需求设定)
Rmax = 100.0
Zmin = -50.0
Zmax = 50.0

@nb.njit()
def in_cylinder(point, Rmax, Zmin, Zmax):
    """
    检查一个点是否在圆柱体空间边界内。
    使用Numba JIT编译，并优化距离计算（避免np.sqrt）。
    point: 单个点的坐标 (x, y, z)
    """
    # 假设point是(x, y, z)数组
    radial_distances_sq = point[0]**2 + point[1]**2
    return (radial_distances_sq <= Rmax ** 2) & (Zmin <= point[2]) & (point[2] <= Zmax)

@nb.njit()
def generate_random_vector(max_magnitude):
    """
    生成一个随机方向和大小的3D向量。
    使用Numba JIT编译。
    """
    direction = np.random.randn(3)
    # 确保方向向量非零，避免除以零
    norm_direction = np.linalg.norm(direction)
    if norm_direction == 0:
        return np.zeros(3) # 或者重新生成
    direction /= norm_direction

    magnitude = np.random.uniform(0, max_magnitude)
    return direction * magnitude

@nb.njit()
def euclidean_distance(vec_a, vec_b):
    """
    计算两个3D向量之间的欧几里得距离。
    使用Numba JIT编译。
    """
    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_neighbors_centers, neighbors_indices, threshold_distance, ignore_idx):
    """
    检查新球心是否与任何潜在邻居重叠。
    使用Numba JIT编译，加速循环和距离计算。
    new_center: 提议的新球心位置
    all_neighbors_centers: 所有球体的当前中心列表
    neighbors_indices: 潜在邻居的索引列表
    threshold_distance: 重叠判断的距离阈值 (2 * r_spheres)
    ignore_idx: 当前移动球体的索引，用于避免与自身比较
    """
    for neighbor_idx in neighbors_indices:
        if neighbor_idx == ignore_idx:
            continue # 忽略自身
        distance = euclidean_distance(new_center, all_neighbors_centers[neighbor_idx])
        if distance < threshold_distance:
            return True # 发现重叠
    return False # 无重叠

def move_spheres_optimized(centers, r_spheres, motion_coef, N_motions):
    """
    优化后的球体随机运动模拟函数。
    centers: 初始球心数组
    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 # 两个球体不重叠的最小距离

    print(f"开始模拟 {n_spheres} 个球体的 {N_motions} 步运动...")

    for step in range(N_motions):
        # 1. 构建KDTree并进行批量邻居查询 (利用多核)
        # 搜索半径应覆盖最大可能的位移和球体直径，以确保找到所有潜在碰撞
        search_radius = overlap_threshold + 2 * motion_magnitude # 考虑球体直径和最大位移
        tree = cKDTree(updated_centers)
        # 使用workers=-1启用所有CPU核心进行并行查询
        potential_neighbors_batch = tree.query_ball_point(updated_centers, search_radius, workers=-1)

        updated_this_step = np.zeros(n_spheres, dtype=bool)

        for i in range(n_spheres):
            # 2. 生成随机位移向量 (Numba加速)
            vector = generate_random_vector(motion_magnitude)
            new_center = updated_centers[i] + vector

            # 3. 检查空间边界 (Numba加速)
            if in_cylinder(new_center, Rmax, Zmin, Zmax):
                # 获取当前球体的潜在邻居索引
                # cKDTree.query_ball_point返回的是列表的列表，需要转换为numpy数组
                neighbors_indices = np.array(potential_neighbors_batch[i])

                # 4. 检查重叠 (Numba加速)
                overlap = any_neighbor_in_range(new_center, updated_centers, neighbors_indices, overlap_threshold, i)

                # 5. 如果没有重叠且在边界内，则更新球心
                if not overlap:
                    updated_centers[i] = new_center
                    updated_this_step[i] = True
            # else:
            #     print(f"球体 {i} 移出边界") # 调试信息，通常在生产代码中移除

        num_updated = np.sum(updated_this_step)
        print(f"步数 {step+1}/{N_motions}: 成功移动 {num_updated}/{n_spheres} 个球体 ({num_updated/n_spheres:.2%})")

    print("模拟完成。")
    return updated_centers

# 示例使用
if __name__ == '__main__':
    # 模拟参数
    num_spheres = 10000 # 示例使用较小数量，百万级别需要更长时间
    sphere_radius = 1.0
    motion_coefficient = 0.1 # 最大位移是半径的10%
    num_motions = 5

    # 初始球心：随机分布在一个圆柱体内，确保不重叠
    # 这是一个简化的初始生成，实际应用中可能需要更复杂的非重叠生成算法
    # 这里我们只是随机生成，不保证初始不重叠，但在move_spheres中会处理重叠
    initial_centers = np.random.rand(num_spheres, 3) * [Rmax, Rmax, Zmax - Zmin]
    initial_centers[:, 0] -= Rmax / 2
    initial_centers[:, 1] -= Rmax / 2
    initial_centers[:, 2] += Zmin

    # 确保初始球心在边界内（如果随机生成可能超出）
    # 这一步可以根据实际需求进行调整，例如拒绝超出边界的初始球心
    valid_indices = [i for i, center in enumerate(initial_centers) if in_cylinder(center, Rmax, Zmin, Zmax)]
    initial_centers = initial_centers[valid_indices[:num_spheres]] # 确保数量不超过num_spheres
    print(f"初始有效球体数量: {len(initial_centers)}")

    # 运行优化后的模拟
    final_centers = move_spheres_optimized(initial_centers, sphere_radius, motion_coefficient, num_motions)

    # 可以进一步分析 final_centers，例如可视化或检查重叠
    print(f"最终球心数据形状: {final_centers.shape}")