今日は何の日。燃えるゴミの日。GPUで圧力計算の並列化。

前回のあらすじ。

3Dヒートマップを、各時間ステップごとに5つのグラフをプロットするように設定します。これにより、各ステップでの圧力分布の変化を3Dで視覚的に確認できます。

燃える炎のコード。

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from mpl_toolkits.mplot3d import Axes3D

# パラメータ設定
num_particles = 1000
num_steps = 5
grid_size = 50
sigma_x = 1.0  # 横方向の広がり
sigma_y = 0.5  # 縦方向の広がり
sigma_z = 1.5  # z方向の広がり

# グリッドを作成
x_grid = np.linspace(-5, 5, grid_size)
y_grid = np.linspace(-5, 5, grid_size)
z_grid = np.linspace(-5, 5, grid_size)
x, y, z = np.meshgrid(x_grid, y_grid, z_grid)
pos = np.vstack([x.ravel(), y.ravel(), z.ravel()]).T

# パーティクルの初期位置を原点に設定
x_initial = np.zeros(num_particles)
y_initial = np.zeros(num_particles)
z_initial = np.zeros(num_particles)

# パーティクルの速度をガウス分布に基づいて生成
vx = np.random.normal(0, sigma_x, num_particles)
vy = np.random.normal(0, sigma_y, num_particles)
vz = np.random.normal(0, sigma_z, num_particles)

# 圧力分布を保存するリスト
pressure_distributions = []

# シミュレーションの実行
for step in range(num_steps):
    # パーティクルの位置を更新
    x_positions = x_initial + vx * step
    y_positions = y_initial + vy * step
    z_positions = z_initial + vz * step

    # 各パーティクルの位置に基づいて圧力分布を計算
    pressure = np.zeros((grid_size, grid_size, grid_size))
    for i in range(num_particles):
        # 各パーティクルがグリッドに与える影響をガウスカーネルで計算
        particle_pos = np.array([x_positions[i], y_positions[i], z_positions[i]])
        rv = multivariate_normal(mean=particle_pos, cov=[[sigma_x**2, 0, 0], [0, sigma_y**2, 0], [0, 0, sigma_z**2]])
        pressure += rv.pdf(pos).reshape(grid_size, grid_size, grid_size)

    pressure_distributions.append(pressure)

# 3Dプロットの設定
fig = plt.figure(figsize=(15, 15))

for step in range(num_steps):
    ax = fig.add_subplot(3, 2, step + 1, projection='3d')
    
    # グリッド座標を生成
    x_plot, y_plot, z_plot = np.meshgrid(np.linspace(-5, 5, grid_size),
                                         np.linspace(-5, 5, grid_size),
                                         np.linspace(-5, 5, grid_size))

    # 圧力の高い部分をプロット
    ax.scatter(x_plot, y_plot, z_plot, c=pressure_distributions[step].ravel(), cmap='hot', marker='o')
    ax.set_title(f'Step {step + 1}')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

plt.tight_layout()
plt.show()

GPU ハイレゾリューション。 num_particles = 5000　　Google corab GPU 計算時間117秒。

grid_size = 100

変更点:
圧力計算の並列化: CuPyでの並列化を活かし、各パーティクルの位置に対してグリッド全体を一度に計算するようにしました。このため、ループの中で各パーティクルに対する計算が一斉に行われる形になります。
cp.expの使用: 距離計算を直接行い、それをcp.expでまとめて計算し、圧力をグリッド全体で更新しています。
これにより、パーティクル数が増えてもGPUの並列処理の力をフルに活かすことができ、全体の計算速度が向上します。

import cupy as cp
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from mpl_toolkits.mplot3d import Axes3D

# パラメータ設定
num_particles = 5000
num_steps = 5
grid_size = 50
sigma_x = 1.0  # 横方向の広がり
sigma_y = 0.5  # 縦方向の広がり
sigma_z = 1.5  # z方向の広がり

# グリッドを作成
x_grid = cp.linspace(-5, 5, grid_size)
y_grid = cp.linspace(-5, 5, grid_size)
z_grid = cp.linspace(-5, 5, grid_size)
x, y, z = cp.meshgrid(x_grid, y_grid, z_grid)
pos = cp.vstack([x.ravel(), y.ravel(), z.ravel()]).T

# パーティクルの初期位置を原点に設定
x_initial = cp.zeros(num_particles)
y_initial = cp.zeros(num_particles)
z_initial = cp.zeros(num_particles)

# パーティクルの速度をガウス分布に基づいて生成
vx = cp.random.normal(0, sigma_x, num_particles)
vy = cp.random.normal(0, sigma_y, num_particles)
vz = cp.random.normal(0, sigma_z, num_particles)

# 圧力分布を保存するリスト
pressure_distributions = []

# シミュレーションの実行
for step in range(num_steps):
    # パーティクルの位置を更新
    x_positions = x_initial + vx * step
    y_positions = y_initial + vy * step
    z_positions = z_initial + vz * step

    # 各パーティクルの位置に基づいて圧力分布を計算
    pressure = cp.zeros((grid_size, grid_size, grid_size))

    # 各パーティクルの位置をグリッドに対してベクトル化して計算
    for i in range(num_particles):
        particle_pos = cp.array([x_positions[i], y_positions[i], z_positions[i]])
        dist = pos - particle_pos  # 各グリッドポイントからパーティクルまでの距離
        pressure += cp.exp(-0.5 * (dist[:, 0]**2 / sigma_x**2 +
                                   dist[:, 1]**2 / sigma_y**2 +
                                   dist[:, 2]**2 / sigma_z**2)).reshape(grid_size, grid_size, grid_size)

    pressure_distributions.append(pressure)

# 5枚目の3Dプロットの設定
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# グリッド座標を生成
x_plot, y_plot, z_plot = cp.meshgrid(cp.linspace(-5, 5, grid_size),
                                     cp.linspace(-5, 5, grid_size),
                                     cp.linspace(-5, 5, grid_size))

# 圧力の高い部分をプロット
ax.scatter(x_plot.get(), y_plot.get(), z_plot.get(), c=pressure_distributions[4].ravel().get(), cmap='hot', marker='o')
ax.set_title('Step 5')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

plt.tight_layout()
plt.show()