simulate nbody problem on a lot of solidified sand called a computer

sim.c

// simulate nbody problem on a lot of solidified sand called a computer
// $ mpicc -o nbody_dist sim.c -lm -O3
// $ mpirun -np 16 ./nbody_dist --steps 2 --n 100000 --frequency 0.01 --info

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>
#include <math.h>
#include <time.h>

struct body {
	float x, y, z;
	float vx, vy, vz;
};

int steps, n, seed = 42;
float frequency = 0.01f;
char *initial_state_file = NULL;
int gen_radius = 10;
int info_flag = 0;

void parse_args(int argc, char **argv)
{
	for (unsigned int i = 1; i < argc; i++) {
		if (strcmp(argv[i], "--steps") == 0 && i + 1 < argc) {
			steps = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--n") == 0 && i + 1 < argc) {
			n = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--seed") == 0 && i + 1 < argc) {
			seed = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--frequency") == 0 && i + 1 < argc) {
			frequency = atof(argv[++i]);
		} else if (strcmp(argv[i], "--initial-state") == 0 && i + 1 < argc) {
			initial_state_file = argv[++i];
		} else if (strcmp(argv[i], "--gen-radius") == 0 && i + 1 < argc) {
			gen_radius = atoi(argv[++i]);
		} else if (strcmp(argv[i], "--info") == 0) {
			info_flag = 1;
		}
	}
}

void load_or_generate_initial_state(struct body *bodies)
{
	if (initial_state_file) {
		// too lazy to implement
	} else {
		srand(seed);
		for (unsigned int i = 0; i < n; i++) {
			bodies[i].x = gen_radius * (2.0f * rand() / RAND_MAX - 1.0f);
			bodies[i].y = gen_radius * (2.0f * rand() / RAND_MAX - 1.0f);
			bodies[i].z = gen_radius * (2.0f * rand() / RAND_MAX - 1.0f);
			bodies[i].vx = bodies[i].vy = bodies[i].vz = 0.0f;
		}
	}
}

void compute_forces(struct body *local_bodies, 
                           struct body *all_bodies, int local_n)
{
	float G = 6.67430e-11; // thanks google
	for (unsigned int i = 0; i < local_n; i++) {
		float fx = 0.0f, fy = 0.0f, fz = 0.0f;
		for (unsigned int j = 0; j < n; j++) {
			if (j != i) {
				float dx = all_bodies[j].x - local_bodies[i].x;
				float dy = all_bodies[j].y - local_bodies[i].y;
				float dz = all_bodies[j].z - local_bodies[i].z;

				float dist = sqrt(dx * dx + dy * dy + dz * dz) + 1.0e-20;
				float force = G / (dist * dist * dist);

				fx += force * dx;
				fy += force * dy;
				fz += force * dz;
			}
		}
		local_bodies[i].vx += fx;
		local_bodies[i].vy += fy;
		local_bodies[i].vz += fz;
	}
}

void update_positions(struct body *bodies, int local_n, float dt)
{
	for (unsigned int i = 0; i < local_n; i++) {
		bodies[i].x += bodies[i].vx * dt;
		bodies[i].y += bodies[i].vy * dt;
		bodies[i].z += bodies[i].vz * dt;
	}
}

void write_to_file(struct body *bodies, int step)
{
	char filename[100];
	sprintf(filename, "nbody_step_%d.bin", step);
	FILE *file = fopen(filename, "wb");
	fwrite(bodies, sizeof(struct body), n, file);
	fclose(file);
}

int main(int argc, char **argv)
{
    int rank, size;
    double start_time, end_time;
    double local_time = 0.0;
    double *all_times = NULL;
    int local_n = 0;
    struct body *all_bodies = NULL;
    struct body *local_bodies = NULL;
	MPI_Init(&argc, &argv);
	
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	MPI_Comm_size(MPI_COMM_WORLD, &size);

	parse_args(argc, argv);

	local_n = n / size;
	
	local_bodies = malloc(local_n * sizeof(struct body));
	all_bodies = malloc(n * sizeof(struct body));

	if (rank == 0)
		load_or_generate_initial_state(all_bodies);

	MPI_Scatter(all_bodies, local_n * sizeof(struct body), MPI_BYTE,
		    local_bodies, local_n * sizeof(struct body), MPI_BYTE, 0, MPI_COMM_WORLD);

	for (unsigned int step = 0; step < steps; ++step) {
		MPI_Bcast(all_bodies, n * sizeof(struct body), MPI_BYTE, 0, MPI_COMM_WORLD);

		start_time = MPI_Wtime();

		compute_forces(local_bodies, all_bodies, local_n);
		update_positions(local_bodies, local_n, frequency);

		end_time = MPI_Wtime();

		MPI_Gather(local_bodies, local_n * sizeof(struct body), MPI_BYTE,
			   all_bodies, local_n * sizeof(struct body), MPI_BYTE, 0, MPI_COMM_WORLD);

		if (rank == 0 && info_flag)
			write_to_file(all_bodies, step);

		if (info_flag) {
            local_time = end_time - start_time;

			if (rank == 0)
				all_times = malloc(size * sizeof(double));

			MPI_Gather(&local_time, 1, MPI_DOUBLE, all_times, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);

			if (rank == 0) {
				printf("step %d\n", step);
				for (unsigned int i = 0; i < size; i++)
					printf("rank %d: %fs\n", i, all_times[i]);
				free(all_times);
			}
		}
	}

	free(all_bodies);
	free(local_bodies);
	MPI_Finalize();

	return 0;
}