268 lines
9.2 KiB
C++
268 lines
9.2 KiB
C++
#include <cmath>
|
||
#include <cstdlib>
|
||
#include <iomanip>
|
||
#include <iostream>
|
||
#include <vector>
|
||
#include <mpi.h>
|
||
|
||
using namespace std;
|
||
|
||
// 物理常量
|
||
const double G = 6.67430e-11; // 引力常数
|
||
const double DT = 0.01; // 时间步长
|
||
const int TMAX = 100; // 总时间步数
|
||
const double mass_scale = 1e24; // 质量缩放因子
|
||
const double dist_scale = 1e8; // 距离缩放因子
|
||
const double vel_scale = 1e3; // 速度缩放因子
|
||
|
||
// 三维向量结构体
|
||
struct Vec3 {
|
||
double x, y, z;
|
||
|
||
Vec3() : x(0), y(0), z(0) {}
|
||
Vec3(double x, double y, double z) : x(x), y(y), z(z) {}
|
||
|
||
Vec3 operator+(const Vec3 &other) const {
|
||
return Vec3(x + other.x, y + other.y, z + other.z);
|
||
}
|
||
|
||
Vec3 operator-(const Vec3 &other) const {
|
||
return Vec3(x - other.x, y - other.y, z - other.z);
|
||
}
|
||
|
||
Vec3 operator*(double scalar) const {
|
||
return Vec3(x * scalar, y * scalar, z * scalar);
|
||
}
|
||
|
||
double magnitude() const {
|
||
return sqrt(x * x + y * y + z * z);
|
||
}
|
||
};
|
||
|
||
// 天体结构体
|
||
struct Body {
|
||
double mass; // 质量
|
||
Vec3 position; // 位置
|
||
Vec3 velocity; // 速度
|
||
};
|
||
|
||
// 初始化天体系统
|
||
void init_bodies(vector<Body> &bodies, int n, bool verbose = false) {
|
||
// 中心天体（类似太阳）
|
||
bodies[0].mass = 1000 * mass_scale;
|
||
bodies[0].position = Vec3(0, 0, 0);
|
||
bodies[0].velocity = Vec3(0, 0, 0);
|
||
|
||
// 其他天体（类似行星）
|
||
for (int i = 1; i < n; i++) {
|
||
bodies[i].mass = (1.0 + i * 0.5) * mass_scale;
|
||
double angle = 2.0 * M_PI * i / n;
|
||
double radius = (1.0 + i * 0.5) * dist_scale;
|
||
|
||
bodies[i].position = Vec3(radius * cos(angle), radius * sin(angle), 0.0);
|
||
|
||
// 给予切向速度以形成轨道
|
||
double orbital_speed = sqrt(G * bodies[0].mass / radius);
|
||
bodies[i].velocity = Vec3(-orbital_speed * sin(angle),
|
||
orbital_speed * cos(angle), 0.0);
|
||
}
|
||
|
||
// 输出初始状态
|
||
if (verbose) {
|
||
cout << fixed << setprecision(6);
|
||
cout << "\n初始状态:" << endl;
|
||
for (int i = 0; i < n; i++) {
|
||
cout << "天体 " << i << ": 质量=" << bodies[i].mass / mass_scale
|
||
<< "e24 kg, "
|
||
<< "位置=(" << bodies[i].position.x / dist_scale << ", "
|
||
<< bodies[i].position.y / dist_scale << ", "
|
||
<< bodies[i].position.z / dist_scale << ")e8 m" << endl;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 计算local_particles中每个物体受到all_particles中所有物体的作用力
|
||
// 并更新local_particles中物体的速度和位置
|
||
void compute_local_forces(vector<Body>& local_particles,
|
||
const vector<Body>& all_particles,
|
||
int local_start) {
|
||
for (size_t i = 0; i < local_particles.size(); i++) {
|
||
Vec3 total_force(0, 0, 0);
|
||
int global_idx = local_start + i;
|
||
|
||
// 计算all_particles中所有物体对local_particles[i]的作用力
|
||
for (size_t j = 0; j < all_particles.size(); j++) {
|
||
// 跳过自己
|
||
if (global_idx == static_cast<int>(j)) continue;
|
||
|
||
// 计算从物体i指向物体j的向量
|
||
Vec3 r_vec = all_particles[j].position - local_particles[i].position;
|
||
double distance = r_vec.magnitude();
|
||
|
||
// 避免除以零
|
||
if (distance < 1e-10) continue;
|
||
|
||
// 计算引力大小
|
||
double force_magnitude = G * local_particles[i].mass * all_particles[j].mass
|
||
/ (distance * distance);
|
||
|
||
// 计算力的方向并累加
|
||
Vec3 force_direction = r_vec / distance;
|
||
total_force = total_force + force_direction * force_magnitude;
|
||
}
|
||
|
||
// 更新local_particles[i]的速度和位置
|
||
Vec3 v_new = local_particles[i].velocity + total_force * DT / local_particles[i].mass;
|
||
Vec3 x_new = local_particles[i].position + v_new * DT;
|
||
|
||
local_particles[i].velocity = v_new;
|
||
local_particles[i].position = x_new;
|
||
}
|
||
}
|
||
|
||
// 获取每个进程负责的天体信息
|
||
void get_rank_info(int rank_id, int bodies_count, int world_size,
|
||
int& send_size, int& send_offset) {
|
||
int particles_per_proc = bodies_count / world_size;
|
||
int remainder = bodies_count % world_size;
|
||
|
||
if (rank_id < remainder) {
|
||
send_size = particles_per_proc + 1;
|
||
send_offset = rank_id * (particles_per_proc + 1);
|
||
} else {
|
||
send_size = particles_per_proc;
|
||
send_offset = rank_id * particles_per_proc + remainder;
|
||
}
|
||
}
|
||
|
||
int main(int argc, char **argv) {
|
||
MPI_Init(&argc, &argv);
|
||
|
||
// 获取进程数量和当前进程rank
|
||
int world_size, world_rank;
|
||
bool verbose = false;
|
||
MPI_Comm_size(MPI_COMM_WORLD, &world_size);
|
||
MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
|
||
|
||
// 从命令行参数获取天体数量
|
||
int n = 4; // 默认4个天体
|
||
if (argc > 1) {
|
||
n = atoi(argv[1]);
|
||
}
|
||
if (argc > 2) {
|
||
verbose = (strcmp(argv[2], "--verbose") == 0 || strcmp(argv[2], "-v") == 0);
|
||
}
|
||
|
||
// 只有rank 0打印初始信息
|
||
if (world_rank == 0) {
|
||
cout << "N体问题并行模拟" << endl;
|
||
cout << "天体数量: " << n << endl;
|
||
cout << "进程数量: " << world_size << endl;
|
||
cout << "时间步长: " << DT << " s" << endl;
|
||
cout << "总步数: " << TMAX << endl;
|
||
cout << "----------------------------------------" << endl;
|
||
}
|
||
|
||
// 定义Body的MPI数据类型
|
||
// Body结构包含: mass(1) + position(3) + velocity(3) = 7个double
|
||
MPI_Datatype MPI_BODY;
|
||
MPI_Type_contiguous(7, MPI_DOUBLE, &MPI_BODY);
|
||
MPI_Type_commit(&MPI_BODY);
|
||
|
||
// 步骤1: 获取分配给本进程的物体的初始信息local_particles
|
||
// 步骤2: 获取应用程序中所有物体的信息all_particles
|
||
|
||
vector<Body> all_particles(n);
|
||
vector<Body> local_particles;
|
||
|
||
// 计算每个进程分配到的物体数量
|
||
int particles_per_proc = n / world_size;
|
||
int remainder = n % world_size;
|
||
|
||
int local_start, local_count;
|
||
if (world_rank < remainder) {
|
||
local_count = particles_per_proc + 1;
|
||
local_start = world_rank * local_count;
|
||
} else {
|
||
local_count = particles_per_proc;
|
||
local_start = world_rank * particles_per_proc + remainder;
|
||
}
|
||
|
||
// Rank 0初始化所有物体
|
||
if (world_rank == 0) {
|
||
init_bodies(all_particles, n, verbose);
|
||
}
|
||
|
||
// 广播所有物体的初始信息到所有进程
|
||
MPI_Bcast(all_particles.data(), n, MPI_BODY, 0, MPI_COMM_WORLD);
|
||
|
||
// 每个进程提取自己负责的物体
|
||
local_particles.resize(local_count);
|
||
for (int i = 0; i < local_count; i++) {
|
||
local_particles[i] = all_particles[local_start + i];
|
||
}
|
||
|
||
if (world_rank == 0) {
|
||
cout << "\n开始模拟..." << endl;
|
||
}
|
||
|
||
// 创建发送和接收缓冲区信息
|
||
vector<int> all_send_size(world_size);
|
||
vector<int> all_send_offset(world_size);
|
||
|
||
for (int r = 0; r < world_size; r++) {
|
||
get_rank_info(r, n, world_size, all_send_size[r], all_send_offset[r]);
|
||
}
|
||
|
||
double start_time = MPI_Wtime();
|
||
vector<Body> send_buf(local_count);
|
||
|
||
// 主循环：N体模拟
|
||
for (int t = 0; t < TMAX; t++) {
|
||
// 计算所有物体对分配给本进程的物体的作用力
|
||
// 并据此更新local_particles的本进程的物体信息
|
||
compute_local_forces(local_particles, all_particles, local_start);
|
||
|
||
// 将本进程信息local_particles保存到发送缓冲区send_buf
|
||
// 同时更新all_particles中的部分信息
|
||
send_buf = local_particles;
|
||
|
||
// 更新all_particles中本进程负责的部分信息
|
||
for (int i = 0; i < local_count; i++) {
|
||
all_particles[local_start + i] = local_particles[i];
|
||
}
|
||
|
||
// 全局通信：同步所有进程的物体信息
|
||
MPI_Allgatherv(send_buf.data(), local_count,
|
||
MPI_BODY, all_particles.data(),
|
||
all_send_size.data(), all_send_offset.data(),
|
||
MPI_BODY, MPI_COMM_WORLD);
|
||
|
||
|
||
// 每10步输出一次状态（仅rank 0）
|
||
if (verbose && (t + 1) % 10 == 0 && world_rank == 0) {
|
||
cout << "时间步 " << t + 1 << ":" << endl;
|
||
for (int i = 0; i < n; i++) {
|
||
cout << " 天体 " << i << ": "
|
||
<< "位置=(" << all_particles[i].position.x / dist_scale << ", "
|
||
<< all_particles[i].position.y / dist_scale << ", "
|
||
<< all_particles[i].position.z / dist_scale << ")e8 m, "
|
||
<< "速度=(" << all_particles[i].velocity.x / vel_scale << ", "
|
||
<< all_particles[i].velocity.y / vel_scale << ", "
|
||
<< all_particles[i].velocity.z / vel_scale << ")e3 m/s" << endl;
|
||
}
|
||
}
|
||
}
|
||
|
||
if (world_rank == 0) {
|
||
cout << "" << endl;
|
||
double end_time = MPI_Wtime();
|
||
cout << "模拟用时: " << end_time - start_time << " 秒" << endl;
|
||
cout << "\n模拟完成!" << endl;
|
||
}
|
||
|
||
MPI_Type_free(&MPI_BODY);
|
||
MPI_Finalize();
|
||
return 0;
|
||
}
|